;;; Copyright (C) Peter McGoron 2024
;;; This program is free software: you can redistribute it and/or modify
;;; it under the terms of the GNU General Public License as published by
;;; the Free Software Foundation, version 3 of the License.
;;; 
;;; This program is distributed in the hope that it will be useful,
;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;;; GNU General Public License for more details.
;;; 
;;; You should have received a copy of the GNU General Public License
;;; along with this program.  If not, see <https://www.gnu.org/licenses/>.

;;; Persistent AVL sets and maps using JOIN.

;;; ;;;;;;;;;;;;;;;;
;;; Nodes, direction
;;; ;;;;;;;;;;;;;;;;

;;; Returns the slot number for
;;; =: The node data
;;; h: the height
;;; <: the left child
;;; >: the right child
(define %set:accessor
  (lambda (sym)
    (cond
      ((eq? sym '=) 0)
      ((eq? sym 'h) 1)
      ((eq? sym '<) 2)
      ((eq? sym '>) 3)
      (else (error "invalid direction")))))

;;; Gets data from node value given accessor symbol.
(define set:get
  (lambda (t sym)
    (let ((acc (%set:accessor sym)))
      (if (null? t)
          (cond
            ((eq? sym 'h) 0)
            (else '()))
          (vector-ref t (%set:accessor sym))))))

(define %set->sexpr
  (lambda (node)
    (if (null? node)
        '()
        (list (list 'data (set:get node '=))
              (list '< (%set->sexpr (set:get node '<)))
              (list '> (%set->sexpr (set:get node '>)))))))

(cond-expand
 (chicken
  (import (chicken process))
  (define (set:display set)
    (call-with-output-pipe
        "dot -T png > /tmp/dot.png"
      (lambda (out)
        (letrec
            ((D
              (lambda (o)
                (display o out)
                (display o)))
             (dump-data
              (lambda (tree)
                (D "\"")
                (D (set:get tree '=))
                (D "\"")))
             (connect
              (lambda (parent child)
                (if (null? child)
                    '()
                    (begin
                      (dump-data parent)
                      (D " -> ")
                      (dump-data child)
                      (D " ;\n")
                      (loop child)))))
             (loop
              (lambda (parent)
                (if (null? parent)
                    '()
                    (begin
                      (connect parent (set:get parent '<))
                      (connect parent (set:get parent '>)))))))
          (D "digraph t {\n")
          (loop set)
          (D "}\n"))))
    (system "imv-wayland /tmp/dot.png")))
 (else
  (define set:display
    (lambda (set)
      (display (%set->sexpr set))
      (newline)))))

;;; Get the height of a node, handling the empty node.
(define %set:height
  (lambda (node)
    (if (null? node)
        0
        (set:get node 'h))))

;;; Get the difference between the heights of two trees.
(define %set:height-diff
  (lambda (t1 t2)
    (- (%set:height t1) (%set:height t2))))

;;; Get the balance factor of a tree.
(define %set:bal
  (lambda (t) (%set:height-diff (set:get t '<)
                                (set:get t '>))))

;;; Set data in node given accessor symbol.
(define %set:set!
  (lambda (node dir x)
    (vector-set! node (%set:accessor dir) x)))

;;; Construct a new tree with data VAL.
(define %set:node
  (lambda (val dir1 node1 dir2 node2)
    (let ((node (vector val (+ 1
                               (max (%set:height node1)
                                    (%set:height node2)))
                        '() '())))
      (%set:set! node dir1 node1)
      (%set:set! node dir2 node2)
      node)))

(define set:node-new-val
  (lambda (node newval)
    (%set:node newval
               '< (set:get node '<)
               '> (set:get node '>))))

(define %set:invdir
  (lambda (dir)
    (cond
      ((eq? dir '<) '>)
      ((eq? dir '>) '<)
      (else (error "invalid direction")))))

;;; ;;;;;;;;;;;;;;
;;; Tree rotations
;;; ;;;;;;;;;;;;;;

;;; Rotate NODE to the left (dir = '>) or right (dir = '<).
(define %set:rotate
  (lambda (node dir)
    (if (null? node)
        #f
        (let ((invdir (%set:invdir dir)))
          (let ((child (set:get node invdir)))
            (let ((to-swap (set:get child dir)))
              (%set:node (set:get child '=)
                         dir (%set:node (set:get node '=)
                                        dir (set:get node dir)
                                        invdir to-swap)
                         invdir (set:get child invdir))))))))

;;; ;;;;;;;;;;;;;;;;;;;
;;; JOIN function for AVL trees.
;;; ;;;;;;;;;;;;;;;;;;;

;;; Handles rebalancing of the tree.
(define %set:join
  (lambda (heavier val lighter heavier-dir)
    (let ((heavy-val (set:get heavier '=))
          (lighter-dir (%set:invdir heavier-dir)))
      (let ((heavy-heavy (set:get heavier heavier-dir))
            (heavy-light (set:get heavier lighter-dir)))
        (if (<= (abs (%set:height-diff heavy-light lighter)) 1)
            (let ((node (%set:node val
                                   heavier-dir heavy-light
                                   lighter-dir lighter)))
              (if (<= (%set:height-diff node heavy-heavy) 1)
                  (%set:node heavy-val
                             heavier-dir heavy-heavy
                             lighter-dir node)
                  (%set:rotate (%set:node heavy-val
                                          heavier-dir heavy-heavy
                                          lighter-dir
                                          (%set:rotate node lighter-dir))
                               heavier-dir)))
            (let ((new-light (%set:join heavy-light val lighter heavier-dir)))
              (let ((node (%set:node heavy-val
                                     heavier-dir heavy-heavy
                                     lighter-dir new-light)))
                (if (<= (abs (%set:bal node)) 1)
                    node
                    (%set:rotate node heavier-dir)))))))))

(define set:join
  (lambda (val dir1 node1 dir2 node2)
    (let ((diff (%set:height-diff node1 node2)))
      (cond
        ((> diff 1) (%set:join node1 val node2 dir1))
        ((< diff -1) (%set:join node2 val node1 dir2))
        (else (%set:node val dir1 node1 dir2 node2))))))

(define set:join2
  (letrec
      ((join2
        (lambda (left right)
          (if (null? left)
              right
              (let ((split-last-tree (split-last left)))
                (set:join (set:get split-last-tree '=)
                          '< (set:get split-last-tree '<)
                          '> right)))))
       (split-last
        (lambda (tree)
          (let ((right (set:get tree '>)))
            (if (null? right)
                tree
                (let ((last (split-last right)))
                  (%set:node (set:get last '=)
                             '< (set:join (set:get tree '=)
                                          '< (set:get tree '<)
                                          '> (set:get last '>))
                             '> '())))))))
    join2))

(define set:split
  (lambda (<=>)
    (letrec
        ((split
          (lambda (set data)
            (if (null? set)
                (%set:node #f '< '() '> '())
                (let ((set-data (set:get set '=)))
                  (let ((dir (<=> data set-data)))
                    (if (eq? dir '=)
                        (set:node-new-val set #t)
                        (let ((new-tree (split (set:get set dir) data))
                              (invdir (%set:invdir dir)))
                          (%set:node (set:get new-tree '=)
                                     dir (set:get new-tree dir)
                                     invdir (set:join set-data
                                                      dir (set:get new-tree invdir)
                                                      invdir (set:get set invdir)))))))))))
      split)))

;;; ;;;;;;;;;;;;;
;;; Set functions
;;; ;;;;;;;;;;;;;

;;; Generate union operation from split operation.
;;; Union prioritizes data in PRIORITY over keys in SECONDARY.
(define set:union
  (lambda (split)
    (letrec
        ((union
          (lambda (priority secondary)
            (cond
              ((null? priority) secondary)
              ((null? secondary) priority)
              (else
               (let ((key (set:get priority '=)))
                 (let ((split-tree (split secondary key)))
                   (set:join key
                             '< (union (set:get priority '<)
                                       (set:get split-tree '<))
                             '> (union (set:get priority '>)
                                       (set:get split-tree '>))))))))))
      union)))

;;; ;;;;;;;;;;;;;;;;;
;;; Element functions
;;; ;;;;;;;;;;;;;;;;;

;;; (SET:IN <=>) generates a search function for comparison function <=>.
;;; (SEARCH TREE DATA) searches TREE for a node that matches DATA.
;;; It will return the node that contains the matched DATA, or '().
(define set:in
  (lambda (<=>)
    (lambda (tree data)
      (letrec
          ((loop
            (lambda (tree)
              (if (null? tree)
                  '()
                  (let ((dir (<=> data (set:get tree '=))))
                    (if (eq? dir '=)
                        tree
                        (loop (set:get tree dir))))))))
        (loop tree)))))

;;; (SET:UPDATE <=>) generates an update function for <=>.
;;; (UPDATE TREE NODE UPDATE) inserts (UPDATE NODE #F) into the tree if
;;; no node comparing equal to NODE is found, and (UPDATE NODE OLD) if
;;; OLD compares equal to NODE.
(define set:update
  (lambda (<=>)
    (lambda (tree node update)
      (letrec
          ((loop
            (lambda (tree)
              (if (null? tree)
                  (update node #f)
                  (let ((dir (<=> (set:get node '=)
                                  (set:get tree '=))))
                    (if (eq? dir '=)
                        (update node tree)
                        (let ((invdir (%set:invdir dir)))
                          (set:join (set:get tree '=)
                                    dir (loop (set:get tree dir))
                                    invdir (set:get tree invdir)))))))))
        (loop tree)))))

;;; (SET:INSERT <=>) generates an insert function for comparison function
;;; <=>.
;;; (INSERT TREE NODE) inserts NODE into TREE. It returns
;;; (CONS NEWTREE FOUND), where FOUND is the node that was replaced by
;;; NODE, and #F otherwise, and NEWTREE is the new root of the tree.
(define set:insert
  (lambda (update)
    (lambda (tree node)
      (let ((found #f))
        (let ((newroot (update tree node
                               (lambda (node oldnode)
                                 (if (not oldnode)
                                     node
                                     (begin
                                       (set! found oldnode)
                                       (set:node-new-val oldnode
                                                         (set:get node '=))))))))
          (cons newroot found))))))

;;; (SET:DELETE <=>) generates a delete function for comparison function
;;; <=>.
;;; (DELETE TREE DATA) deletes a node from TREE that compares equal to
;;; DATA. The function returns (CONS NEWTREE FOUND), where FOUND is the
;;; deleted node, or NULL if not found, and NEWTREE is the root of the new
;;; tree.
(define set:delete
  (lambda (<=>)
    (lambda (tree node)
      (let ((found '()))
        (letrec
            ((loop
              (lambda (tree)
                (if (null? tree)
                    node
                    (let ((dir (<=> (set:get node '=)
                                    (set:get tree '=))))
                      (if (eq? dir '=)
                          (begin
                            (set! found tree)
                            (set:join2 (set:get tree '<)
                                       (set:get tree '>)))
                          (let ((invdir (%set:invdir dir)))
                            (set:join (set:get tree '=)
                                      dir (loop (set:get tree dir))
                                      invdir (set:get tree invdir)))))))))
          (let ((newtree (loop tree)))
            (cons newtree found)))))))

;;; ;;;;;;;;;;;;;;;;;;;;;;;
;;; Converting sets to maps
;;; ;;;;;;;;;;;;;;;;;;;;;;;

;;; Convert a <=> for sets to one for maps.
(define set:<=>-to-map
  (lambda (<=>)
    (lambda (x y)
      (<=> (car x) (car y)))))

(define map:node-new-val
  (lambda (node newkey newval)
    (set:node-new-val node
                      (cons newkey newval))))

(define map:empty-node
  (lambda (key val)
    (%set:node (cons key val)
               '< '() '> '())))

(define map:key
  (lambda (node)
    (car (set:get node '=))))

(define map:val
  (lambda (node)
    (cdr (set:get node '=))))

(define map:update
  (lambda (%update-recursive)
    (lambda (tree key update)
      (%update-recursive tree (map:empty-node key '())
                         update))))

(define map:insert
  (lambda (%update-recursive)
    (let ((insert (set:insert %update-recursive)))
      (lambda (tree key val)
        (insert tree (map:empty-node key val))))))

(define map:search
  (lambda (<=>)
    (let ((search (set:in <=>)))
      (lambda (tree key)
        (search tree (cons key '()))))))

(define map:delete
  (lambda (<=>)
    (let ((delete (set:delete <=>)))
      (lambda (tree key)
        (delete tree (map:empty-node key '()))))))

(define map:split
  (lambda (<=>) (set:split <=>)))

(define map:union
  (lambda (split) (set:union split)))

;;; ;;;;;;;;;;;
;;; For strings
;;; ;;;;;;;;;;;

(cond-expand
 ((not miniscm-unslisp)
  (define (string<=> x y)
    (cond
      ((string<? x y) '<)
      ((string>? x y) '>)
      (else '=))))
 (else #f))

(cond-expand
 ((and (not miniscm-unslisp) (not r7rs))
  (define (list-set! lst n val)
    (if (= n 0)
        (set-car! lst val)
        (list-set! (cdr lst) (- n 1) val))))
 (else #f))

(define map:string<=> (set:<=>-to-map string<=>))
(define %smap:update (set:update map:string<=>))

(define smap:update (map:update %smap:update))
(define smap:insert (map:insert %smap:update))
(define smap:search (map:search  map:string<=>))
(define smap:delete (map:delete map:string<=>))

(define %smap:split (map:split map:string<=>))
(define smap:union (map:union %smap:split))

;;; ;;;;;
;;; Tests
;;; ;;;;;

;;; LST is a list of elements of the form
;;; (KEY VAL ALREADY-IN)
;;; where ALREADY-IN is #F for an element not in the set, or the value
;;; that should be in the set.

(define %set:operate-all
  (lambda (f tree lst)
    (if (null? lst)
        tree
        (let ((key (list-ref (car lst) 0))
              (val (list-ref (car lst) 1))
              (already-in (list-ref (car lst) 2)))
          (let ((insert-return (f tree key val)))
            (cond
              ((and already-in (null? (cdr insert-return)))
               "should have been found")
              ((and already-in (not (equal? already-in
                                            (map:val (cdr insert-return)))))
               "found is not correct")
              (else (%set:operate-all f (car insert-return) (cdr lst)))))))))

(define %set:insert-all
  (lambda (tree lst)
    (%set:operate-all smap:insert tree lst)))

(define %set:search-all
  (lambda (tree lst)
    (%set:operate-all (lambda (tree key _)
                        (let ((search-res (smap:search tree key)))
                          (cons tree search-res))) tree lst)))

(define %set:delete-all
  (lambda (tree lst)
    (%set:operate-all (lambda (tree key _)
                        (smap:delete tree key)) tree lst)))

(define %set:tests
  (list
   (cons "rotate right"
         (lambda ()
           (let ((right (%set:rotate (%set:node 1
                                                '< (%set:node 2
                                                              '< (%set:node 3
                                                                            '< '()
                                                                            '> '())
                                                              '> (%set:node 4
                                                                            '< '()
                                                                            '> '()))
                                                '> (%set:node 5 '< '() '> '()))
                                     '>)))
             (cond
               ((not (eqv? (set:get right '=) 2)) "bad parent")
               ((not (eqv? (set:get (set:get right '>) '=) 1)) "bad right child")
               ((not (eqv? (set:get (set:get right '<) '=) 3)) "bad left child")
               ((not (eqv? (set:get (set:get (set:get right '>) '>) '=) 5))
                "bad right child of right child")
               ((not (eqv? (set:get (set:get (set:get right '>) '<) '=) 4))
                "bad left child of right child")
               (else #t)))))
   (cons "rotate left"
         (lambda ()
           (let ((right (%set:rotate (%set:node 1
                                                '> (%set:node 2
                                                              '< (%set:node 3
                                                                            '< '()
                                                                            '> '())
                                                              '> (%set:node 4
                                                                            '< '()
                                                                            '> '()))
                                                '< (%set:node 5 '< '() '> '()))
                                     '<)))
             (cond
               ((not (eqv? (set:get right '=) 2)) "bad parent")
               ((not (eqv? (set:get (set:get right '>) '=) 4)) "bad right child")
               ((not (eqv? (set:get (set:get right '<) '=) 1)) "bad left child")
               ((not (eqv? (set:get (set:get (set:get right '<) '>) '=) 3))
                "bad right child of left child")
               ((not (eqv? (set:get (set:get (set:get right '<) '<) '=) 5))
                "bad left child of left child")
               (else #t)))))
   (cons "insert then delete"
         (lambda ()
           (let ((insert-return (smap:insert '() (string #\a) 5)))
             (cond
               ((not (pair? insert-return)) "invalid insert return")
               ((cdr insert-return) "string found in empty tree")
               (else
                (let ((tree (car insert-return)))
                  (let ((found (smap:search tree (string #\a))))
                    (cond
                      ((null? found) "string not in tree")
                      ((not (equal? (map:key tree) (string #\a)))
                       "returned key not equal to a")
                      ((not (equal? (map:val tree) 5))
                       "returned value not equal to 5")
                      (else
                       (let ((delete-return (smap:delete tree (string #\a))))
                         (cond
                           ((not (pair? delete-return))
                            "invalid delete return")
                           ((not (cdr delete-return)) "string not found")
                           ((not (eqv? (car delete-return) '()))
                            "returned tree not null")
                           (else #t))))))))))))
   (cons "insert a few unique then delete"
         (lambda ()
           (let ((to-insert (list
                             (list (string #\a #\b #\c) 1 #f)
                             (list (string #\a #\b #\d) 2 #f)
                             (list (string #\d #\e #\f) 3 #f)
                             (list (string #\1 #\2 #\3 #\a #\C) 4 #f)
                             (list (string #\q #\w #\e) 5 #f)
                             (list (string #\1 #\2 #\3) 5 #f)
                             (list (string #\l #\i #\s #\p) 6 #f)
                             (list (string #\c) 7 #f)
                             (list (string #\s #\c #\m) 8 #f)
                             (list (string #\a #\l #\g #\o #\l) 9 #f)
                             (list (string #\a #\s #\m) 10 #f)
                             (list (string #\4) 11 #f)
                             (list (string #\a #\s #\m #\e) 12 #f))))
             (display "insert all") (newline)
             (let ((tree (%set:insert-all '() to-insert)))
               (if (string? tree)
                   tree
                   (begin
                     (for-each (lambda (x)
                                 (list-set! x 2 (list-ref x 1)))
                               to-insert)
                     (display "search all") (newline)
                     (let ((res (%set:search-all tree to-insert)))
                       (if (string? res)
                           res
                           (begin
                             (display "delete all") (newline)
                             (let ((tree (%set:delete-all tree to-insert)))
                               (cond
                                 ((string? tree) tree)
                                 ((not (null? tree)) "did not delete everything")
                                 (else #t))))))))))))
   (cons "insert a few, update"
         (lambda ()
           (let ((tree (%set:insert-all '()
                                        (list
                                         (list (string #\a #\b #\c #\d) 1 #f)
                                         (list (string #\e #\f #\g #\h) 2 #f)
                                         (list (string #\1 #\4 #\2 #\9 #\3) 3 #f)
                                         (list (string #\a #\b #\c #\d #\e) 4 #f)))))
             (if (string? tree)
                 tree
                 (let ((tree (smap:update tree
                                          (string #\a #\b #\c #\d #\e)
                                          (lambda (node oldnode)
                                            (map:node-new-val oldnode
                                                              (string #\a #\b #\c #\d #\e)
                                                              10)))))
                   (let ((res (%set:search-all tree
                                               (list
                                                (list (string #\a #\b #\c #\d) 1 1)
                                                (list (string #\e #\f #\g #\h) 2 2)
                                                (list (string #\1 #\4 #\2 #\9 #\3) 3 3)
                                                (list (string #\a #\b #\c #\d #\e) 10 10)))))
                     (if (string? res)
                         res
                         #t)))))))
   (cons "union a few"
         (lambda ()
           (let ((tree1 (%set:insert-all '()
                                         (list
                                          (list (string #\a) 1 #f)
                                          (list (string #\b) 2 #f)
                                          (list (string #\c) 3 #f))))
                 (tree2 (%set:insert-all '()
                                         (list
                                          (list (string #\c) 4 #f)
                                          (list (string #\d) 5 #f)
                                          (list (string #\e) 6 #f)))))
             (let ((tree (smap:union tree1 tree2))
                   (to-search (list
                               (list (string #\a) 1 1)
                               (list (string #\b) 2 2)
                               (list (string #\c) 3 3)
                               (list (string #\d) 5 5)
                               (list (string #\e) 6 6))))
               (not (string? (%set:search-all tree to-search)))))))))