;;; ;;;;;;;;;;;;;;;;;;;
;;; Definition of nodes and functions to calculate values for nodes.
;;; ;;;;;;;;;;;;;;;;;;;

(define-type node-type (or (struct <wb-tree>) null))

(: %wb-tree-node (* fixnum node-type node-type --> (struct <wb-tree>)))
(: get-data ((struct <wb-tree>) --> *))
(: %get-weight ((struct <wb-tree>) --> fixnum))
(: get-left ((struct <wb-tree>) --> node-type))
(: get-right ((struct <wb-tree>) --> node-type))

(define-record-type <wb-tree>
  (%wb-tree-node data weight left right)
  non-null-wb-tree-node?
  (data get-data)
  (weight %get-weight)
  (left get-left)
  (right get-right))

(: wb-tree-node? (* -> boolean : node-type))
(define (wb-tree-node? x)
  (or (null? x) (non-null-wb-tree-node? x)))

(: get-weight (node-type --> fixnum))
(define (get-weight node)
  ;; Get the stored size of a node.
  (cond
    ((null? node) 1)
    (else (%get-weight node))))

(: get-size (node-type --> fixnum))
(define (get-size node)
  (- (get-weight node) 1))

(: fixnum-calculate-weight (fixnum fixnum --> fixnum))
(define (fixnum-calculate-weight left-weight right-weight)
  ;; Calculate the weight of a node given the weight of its children.
  (+ left-weight right-weight))

(: calculate-weight (node-type node-type --> fixnum))
(define (calculate-weight left right)
  ;; Calculate the weight of a node that has children `left` and `right`.
  (fixnum-calculate-weight (get-weight left) (get-weight right)))

(: wb-tree-node (* node-type node-type --> (struct <wb-tree>)))
(define (wb-tree-node data left right)
  ;; Construct a node with `data`, `left`, and `right`, with the correct
  ;; weight.
  (%wb-tree-node data (calculate-weight left right) left right))

(: balanced-as-child? (fixnum fixnum --> boolean))
(define (balanced-as-child? child-weight node-weight)
  ;; Determine if child would be weight-balanced if its parent was a node
  ;; with weight `node-weight`.
  (let ((alpha #e0.29))
    (>= child-weight (* alpha node-weight))))

(: fixnum-would-be-balanced? (fixnum fixnum --> boolean))
(define (fixnum-would-be-balanced? left-weight right-weight)
  ;; Determine if the two weights would be balanced if placed into a node.
  (let ((size (+ left-weight right-weight)))
    (and (balanced-as-child? left-weight size)
         (balanced-as-child? right-weight size))))

(: would-be-balanced? (node-type node-type --> boolean))
(define (would-be-balanced? left right)
  ;; Determine if the two nodes would be balanced if placed into a node.
  (fixnum-would-be-balanced? (get-weight left) (get-weight right))) 

(: heavy<=> (node-type node-type --> fixnum))
(define (heavy<=> left right)
  ;; Return 1 if right > left, -1 if left < right, and 0 if left = right
  ;; weightwise.
  (let ((left (get-weight left))
        (right (get-weight right)))
    (cond
      ((< left right) -1)
      ((> left right) 1)
      (else 0))))

(: balanced? (node-type --> boolean))
(define (balanced? node)
  ;; Recursively check if node is weight balanced.
  (cond
    ((null? node) #t)
    (else (let ((left (get-left node))
                (right (get-right node))
                (weight (get-weight node)))
            (and (balanced? left) (balanced? right)
                 (balanced-as-child? (get-weight left) weight)
                 (balanced-as-child? (get-weight right) weight))))))

;;; ;;;;;;;;;;;;;;;;;;;;;;;;;
;;; Convert in-order vectors to ordered trees.
;;; ;;;;;;;;;;;;;;;;;;;;;;;;;

(: in-order-vector->node (vector --> node-type))
(define (in-order-vector->node vec)
  (define (divide left right)
    (if (< right left)
        '()
        (let ((midpoint (floor (/ (+ left right) 2))))
          (wb-tree-node (vector-ref vec midpoint)
                        (divide left (- midpoint 1))
                        (divide (+ midpoint 1) right)))))
  (divide 0 (- (vector-length vec) 1)))

(: node->in-order-list (node-type --> list))
(define (node->in-order-list node)
  (if (null? node)
      '()
      (append (node->in-order-list (get-left node))
              (list (get-data node))
              (node->in-order-list (get-right node)))))

;;; ;;;;;;;;;;;;;;;;;;;;
;;; Fundamental tree operations
;;; 
;;; These macros are used to automatically generate symmetric cases for
;;; tree operations. They work with a trick with the syntax-rules
;;; matcher where it will match literal strings.
;;; 
;;; Since syntax-rules is not eager it cannot calculate the inverse
;;; direction, so both must be supplied. This is fine for most cases,
;;; and helps make the algorithm more explicit.
;;; ;;;;;;;;;;;;;;;;;;;;

(define-syntax s-make-node
  ;; Create a node with these directions.
  (syntax-rules ()
    ((_ data ("<" left) (">" right))
     (wb-tree-node data left right))
    ((_ data (">" right) ("<" left))
     (s-make-node data ("<" left) (">" right)))))

(define-syntax with-node
  (syntax-rules ()
    ((with-node (%node data ("<" left) (">" right)) body ...)
     (let* ((node %node)
            (left (get-left node))
            (right (get-right node))
            (data (get-data node)))
       body ...))
    ((with-node (%node data (">" right) ("<" left)) body ...)
     (with-node (%node data ("<" left) (">" right)) body ...))))

(define-syntax s-get
  (syntax-rules ()
    ((s-get "<" node) (get-left node))
    ((s-get ">" node) (get-right node))))

(define-syntax s-rotate
  ;; Generate rotation based on direction. Rotations are:
  ;; 
  ;;     A                  C
  ;;    / \                / \
  ;;   B   C      ->      A   E
  ;;      / \            / \
  ;;     D   E          B   D
  ;; 
  ;;     A                  C
  ;;    / \                / \
  ;;   C   B      ->      E   A
  ;;  / \                    / \
  ;; E   D                  D   B
  (syntax-rules ()
    ((_ dir invdir A)
     (with-node (A A-data (dir B) (invdir C))
       (with-node (C C-data (dir D) (invdir E))
         (s-make-node C-data
                      (dir (s-make-node A-data (dir B) (invdir D)))
                      (invdir E)))))))

(define-syntax s-join
  ;; Generate a macro that traverses `invdir` to make a balanced tree with
  ;; `dir` in it and `data` in the middle.
  (syntax-rules ()
    ((s-join %data (dir init-in-dir) (invdir init-in-invdir))
     (let ((in-dir init-in-dir)
           (%data data))
       (let join ((in-invdir init-in-invdir))
         (if (would-be-balanced? in-invdir in-dir)
             (s-make-node data
                          (invdir in-invdir)
                          (dir in-dir))
             (with-node (in-invdir invdir-data
                                   (dir dir-in-invdir)
                                   (invdir invdir-in-invdir))
               (let ((new-dir (join dir-in-invdir)))
                 (if (would-be-balanced? invdir-in-invdir new-dir)
                     (s-make-node invdir-data
                                  (invdir invdir-in-invdir)
                                  (dir new-dir))
                     (with-node (new-dir _
                                         (dir dir-in-new-dir)
                                         (invdir invdir-in-new-dir))
                       (if (and (would-be-balanced? invdir-in-invdir
                                                    invdir-in-new-dir)
                                (fixnum-would-be-balanced?
                                 (+ (get-weight invdir-in-invdir)
                                    (get-weight invdir-in-new-dir))
                                 (get-weight dir-in-new-dir)))
                           (s-rotate invdir
                                     dir
                                     (s-make-node invdir-data
                                                  (invdir invdir-in-invdir)
                                                  (dir new-dir)))
                           (s-rotate invdir
                                     dir
                                     (s-make-node invdir-data
                                                  (invdir invdir-in-invdir)
                                                  (dir (s-rotate dir invdir new-dir)))))))))))))))

(: join (* node-type node-type --> (struct <wb-tree>)))
(define (join data left right)
  (let ((dir (heavy<=> left right)))
    (cond
      ((positive? dir) (s-join data (">" right) ("<" left)))
      ((negative? dir) (s-join data ("<" left) (">" right)))
      (else (wb-tree-node data left right)))))

(: join2 (node-type node-type --> node-type))
(define (join2 left right)
  (define split-last
    (the ((struct <wb-tree>) --> node-type *)
         (lambda (tree)
           (with-node (tree data ("<" left) (">" right))
             (if (null? right)
                 (values left data)
                 (let-values (((new-right new-data)
                              (split-last right)))
                   (values (join data left new-right)
                           new-data)))))))
  (if (null? left)
      right
      (let-values (((new-left new-data)
                    (split-last left)))
        (join new-data new-left right))))

;;; XXX: The comparator library does not export the struct type for
;;; the comparator.
(: split (* node-type * (-> *) --> node-type * node-type))
(define (split cmp tree key default)
  (let split ((tree tree))
    (if (null? tree)
        (values '() (default) '())
        (with-node (tree data ("<" left) (">" right))
          (comparator-if<=> cmp key data
            (let-values (((new-left bool new-right) (split left)))
              (values new-left bool (join data new-right right)))
            (values left data right)
            (let-values (((new-left bool new-right) (split right)))
              (values (join data left new-left) bool new-right)))))))

;;; ;;;;;;;;;;;;;;;;;;;;;;;;
;;; Derived tree operations
;;; ;;;;;;;;;;;;;;;;;;;;;;;;

(define search
  (case-lambda
    ((cmp key tree) (search cmp key tree (lambda () #f)))
    ((cmp key tree default)
     (let search ((tree tree))
       (if (null? tree)
           (default)
           (with-node (tree data ("<" left) (">" right))
             (comparator-if<=> cmp key data
               (search left)
               data
               (search right))))))))

(define the-sentinel-value
  ;; A dummy list allocated so that it is not `eq?` to any other object.
  ;; Used as the return value for internal set functions.
  (cons #f '()))
(define return-sentinel (lambda () the-sentinel-value))
(define (sentinel? x) (eq? x the-sentinel-value))

(: union (* node-type node-type --> node-type))
(define (union cmp left right)
  (let union ((left left)
              (right right))
    (cond
      ((null? left) right)
      ((null? right) left)
      (else (with-node (right right-data
                              ("<" left-of-right)
                              (">" right-of-right))
              (let-values (((new-left _ new-right)
                            (split cmp left right-data return-sentinel)))
                (join right-data
                      (union new-left left-of-right)
                      (union new-right right-of-right))))))))

(: intersection (* node-type node-type --> node-type))
(define (intersection cmp left right)
  (let intersection ((left left)
                     (right right))
    (if (or (null? left) (null? right))
        '()
        (with-node (right right-data
                          ("<" left-of-right)
                          (">" right-of-right))
          (let-values (((new-left new-key new-right)
                        (split cmp left right-data return-sentinel)))
            (let ((final-left (intersection new-left left-of-right))
                  (final-right (intersection new-right right-of-right)))
              (if (sentinel? new-key)
                  (join2 final-left final-right) ; right-data not found
                  (join new-key final-left final-right))))))))

(: difference (* node-type node-type --> node-type))
(define (difference cmp left right)
  (let difference ((left left)
                   (right right))
    (cond
      ((null? left) '())
      ((null? right) left)
      (else (with-node (right right-data
                              ("<" left-of-right)
                              (">" right-of-right))
              (let-values (((new-left new-key new-right)
                            (split cmp left right-data return-sentinel)))
                (join2 (difference new-left left-of-right)
                       (difference new-right right-of-right))))))))

;;; ;;;;;;;;;;;;;;;;;;;;
;;; Single value operations
;;; ;;;;;;;;;;;;;;;;;;;;

(: update (* node-type * (-> *) (-> node-type) --> node-type))
(define (update cmp set to-search on-found on-not-found)
  (let update ((set set))
    (if (null? set)
        (on-not-found)
        (with-node (set data ("<" left) (">" right))
          (comparator-if<=> cmp to-search data
            (join data (update left) right)
            (wb-tree-node (on-found data) left right)
            (join data left (update right)))))))

(: delete (* * node-type --> node-type))
(define (delete cmp to-search set)
  (let delete ((set set))
    (if (null? set)
        '()
        (with-node (set data ("<" left) (">" right))
          (comparator-if<=> cmp to-search data
            (join data (delete left) right)
            (join2 left right)
            (join data left (delete right)))))))

;;; ;;;;;;;;;;;;;;;;;
;;; Generic tree functions
;;; ;;;;;;;;;;;;;;;;;

(: every (* node-type --> *))
(define (every predicate? tree)
  (if (null? tree)
      #t
      (with-node (tree data ("<" left) (">" right))
        (and (predicate? data)
             (every predicate? left)
             (every predicate? right)))))