; -*- Scheme -*-
;
; $Id: btree.scm,v 1.1 1998/03/16 07:57:39 foner Exp $

;+file-summary
; A balanced binary tree package.
; Written by cph@altdorf.ai.mit.edu
; Based on :-
; Balanced Trees
;  Knuth, Donald E., "The Art of Computer Programming",
;  volume 3, "Sorting and Searching",
;  section 6.2.3, "Balanced Trees".
;
; Modified by Stephen J. Bevan <bevan@cs.man.ac.uk> to include a few
; comments and also make the names acceptable to Aubry Jaffer's
; Scheme.
;-file-summary

(define (make-node wrapped-key left-link right-link balance-factor)
  (vector wrapped-key left-link right-link balance-factor))

(define (get-wrapped-key node) (vector-ref node 0))
(define (set-wrapped-key! node key) (vector-set! node 0 key))

(define (get-left-link node) (vector-ref node 1))
(define (set-left-link! node link) (vector-set! node 1 link))

(define (get-right-link node) (vector-ref node 2))
(define (set-right-link! node link) (vector-set! node 2 link))

(define (get-balance-factor node) (vector-ref node 3))
(define (set-balance-factor! node b) (vector-set! node 3 b))

(define balanced 'BALANCED)		; Knuth's 0
(define tipped-left 'TIPPED-LEFT)	; Knuth's -1
(define tipped-right 'TIPPED-RIGHT)	; Knuth's +1

(define left-d (vector tipped-left tipped-right 1 2))
(define right-d (vector tipped-right tipped-left 2 1))

(define (btree::-d d)
  (cond ((eq? d left-d) right-d)
	((eq? d right-d) left-d)
	(else (error '-d "-D: Bad argument" d))))

(define (btree::+a d)
  (vector-ref d 0))

(define (btree::-a d)
  (vector-ref d 1))

(define (get-link+ p d)
  (vector-ref p (vector-ref d 2)))

(define (get-link- p d)
  (vector-ref p (vector-ref d 3)))

(define (set-link+! p d l)
  (vector-set! p (vector-ref d 2) l))

(define (set-link-! p d l)
  (vector-set! p (vector-ref d 3) l))

(define (cons-path p d path)
  (list p d path))

(define (path-components path receiver)
  (apply receiver path))

;+fs
; This generates a new (empty) balanced binary tree.
;-
(define (make-btree)
  (make-node 'DUMMY-KEY '() '() balanced))

;+fs
; Inserts `k' into the `btree' using '<' as the ordering relation.
; `unwrap-key' should be a function taking one argument, the key `k'.
;   If the key is a structured item, this should return the part of it
;   that is to be used for the comparison.
;   If the key is not a structured item, then this should just return the
;   key.
; `wrap-key' should be a function taking one argument, the key `k'
;   If the key is part of the data to be stored, this should insert `k' into
;   the data to be stored and return this.
;   If the key is not part of the data, then this should just return the
;   data to be stored.
; `if-found' should be a function taking one argument (the key with
;   wrap-key applied to it).
; `if-not-found' should be a function taking one argument (the key with
;   wrap-key applied to it).
;-
(define (btree-insert! btree < unwrap-key k wrap-key if-found if-not-found)
  (let ((p (get-right-link btree)))
    (if (null? p)
	(let ((wk (wrap-key k)))
	  (set-right-link! btree (make-node wk '() '() balanced))
	  (if-not-found wk))
	(let search ((t btree) (set-s-link! set-right-link!) (s p) (p p))
	  (define (move-once set-link! q)
	    (cond ((null? q)
		   (let ((wk (wrap-key k)))
		     (let ((q (make-node wk '() '() balanced)))
		       (set-link! p q)
		       (let ((d (if (< k (unwrap-key (get-wrapped-key s)))
				    left-d
				    right-d)))
			 (let adjust-balance-factors! ((p (get-link+ s d)))
			   (cond ((eq? p q) 'DONE)
				 ((< k (unwrap-key (get-wrapped-key p)))
				  (set-balance-factor! p tipped-left)
				  (adjust-balance-factors!
				   (get-left-link p)))
				 (else
				  (set-balance-factor! p tipped-right)
				  (adjust-balance-factors!
				   (get-right-link p)))))
			 (cond ((eq? (get-balance-factor s) balanced)
				(set-balance-factor! s (btree::+a d)))
			       ((eq? (get-balance-factor s) (btree::-a d))
				(set-balance-factor! s balanced))
			       (else
				(rebalance! s d
				  (lambda (new-s)
				    (set-s-link! t new-s))
				  (lambda (new-s)
				    new-s
				    (error "Tree shouldn't be same height!"
					   'BTREE-INSERT!)))))))
		     (if-not-found wk)))
		  ((eq? (get-balance-factor q) balanced)
		   (search t set-s-link! s q))
		  (else
		   (search p set-link! q q))))
	  (let ((kp (unwrap-key (get-wrapped-key p))))
	    (cond ((< k kp)
		   (move-once set-left-link! (get-left-link p)))
		  ((< kp k)
		   (move-once set-right-link! (get-right-link p)))
		  (else
		   (if-found (get-wrapped-key p)))))))))

(define (btree-delete! btree < unwrap-key k if-found if-not-found)
  (let loop ((p (get-right-link btree))
	     (path (cons-path btree right-d '())))
    (if (null? p)
	(if-not-found k)
	(let ((kp (unwrap-key (get-wrapped-key p))))
	  (cond ((< k kp)
		 (loop (get-left-link p)	 
		       (cons-path p left-d path)))
		((< kp k)
		 (loop (get-right-link p)
		       (cons-path p right-d path)))
		(else
		 (let ((result (get-wrapped-key p)))
		   (cond ((null? (get-left-link p))
			  (replace-node! path (get-right-link p)))
			 ((null? (get-right-link p))
			  (replace-node! path (get-left-link p)))
			 (else
			  (set-wrapped-key!
			   p
			   (get-wrapped-key
			    (remove-successor! (get-right-link p)
					       (cons-path p right-d path))))))
		   (if-found result))))))))

;+fs
; Looks for the item `k' in the balanced tree `btree' using the ordering
; relation '<'.
; `unwrap-key' - a function taking one argument which extracts the actual
;   value to be used in the comparison from the data in the tree.
;   `k' is already assumed to be un-wrapped.
; `if-found' - a function taking one argument, the item that was found.
; `if-not-found' - a function taking one argument, the key `k'.
;-
(define (btree-lookup btree < unwrap-key k if-found if-not-found)
  (let loop ((p (get-right-link btree)))
    (if (null? p)
	(if-not-found k)
	(let ((kp (unwrap-key (get-wrapped-key p))))
	  (cond ((< k kp)
		 (loop (get-left-link p)))
		((< kp k)
		 (loop (get-right-link p)))
		(else
		 (if-found (get-wrapped-key p))))))))

;+fs
; Returns all the elements of the tree as a list.
;-
(define (btree-fringe btree)
  (let loop ((p (get-right-link btree)) (tail '()))
    (if (null? p)
	tail
	(loop (get-left-link p)
	      (cons (get-wrapped-key p)
		    (loop (get-right-link p) tail))))))

(define (remove-successor! p path)
  (if (null? (get-left-link p))
      (begin (replace-node! path (get-right-link p)) p)
      (remove-successor! (get-left-link p)
			 (cons-path p left-d path))))

(define (replace-node! path new-node)
  (path-components path
    (lambda (pl-1 dl-1 rest)
      (set-link+! pl-1 dl-1 new-node)
      (adjust-balance-factors! pl-1 dl-1 rest))))

(define (adjust-balance-factors! pk dk path)
  (cond ((null? path) 'DONE)
	((eq? (get-balance-factor pk) balanced)
	 (set-balance-factor! pk (btree::-a dk)))
	(else
	 (path-components path
	   (lambda (pk-1 dk-1 rest)
	     (if (eq? (get-balance-factor pk) (btree::+a dk))
		 (begin (set-balance-factor! pk balanced)
			(adjust-balance-factors! pk-1 dk-1 rest))
		 (rebalance! pk (btree::-d dk)
			     (lambda (new-pk)
			       (set-link+! pk-1 dk-1 new-pk)
			       (adjust-balance-factors! pk-1 dk-1 rest))
			     (lambda (new-pk)
			       (set-link+! pk-1 dk-1 new-pk)))))))))

(define (rebalance! A d if-shorter if-same-height)
  (let ((B (get-link+ A d)))
    (define (case-1)
      (set-link+! A d (get-link- B d))
      (set-balance-factor! A balanced)
      (set-link-! B d A)
      (set-balance-factor! B balanced)
      (if-shorter B))

    (define (case-2 X)
      (set-link-! B d (get-link+ X d))
      (set-link+! X d B)
      (set-link+! A d (get-link- X d))
      (set-link-! X d A)
      (cond ((eq? (get-balance-factor X) balanced)
	     (set-balance-factor! A balanced)
	     (set-balance-factor! B balanced))
	    ((eq? (get-balance-factor X) (btree::+a d))
	     (set-balance-factor! A (btree::-a d))
	     (set-balance-factor! B balanced))
	    (else
	     (set-balance-factor! A balanced)
	     (set-balance-factor! B (btree::+a d))))
      (set-balance-factor! X balanced)
      (if-shorter X))

    (define (case-3)
      (set-link+! A d (get-link- B d))
      (set-balance-factor! A (btree::+a d))

      (set-link-! B d A)
      (set-balance-factor! B (btree::-a d))
      (if-same-height B))

    (cond ((eq? (get-balance-factor B) (btree::+a d))
	   (case-1))
	  ((eq? (get-balance-factor B) (btree::-a d))
	   (case-2 (get-link- B d)))
	  (else
	   (case-3)))))
