From 5fc654d57ead71636746ea5d9c57aaf1f6d0c5da Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dominik=20Pant=C5=AF=C4=8Dek?=
 <dominik.pantucek@trustica.cz>
Date: Thu, 6 Jul 2023 22:00:41 +0200
Subject: [PATCH] Implement ref, contains? and set.

---
 src/util-bst.scm | 123 ++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 112 insertions(+), 11 deletions(-)

diff --git a/src/util-bst.scm b/src/util-bst.scm
index 25559db..03b5cfc 100644
--- a/src/util-bst.scm
+++ b/src/util-bst.scm
@@ -11,10 +11,10 @@
   bst?
 
   bst-empty?
-  ;;bst-ref
-  ;;bst-contains?
+  bst-ref
+  bst-contains?
 
-  ;;bst-set
+  bst-set
   ;;bst-remove ;;
 
   ;;bst->kvv
@@ -27,12 +27,45 @@
   )
 
  (import scheme
+	 (chicken condition)
 	 util-tag
 	 testing)
 
- ;; Unique tags
+ ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+ ;; 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 TAG-KVV (make-tag kvv))
 
  (define/doc (make-bst subtag EQ? <?)
    ("Creates empty BST with given comparators")
@@ -52,6 +85,26 @@
  (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)
@@ -63,13 +116,61 @@
    (not (bst-root bst)))
 
  ;; Wrapper to setup comparators
- (define-syntax define-bst-proc
+ (define-syntax let-comparators
    (syntax-rules ()
-     ((_ EQ? <? (name bst . args) expr ...)
-      (define (name bst . args)
-	(let ((EQ? (bst-eq? bst))
-	      (<? (bst-<? bst)))
-	  expr ...)))))
+     ((_ (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!)