set: add update

1 files changed, 150 insertions, 79 deletions

diff --git a/set.scm b/set.scm
index c870af9..94a6e1d 100644
--- a/set.scm
+++ b/set.scm
@@ -32,7 +32,7 @@
       (else (error "invalid direction")))))
 
 ;;; Gets data from node value given accessor symbol.
-(define %set:get
+(define set:get
   (lambda (t sym)
     (vector-ref t (%set:accessor sym))))
 
@@ -40,16 +40,16 @@
   (lambda (node)
     (if (null? node)
         '()
-        (list (list 'data (%set:get node '=))
-              (list '< (%set->sexpr (%set:get node '<)))
-              (list '> (%set->sexpr (%set:get node '>)))))))
+        (list (list 'data (set:get node '=))
+              (list '< (%set->sexpr (set:get node '<)))
+              (list '> (%set->sexpr (set:get node '>)))))))
 
 ;;; Get the height of a node, handling the empty node.
 (define %set:height
   (lambda (node)
     (if (null? node)
         0
-        (%set:get node 'h))))
+        (set:get node 'h))))
 
 ;;; Get the difference between the heights of two trees.
 (define %set:height-diff
@@ -58,8 +58,8 @@
 
 ;;; Get the balance factor of a tree.
 (define %set:bal
-  (lambda (t) (%set:height-diff (%set:get t '<)
-                                (%set:get t '>))))
+  (lambda (t) (%set:height-diff (set:get t '<)
+                                (set:get t '>))))
 
 ;;; Set data in node given accessor symbol.
 (define %set:set!
@@ -77,6 +77,12 @@
       (%set:set! node dir2 node2)
       node)))
 
+(define set:node-new-val
+  (lambda (node newval)
+    (%set:node newval
+               '< (set:get node '<)
+               '> (set:get node '>))))
+
 (define %set:invdir
   (lambda (dir)
     (cond
@@ -94,13 +100,13 @@
     (if (null? node)
         #f
         (let ((invdir (%set:invdir dir)))
-          (let ((child (%set:get node invdir)))
-            (let ((to-swap (%set:get child dir)))
-              (%set:node (%set:get child '=)
-                         dir (%set:node (%set:get node '=)
-                                        dir (%set:get node dir)
+          (let ((child (set:get node invdir)))
+            (let ((to-swap (set:get child dir)))
+              (%set:node (set:get child '=)
+                         dir (%set:node (set:get node '=)
+                                        dir (set:get node dir)
                                         invdir to-swap)
-                         invdir (%set:get child invdir))))))))
+                         invdir (set:get child invdir))))))))
 
 ;;; ;;;;;;;;;;;;;;;;;;;
 ;;; JOIN function for AVL trees.
@@ -109,10 +115,10 @@
 ;;; Handles rebalancing of the tree.
 (define %set:join
   (lambda (heavier val lighter heavier-dir)
-    (let ((heavy-val (%set:get heavier '=))
+    (let ((heavy-val (set:get heavier '=))
           (lighter-dir (%set:invdir heavier-dir)))
-      (let ((heavy-heavy (%set:get heavier heavier-dir))
-            (heavy-light (%set:get heavier lighter-dir)))
+      (let ((heavy-heavy (set:get heavier heavier-dir))
+            (heavy-light (set:get heavier lighter-dir)))
         (if (<= (abs (%set:height-diff heavy-light lighter)) 1)
             (let ((node (%set:node val
                                    heavier-dir heavy-light
@@ -149,19 +155,19 @@
           (if (null? left)
               right
               (let ((split-last-tree (split-last left)))
-                (set:join (%set:get split-last-tree '=)
-                          '< (%set:get split-last-tree '<)
+                (set:join (set:get split-last-tree '=)
+                          '< (set:get split-last-tree '<)
                           '> right)))))
        (split-last
         (lambda (tree)
-          (let ((right (%set:get tree '>)))
+          (let ((right (set:get tree '>)))
             (if (null? right)
                 tree
                 (let ((last (split-last right)))
-                  (%set:node (%set:get last '=)
-                             (set:join (%set:get tree '=)
-                                       '< (%set:get tree '<)
-                                       '> (%set:get last '>)))))))))
+                  (%set:node (set:get last '=)
+                             (set:join (set:get tree '=)
+                                       '< (set:get tree '<)
+                                       '> (set:get last '>)))))))))
     join2))
 
 ;;; ;;;;;;;;;;;;;;;;;
@@ -179,10 +185,32 @@
             (lambda (tree)
               (if (null? tree)
                   #f
-                  (let ((dir (<=> (%set:get tree '=) data)))
+                  (let ((dir (<=> (set:get tree '=) data)))
                     (if (eq? dir '=)
                         tree
-                        (loop (%set:get tree dir))))))))
+                        (loop (set:get tree dir))))))))
+        (loop tree)))))
+
+;;; (SET:UPDATE <=>) generates an update function for <=>.
+;;; (UPDATE TREE NODE UPDATE) inserts (UPDATE NODE #F) into the tree if
+;;; no node comparing equal to NODE is found, and (UPDATE NODE OLD) if
+;;; OLD compares equal to NODE.
+(define set:update
+  (lambda (<=>)
+    (lambda (tree node update)
+      (letrec
+          ((loop
+            (lambda (tree)
+              (if (null? tree)
+                  (update node #f)
+                  (let ((dir (<=> (set:get tree '=)
+                                  (set:get node '=))))
+                    (if (eq? dir '=)
+                        (update node tree)
+                        (let ((invdir (%set:invdir dir)))
+                          (set:join (set:get tree '=)
+                                    dir (loop (set:get tree dir))
+                                    invdir (set:get tree invdir)))))))))
         (loop tree)))))
 
 ;;; (SET:INSERT <=>) generates an insert function for comparison function
@@ -191,27 +219,18 @@
 ;;; (CONS NEWTREE FOUND), where FOUND is the node that was replaced by
 ;;; NODE, and #F otherwise, and NEWTREE is the new root of the tree.
 (define set:insert
-  (lambda (<=>)
+  (lambda (update)
     (lambda (tree node)
       (let ((found #f))
-        (letrec
-            ((loop
-              (lambda (tree)
-                (if (null? tree)
-                    node
-                    (let ((dir (<=> (%set:get tree '=)
-                                    (%set:get node '=))))
-                      (if (eq? dir '=)
-                          (begin
-                            (set! found tree)
-                            (%set:node (%set:get node '=)
-                                       '< left '> right))
-                          (let ((invdir (%set:invdir dir)))
-                            (set:join (%set:get tree '=)
-                                      dir (loop (%set:get tree dir))
-                                      invdir (%set:get tree invdir)))))))))
-          (let ((newtree (loop tree)))
-            (cons newtree found)))))))
+        (let ((newroot (update tree node
+                               (lambda (node oldnode)
+                                 (if (not oldnode)
+                                     node
+                                     (begin
+                                       (set! found oldnode)
+                                       (set:node-new-val oldnode
+                                                         (set:get node '=))))))))
+          (cons newroot found))))))
 
 ;;; (SET:DELETE <=>) generates a delete function for comparison function
 ;;; <=>.
@@ -228,20 +247,48 @@
               (lambda (tree)
                 (if (null? tree)
                     node
-                    (let ((dir (<=> (%set:get tree '=)
-                                    (%set:get node '=))))
+                    (let ((dir (<=> (set:get tree '=)
+                                    (set:get node '=))))
                       (if (eq? dir '=)
                           (begin
                             (set! found tree)
-                            (set:join2 (%set:get tree '<)
-                                       (%set:get tree '>)))
+                            (set:join2 (set:get tree '<)
+                                       (set:get tree '>)))
                           (let ((invdir (%set:invdir dir)))
-                            (set:join (%set:get tree '=)
-                                      dir (loop (%set:get tree dir))
-                                      invdir (%set:get tree invdir)))))))))
+                            (set:join (set:get tree '=)
+                                      dir (loop (set:get tree dir))
+                                      invdir (set:get tree invdir)))))))))
           (let ((newtree (loop tree)))
             (cons newtree found)))))))
 
+;;; ;;;;;;;;;;;;;;;;;;;;;;;
+;;; Converting sets to maps
+;;; ;;;;;;;;;;;;;;;;;;;;;;;
+
+;;; Convert a <=> for sets to one for maps.
+(define set:<=>-to-map
+  (lambda (<=>)
+    (lambda (x y)
+      (<=> (car x) (car y)))))
+
+(define map:node-new-val
+  (lambda (node newkey newval)
+    (set:node-new-val node
+                      (cons newkey newval))))
+
+(define map:empty-node
+  (lambda (key val)
+    (%set:node (cons key val)
+               '< '() '> '())))
+
+(define map:key
+  (lambda (node)
+    (car (set:get node '=))))
+
+(define map:val
+  (lambda (node)
+    (cdr (set:get node '=))))
+
 ;;; ;;;;;;;;;;;
 ;;; For strings
 ;;; ;;;;;;;;;;;
@@ -261,11 +308,18 @@
         (set-car! lst val)
         (list-set! (cdr lst) (- n 1) val)))))
 
-(define map:string<=>
-  (lambda (x y)
-    (string<=> (car x) (car y))))
+(define map:string<=> (set:<=>-to-map string<=>))
+
+(define %smap:update (set:update map:string<=>))
+(define %smap:insert (set:insert %smap:update))
 
-(define %smap:insert (set:insert map:string<=>))
+;;; (SMAP:UPDATE TREE KEY UPDATE) updates NODE in TREE with
+;;; (UPDATE NODE-WITH-KEY NODE), where NODE-WITH-KEY is an empty node
+;;; with the key KEY, and (UPDATE NODE-WITH-KEY #F) if no NODE is
+;;; found.
+(define smap:update
+  (lambda (tree key update)
+    (%smap:update tree (map:empty-node key '()) update)))
 
 ;;; (SMAP:INSERT TREE KEY VAL) inserts (CONS KEY VAL) into TREE, and
 ;;; returns (CONS NEWROOT FOUND), where NEWROOT is the new root of
@@ -273,9 +327,7 @@
 ;;; or the matching element if found.
 (define smap:insert
   (lambda (tree key val)
-    (%smap:insert tree (%set:node (cons key val)
-                                  '< '()
-                                  '> '()))))
+    (%smap:insert tree (map:empty-node key val))))
 
 (define %smap:search (set:in map:string<=>))
 
@@ -287,13 +339,7 @@
 (define %smap:delete (set:delete map:string<=>))
 (define smap:delete
   (lambda (tree key)
-    (%smap:delete tree (%set:node (cons key '())
-                                  '< '() '> '()))))
-
-(define smap:key
-  (lambda (node) (car (%set:get node '=))))
-(define smap:val
-  (lambda (node) (cdr (%set:get node '=))))
+    (%smap:delete tree (map:empty-node key '()))))
 
 ;;; ;;;;;
 ;;; Tests
@@ -316,7 +362,7 @@
               ((and already-in (not (cdr insert-return)))
                "should have been found")
               ((and already-in (not (equal? already-in
-                                            (smap:val (cdr insert-return)))))
+                                            (map:val (cdr insert-return)))))
                "found is not correct")
               (else (%set:operate-all f (car insert-return) (cdr lst)))))))))
 
@@ -350,12 +396,12 @@
                                                 '> (%set:node 5 '< '() '> '()))
                                      '>)))
              (cond
-               ((not (eqv? (%set:get right '=) 2)) "bad parent")
-               ((not (eqv? (%set:get (%set:get right '>) '=) 1)) "bad right child")
-               ((not (eqv? (%set:get (%set:get right '<) '=) 3)) "bad left child")
-               ((not (eqv? (%set:get (%set:get (%set:get right '>) '>) '=) 5))
+               ((not (eqv? (set:get right '=) 2)) "bad parent")
+               ((not (eqv? (set:get (set:get right '>) '=) 1)) "bad right child")
+               ((not (eqv? (set:get (set:get right '<) '=) 3)) "bad left child")
+               ((not (eqv? (set:get (set:get (set:get right '>) '>) '=) 5))
                 "bad right child of right child")
-               ((not (eqv? (%set:get (%set:get (%set:get right '>) '<) '=) 4))
+               ((not (eqv? (set:get (set:get (set:get right '>) '<) '=) 4))
                 "bad left child of right child")
                (else #t)))))
    (cons "rotate left"
@@ -371,12 +417,12 @@
                                                 '< (%set:node 5 '< '() '> '()))
                                      '<)))
              (cond
-               ((not (eqv? (%set:get right '=) 2)) "bad parent")
-               ((not (eqv? (%set:get (%set:get right '>) '=) 4)) "bad right child")
-               ((not (eqv? (%set:get (%set:get right '<) '=) 1)) "bad left child")
-               ((not (eqv? (%set:get (%set:get (%set:get right '<) '>) '=) 3))
+               ((not (eqv? (set:get right '=) 2)) "bad parent")
+               ((not (eqv? (set:get (set:get right '>) '=) 4)) "bad right child")
+               ((not (eqv? (set:get (set:get right '<) '=) 1)) "bad left child")
+               ((not (eqv? (set:get (set:get (set:get right '<) '>) '=) 3))
                 "bad right child of left child")
-               ((not (eqv? (%set:get (%set:get (%set:get right '<) '<) '=) 5))
+               ((not (eqv? (set:get (set:get (set:get right '<) '<) '=) 5))
                 "bad left child of left child")
                (else #t)))))
    (cons "insert then delete"
@@ -390,9 +436,9 @@
                   (let ((found (smap:search tree (string #\a))))
                     (cond
                       ((not found) "string not in tree")
-                      ((not (equal? (smap:key tree) (string #\a)))
+                      ((not (equal? (map:key tree) (string #\a)))
                        "returned key not equal to a")
-                      ((not (equal? (smap:val tree) 5))
+                      ((not (equal? (map:val tree) 5))
                        "returned value not equal to 5")
                       (else
                        (let ((delete-return (smap:delete tree (string #\a))))
@@ -428,4 +474,29 @@
                                (cond
                                  ((string? tree) tree)
                                  ((not (null? tree)) "did not delete everything")
-                                 (else #t))))))))))))))
+                                 (else #t))))))))))))
+   (cons "insert a few, update"
+         (lambda ()
+           (let ((tree (%set:insert-all '()
+                                        (list
+                                         (list (string #\a #\b #\c #\d) 1 #f)
+                                         (list (string #\e #\f #\g #\h) 2 #f)
+                                         (list (string #\1 #\4 #\2 #\9 #\3) 3 #f)
+                                         (list (string #\a #\b #\c #\d #\e) 4 #f)))))
+             (if (string? tree)
+                 tree
+                 (let ((tree (smap:update tree
+                                          (string #\a #\b #\c #\d #\e)
+                                          (lambda (node oldnode)
+                                            (map:node-new-val oldnode
+                                                              (string #\a #\b #\c #\d #\e)
+                                                              10)))))
+                   (let ((res (%set:search-all tree
+                                               (list
+                                                (list (string #\a #\b #\c #\d) 1 1)
+                                                (list (string #\e #\f #\g #\h) 2 2)
+                                                (list (string #\1 #\4 #\2 #\9 #\3) 3 3)
+                                                (list (string #\a #\b #\c #\d #\e) 10 10)))))
+                     (if (string? res)
+                         res
+                         #t)))))))))