; Copyright (c) 1993, 1994 Richard Kelsey and Jonathan Rees.  See file COPYING.

;Date: Thu, 4 Nov 93 13:30:46 EST
;To: jar@ai.mit.edu
;Subject: binary search trees
;From: kelsey@research.nj.nec.com
;
;
;For no particular reason I implemented balanced binary search
;trees as another random data structure to go in big.  The only
;things it uses from BIG-SCHEME are DEFINE-RECORD-TYPE and
;RECEIVE.
;
;(define-interface search-tree-interface
;  (export make-search-tree
;          search-tree-ref
;          search-tree-set!
;          search-tree-modify!
;          search-tree-max
;          search-tree-min
;          walk-search-tree))
;
;(define-structure search-tree search-tree-signature
;  (open big-scheme scheme)
;  (files (big search-tree)))


; Red-Black binary search trees as described in Introduction to Algorithms
; by Cormen, Leiserson, and Rivest.
;
; (make-search-tree key-= key-<)      -> tree
;
; (search-tree-ref tree key)          -> value
;
; (search-tree-set! tree key value)
;
; (search-tree-modify! tree key proc)
;   == (search-tree-set! tree key (proc (search-tree-ref tree key)))
;
; (search-tree-max tree)              -> key + value
;
; (search-tree-min tree)              -> key + value
;
; (walk-search-tree proc tree)
;   applies PROC in order to all key + value pairs with a non-#F value

(define-record-type tree
  (lookup
   nil)        ; node marker for missing leaf nodes (used in REALLY-DELETE!)
  ((root #f)))

(define (make-search-tree = <)
  (let ((nil (make-node #f #f #f)))
    (set-node-red?! nil #f)
    (tree-maker (make-lookup = <) nil)))

(define-record-type node
  ((key)
   (value)
   (parent))         ; #F for the root node
  ((red? #t)
   (left  #f)
   (right #f)))

(define make-node node-maker)

(define-record-discloser type/node
  (lambda (node)
    (list 'node (node-key node))))

(define (search-tree-ref tree key)
  (receive (node parent left?)
      ((tree-lookup tree) tree key)
    (if node
	(node-value node)
	#f)))

(define (search-tree-set! tree key value)
  (search-tree-modify! tree key (lambda (ignore) value)))

(define (search-tree-modify! tree key proc)
  (receive (node parent left?)
      ((tree-lookup tree) tree key)
    (let ((new-value (proc (if node (node-value node) #f))))
      (cond ((and node new-value)
	     (set-node-value! node new-value))
	    (new-value
	     (really-insert! tree parent left? (make-node key new-value parent)))
	    (node
	     (really-delete! tree node))))))

(define (search-tree-max tree)
  (let ((node (tree-root tree)))
    (if node
	(let loop ((node node))
	  (cond ((node-right node)
		 => loop)
		(else
		 (values (node-key node) (node-value node)))))
	(values #f #f))))

(define (search-tree-min tree)
  (let ((node (tree-root tree)))
    (if node
	(let loop ((node node))
	  (cond ((node-left node)
		 => loop)
		(else
		 (values (node-key node) (node-value node)))))
	(values #f #f))))

(define (walk-search-tree proc tree)
  (let recur ((node (tree-root tree)))
    (cond (node
	   (recur (node-left node))
	   (proc (node-key node) (node-value node))
	   (recur (node-right node))))))

(define (make-lookup = <)
  (lambda (tree key)
    (let loop ((node (tree-root tree))
	       (parent #f)
	       (left? #f))
      (cond ((not node)
	     (values #f parent left?))
	    ((= (node-key node) key)
	     (values node #f #f))
	    ((< key (node-key node))
	     (loop (node-left node) node #t))
	    (else
	     (loop (node-right node) node #f))))))

; Parameterized node access

(define (node-child node left?)
  (if left?
      (node-left node)
      (node-right node)))

(define (set-node-child! node left? child)
  (if left?
      (set-node-left! node child)
      (set-node-right! node child)))

; Empty leaf slots are considered black.

(define (node-black? node)
  (not (and node (node-red? node))))

; The next node (used in REALLY-DELETE!)

(define (successor node)
  (cond ((node-right node)
	 => (lambda (node)
	      (let loop ((node node))
		(cond ((node-left node)
		       => loop)
		      (else node)))))
	(else
	 (let loop ((node node) (parent (node-parent node)))
	   (if (and parent
		    (eq? node (node-right parent)))
	       (loop parent (node-parent parent))
	       parent)))))

(define (really-insert! tree parent left? node)
  (if (not parent)
      (set-tree-root! tree node)
      (set-node-child! parent left? node))
  (fixup-insertion! node tree))

(define (fixup-insertion! node tree)
  (let loop ((node node))
    (let ((parent (node-parent node)))
      (if (and parent (node-red? parent))
	  (let* ((grand (node-parent parent))
		 (left? (eq? parent (node-left grand)))
		 (y (node-child grand (not left?))))
	    (cond ((node-black? y)
		   (let* ((node (cond ((eq? node (node-child parent (not left?)))
				       (rotate! parent left? tree)
				       parent)
				      (else node)))
			  (parent (node-parent node))
			  (grand (node-parent parent)))
		     (set-node-red?! parent #f)
		     (set-node-red?! grand  #t)
		     (rotate! grand (not left?) tree)
		     (loop node)))
		  (else
		   (set-node-red?! parent #f)
		   (set-node-red?! y      #f)
		   (set-node-red?! grand  #t)
		   (loop grand)))))))
  (set-node-red?! (tree-root tree) #f))

;        A                                 B
;       / \    =(rotate! A #f tree)=>     / \
;      B   k                             i   A
;     / \      <=(rotate! B #t tree)=       / \
;    i   j                                 j   k

(define (rotate! node left? tree)
  (let* ((y (node-child node (not left?)))
	 (y-left (node-child y left?))
	 (parent (node-parent node)))
    (set-node-child! node (not left?) y-left)
    (if y-left
	(set-node-parent! y-left node))
    (replace! parent y node tree)
    (set-node-child! y left? node)
    (set-node-parent! node y)))
		
; Replace CHILD (of PARENT) with NEW-CHILD

(define (replace! parent new-child child tree)
    (set-node-parent! new-child parent)
    (cond ((eq? child (tree-root tree))
	   (set-tree-root! tree new-child))
	  ((eq? child (node-left parent))
	   (set-node-left! parent new-child))
	  (else
	   (set-node-right! parent new-child))))

(define (really-delete! tree node)
  (let* ((y (cond ((or (not (node-left node))
		       (not (node-right node)))
		   node)
		  (else
		   (let ((y (successor node)))
		     (set-node-key!   node (node-key y))
		     (set-node-value! node (node-value y))
		     y))))
	 (x (or (node-left y)
		(node-right y)
		(let ((x (tree-nil tree)))
		  (set-node-right! y x)
		  x)))
	 (parent (node-parent y)))
    (replace! parent x y tree)
    (if (not (node-red? y))
	(fixup-delete! x tree))
    (let ((nil (tree-nil tree)))
      (cond ((node-parent nil)
	     => (lambda (p)
		  (if (eq? (node-right p) nil)
		      (set-node-right! p #f)
		      (set-node-left! p #f))
		  (set-node-parent! (tree-nil tree) #f)))
	    ((eq? nil (tree-root tree))
	     (set-tree-root! tree #f))))))

(define (fixup-delete! x tree)
  (let loop ((x x))
    (if (or (eq? x (tree-root tree))
	    (node-red? x))
	(set-node-red?! x #f)
	(let* ((parent (node-parent x))
	       (left? (eq? x (node-left parent)))
	       (w (node-child parent (not left?)))
	       (w (cond ((node-red? w)
			 (set-node-red?! w #f)
			 (set-node-red?! parent #t)
			 (rotate! parent left? tree)
			 (node-child (node-parent x) (not left?)))
			(else
			 w))))
	  (cond ((and (node-black? (node-left w))
		      (node-black? (node-right w)))
		 (set-node-red?! w #t)
		 (loop (node-parent x)))
		(else
		 (let ((w (cond ((node-black? (node-child w (not left?)))
				 (set-node-red?! (node-child w left?) #f)
				 (set-node-red?! w #t)
				 (rotate! w (not left?) tree)
				 (node-child (node-parent x) (not left?)))
				(else
				 w))))
		   (let ((parent (node-parent x)))
		     (set-node-red?! w (node-red? parent))
		     (set-node-red?! parent #f)
		     (set-node-red?! (node-child w (not left?)) #f)
		     (rotate! parent left? tree)
		     (set-node-red?! (tree-root tree) #f)))))))))
		 
; Verify that the coloring is correct
;
;(define (okay-tree? tree)
;  (receive (okay? red? count)
;      (let recur ((node (tree-root tree)))
;        (if (not node)
;            (values #t #f 0)
;            (receive (l-ok? l-r? l-c)
;                (recur (node-left node))
;              (receive (r-ok? r-r? r-c)
;                  (recur (node-right node))
;                (values (and l-ok?
;                             r-ok?
;                             (not (and (node-red? node)
;                                       (or l-r? r-r?)))
;                             (= l-c r-c))
;                        (node-red? node)
;                        (if (node-red? node)
;                            l-c
;                            (+ l-c 1)))))))
;    okay?))
;
;
;(define (walk-sequences proc list)
;  (let recur ((list list) (r '()))
;    (if (null? list)
;        (proc (reverse r))
;        (let loop ((list list) (done '()))
;          (if (not (null? list))
;              (let ((next (car list)))
;                (recur (append (reverse done) (cdr list)) (cons next r))
;                (loop (cdr list) (cons next done))))))))
;           
;(define (tree-test n)
;  (let ((iota (do ((i n (- i 1))
;                   (l '() (cons i l)))
;                  ((<= i 0) l))))
;    (walk-sequences (lambda (in)
;                      (walk-sequences (lambda (out)
;                                        (do-tree-test in out))
;                                      iota))
;                    iota)
;    #t))
;
;(define (do-tree-test in out)
;  (let ((tree (make-search-tree = <)))
;    (for-each (lambda (i)
;                (search-tree-set! tree i (- 0 i)))
;              in)
;    (if (not (okay-tree? tree))
;        (breakpoint "tree ~S is not okay" in))
;    (if (not (tree-ordered? tree (length in)))
;        (breakpoint "tree ~S is not ordered" in))
;    (for-each (lambda (i)
;                (if (not (= (search-tree-ref tree i) (- 0 i)))
;                    (breakpoint "looking up ~S in ~S lost" i in)))
;              in)
;    (do ((o out (cdr o)))
;        ((null? o))
;      (search-tree-set! tree (car o) #f)
;      (if (not (okay-tree? tree))
;          (breakpoint "tree ~S is not okay after deletions ~S" in out)))))
;
;(define (tree-ordered? tree count)
;  (let ((l '()))
;    (walk-search-tree (lambda (key value)
;                        (set! l (cons (cons key value) l)))
;                      tree)
;    (let loop ((l l) (n count))
;      (cond ((null? l)
;             (= n 0))
;            ((and (= (caar l) n)
;                  (= (cdar l) (- 0 n)))
;             (loop (cdr l) (- n 1)))
;            (else #f)))))
;
;(define (do-tests tester)
;  (do ((i 0 (+ i 1)))
;      (#f)
;    (tester i)
;    (format #t " done with ~D~%" i)))
;
;(define (another-test n)
;  (let ((iota (do ((i n (- i 1))
;                   (l '() (cons i l)))
;                  ((<= i 0) l))))
;    (walk-sequences (lambda (in)
;                      (do ((i 1 (+ i 1)))
;                          ((> i n))
;                        (let ((tree (make-search-tree = <)))
;                          (for-each (lambda (i)
;                                      (search-tree-set! tree i (- 0 i)))
;                                    in)
;                          (if (not (okay-tree? tree))
;                              (breakpoint "tree ~S is not okay" in))
;                          (if (not (tree-ordered? tree (length in)))
;                              (breakpoint "tree ~S is not ordered" in))
;                          (for-each (lambda (i)
;                                      (if (not (= (search-tree-ref tree i) (- 0 i)))
;                                          (breakpoint "looking up ~S in ~S lost" i in)))
;                                    in)
;                          (search-tree-set! tree i #f)
;                          (if (not (okay-tree? tree))
;                              (breakpoint "tree ~S is not okay after deletion ~S"
;                                          in i))
;                          (for-each (lambda (j)
;                                      (let ((ref (search-tree-ref tree j)))
;                                        (if (not (eq? ref (if (= j i) #f (- 0 j))))
;                                            (breakpoint "looking up ~S in ~S lost" i in))))
;                                    in))))
;                    iota)))