;;; 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)
    (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 '>)))))))

;;; 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: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 (<= (abs (%set:bal node)) 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))

;;; ;;;;;;;;;;;;;;;;;
;;; 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 #F.
(define set:in
  (lambda (<=>)
    (lambda (tree data)
      (letrec
          ((loop
            (lambda (tree)
              (if (null? tree)
                  #f
                  (let ((dir (<=> (%set:get tree '=) data)))
                    (if (eq? dir '=)
                        tree
                        (loop (%set:get tree dir))))))))
        (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 (<=>)
    (lambda (tree node)
      (let ((found #f))
        (letrec
            ((loop
              (lambda (tree)
                (if (null? tree)
                    node
                    (let ((dir (<=> (%set:get tree '=)
                                    (%set:get node '=))))
                      (if (eq? dir '=)
                          (begin
                            (set! found tree)
                            (%set:node (%set:get node '=)
                                       '< left '> right))
                          (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)))))))

;;; (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 #F if not found, and NEWTREE is the root of the new
;;; tree.
(define set:delete
  (lambda (<=>)
    (lambda (tree node)
      (let ((found #f))
        (letrec
            ((loop
              (lambda (tree)
                (if (null? tree)
                    node
                    (let ((dir (<=> (%set:get tree '=)
                                    (%set:get node '=))))
                      (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)))))))

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

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

(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)))))

(define map:string<=>
  (lambda (x y)
    (string<=> (car x) (car y))))

(define %smap:insert (set:insert map:string<=>))

;;; (SMAP:INSERT TREE KEY VAL) inserts (CONS KEY VAL) into TREE, and
;;; returns (CONS NEWROOT FOUND), where NEWROOT is the new root of
;;; the tree, and FOUND is #F if no element matching KEY was found,
;;; or the matching element if found.
(define smap:insert
  (lambda (tree key val)
    (%smap:insert tree (%set:node (cons key val)
                                  '< '()
                                  '> '()))))

(define %smap:search (set:in map:string<=>))

;;; (SMAP:SEARCH TREE KEY)
(define smap:search
  (lambda (tree key)
    (%smap:search tree (cons key '()))))

(define %smap:delete (set:delete map:string<=>))
(define smap:delete
  (lambda (tree key)
    (%smap:delete tree (%set:node (cons key '())
                                  '< '() '> '()))))

(define smap:key
  (lambda (node) (car (%set:get node '=))))
(define smap:val
  (lambda (node) (cdr (%set:get node '=))))

;;; ;;;;;
;;; 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 (not (cdr insert-return)))
               "should have been found")
              ((and already-in (not (equal? already-in
                                            (smap: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
                      ((not found) "string not in tree")
                      ((not (equal? (smap:key tree) (string #\a)))
                       "returned key not equal to a")
                      ((not (equal? (smap: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))))
             (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))))))))))))))