
;; Two variants of red-black trees.

;; In the expunge variant, each node stores the count of nodes in its
;; left branch. A nodes' value corresponds to a pre-delete index, and
;; subtracting the right-branch count produces a post-delete index
;; (i.e., as if all deleteions for earlier pre-delete indices happened
;; first).  Inserting into the tree provides a post-delete index, so
;; the insert must use the left-branch count to compare the node's
;; pre-delete index to the given post-delete index. Rotations to
;; maintain the red-black property can locally update the left-branch
;; count stored in the rotated nodes. It's not possible to insert the
;; same pre-delete index twice, since insertion uses the post-delete
;; index.

;; In the fetch variant, each node's value is a list, where the first
;; number in the list is the key. But for a right branch, all keys are
;; decremented by the node's key (recursively). This allows a ln-time
;; shift operation when a message is expunged.

(module rbtree mzscheme
  (provide new-tree tree-empty?
	   expunge-insert! expunge-tree->list
	   fetch-insert! fetch-find fetch-delete! fetch-shift! fetch-tree->list)

  (define-struct tree (v red? left-count left right parent) (make-inspector))

  (define (new-tree)
    (make-tree 'pre-root #f 0 #f #f #f))

  (define (tree-empty? t)
    (not (tree-left t)))

  (define (k+ a b)
    (cons (+ (car a) (if (number? b) b (car b)))
	  (cdr a)))
  (define (k- a b)
    (cons (- (car a) (if (number? b) b (car b)))
	  (cdr a)))
  (define kv car)

  (define (mk-insert sort-to-left? sort=? right+ 
		     left-insert-adjust!
		     left-rotate-adjust! right-rotate-adjust!)
    (define-values (rotate-left! rotate-right!)
      (let ([mk
	     (lambda (tree-west tree-east set-tree-west! set-tree-east! adj-count!)
	       (lambda (t)
		 (let ([old-east (tree-east t)])
		   (let ([r (tree-west old-east)])
		     (set-tree-east! t r)
		     (when r
		       (set-tree-parent! r t)))
		   (let ([p (tree-parent t)])
		     (set-tree-parent! old-east p)
		     (if (eq? t (tree-left p))
			 (set-tree-left! p old-east)
			 (set-tree-right! p old-east)))
		   (set-tree-west! old-east t)
		   (set-tree-parent! t old-east)
		   (adj-count! t old-east))))])
	(values (mk tree-left tree-right set-tree-left! set-tree-right!
		    left-rotate-adjust!)
		(mk tree-right tree-left set-tree-right! set-tree-left!
		    right-rotate-adjust!))))

    (values
     ;; insert
     (lambda (pre-root n)
       (let ([new
	      ;; Insert:
	      (let loop ([t (tree-left pre-root)] 
			 [n n] 
			 [parent pre-root] 
			 [set-child! (lambda (t v)
				       (set-tree-left! pre-root v))])
		(cond
		 [(not t) (let ([new (make-tree n #t 0 #f #f parent)])
			    (set-child! parent new)
			    new)]
		 [(sort=? n t)
		  (set-tree-v! t n)
		  pre-root]
		 [(sort-to-left? n t)
		  (left-insert-adjust! t)
		  (loop (tree-left t) n t set-tree-left!)]
		 [else
		  (loop (tree-right t) (right+ n t) t set-tree-right!)]))])
	 ;; Restore red-black property:
	 (let loop ([v new])
	   (let ([p (tree-parent v)])
	     (when (and p (tree-red? p))
	       (let ([gp (tree-parent p)])
		 (let-values ([(tree-west tree-east rotate-west! rotate-east!)
			       (if (eq? p (tree-left gp))
				   (values tree-left tree-right rotate-left! rotate-right!)
				   (values tree-right tree-left rotate-right! rotate-left!))])
		   (let ([uncle (tree-east (tree-parent p))])
		     (if (and uncle (tree-red? uncle))
			 (begin
			   (set-tree-red?! p #f)
			   (set-tree-red?! uncle #f)
			   (set-tree-red?! gp #t)
			   (loop gp))
			 (let ([finish (lambda (v)
					 (let* ([p (tree-parent v)]
						[gp (tree-parent p)])
					   (set-tree-red?! p #f)
					   (set-tree-red?! gp #t)
					   (rotate-east! gp)
					   (loop gp)))])
			   (if (eq? v (tree-east p))
			       (begin
				 (rotate-west! p)
				 (finish p))
			       (finish v))))))))))
	 (set-tree-red?! (tree-left pre-root) #f)))

     ;; delete (fetch only)
     (lambda (pre-root n)
       (let ([orig-t (fetch-find-node pre-root n)])
	 (when orig-t
	   ;; Delete note t if it has at most one child.
	   ;; Otherwise, move a leaf's data to here, and
	   ;; delete the leaf.
	   (let ([t (if (and (tree-left orig-t)
			     (tree-right orig-t))
			(let loop ([t (tree-right orig-t)])
			  (if (tree-left t)
			      (loop (tree-left t))
			      t))
			orig-t)])
	     (unless (eq? t orig-t)
	       ;; Swap out:
	       (let ([delta (kv (tree-v t))])
		 (set-tree-v! orig-t (k+ (tree-v t) (tree-v orig-t)))
		 (let loop ([c (tree-right orig-t)])
		   (when c
		     (set-tree-v! c (k- (tree-v c) delta))
		     (loop (tree-left c))))))
	     ;; Now we can delete t:
	     (let ([child-t (or (tree-left t)
				(tree-right t))]
		   [p (tree-parent t)])
	       (when child-t
		 (set-tree-parent! child-t p)
		 ;; Adjust relative index of left spine of the
		 ;;  right branch (in the case that there was only
		 ;;  a right branch)
		 (let loop ([c (tree-right t)])
		   (when c
		     (set-tree-v! c (k+ (tree-v c) (tree-v t)))
		     (loop (tree-left c)))))
	       (if (eq? (tree-left p) t) 
		   (set-tree-left! p child-t)
		   (set-tree-right! p child-t))
	       ;; Restore red-black property:
	       (when (not (tree-red? t))
		 (let loop ([c child-t] [p p])
		   (cond
		    [(and c (tree-red? c)) (set-tree-red?! c #f)]
		    [(tree-parent p)
		     (let-values ([(tree-west tree-east rotate-west! rotate-east!)
				   (if (eq? c (tree-left p))
				       (values tree-left tree-right rotate-left! rotate-right!)
				       (values tree-right tree-left rotate-right! rotate-left!))])
		       (let ([sibling (tree-east p)])
			 (let ([z (if (tree-red? sibling)
				      (begin
					(set-tree-red?! sibling #f)
					(set-tree-red?! p #t)
					(rotate-west! p)
					(tree-east p))
				      sibling)])
			   (if (not (or (and (tree-west z)
					     (tree-red? (tree-west z)))
					(and (tree-east z)
					     (tree-red? (tree-east z)))))
			       (begin
				 (set-tree-red?! z #t)
				 (loop p (tree-parent p)))
			       (let ([w (if (not (and (tree-east z)
						      (tree-red? (tree-east z))))
					    (begin
					      (set-tree-red?! (tree-west z) #f)
					      (set-tree-red?! z #t)
					      (rotate-east! z)
					      (tree-east p))
					    z)])
				 (set-tree-red?! w (tree-red? p))
				 (set-tree-red?! p #f)
				 (set-tree-red?! (tree-east w) #f)
				 (rotate-west! p))))))]))))))))))

  (define-values (expunge-insert! ---)
    (mk-insert 
     ;; sort-to-left?
     (lambda (n t)
       ((+ n (tree-left-count t)) . < . (tree-v t)))
     ;; sort=?
     (lambda (n t) #f)
     ;; right+
     (lambda (n t)
       (+ n 1 (tree-left-count t)))
     ;; left-insert-adjust!
     (lambda (t)
       (set-tree-left-count! t (add1 (tree-left-count t))))
     ;; left-rotate-adjust!
     (lambda (t old-right)
       (set-tree-left-count! old-right (+ 1 
					  (tree-left-count old-right)
					  (tree-left-count t))))
     ;; right-rotate-adjust!
     (lambda (t old-left)
       (set-tree-left-count! t (- (tree-left-count t)
				  (tree-left-count old-left)
				  1)))))

  (define-values (fetch-insert! fetch-delete!)
    (mk-insert
     ;; sort-to-left?
     (lambda (n t)
       ((kv n) . < . (kv (tree-v t))))
     ;; sort=?
     (lambda (n t)
       (= (kv n) (kv (tree-v t))))
     ;; right+
     (lambda (n t)
       (k- n (tree-v t)))
     ;; left-insert-adjust!
     void
     ;; left-rotate-adjust!
     (lambda (t old-right)
       (set-tree-v! old-right (k+ (tree-v old-right)
				  (tree-v t))))
     ;; right-rotate-adjust!
     (lambda (t old-left)
       (set-tree-v! t (k- (tree-v t)
			  (tree-v old-left))))))
  
  (define (expunge-tree->list pre-root)
    (let loop ([t (tree-left pre-root)])
      (if t
	  (append (loop (tree-left t))
		  (list (tree-v t))
		  (loop (tree-right t)))
	  null)))

  (define (fetch-find-node pre-root n)
    (let loop ([t (tree-left pre-root)]
	       [n n])
      (and t
	   (cond
	    [(= n (kv (tree-v t))) t]
	    [(< n (kv (tree-v t))) (loop (tree-left t) n)]
	    [else (loop (tree-right t) (- n (kv (tree-v t))))]))))

  (define (fetch-find pre-root n)
    (let ([t (fetch-find-node pre-root n)])
      (and t (tree-v t))))

  (define (fetch-shift! pre-root n)
    (fetch-delete! pre-root n)
    (let loop ([t (tree-left pre-root)]
	       [n n])
      (when t
	(if (n . < . (kv (tree-v t)))
	    (begin
	      (set-tree-v! t (k- (tree-v t) 1))
	      (loop (tree-left t) n))
	    (loop (tree-right t)
		  (- n (kv (tree-v t))))))))

  (define (fetch-tree->list pre-root)
    (let loop ([t (tree-left pre-root)][d 0])
      (if t
	  (append (loop (tree-left t) d)
		  (list (k+ (tree-v t) d))
		  (loop (tree-right t) (+ d (kv (tree-v t)))))
	  null))))

#|

Tests:

(require rbtree)
(require (lib "pretty.ss"))
(print-struct #t)

(define t (new-tree))
(define (test n v)
  (expunge-insert! t n)
  (unless (equal? (expunge-tree->list t) v)
    (error 'bad "~s != ~s" (tree->list t) v)))

(test 12 '(12))
(test 8 '(8 12))
(test 1 '(1 8 12))

(define t (new-tree))

(test 10 '(10))
(test 8 '(8 10))
(test 10 '(8 10 12))
(test 8 '(8 9 10 12))
(test 8 '(8 9 10 11 12))
(test 100 '(8 9 10 11 12 105))
(test 1 '(1 8 9 10 11 12 105))
(test 105 '(1 8 9 10 11 12 105 112))
(test 99 '(1 8 9 10 11 12 105 106 112))
(test 2 '(1 3 8 9 10 11 12 105 106 112))
(test 6 '(1 3 8 9 10 11 12 13 105 106 112))
(test 5 '(1 3 7 8 9 10 11 12 13 105 106 112))
(test 15 '(1 3 7 8 9 10 11 12 13 24 105 106 112))
(test 15 '(1 3 7 8 9 10 11 12 13 24 25 105 106 112))
(test 15 '(1 3 7 8 9 10 11 12 13 24 25 26 105 106 112))

(define t (new-tree))
(define (test n v)
  (cond
   [(< n 0) (fetch-delete! t (- n))]
   [(inexact? n) (fetch-shift! t (inexact->exact n))]
   [else (fetch-insert! t (list n))])
  (printf "Check ~a~n" v)
  (let ([v (map list v)])
    (unless (equal? (fetch-tree->list t) v)
      (error 'bad "~s != ~s" (fetch-tree->list t) v))))

(test 10 '(10))
(test 12 '(10 12))
(test 8 '(8 10 12))
(test 10 '(8 10 12))
(test -10 '(8 12))
(test -10 '(8 12))
(test 10.0 '(8 11))
(test 100.0 '(8 11))
(test 5.0 '(7 10))
(test 1 '(1 7 10))
(test 2 '(1 2 7 10))
(test 3 '(1 2 3 7 10))
(test 4 '(1 2 3 4 7 10))
(test 5 '(1 2 3 4 5 7 10))
(test 6 '(1 2 3 4 5 6 7 10))
(test -6 '(1 2 3 4 5 7 10))
(test -5 '(1 2 3 4 7 10))
(test -4 '(1 2 3 7 10))
(test -3 '(1 2 7 10))
(test -2 '(1 7 10))
(test -7 '(1 10))
(test -1 '(10))
(test -10 '())

(define (in-all-positions n l)
  (if (null? l)
      (list (list n))
      (cons
       (cons n l)
       (map (lambda (r) (cons (car l) r))
	    (in-all-positions n (cdr l))))))

(define (permutations l)
  (if (or (null? l)
	  (null? (cdr l)))
      (list l)
      (apply
       append
       (map (lambda (lol)
	      (in-all-positions (car l) lol))
	    (permutations (cdr l))))))

(define perms (permutations '(1 2 3 4 5 6 7 8)))

(for-each (lambda (l)
	    (let ([t (new-tree)])
	      (for-each (lambda (i)
			  (fetch-insert! t (list i)))
			l)
	      (unless (equal? (fetch-tree->list t) '((1) (2) (3) (4) (5) (6) (7) (8)))
		(error 'perms "bad: ~a" l))
	      (for-each (lambda (i)
			  (fetch-delete! t i))
			l)
	      (unless (equal? (fetch-tree->list t) '())
		(error 'perms "remove bad: ~a" l))))
	  perms)

|#
