(declare (unit util-bst))

(import duck)

(module*
 util-bst
 #:doc ("Binary Search Tree implementation")
 (
  make-bst
  bst?

  bst-EQ?

  bst-count

  bst-empty?
  bst-ref
  bst-contains?

  bst-set
  bst-remove

  bst-keys

  bst->kvv
  kvv->bst
  kvv-filter
  
  bst-balance

  bst-find-pair
  bst-filter-pairs

  bst-map-list
  bst-map-bst

  list->bst

  bst-update

  bst-filter
  bst-reduce

  bst-equal?

  util-bst-tests!
  )

 (import scheme
	 (chicken condition)
	 (chicken sort)
	 util-tag
	 testing
	 racket-kwargs)

 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 ;; Node

 ;; Creates BST node with no children
 (define (make-bst-node key value)
   (cons (cons key value)
	 (cons #f #f)))

 ;; Read-only accessors to BST node
 (define bst-node-kv car)
 (define bst-node-key caar)
 (define bst-node-value cdar)
 (define bst-node-left cadr)
 (define bst-node-right cddr)

 ;; Returns BST node with updated node value
 (define (set-bst-node-value n v)
   (cons (cons (bst-node-key n) v)
	 (cdr n)))

 ;; Updates BST node left child
 (define (set-bst-node-left n l)
   (cons (car n)
	 (cons l (bst-node-right n))))

 ;; Updates BST node right child
 (define (set-bst-node-right n r)
   (cons (car n)
	 (cons (bst-node-left n) r)))

 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 ;; BST

 ;; Unique tag
 (define TAG-BST (make-tag bst))

 (define/doc (make-bst subtag EQ? <?)
   ("Creates empty BST with given comparators")
   (cons TAG-BST
	 (cons (cons #f 0)
	       (cons subtag
		     (cons EQ? <?)))))

 ;; BST accessors
 (define bst-tag car)
 (define bst-root+count cadr)
 (define bst-root caadr)
 (define bst-count cdadr)
 (define bst-type cddr)
 (define bst-subtag caddr)
 (define bst-comparators cdddr)
 (define bst-EQ? cadddr)
 (define bst-<? cddddr)

 ;; Update BST root node
 (define (set-bst-root bst root)
   (cons (car bst)
	 (cons (cons root
		     (bst-count bst))
	       (cddr bst))))

 ;; Update BST count
 (define (set-bst-count bst count)
   (cons (car bst)
	 (cons (cons (bst-root bst)
		     count)
	       (cddr bst))))

 ;; Update BST root node and count
 (define (set-bst-root+count bst root count)
   (cons (car bst)
	 (cons (cons root count)
	       (cddr bst))))

 (define/doc ((bst? subtag) v)
   ("Curried predicate for particular bst type.")
   (and (pair? v)
	(eq? (bst-tag v) TAG-BST)
	(eq? (bst-subtag v) subtag)))

 (define/doc (bst-empty? bst)
   ("Returns #t if given BST is empty.")
   (not (bst-root bst)))

 ;; Wrapper to setup comparators
 (define-syntax let-comparators
   (syntax-rules ()
     ((_ (EQ? <? bst) expr ...)
      (let ((EQ? (bst-EQ? bst))
	    (<? (bst-<? bst)))
	expr ...))))

 (define/doc (bst-ref bst k . vs)
   ("Retrieves value associated with given key.")
   (let-comparators
    (EQ? <? bst)
    (let loop ((n (bst-root bst)))
      (if n
	  (let ((nk (bst-node-key n)))
	    (if (EQ? k nk)
		(bst-node-value n)
		(loop (if (<? k nk)
			  (bst-node-left n)
			  (bst-node-right n)))))
	  (if (null? vs)
	      (error 'bst-ref "Key does not exist" k)
	      (car vs))))))

 (define/doc (bst-contains? bst k)
   ("Predicate for key existence in BST.")
   (handle-exceptions
    ex
    #f
    (let ()
      (bst-ref bst k)
      #t)))

 (define/doc (bst-set bst k v)
   ("Sets given key to given value and updates count if needed.")
   (let-comparators
    (EQ? <? bst)
    (let-values (((new-root count-add)
		  (let loop ((n (bst-root bst)))
		    (if n
			(let ((nk (bst-node-key n)))
			  (if (EQ? k nk)
			      (values (set-bst-node-value n v) 0)
			      (if (<? k nk)
				  (let-values (((new-left-node count-add)
						(loop (bst-node-left n))))
				    (values (set-bst-node-left n new-left-node)
					    count-add))
				  (let-values (((new-right-node count-add)
						(loop (bst-node-right n))))
				    (values (set-bst-node-right n new-right-node)
					    count-add)))))
			(values (make-bst-node k v) 1)))))
      (set-bst-root+count bst
			  new-root
			  (+ (bst-count bst) count-add)))))

 (define/doc (bst-remove bst k . nos)
   ("Removes given key from the BST.")
   (let-comparators
    (EQ? <? bst)
    (let-values
	(((new-root sub-count)
	  (let loop ((n (bst-root bst)))
	    (if n
		(let ((nk (bst-node-key n)))
		  (if (EQ? nk k)
		      (values (if (bst-node-left n)
				  (if (bst-node-right n)
				      (let aloop ((an (bst-node-right n)))
					(if an
					    (set-bst-node-left an
							       (aloop (bst-node-left an)))
					    (bst-node-left n)))
				      (bst-node-left n))
				  (if (bst-node-right n)
				      (bst-node-right n)
				      #f))
			      1)
		      (if (<? k nk)
			  (let-values (((new-left-node sub-count)
					(loop (bst-node-left n))))
			    (values (set-bst-node-left n new-left-node)
				    sub-count))
			  (let-values (((new-right-node sub-count)
					(loop (bst-node-right n))))
			    (values (set-bst-node-right n new-right-node)
				    sub-count)))))
		(if (and (not (null? nos))
			 (car nos))
		    (values #f 0)
		    (error 'bst-remove "Key does not exist" k))))))
      (if (eq? sub-count 0)
	  bst
	  (set-bst-root+count bst
			      new-root
			      (- (bst-count bst) sub-count))))))

 ;; Unique tag for static vector representation that retains type info
 (define TAG-KVV (make-tag kvv))

 ;; Iterates over given BST and passes the KV pairs to given proc
 (define (bst-iter-kv bst proc)
   (let-comparators
    (EQ? <? bst)
    (let loop ((n (bst-root bst)))
      (when n
	(loop (bst-node-left n))
	(proc (bst-node-kv n))
	(loop (bst-node-right n))))))

 ;; Used for iterating over multiple BSTs simultaneously
 (define (bst-kv-iterator bst)
   (let* ((break #f)
	  (resume #f)
	  (yield (lambda (val)
		   (call/cc
		    (lambda (r)
		      (set! resume r)
		      (break val))))))
     (lambda ()
       (call/cc
	(lambda (cc)
	  (set! break cc)
	  (if resume
	      (resume '())
	      (bst-iter-kv bst yield))
	  #f)))))

 (define/doc (bst-keys bst)
   ("Returns all the keys contained in given dictionary.")
   (let ((res '()))
     (bst-iter-kv bst
		  (lambda (kv)
		    (set! res (cons (car kv) res))))
     (reverse res)))

 ;; Converts BST to KVV
 (define (bst->kvv bst)
   (let ((vec (make-vector (bst-count bst)))
	 (idx 0))
     (bst-iter-kv bst
		  (lambda (kv)
		    (vector-set! vec idx kv)
		    (set! idx (add1 idx))))
     (cons TAG-KVV
	   (cons vec
		 (cddr bst)))))

 ;; Converts a vector of key-value cons cells to BST dictionary
 (define (kvv->bst kvv)
   (let* ((vec (cadr kvv))
	  (count (vector-length vec)))
     (cons TAG-BST
	   (cons (cons (let loop ((s 0)
				  (e count))
			 (if (= s e)
			     #f
			     (let* ((c (- e s))
				    (h (quotient c 2))
				    (m (+ s h))
				    (kv (vector-ref vec m)))
			       (cons kv
				     (cons (loop s m)
					   (loop (add1 m) e))))))
		       count)
		 (cddr kvv)))))

 ;; Returns KVV that contains only KV pairs matching pred?
 (define (kvv-filter kvv pred?)
   (cons (car kvv)
	 (cons (apply vector
		      (let ((vec (cadr kvv)))
			(let loop ((idx (sub1 (vector-length vec)))
				   (res '()))
			  (if (< idx 0)
			      res
			      (let ((kv (vector-ref vec idx)))
				(loop (if (pred? (car kv) (cdr kv))
					  (cons kv res)
					  res)
				      (sub1 idx)))))))
	       (cddr kvv))))

 (define/doc (bst-balance bst)
   ("Unconditionally balances the BST.")
   (kvv->bst
    (bst->kvv bst)))

 (define/doc (bst-find-pair bst pred?)
   ("Finds pair by predicate that accepts both key and value.")
   (call/cc
    (lambda (cc)
      (bst-iter-kv bst
		   (lambda (kv)
		     (when (pred? (car kv) (cdr kv))
		       (cc kv))))
      #f)))

 (define/doc (bst-filter-pairs bst pred?)
   ("Returns a list of key-value pairs matching predicate.")
   (let ((res '()))
     (bst-iter-kv bst
		  (lambda (kv)
		    (when (pred? (car kv) (cdr kv))
		      (set! res (cons kv res)))))
     (reverse res)))

 (define/doc (bst-map-list bst proc)
   ("Returns arbitrary list created by mapping all elements.")
   (let ((res '()))
     (bst-iter-kv bst
		  (lambda (kv)
		    (set! res
			  (cons res
				(proc (car kv) (cdr kv))))))
     (reverse res)))

 (define/doc (bst-map-bst bst proc)
   ("Returns a new dictionary with all values processed (keys are left intact).")
   (set-bst-root bst
		 (let loop ((n (bst-root bst)))
		   (if n
		       (cons (cons (caar n) (proc (caar n) (cdar n)))
			     (cons (loop (cadr n))
				   (loop (cddr n))))
		       #f))))

 (define/doc (list->bst lst subtag EQ? <?)
   ("Converts list of pairs into BST dictionary.")
   (kvv->bst
    (cons TAG-KVV
	  (cons (list->vector
		 (sort lst (lambda (akv bkv)
			     (<? (car akv) (car bkv)))))
		(cons subtag
		      (cons EQ? <?))))))

 (define*/doc (bst-update bst k proc (v #f))
   ("Functional update with optional default value (defaults to #f).")
   (let-comparators
    (EQ? <? bst)
    (let-values	(((new-root add-count)
		  (let loop ((n (bst-root bst)))
		    (if n
			(let ((nk (bst-node-key n)))
			  (if (EQ? k nk)
			      (values (set-bst-node-value n (proc (bst-node-value n)))
				      0)
			      (if (<? k nk)
				  (let-values (((new-left-node add-count)
						(loop (bst-node-left n))))
				    (values (set-bst-node-left n new-left-node)
					    add-count))
				  (let-values (((new-right-node add-count)
						(loop (bst-node-right n))))
				    (values (set-bst-node-right n new-right-node)
					    add-count)))))
			(values (make-bst-node k (proc v)) 1)))))
      (set-bst-root+count bst
			  new-root
			  (+ (bst-count bst) add-count)))))

 (define/doc (bst-filter bst pred?)
   ("Returns a BST with only KV pairs matching the predicate which must
accept two arguments.")
   (kvv->bst
    (cons TAG-KVV
	  (cons (apply vector (bst-filter-pairs bst pred?))
		(cddr bst)))))

 (define/doc (bst-reduce bst proc init)
   ("Like generic reduce, the proc gets accumulator, key and value
arguments.")
   (let ((acc init))
     (bst-iter-kv bst
		  (lambda (kv)
		    (set! acc (proc acc (car kv) (cdr kv)))))
     acc))

 ;; Returns true if these are compatible BSTs
 (define (bst-compat? b1 b2)
   (and (eq? (bst-tag b1)
	     (bst-tag b2))
	(eq? (bst-subtag b1)
	     (bst-subtag b2))
	(eq? (bst-EQ? b1)
	     (bst-EQ? b2))
	(eq? (bst-<? b1)
	     (bst-<? b2))))

 (define*/doc (bst-equal? b1 b2 (equality? equal?))
   ("Returns true if both BSTs contain the same keys and values.")
   (if (bst-compat? b1 b2)
       (let-comparators
	(EQ? _ b1)
	(let ((g1 (bst-kv-iterator b1))
	      (g2 (bst-kv-iterator b2)))
	  (let loop ()
	    (let ((kv1 (g1))
		  (kv2 (g2)))
	      (if (and (not kv1)
		       (not kv2))
		  #t
		  (if (and (EQ? (car kv1)
				(car kv2))
			   (equality? (cdr kv1)
				      (cdr kv2)))
		      (loop)
		      #f))))))
       #f))

 ;; Module self-tests
 (define (util-bst-tests!)
   (run-tests
    util-bst
    (test-equal? make-bst
		 (make-bst 'fixnum eq? <)
		 `(,TAG-BST
		   . ((#f . 0)
		      . (fixnum . (,eq? . ,<)))))
    (test-true bst?
	       ((bst? 'fixnum) (make-bst 'fixnum eq? <)))
    (test-false bst?
		((bst? 'fixnum) (make-bst 'string eq? <)))
    (test-false bst?
		((bst? 'fixnum) "string"))
    (test-true bst-empty?
	       (bst-empty? (make-bst 'fixnum eq? <)))
    (test-true bst-contains?
	       (bst-contains?
		(bst-set
		 (make-bst 'fixnum eq? <)
		 1 2) 1))
    (test-false bst-contains?
		(bst-contains? (make-bst 'fixnum eq? <) 1))
    (test-equal? bst-count
		 (bst-count
		  (make-bst 'fixnum eq? <))
		 0)
    (test-equal? bst-set
		 (bst-count
		  (bst-set
		   (make-bst 'fixnum eq? <)
		   1 2))
		 1)
    (test-true bst-remove
	       (bst-empty?
		(bst-remove
		 (bst-set
		  (make-bst 'fixnum eq? <)
		  1 2) 1)))
    (test-equal? bst-balance
		 (bst-count
		  (bst-balance
		   (bst-set
		    (make-bst 'fixnum eq? <)
		    1 2)))
		 1)
    ))

 )