;;
;; Red-black tree
;;
;; An implementation of an ordered dictionary data structure, based
;; on red-black trees.
;;
;; This code is based on the SML/NJ library implementation of
;; red-black trees, which is in turn based on Chris Okasaki's
;; implementation of red-black trees.  The delete function is based on
;; the description in Cormen, Leiserson, and Rivest.
;;
;; Some helper code was borrowed from treap.scm by Oleg Kiselyov.
;;
;;
;; Copyright 2007-2010 Ivan Raikov and the Okinawa Institute of
;; Science and Technology.
;;
;;
;; 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, either version 3 of the
;; License, or (at your option) any later version.
;;
;; 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.
;;
;; A full copy of the GPL license can be found at
;; <http://www.gnu.org/licenses/>.
;;
;;
;; TODO: Add the linear-time tree construction code from the
;; paper _Constructing red-black trees_ by Hinze.
;;

(module rb-tree
	
  (make-rb-tree)

  (import scheme chicken data-structures)
  
  (require-extension srfi-1 datatype matchable)


;;
;; A red-black tree should satisfy the following two invariants:
;;
;;   Red Invariant: each red node has a black parent.
;;
;;   Black Condition: each path from the root to an empty node has the
;;     same number of black nodes (the tree's black height).
;;
;; The Red condition implies that the root is always black and the Black
;; condition implies that any node with only one child will be black and
;; its child will be a red leaf.
;;
;;
;; The red-black tree object is created by procedure make-rb-tree, the
;; only user-visible function defined in this library: 
;;
;;  make-rb-tree:: KEY-COMPARE-PROC -> RB-TREE
;; 
;;  where KEY-COMPARE-PROC is a user-supplied function
;;
;;  KEY-COMPARE-PROC:: key1 key2 -> INTEGER
;; 
;;  that takes two keys and returns a negative, positive, or zero
;;  number depending on how the first key compares to the second.
;;
;; The red-black tree object responds to the following messages
;;
;;	'get 
;;
;;		returns a procedure LAMBDA KEY . DEFAULT-CLAUSE which
;;		searches the red-black tree for an association with a
;;		given KEY, and returns a (key . value) pair of the
;;		found association.  If an association with the KEY
;;		cannot be located in the red-black tree, the PROC
;;		returns the result of evaluating the DEFAULT-CLAUSE.
;;		If the default clause is omitted, an error is
;;		signalled.  The KEY must be comparable to the keys in
;;		the red-black tree by a key-compare predicate (which
;;		has been specified when the red-black tree was
;;		created)
;;
;;	'get-min
;;
;;		returns a (key . value) pair for an association in the
;;		red-black tree with the smallest key. If the red-black
;;		tree is empty, an error is signalled.
;;
;;	'delete-min!
;;
;;		removes the min key and the corresponding association
;;		from the red-black tree. Returns a (key . value) pair
;;		of the removed association.  If the red-black tree is
;;		empty, an error is signalled.
;;
;;	'get-max
;;
;;		returns a (key . value) pair for an association in the
;;		red-black tree with the largest key. If the red-black
;;		tree is empty, an error is signalled.
;;
;;	'delete-max!
;;
;;		removes the max key and the corresponding association
;;		from the red-black tree. Returns a (key . value) pair
;;		of the removed association.  If the red-black tree is
;;		empty, an error is signalled.
;;
;;	empty?
;;		returns #t if the red-black tree is empty
;;
;;	size
;;
;;		returns the size (the number of associations) in the
;;		red-black tree
;;
;;	depth
;;
;;		returns the depth of the tree. It requires the
;;		complete traversal of the tree, so use sparingly
;;
;;	clear!
;;
;;		removes all associations from the red-black tree (thus
;;		making it empty)
;;
;;	'put!
;;
;;		returns a procedure LAMBDA KEY VALUE which, given a
;;		KEY and a VALUE, adds the corresponding association to
;;		the red-black tree. If an association with the same
;;		KEY already exists, its value is replaced with the
;;		VALUE (and the old (key . value) association is
;;		returned). Otherwise, the return value is #f.
;;
;;	'delete!
;;
;;		returns a procedure LAMBDA KEY . DEFAULT-CLAUSE which
;;		searches the red-black tree for an association with a
;;		given KEY, deletes it, and returns a (key . value)
;;		pair of the found and deleted association.  If an
;;		association with the KEY cannot be located in the
;;		red-black tree, the PROC returns the result of
;;		evaluating the DEFAULT-CLAUSE.  If the default clause
;;		is omitted, an error is signalled.
;;
;;	for-each-ascending
;;
;;		returns a procedure LAMBDA PROC that will apply the
;;		given procedure PROC to each (key . value) association
;;		of the red-black tree, from the one with the smallest
;;		key all the way to the one with the max key, in an
;;		ascending order of keys.  The red-black tree must not
;;		be empty.
;;
;;	for-each-descending
;;
;;		returns a procedure LAMBDA PROC that will apply the
;;		given procedure PROC to each (key . value) association
;;		of the red-black tree, in the descending order of
;;		keys.  The red-black tree must not be empty.
;;


(define (rb-tree:error x . rest)
  (let ((port (open-output-string)))
    (let loop ((objs (cons x rest)))
      (if (null? objs)
	  (begin
	    (newline port)
	    (error 'rb-tree (get-output-string port)))
	  (begin (display (car objs) port)
		 (display " " port)
		 (loop (cdr objs)))))))

(define R 'Red)
(define B 'Black)

(define (color? x) (or (eq? x 'Red) (eq? x 'Black)))

(define-datatype tree tree?
  (Empty)
  (Tree  (color color?) (left tree?) (key identity) (value identity) (right tree?)))

(define-datatype zipper zipper?
  (Top) 
  (Left (color color?) (key identity) (value identity) (tree tree?) (zipper zipper?))
  (Right (color color?) (tree tree?) (key identity) (value identity) (zipper zipper?)))


(define (tree-tag x)
  (cases tree x
	 (Empty () 'Empty)
	 (Tree (c l k v r) 'Tree)))

(define-record-printer (tree x out)
  (cases tree x 
	 (Empty () (display "#(Empty)" out))
	 (Tree (c l k v r)
	       (display "#(Tree " out)
	       (display (conc c " ") out) 
	       (display (tree-tag l) out)
	       (display (conc " " k ":" v " ") out)
	       (display (tree-tag r) out)
	       (display ")" out))))

;;
;; This macro was borrowed from treap.scm by Oleg Kiselyov
;;
(define-syntax dispatch-on-key 
  (lambda (x r c)
    (let ((key-compare (second x)) 
	  (key         (third x)) 	  
	  (node-key    (fourth x))
	  (on-less     (fifth x)) 
	  (on-equal    (sixth x)) 
	  (on-greater  (seventh x)))
      (let ((%let   (r 'let))
	    (%cond  (r 'cond))
	    (%else  (r 'else))
	    (%zero?  (r 'zero?))
	    (%positive?  (r 'positive?))
	    (result (r 'result)))
	`(,%let ((,result (,key-compare ,key ,node-key )))
		(,%cond
		 ((,%zero? ,result)     ,on-equal)
		 ((,%positive? ,result) ,on-greater)
		 (,%else              ,on-less)))))))


(define (make-rb-tree key-compare #!key 
		      (insert-key-compare key-compare)
		      (delete-key-compare key-compare) )

  (let ((root (Empty)) (size 0))

    (define (make-rb-tree-dispatcher root size)

    ;; Adds a new association to the tree (or replaces the old one if
    ;; existed). Returns the (key . value) pair of the old
    ;; association, or #f if a new association was really added
    (define (insert root key value)
      (define (ins root)
	(cases tree root
	       (Empty ()  (values #f (Tree R (Empty) key value (Empty))))
	       (Tree (color a yk y b)  
		     (dispatch-on-key insert-key-compare 
		      key yk 
		      ;; Case 1: key < yk
		      (match a
			     (($ tree 'Tree 'Red c zk z d)
			      (dispatch-on-key insert-key-compare
			       key zk
			       ;; Case 1.1: key < zk
			       (let-values (((found? c1) (ins c)))
				 (values found?
					 (match c1
						(($ tree 'Tree 'Red e wk w f)
						 (Tree R (Tree B e wk w f) zk z (Tree B d yk y b)))
						(else  (Tree B (Tree R c1 zk z d) yk y b)))))
			       ;; Case 1.2: key = zk
			       (values a (Tree color (Tree R c key value d) yk y b))
			       ;; Case 1.3: key > zk
			       (let-values (((found? d1) (ins d)))
				 (values found?
					 (match d1
						(($ tree 'Tree 'Red e wk w f)
						 (Tree R (Tree B c zk z e)  wk  w  (Tree B f yk y b)))
						(else (Tree B (Tree R c zk z d1) yk y b)))))))
			     (else  (let-values (((found? a1)  (ins a)))
				      (values found? (Tree B a1 yk y b)))))
		      ;; Case 2: key  = yk
		      (values root (Tree color a key value b))
		      ;; Case 3: key  > yk
		      (match b
			     (($ tree 'Tree 'Red c zk z d)
			      (dispatch-on-key insert-key-compare
			       key zk
			       ;; Case 3.1: key < zk
			       (let-values (((found? c1) (ins c)))
				 (values found?
					 (match c1
						(($ tree 'Tree 'Red e wk w f)
						 (Tree R (Tree B a yk y e)  wk  w (Tree B f zk z d)))
						(else (Tree B a yk y (Tree R c1 zk z d))))))
			       ;; Case 3.2: key = zk
			       (values b (Tree color a yk y (Tree R c key value d)))
			       ;; Case 3.3: key > zk
			       (let-values (((found? d1) (ins d)))
				 (values found?
					 (match d1
						(($ tree 'Tree 'Red e wk w f)
						 (Tree R (Tree B a yk y c)  zk z (Tree B e wk w f)))
						(else (Tree B a yk y (Tree R c zk z d1))))))))
			     (else (let-values (((found? b1) (ins b)))
				     (values found? (Tree B a yk y b1)))))))))

      (ins root))


    ;; Looks for an item: Given a key, returns the corresponding (key
    ;; . value) association or #f if the tree does not contain an
    ;; association with that key.
    (define (find-assoc key)
      (define (find root)
	(cases tree root 
	       (Empty ()  #f)
	       (Tree (c a yk y b)
		     (dispatch-on-key key-compare 
		      key yk (find a) (cons yk y) (find b)))))
      (find root))

    ;; Finds an association with a given key, and deletes it.  Returns
    ;; the (key . value) pair of the deleted association, or #f if it
    ;; couldn't be found
    (define (delete root key)

      (define (zip zipper tree)
	(match (cons zipper tree)
	       ((($ zipper 'Top) . a)  tree)
	       ((($ zipper 'Left color xk x b z) . a)   (zip z (Tree color a xk x b)))
	       ((($ zipper 'Right color a xk x z) . b)  (zip z (Tree color a xk x b)))))

      ;; bbZip propagates a black deficit up the tree until either
      ;; the top is reached, or the deficit can be covered.  It
      ;; returns a boolean that is true if there is still a deficit
      ;; and the zipped tree.
     (define (bbZip zipper tree)
	(match (cons zipper tree)
	       ((($ zipper 'Top) . a)  (cons #t a))
	       ;; case 1L 
	       ((($ zipper 'Left 'Black xk x ($ tree 'Tree 'Red c yk y d) z) . a)
		(bbZip (Left R xk x c (Left B yk y d z)) a))
	       ;; case 3L 
	       ((($ zipper 'Left color xk x ($ tree 'Tree 'Black ($ tree 'Tree 'Red c yk y d) wk w e) z) . a)
		(bbZip (Left color xk x (Tree B c yk y (Tree R d wk w e)) z) a))
	       ;; case 4L 
	       ((($ zipper 'Left color xk x ($ tree 'Tree 'Black c yk y ($ tree 'Tree 'Red d wk w e)) z) . a)
		(cons #f (zip z (Tree color (Tree B a xk x c) yk y (Tree B d wk w e)))))
	       ;; case 2L 
	       ((($ zipper 'Left 'Red xk x ($ tree 'Tree 'Black c yk y d) z) . a)
		(cons #f (zip z (Tree B a xk x (Tree R c yk y d)))))
	       ;; case 2L 
	       ((($ zipper 'Left 'Black xk x ($ tree 'Tree 'Black c yk y d) z) . a)
		(bbZip z (Tree B a xk x (Tree R c yk y d))))
	       ;; case 1R 
	       ((($ zipper 'Right color ($ tree 'Tree 'Red c yk y d) xk x z) . b) 
		(bbZip (Right R d xk x (Right B c yk y z)) b))
	       ;; case 3R 
	       ((($ zipper 'Right color ($ tree 'Tree 'Black ($ tree 'Tree 'Red c wk w d) yk y e) xk x z) . b) 
		(bbZip (Right color (Tree B c wk w (Tree R d yk y e)) xk x z) b))
	       ;; case 4R
	       ((($ zipper 'Right color ($ tree 'Tree 'Black c yk y ($ tree 'Tree 'Red d wk w e)) xk x z) . b) 
		(cons #f (zip z (Tree color c yk y (Tree B (Tree R d wk w e) xk x b)))))
	       ;; case 2R 
	       ((($ zipper 'Right 'Red ($ tree 'Tree 'Black c yk y d) xk x z) . b) 
		(cons #f (zip z (Tree B (Tree R c yk y d) xk x b))))
	       ;; case 2R 
	       ((($ zipper 'Right 'Black ($ tree 'Tree 'Black c yk y d) xk x z) . b) 
		(bbZip z (Tree B (Tree R c yk y d) xk x b)))
	       (else   (cons #f (zip zipper tree)))))

      (define (delMin tree z)
	(match tree
	       (($ tree 'Tree 'Red ($ tree 'Empty) yk y b) 
		(values yk y (cons #f (zip z b))))
	       (($ tree 'Tree 'Black ($ tree Empty) yk y b) 
		(values yk y (bbZip z b)))
	       (($ tree 'Tree color a yk y b) 
		(delMin a (Left color yk y b z)))
	       (($ tree 'Empty) (rb-tree:error 'delete! "invalid tree"))))

      (define (join color a b z)
	(match (list color a b)
	       (( 'Red ($ tree 'Empty) ($ tree 'Empty))  
		(zip z (Empty)))
	       (( _ a ($ tree 'Empty))  
		(cdr  (bbZip z a)))
	       (( _ ($ tree 'Empty) b)
		(cdr  (bbZip z b)))
	       (( color a b)
		(let-values (((xk x b)  (delMin b (Top))))
		    (match b
			   ((#t . b1)  (cdr  (bbZip z (Tree color a xk x b1))))
			   ((#f . b1)  (zip z (Tree color a xk x b1))))))))


      (define (del tree z)
	(match tree 
	       (($ tree 'Empty)  #f)
	       (($ tree 'Tree color a yk y b)  
		(dispatch-on-key delete-key-compare 
		 key yk 
		 (del a (Left color yk y b z))
		 (cons (cons yk y) (join color a b z))
		 (del b (Right color a yk y z))))))

      (del root (Top)))

    (define (delete! key)
      (let ((item+tree  (delete root key)))
	(and item+tree
	     (begin
	       (set! root (cdr item+tree))
	       (set! size (- size 1))
	       (car item+tree)))))


    (define (get-min)
      (define (f root)
	(match root
	       (($ tree 'Empty)  #f)
	       (($ tree 'Tree _ _ ($ tree 'Empty) xk x _)  (cons xk x))
	       (($ tree 'Tree _ a _ _ _)   (f a))))
      (f root))

    (define (get-max)
      (define (f root)
	(match root
	       (($ tree 'Empty)  #f)
	       (($ tree 'Tree _ _ xk x ($ tree 'Empty))  (cons xk x))
	       (($ tree 'Tree _ _ _ _ b)   (f b))))
      (f root))


    (define (fold-limit p f init) 
      (define (foldf tree ax)
	(match tree
	       (($ tree 'Empty)  ax)
	       (($ tree 'Tree _ a _ x b)  
		(if (p ax) ax (foldf b (f x (foldf a ax)))))))
      (foldf root init))

    (define (fold-right-limit p f init) 
      (define (foldf tree ax)
	(match tree
	       (($ tree 'Empty)  ax)
	       (($ tree 'Tree _ a _ x b)  
		(if (p ax) ax (foldf a (f x (foldf b ax)))))))
      (foldf root init))

    (define (fold-partial p f init) 
      (define (foldf tree ax)
	(match tree
	       (($ tree 'Empty)  ax)
	       (($ tree 'Tree _ a _ x b)  
		(if (p x) (foldf b (f x (foldf a ax))) ax))))
      (foldf root init))

    (define (foldi-partial p f init) 
      (define (foldf tree ax)
	(match tree
	       (($ tree 'Empty)  ax)
	       (($ tree 'Tree _ a xk x b)  
		(if (p xk x) (foldf b (f xk x (foldf a ax))) ax))))
      (foldf root init))

    (define (fold-right-partial p f init) 
      (define (foldf tree ax)
	(match tree
	       (($ tree 'Empty)  ax)
	       (($ tree 'Tree _ a _ x b)  
		(if (p x) (foldf a (f x (foldf b ax))) ax))))
      (foldf root init))

    (define (foldi-right-partial p f init) 
      (define (foldf tree ax)
	(match tree
	       (($ tree 'Empty)  ax)
	       (($ tree 'Tree _ a xk x b)  
		(if (p xk x) (foldf a (f xk x (foldf b ax))) ax))))
      (foldf root init))


    (define (fold f init)
      (define (foldf tree ax)
	(match tree
	       (($ tree 'Empty)  ax)
	       (($ tree 'Tree _ a _ x b)  (foldf b (f x (foldf a ax))))))
      (foldf root init))

    (define (foldi f init)
      (define (foldf tree ax)
	(match tree
	       (($ tree 'Empty)  ax)
	       (($ tree 'Tree _ a xk x b)  (foldf b (f xk x (foldf a ax))))))
      (foldf root init))

    (define (fold-right f init)
      (define (foldf tree ax)
	(match tree
	       (($ tree 'Empty)  ax)
	       (($ tree 'Tree _ a _ x b)  (foldf a (f x (foldf b ax))))))
      (foldf root init))

    (define (foldi-right f init)
      (define (foldf tree ax)
	(match tree
	       (($ tree 'Empty)  ax)
	       (($ tree 'Tree _ a xk x b)  (foldf a (f xk x (foldf b ax))))))
      (foldf root init))


    (define (get-depth)
      (let loop ((node root) (level 0))
	(match node 
	       (($ tree 'Empty)  level)
	       (($ tree 'Tree _ a _ _ b)  (max (loop a (+ 1 level))
						(loop b (+ 1 level)))))))

    ;; Returns an ordered list of the keys in the tree
    (define (list-keys)
      (foldi-right (lambda (k x l) (cons k l)) (list)))

    ;; Returns an ordered list of the (key . item) pairs in the tree
    (define (list-items)
      (foldi-right (lambda (k x l) (cons (cons k x) l)) (list)))

    (define (for-each-ascending f)
      (define (appf tree)
	(match tree
	       (($ tree 'Empty)  (void))
	       (($ tree 'Tree _ a k x b)  (begin (appf a) (f (cons k x)) (appf b)))))
      (appf root))

    (define (for-each-descending f)
      (define (appf tree)
	(match tree
	       (($ tree 'Empty)  (void))
	       (($ tree 'Tree _ a k x b)  (begin (appf b) (f (cons k x)) (appf a)))))
      (appf root))

    (define (map f)
      (define (mapf tree)
	(match tree
	       (($ tree 'Empty)  (Empty))
	       (($ tree 'Tree color a xk x b)  
		(Tree color (mapf a) xk (f x) (mapf b)))))
      (make-rb-tree-dispatcher (mapf root) size))

    (define (mapi f)
      (define (mapf tree)
	(match tree
	       (($ tree 'Empty)   (Empty))
	       (($ tree 'Tree color a xk x b)  
		 (Tree color (mapf a) xk (f xk x) (mapf b)))))
      (make-rb-tree-dispatcher  (mapf root) size))


    (define (apply-default-clause label key default-clause)
      (cond
        ((null? default-clause)
          (rb-tree:error label "key " key " was not found in the tree"))
        ((pair? (cdr default-clause))
          (rb-tree:error label "default argument must be a single clause"))
        ((procedure? (car default-clause)) ((car default-clause)))
        (else (car default-clause))))
    
    ;; Dispatcher
    (lambda (selector)
      (case selector
        ((get)
          (lambda (key . default-clause)
            (or (find-assoc key) (apply-default-clause 'get key default-clause))))

        ((delete!)
          (lambda (key . default-clause)
            (or (delete! key)
		(apply-default-clause 'delete! key default-clause))))

        ((delete)
          (lambda (key . default-clause)
            (or (let ((item+tree  (delete root key)))
		  (and item+tree
		       (make-rb-tree-dispatcher  (cdr item+tree) 
						 (if (car item+tree) (- size 1) size))))
		(apply-default-clause 'delete key default-clause))))

        ((put!) 
	 (lambda (key value)
	   (let-values (((found? new-root)  (insert root key value)))
		       (set! root new-root)
		       (if (not found?)  (set! size (+ 1 size)))
		       found?)))

        ((put) 
	 (lambda (key value)
	   (let-values (((found? new-root)  (insert root key value)))
		       (make-rb-tree-dispatcher  new-root (if (not found?) (+ 1 size) size)))))


        ((get-min) (get-min))
        ((get-max) (get-max))

        ((delete-min!)   (delete! (car (get-min))))
        ((delete-max!)   (delete! (car (get-max))))

        ((empty?)  (cases tree root
			  (Empty () #t)
			  (else #f)))
        ((size)    size)

        ((depth)   (get-depth))

        ((clear!)  (begin
		     (set! root (Empty)) 
		     (set! size 0)))

        ((for-each-ascending)   for-each-ascending)
        ((for-each-descending)  for-each-descending)

	((list-keys)            (list-keys))
	((list-items)           (list-items))
	
	((map)                  map)
	((mapi)                 mapi)

	((fold)                fold)
	((foldi)               foldi)
	((fold-right)          fold-right)
	((foldi-right)         foldi-right)

	((fold-partial)        fold-partial)
	((foldi-partial)       foldi-partial)
	((fold-right-partial)          fold-right-partial)
	((foldi-right-partial)         foldi-right-partial)

	((fold-limit)          fold-limit)
	((fold-right-limit)    fold-right-limit)

        (else
          (rb-tree:error 'selector "unknown message " selector " sent to a red-black tree")))))

  (make-rb-tree-dispatcher root size)))

)