;;
;; util-dict-bst.scm
;;
;; Simple dictionary implementation using BST backend. Only numbers
;; supported as keys.
;;
;; ISC License
;;
;; Copyright 2023 Brmlab, z.s.
;; Dominik Pantůček <dominik.pantucek@trustica.cz>
;;
;; Permission to use, copy, modify, and/or distribute this software
;; for any purpose with or without fee is hereby granted, provided
;; that the above copyright notice and this permission notice appear
;; in all copies.
;; 
;; THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL
;; WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
;; WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE
;; AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR
;; CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS
;; OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
;; NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
;; CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
;;

(declare (unit util-dict-bst))

(module
 util-dict-bst
 (
  make-bdict

  bdict?

  bdict-empty?
  bdict-contains?
  bdict-ref
  
  bdict-set
  bdict-remove
  bdict-update

  bdict-find-pair
  bdict-find-value
  bdict-find-key

  bdict-reduce

  bdict-keys
  bdict-values

  bdict-balance

  bdict-filter-pairs
  bdict-filter-keys
  bdict-filter-values

  list->bdict

  bdict-map-list
  bdict-map-dict
  
  bdict-tests!
  )
 (import scheme
	 (chicken base)
	 util-tag
	 testing)

 ;; Unique tag
 (define TAG-BDICT (make-tag dict-bst))

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

 ;; Creates empty BST dictionary
 (define (make-bdict)
   (cons TAG-BDICT #f))

 ;; Accessor for root
 (define bdict-root cdr)

 ;; Updates BST dictionary root node
 (define (set-bdict-root d r)
   (cons (car d) r))

 ;; Checks whether given value is BST dictionary
 (define (bdict? v)
   (and (pair? v)
	(eq? (car v) TAG-BDICT)))

 ;; Returns true if given dictionary is empty
 (define (bdict-empty? d)
   (and (bdict? d)
	(not (bdict-root d))))

 ;; Returns true if dictionary contains given key
 (define (bdict-contains? d k)
   (let loop ((n (bdict-root d)))
     (if n
	 (let ((nk (bst-node-key n)))
	   (if (eq? k nk)
	       #t
	       (loop (if (< k nk)
			 (bst-node-left n)
			 (bst-node-right n)))))
	 #f)))

 ;; Returns the value associated with given key
 (define (bdict-ref d k . vs)
   (let loop ((n (bdict-root d)))
     (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 'bdict-ref "Key does not exist" k)
	     (car vs)))))

 ;; Returns a dictionary with given key set to given value
 (define (bdict-set d k v)
   (set-bdict-root
    d
    (let loop ((n (bdict-root d)))
      (if n
	  (let ((nk (bst-node-key n)))
	    (if (eq? k nk)
		(set-bst-node-value n v)
		(if (< k nk)
		    (set-bst-node-left n (loop (bst-node-left n)))
		    (set-bst-node-right n (loop (bst-node-right n))))))
	  (make-bst-node k v)))))

 ;; Functional update with optional default value (defaults to #f)
 (define (bdict-update d k p . vs)
   (let ((v (if (null? vs)
		#f
		(car vs))))
     (set-bdict-root
      d
      (let loop ((n (bdict-root d)))
	(if n
	    (let ((nk (bst-node-key n)))
	      (if (eq? k nk)
		  (set-bst-node-value n (p (bst-node-value n)))
		  (if (< k nk)
		      (set-bst-node-left n (loop (bst-node-left n)))
		      (set-bst-node-right n (loop (bst-node-right n))))))
	    (make-bst-node k (p v)))))))

 ;; Finds key-value pair based on predicate - linear search from left
 ;; to right
 (define (bdict-find-pair d p?)
   (let loop ((n (bdict-root d)))
     (if n
	 (let ((l (loop (bst-node-left n))))
	   (or l
	       (if (p? (bst-node-key n) (bst-node-value n))
		   (car n)
		   (loop (bst-node-right n)))))
	 #f)))

 ;; Finds value based on predicate
 (define (bdict-find-value d p?)
   (let ((p (bdict-find-pair d p?)))
     (if p
	 (cdr p)
	 #f)))

 ;; Finds key based on predicate
 (define (bdict-find-key d p?)
   (let ((p (bdict-find-pair d p?)))
     (if p
	 (car p)
	 #f)))

 ;; Returns a dictionary with given key removed, if last argument is
 ;; #t, allows "removing" non-existent keys
 (define (bdict-remove d k . nos)
   (let ((new-root
	  (let loop ((n (bdict-root d)))
	    (if n
		(let ((nk (bst-node-key n)))
		  (if (eq? nk k)
		      (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))
		      (if (< k nk)
			  (set-bst-node-left n (loop (bst-node-left n)))
			  (set-bst-node-right n (loop (bst-node-right n))))))
		(if (and (not (null? nos))
			 (car nos))
		    #f
		    (error 'bdict-remove "Key does not exist" k))))))
     (if new-root
	 (set-bdict-root d new-root)
	 d)))

 ;; Reduce over key-value pairs
 (define (bdict-reduce init proc bd)
   (let loop ((n (bdict-root bd))
	      (acc init))
     (if n
	 (loop (bst-node-right n)
	       (proc (loop (bst-node-left n) acc)
		     (bst-node-key n) (bst-node-value n)))
	 acc)))

 ;; Returns all the keys contained in given dictionary
 (define (bdict-keys d)
   (reverse
    (bdict-reduce '() (lambda (a k v) (cons k a)) d)))

 ;; Returns only values
 (define (bdict-values d)
   (reverse
    (bdict-reduce '() (lambda (a k v) (cons v a)) d)))

 ;; Converts to vector of key-value cons cells
 (define (bdict->vector d)
   (apply
    vector
    (reverse
     (bdict-reduce '() (lambda (a k v) (cons (cons k v) a)) d))))

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

 ;; Returns optimally-balanced version of given dictionary
 (define (bdict-balance d)
   (vector->bdict
    (bdict->vector d)))

 ;; Returns a list of key-value pairs matching predicate
 (define (bdict-filter-pairs d p)
   (bdict-reduce
    '()
    (lambda (a k v)
      (let ((e (p k v)))
	(if e
	    (cons (cons k v) a)
	    a)))
    d))

 ;; Returns a list of keys pairs matching predicate
 (define (bdict-filter-keys d p)
   (bdict-reduce
    '()
    (lambda (a k v)
      (let ((e (p k v)))
	(if e
	    (cons k a)
	    a)))
    d))

 ;; Returns a list of keys pairs matching predicate
 (define (bdict-filter-values d p)
   (bdict-reduce
    '()
    (lambda (a k v)
      (let ((e (p k v)))
	(if e
	    (cons v a)
	    a)))
    d))

 ;; Converts list of pairs into BST dictionary
 (define (list->bdict l)
   (vector->bdict
    (apply vector l)))

 ;; Returns arbitrary list created by mapping all elements
 (define (bdict-map-list d p)
   (reverse
    (bdict-reduce
     '()
     (lambda (a k v)
       (cons (p k v) a))
     d)))

 ;; Returns new dictionary with all values processed (keys are left
 ;; intact)
 (define (bdict-map-dict d p)
   (set-bdict-root
    d
    (let loop ((n (bdict-root d)))
      (if n
	  (cons (cons (caar n) (p (caar n) (cdar n)))
		(cons (loop (cadr n))
		      (loop (cddr n))))
	  #f))))

 ;; Performs module self-tests
 (define (bdict-tests!)
   (run-tests
    bdict
    (test-equal? make-bdict (make-bdict) `(,TAG-BDICT . #f))
    (test-true bdict-empty?
	       (bdict-empty? (make-bdict)))
    (test-false bdict-set/bdict-empty?
		(bdict-empty?
		 (bdict-set
		  (make-bdict)
		  1 "Hello")))
    (test-true bdict-set/bdict-contains?
	       (bdict-contains?
		(bdict-set
		 (make-bdict)
		 1 "Hello")
		1))
    (test-false bdict-set/bdict-contains?
		(bdict-contains?
		 (bdict-set
		  (make-bdict)
		  1 "Hello")
		 2))
    (test-true bdict-set/bdict-contains?
	       (bdict-contains?
		(bdict-set
		 (bdict-set
		  (make-bdict)
		  1 "Hello")
		 2 "World")
		1))
    (test-equal? bdict-keys
		 (bdict-keys
		  (bdict-set
		   (bdict-set
		    (bdict-set
		     (make-bdict)
		     23 "Hello")
		    666 "World")
		   7 "BST"))
		 '(7 23 666))
    (test-false bdict-remove
		(bdict-contains?
		 (bdict-remove
		  (bdict-set
		   (bdict-set
		    (bdict-set
		     (make-bdict)
		     23 "Hello")
		    666 "World")
		   7 "BST")
		  23)
		 23))
    (test-equal? bdict-balance
		 (bdict-keys
		  (bdict-balance
		   (bdict-set
		    (bdict-set
		     (bdict-set
		      (make-bdict)
		      1 "Hello")
		     2 "World")
		    3 "BST")))
		 '(1 2 3))
    (test-equal? bdict-values
		 (bdict-values
		  (bdict-balance
		   (bdict-set
		    (bdict-set
		     (bdict-set
		      (make-bdict)
		      1 "Hello")
		     2 "World")
		    3 "BST")))
		 '("Hello" "World" "BST"))
    (test-exn bdict-ref
	      (bdict-ref (make-bdict) 1))
    (test-equal? bdict-ref
		 (bdict-ref
		  (bdict-set
		   (bdict-set
		    (bdict-set
		     (make-bdict)
		     1 "Hello")
		    2 "World")
		   3 "BST")
		  2)
		 "World")
    (test-equal? bdict-update
		 (bdict-ref
		  (bdict-update
		   (bdict-set
		    (bdict-set
		     (bdict-set
		      (make-bdict)
		      1 "Hello")
		     2 "World")
		    3 "BST")
		   2
		   (lambda (v)
		     "Scheme"))
		  2)
		 "Scheme")
    (test-equal? bdict-find-pair
		 (bdict-find-pair
		  (bdict-set
		   (bdict-set
		    (bdict-set
		     (make-bdict)
		     1 "Hello")
		    2 "World")
		   3 "BST")
		  (lambda (k v)
		    (equal? v "World")))
		 (cons 2 "World"))
    ))

 )