
;; Implements a red-black tree with relative indexing along right
;; splines. This allows the usual O(log(n)) operations, plus a
;; O(log(n)) shift operation.

;; (This is the same data structure as used for lines by MrEd's text%
;; class, but that one is implemented in C++.)

(unit/sig relative-btree^
  (import)
  (rename (create-btree make-btree))
  
  (define-struct btree (root))

  (define-struct node (pos data parent left right color))

  (define (adjust-offsets n new-child)
    (when new-child
      (set-node-pos! new-child (- (node-pos new-child)
				  (node-pos n)))))

  (define (deadjust-offsets n old-child)
    (when old-child
      (set-node-pos! old-child (+ (node-pos old-child)
				  (node-pos n)))))

  (define (rotate-left n btree)
    (let ([old-right (node-right n)])
      (deadjust-offsets n old-right)
      
      (let ([r (node-left old-right)])
	(set-node-right! n r)
	(when r
	  (set-node-parent! r n)))
  
      (let ([p (node-parent n)])
	(set-node-parent! old-right p)
	(cond
	 [(not p) (set-btree-root! btree old-right)]
	 [(eq? n (node-left p)) (set-node-left! p old-right)]
	 [else                 (set-node-right! p old-right)]))

      (set-node-left! old-right n)
      (set-node-parent! n old-right)))

  (define (rotate-right n btree)
    (let ([old-left (node-left n)])
      (adjust-offsets old-left n)
      
      (let ([l (node-right old-left)])
	(set-node-left! n l)
	(when l
	  (set-node-parent! l n)))
  
      (let ([p (node-parent n)])
	(set-node-parent! old-left p)
	(cond
	 [(not p) (set-btree-root! btree old-left)]
	 [(eq? n (node-left p)) (set-node-left! p old-left)]
	 [else                 (set-node-right! p old-left)]))

      (set-node-right! old-left n)
      (set-node-parent! n old-left)))


  (define (insert before? n btree pos data)
    (let ([new (make-node pos data #f #f #f 'black)])
      (if (not (btree-root btree))
	  (set-btree-root! btree new)

	  (begin

	    (set-node-color! new 'red)
	    
	    ; Insert into tree
	    (if before?

		(if (not (node-left n))
		    (begin
		      (set-node-left! n new)
		      (set-node-parent! new n))
		    
		    (let loop ([node (node-left n)])
		      (if (node-right node)
			  (loop (node-right node))
			  (begin
			    (set-node-right! node new)
			    (set-node-parent! new node)))))

		(if (not (node-right n))
		    (begin
		      (set-node-right! n new)
		      (set-node-parent! new n))

		    (let loop ([node (node-right n)])
		      (if (node-left node)
			  (loop (node-left node))
			  (begin
			    (set-node-left! node new)
			    (set-node-parent! new node))))))

	    ; Make value in new node relative to right-hand parents
	    (let loop ([node new])
	      (let ([p (node-parent node)])
		(when p
		  (when (eq? node (node-right p))
		    (adjust-offsets p new))
		  (loop p))))

	    ; Balance tree
	    (let loop ([node new])
	      (let ([p (node-parent node)])
		(when (and (not (eq? node (btree-root btree)))
			   (eq? 'red (node-color p)))
		  (let* ([recolor-k
			  (lambda (y)
			    (set-node-color! p 'black)
			    (set-node-color! y 'black)
			    (let ([pp (node-parent p)])
			      (set-node-color! pp 'red)
			      (loop pp)))]
			 [rotate-k
			  (lambda (rotate node)
			    (let ([p (node-parent node)])
			      (set-node-color! p 'black)
			      (let ([pp (node-parent p)])
				(set-node-color! pp 'red)
				(rotate pp btree)
				(loop pp))))]
			 [k
			  (lambda (node-y long-rotate always-rotate)
			    (let ([y (node-y (node-parent p))])
			      (if (and y (eq? 'red (node-color y)))
				  (recolor-k y)
				  (let ([k (lambda (node)
					     (rotate-k always-rotate node))])
				    (if (eq? node (node-y p))
					(begin
					  (long-rotate p btree)
					  (k p))
					(k node))))))])
		    (if (eq? p (node-left (node-parent p)))
			(k node-right rotate-left rotate-right)
			(k node-left rotate-right rotate-left))))))
	    
	    (set-node-color! (btree-root btree) 'black)))))

  (define (find-following-node btree pos)
    (let ([root (btree-root btree)])
      (let loop ([n root]
		 [so-far root]
		 [so-far-pos (and root (node-pos root))]
		 [v 0])
	(if (not n) 
	    (values so-far so-far-pos)
	    (let ([npos (+ (node-pos n) v)])
	      (cond
	       [(<= pos npos)
		(loop (node-left n) n npos v)]
	       [(or (not so-far-pos)
		    (> npos so-far-pos))
		(loop (node-right n) n npos npos)]
	       [else
		(loop (node-right n) so-far so-far-pos npos)]))))))
  
  (define (create-btree)
    (make-btree #f))

  (define (btree-get btree pos)
    (let-values ([(n npos) (find-following-node btree pos)])
      (and n
	   (= npos pos)
	   (node-data n))))

  (define (btree-put! btree pos data)
    (let-values ([(n npos) (find-following-node btree pos)])
      (if (and n (= npos pos))
	  (set-node-data! n data)
	  (insert (and n (< pos npos))
		  n btree pos data))))
  
  (define (btree-shift! btree start delta)
    (let loop ([n (btree-root btree)]
	       [v 0])
      (when n
	(let ([npos (node-pos n)])
	  (cond
	   [(< start (+ v npos))
	    (set-node-pos! n (+ npos delta))
	    (loop (node-left n) v)]
	   [else
	    (loop (node-right n) (+ v npos))])))))
  
  (define (btree-for-each btree f)
    (when (btree-root btree)
      (let loop ([n (btree-root btree)]
		 [v 0])
	(when (node-left n)
	  (loop (node-left n) v))
	(f (+ v (node-pos n)) (node-data n))
	(when (node-right n)
	  (loop (node-right n)
		(+ v (node-pos n)))))))

  (define (btree-map btree f)
    (let ([start (cons #f null)])
      (let loop ([n (btree-root btree)]
		 [l start]
		 [v 0])
	(if (not n)
	    l
	    (let ([pre (loop (node-left n) l v)]
		  [here (cons (f (+ v (node-pos n))
				 (node-data n))
			      null)])
	      (set-cdr! pre here)
	      (loop (node-right n)
		    here
		    (+ v (node-pos n))))))
      (cdr start))))
