;;; 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 (%avl: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)))))))))

;;; JOIN takes two trees and the data for a node, such that all nodes
;;; in LEFT are less than or equal to VAL, and all nodes in RIGHT are
;;; greater than or equal to VAL, and outputs a balanced binary search
;;; tree.
(define set:join
  (lambda (left val right)
    (let ((diff (%set:diff left right)))
      (cond
        ((> diff 1) (%set:join left val right '<))
        ((< diff -1) (%set:join right val left '>))
        (else (set:node val '< left '> right))))))

(define set:join2
  (letrec
      ((join2
        (lambda (left right)
          (if (null? left)
              right
              (let ((split-last-tree (split-last left)))
                (set:join (bst:get split-last-tree '<)
                          (bst:get split-last-tree '=)
                          right)))))
       (split-last
        (lambda (tree)
          (let ((right (set:get tree '>)))
            (if (null? right)
                tree
                (let ((last (split-last right)))
                  (bst:node (set:get last '=)
                            (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 '=)))
                          (left (%set:get tree '<))
                          (right (%set:get tree '>)))
                      (if (eq? dir '=)
                          (begin
                            (set! found tree)
                            (set:node (set:get node '=)
                                      '< left '> right))
                          (join (set:node (%set:get tree '=)
                                          dir (loop left)
                                          (set:invdir dir) right))))))))
          (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 '=)))
                          (left (%set:get tree '<))
                          (right (%set:get tree '>)))
                      (if (eq? dir '=)
                          (begin
                            (set! found tree)
                            (set:join2 left right))
                          (join (set:node (%set:get tree '=)
                                          dir (loop left)
                                          (set:invdir dir) right))))))))
          (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 '=)))))

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

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

(define %smap:search (set:in map:string<=>))
(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 '=))))

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