(declare (unit util-bst))

(import duck)

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

  bst-empty?
  bst-ref
  bst-contains?

  bst-set
  ;;bst-remove ;;

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

  util-bst-tests!
  )

 (import scheme
	 (chicken condition)
	 util-tag
	 testing)

 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 ;; 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)))))

 ;; 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-count
		 (bst-count
		  (bst-set
		   (make-bst 'fixnum eq? <)
		   1 2))
		 1)
    ))

 )

(import util-bst)
(util-bst-tests!)