Change APIs and add union procedure.

1 files changed, 161 insertions, 63 deletions

diff --git a/lib/tokyo/tojo/map/avl.ss b/lib/tokyo/tojo/map/avl.ss
index 5ae4c85..59a2deb 100644
--- a/lib/tokyo/tojo/map/avl.ss
+++ b/lib/tokyo/tojo/map/avl.ss
@@ -1,27 +1,46 @@
 #!r6rs
 (library (tokyo tojo map avl)
   (export map? empty? empty
-          insert search (rename (avl:remove remove))
+          insert search
+          (rename (avl:remove remove))
           size
           (rename (avl:map map)
+                  (avl:map/key map/key)
                   (avl:filter filter)
+                  (avl:filter/key filter/key)
                   (avl:partition partition)
+                  (avl:partition/key partition/key)
                   (avl:for-each for-each)
+                  (avl:for-each/key for-each/key)
                   (avl:fold-left fold-left)
-                  (avl:fold-right fold-right)))
+                  (avl:fold-left/key fold-left/key)
+                  (avl:fold-right fold-right)
+                  (avl:fold-right/key fold-right/key))
+          union union/key)
   (import (rnrs))
 
   (define map?
     (lambda (x)
       (or (empty? x) (node? x))))
+
+  ;; (define empty (list 'avl:empty))
+  ;; (define empty? (lambda (x) (eq? x empty)))
+  
+  ;; (define node (list 'node))
+  ;; (define node? (lambda (x)
+  ;;                 (and (pair? x)
+  ;;                      (eq? (car x) node))))
+  ;; (define make-node
+  ;;   (lambda (kv l r h)
+  ;;     (list node kv l r h)))
+  
+  ;; (define key&value (lambda (node) (cadr node)))
+  ;; (define left (lambda (node) (caddr node)))
+  ;; (define right (lambda (node) (cadddr node)))
+  ;; (define %height (lambda (node) (cadddr (cdr node))))
   
   (define-record-type (avl:empty make-empty empty?)
     (opaque #t))
-  
-  (define error-not-avl-tree
-    (lambda (name x)
-      (unless (map? x)
-        (assertion-violation name "not avl tree" x))))
 
   (define empty (make-empty))
 
@@ -32,6 +51,11 @@
             (immutable h %height))
     (opaque #t))
 
+  (define error-not-map-tree
+    (lambda (name x)
+      (unless (map? x)
+        (assertion-violation name "not map tree" x))))
+
   (define height
     (lambda (tr)
       (if (empty? tr)
@@ -145,16 +169,18 @@
   (define insert
     (case-lambda
       [(<? =? tr k v)
-       (error-not-avl-tree 'insert tr)
+       (error-not-map-tree 'insert tr)
        (let f ([tr tr])
          (if (not (or (node? tr) (empty? tr)))
-             (assertion-violation 'insert "not avl tree" tr))
+             (assertion-violation 'insert "not map tree" tr))
          (if (empty? tr)
              (make-node (cons k v) empty empty 0)
              (let ([kv (key&value tr)])
                (cond
                 [(=? k (car kv))
-                 (make-node (cons k v) (left tr) (right tr) (height tr))]
+                 (make-node (cons k v)
+                            (left tr) (right tr)
+                            (height tr))]
                 [(<? k (car kv))
                  (let ([l (f (left tr))]
                        [r (right tr)])
@@ -173,7 +199,7 @@
   (define search
     (case-lambda
       [(<? =? tr k)
-       (error-not-avl-tree 'search tr)
+       (error-not-map-tree 'search tr)
        (let f ([tr tr])
          (if (empty? tr)
              #f
@@ -195,22 +221,23 @@
                                 (height+ (left node) tr)))
                     m)))))
 
-  (define rem
-    (lambda (node)
-      (cond
-       [(empty? (left node)) (right node)]
-       [(empty? (right node)) (left node)]
-       [else
-        (let-values ([(tr kv) (node-max (left node))])
-          (balance
-           (make-node kv
-                      tr (right node)
-                      (height+ tr (right node)))))])))
+  
 
   (define avl:remove
     (case-lambda
       [(<? =? tr k)
-       (error-not-avl-tree 'remove tr)
+       (define rem
+         (lambda (node)
+           (cond
+            [(empty? (left node)) (right node)]
+            [(empty? (right node)) (left node)]
+            [else
+             (let-values ([(tr kv) (node-max (left node))])
+               (balance
+                (make-node kv
+                           tr (right node)
+                           (height+ tr (right node)))))])))
+       (error-not-map-tree 'remove tr)
        (let f ([tr tr])
          (if (empty? tr)
              empty
@@ -236,77 +263,148 @@
   
   (define avl:map
     (lambda (f tr)
-      (error-not-avl-tree 'map tr)
+      (error-not-map-tree 'map tr)
+      (avl:map/key (lambda (k v) (f v)) tr)))
+  
+  (define avl:map/key
+    (lambda (f tr)
+      (error-not-map-tree 'map/key tr)
       (let g ([tr tr])
         (if (empty? tr)
             empty
             (let ([kv (key&value tr)])
-              (make-node (cons (car kv) (f kv))
+              (make-node (cons (car kv) (f (car kv)
+                                           (cdr kv)))
                          (g (left tr))
                          (g (right tr))
                          (height tr)))))))
-
+  
   (define avl:for-each
     (lambda (f tr)
-      (error-not-avl-tree 'for-each tr)
-      (let g [(tr tr)]
+      (error-not-map-tree 'for-each tr)
+      (avl:for-each/key (lambda (k v) (f v)) tr)))
+  
+  (define avl:for-each/key
+    (lambda (f tr)
+      (error-not-map-tree 'for-each/key tr)
+      (let g ([tr tr])
         (if (empty? tr)
             (if #f #f)
-            (begin
+            (let ([kv (key&value tr)])
               (g (left tr))
-              (f (key&value tr))
+              (f (car kv) (cdr kv))
               (g (right tr)))))))
-
-  (define fold
-    (lambda (f init tr)
-      (if (empty? tr)
-          init
-          (f (fold f init (left tr))
-             (key&value tr)
-             (fold f init (right tr))))))
-
+  
   (define avl:fold-left
     (lambda (f init tr)
-      (error-not-avl-tree 'fold-left tr)
-      (let g [[acc init] (tr tr)]
+      (error-not-map-tree 'fold-left tr)
+      (avl:fold-left/key (lambda (acc k v) (f acc v))
+                         init tr)))
+  
+  (define avl:fold-left/key
+    (lambda (f init tr)
+      (error-not-map-tree 'fold-left/key tr)
+      (let g ([acc init] [tr tr])
         (if (empty? tr)
             acc
-            (g (f (g acc (left tr))
-                  (key&value tr))
-               (right tr))))))
+            (let ([kv (key&value tr)])
+              (g (f (g acc (left tr))
+                    (car kv)
+                    (cdr kv))
+                 (right tr)))))))
 
   (define avl:fold-right
     (lambda (f init tr)
-      (error-not-avl-tree 'fold-right tr)
-      (let g [[acc init] (tr tr)]
+      (error-not-map-tree 'fold-right tr)
+      (avl:fold-right/key (lambda (k v acc) (f v acc))
+                          init tr)))
+  
+  (define avl:fold-right/key
+    (lambda (f init tr)
+      (error-not-map-tree 'fold-right/key tr)
+      (let g ([acc init] [tr tr])
         (if (empty? tr)
             acc
-            (g (f (key&value tr)
-                  (g acc (right tr)))
-               (left tr))))))
-
+            (let ([kv (key&value tr)])
+              (g (f (car kv) (cdr kv)
+                    (g acc (right tr)))
+                 (left tr)))))))
+  
   (define size
     (lambda (tr)
-      (error-not-avl-tree 'size tr)
+      (error-not-map-tree 'size tr)
       (avl:fold-left (lambda (acc v) (+ acc 1)) 0 tr)))
 
   (define avl:filter
     (lambda (p? tr)
-      (error-not-avl-tree 'filter tr)
-      (avl:for-each (lambda (kv)
-                      (if (not (p? kv))
-                          (set! tr (avl:remove tr (car kv)))))
-                    tr)
+      (error-not-map-tree 'filter tr)
+      (avl:filter/key (lambda (k v) (p? v)) tr)))
+  
+  (define avl:filter/key
+    (lambda (p? tr)
+      (error-not-map-tree 'filter/key tr)
+      (avl:for-each/key (lambda (k v)
+                          (if (not (p? k v))
+                              (set! tr (avl:remove tr k))))
+                        tr)
       tr))
   
   (define avl:partition
     (lambda (p? tr)
-      (error-not-avl-tree 'filter tr)
+      (error-not-map-tree 'filter tr)
+      (avl:partition/key (lambda (k v) (p? v)) tr)))
+
+  (define avl:partition/key
+    (lambda (p? tr)
+      (error-not-map-tree 'filter tr)
       (let ([t tr]
             [f tr])
-        (avl:for-each (lambda (kv)
-                        (if (p? kv)
-                            (set! f (avl:remove f (car kv)))
-                            (set! t (avl:remove t (car kv)))))
-                      tr)
-        (values t f)))))
+        (avl:for-each/key (lambda (k v)
+                            (if (p? k v)
+                                (set! f (avl:remove f k))
+                                (set! t (avl:remove t k))))
+                          tr)
+        (values t f))))
+
+  (define union
+    (case-lambda
+      [(f m1 m2)
+       (union < = f m1 m2)]
+      [(<? =? f m1 m2)
+       (union/key <? =? (lambda (k v1 v2) (f v1 v2)) m1 m2)]))
+
+  (define union/key
+    (case-lambda
+      [(f m1 m2)
+       (union < = f m1 m2)]
+      [(<? =? f m1 m2)
+       (let ([l1 (avl:fold-right/key (lambda (k v acc)
+                                       (cons (cons k v) acc))
+                                     '() m1)]
+             [l2 (avl:fold-right/key (lambda (k v acc)
+                                       (cons (cons k v) acc))
+                                     '() m2)])
+         (let loop ([l1 l1] [l2 l2] [acc empty])
+           (cond [(null? l1)
+                  (fold-left
+                   (lambda (acc p) (insert <? =? acc
+                                           (car p) (cdr p)))
+                   acc l2)]
+                 [(null? l2)
+                  (fold-left (lambda (acc p)
+                               (insert <? =? acc
+                                       (car p) (cdr p)))
+                             acc l1)]
+                 [(=? (caar l1) (caar l2))
+                  (loop (cdr l1) (cdr l2)
+                        (insert <? =? acc
+                                (caar l1)
+                                (f (caar l1) (cdar l1) (cdar l2))))]
+                 [(<? (caar l1) (caar l2))
+                  (loop (cdr l1) l2
+                        (insert <? =? acc
+                                (caar l1) (cdar l1)))]
+                 [else
+                  (loop l1 (cdr l2)
+                        (insert <? =? acc
+                                (caar l2) (cdar l2)))])))])))