;;;; Copyright (c) 2024, Peter McGoron
;;;;
;;;; Redistribution and use in source and binary forms, with or without
;;;; modification, are permitted provided that the following conditions
;;;; are met:
;;;;
;;;; 1) Redistributions of source code must retain the above copyright
;;;;    notice, this list of conditions and the following disclaimer.
;;;; 2) Redistributions in binary form must reproduce the above copyright
;;;;    notice, this list of conditions and the following disclaimer
;;;;    in the documentation and/or other materials provided with the
;;;;    distribution.
;;;;
;;;; THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
;;;; "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
;;;; LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
;;;; FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
;;;; COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
;;;; INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
;;;; BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
;;;; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
;;;; CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
;;;; LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
;;;; ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
;;;; POSSIBILITY OF SUCH DAMAGE.
;;;; 
;;;; This implements functional balanced binary search trees (AA Trees).
;;;; 
;;;; (NEW <=>) makes a new AA Tree root. 
;;;; (<=> KEY1 KEY2) must return one of the symbols in '(< = >), which
;;;; denote that KEY1 is less than, equivalent to, or greater than KEY2.
;;;; 
;;;; (SEARCH TREE KEY) searches the tree for the first node equivalent
;;;; to KEY. If successful, it returns (VALUES 'FOUND KEY VAL). Otherwise
;;;; it returns (VALUES 'NOT-FOUND '() '()).
;;;; 
;;;; (INSERT TREE KEY VAL) inserts KEY into TREE with value VAL. If
;;;; a key equivalent to KEY is already in TREE, the old key is replaced with
;;;; KEY and the old value is replaced with VAL, and the function returns
;;;; (VALUES 'FOUND OLDKEY OLDVAL). If the value is not found, then
;;;; the function returns (VALUES 'NOT-FOUND OLDKEY OLDVAL).
;;;; 
;;;; (DELETE TREE KEY) deletes the node equivalent to KEY, if it exists.
;;;; It returns 

(define-namespace aatree
  (begin
    (define-record-type :aatree-node
      (new-node key value left right level)
      node?
      (key get-key key-set!)
      (value get-value value-set!)
      (left get-left left-set!)
      (right get-right right-set!)
      (level get-level-raw level-set!))
    (define-record-type :aatree
      (%aatree <=> root)
      aatree?
      (<=> get-<=>)
      (root get-root set-root!))
    (define leaf '())
    (define (new <=>) (%aatree <=> leaf))
    (define leaf? null?)
    (define (get-level r)
      (if (leaf? r)
          0
          (get-level-raw r)))
    ; Option-like accessors
    (define (maybe-right t)
      (if (null? t)
          leaf
          (get-right t)))
    (define (maybe-left t)
      (if (null? t)
          leaf
          (get-left t)))
    ; right rotation
    ;        a        b
    ;       / \      / \
    ;      b   c -> d   a
    ;     / \          / \
    ;    d   e        e   c
    (define (skew A)
      (let* ((B (maybe-left A))
             (E (maybe-right B)))
        (if (and (not (leaf? A))
                 (eq? (get-level B)
                      (get-level A)))
            (begin
              (right-set! B A)
              (left-set! A E)
              B)
            A)))
    ; left rotation
    ;     a            c
    ;    / \          / \
    ;   b   c   ->   a   e
    ;      / \      / \
    ;     d   e    b   d
    ;   
    (define (split A)
      (let* ((C (maybe-right A))
             (E (maybe-right C))
             (D (maybe-left C)))
        (if (and (not (leaf? A))
                 (not (leaf? C))
                 (eq? (get-level E)
                      (get-level A)))
            (begin
              (left-set! C A)
              (right-set! A D)
              (level-set! C (+ (get-level C) 1))
              C)
            A)))
    (define (search* <=> tree key)
      (if (leaf? tree)
          (values 'not-found '() '())
          (let ((nodekey (get-key tree)))
            (case (<=> key nodekey)
              ((<) (search* <=> (get-left tree) key))
              ((>) (search* <=> (get-right tree) key))
              ((=) (values 'found nodekey (get-value tree)))))))
    (define (search tree key) (search* (get-<=> tree) (get-root tree) key))
    (define (insert* <=> node key val)
      (if (leaf? node)
          (values (new-node key val '() '() 1) 'not-found '() '())
          (case (<=> key (get-key node))
            ((=) (let ((oldval (get-value node)))
                   (value-set! node val)
                   (values node 'found (get-key node) oldval)))
            ((<) (let-values (((newnode . rest)
                               (insert* <=> (get-left node) key val)))
                   (left-set! node newnode)
                   (apply values (cons (split (skew node)) rest))))
            ((>) (let-values (((newnode . rest)
                               (insert* <=> (get-right node) key val)))
                   (right-set! node newnode)
                   (apply values (cons (split (skew node)) rest))))
            (else (error "comparision must return <, =, or >")))))
    (define (insert tree key val)
      (let-values (((new-root . rest) (insert* (get-<=> tree) (get-root tree) key val)))
        (set-root! tree new-root)
        (apply values rest)))
    (define (delete* <=> tree key)
      (if (leaf? tree)
          (values tree 'not-found '() '())
          (let ((process (lambda (t)
                           (if (leaf? t)
                               t
                               (let* ((level (get-level t))
                                      (level-l (get-level (get-left t)))
                                      (level-r (get-level (get-right t)))
                                      (new-level (- level 1)))
                                 (if (or (< level-l new-level)
                                         (< level-r new-level))
                                     (begin
                                       (if (> level-r new-level)
                                           (level-set! (get-right t)
                                                       (min level-r new-level)))
                                       (level-set! t new-level)
                                       (set! t (skew t))
                                       (right-set! t (skew (get-right t)))
                                       (right-set! (get-right t)
                                                   (skew (get-right (get-right t))))
                                       (set! t (split t))
                                       (right-set! t (split (get-right t)))))
                                 t)))))
            (case (<=> key (get-key tree))
              ((<) (let-values (((newnode . rest)
                                 (delete* <=> (get-left tree) key)))
                     (left-set! tree newnode)
                     (apply values (cons (process tree) rest))))
              ((>) (let-values (((newnode . rest)
                                 (delete* <=> (get-right tree) key)))
                     (right-set! tree newnode)
                     (apply values (cons (process tree) rest))))
              ((=) (letrec
                       ((del-min (lambda (t)
                                   (if (leaf? (get-left t))
                                       (values (get-right t) 'found (get-key t)
                                               (get-value t))
                                       (let-values (((newnode . rest)
                                                     (del-min (get-left t))))
                                         (if (leaf? (get-right t))
                                             (error "imbalanced tree"))
                                         (left-set! t newnode)
                                         (apply values
                                                (cons (process t) rest)))))))
                     (if (leaf? (get-right tree))
                         (if (not (leaf? (get-left tree)))
                             (error "imbalanced tree")
                             (values '() 'found (get-key tree) (get-value tree)))
                         (let-values (((newnode status key value)
                                       (del-min (get-right tree)))
                                      ((oldvalue) (get-value tree))
                                      ((oldkey) (get-key tree)))
                           (key-set! tree key)
                           (value-set! tree value)
                           (right-set! tree newnode)
                           (values (process tree) status oldkey oldvalue)))))))))
    (define (delete tree key)
      (let-values (((new-root . rest) (delete* (get-<=> tree)
                                               (get-root tree)
                                               key)))
        (set-root! tree new-root)
        (apply values rest)))
    (define (aatree->alist* node alist)
      (if (leaf? node)
          alist
          (let ((alist-right (aatree->alist* (get-right node) alist)))
            (aatree->alist* (get-left node)
                            (cons (cons (get-key node)
                                        (get-value node))
                                  alist-right)))))
    (define (aatree->alist tree) (aatree->alist* (get-root tree) '()))
    (define (alist->aatree* node <=> alist)
      (if (null? alist)
          node
          (let ((pair (car alist)))
            (let-values (((node . _)
                          (insert node <=>
                                  (car pair)
                                  (cdr pair))))
              (alist->aatree* node (cdr alist) <=>)))))
    (define (alist->aatree tree)
      (alist->aatree* (get-root tree) (get-<=> tree) '())))
  (export
   aatree?
   new
   aatree->alist alist->aatree
   search insert delete))