set operations

1 files changed, 69 insertions, 2 deletions

diff --git a/mcgoron.weight-balanced-trees.internal.scm b/mcgoron.weight-balanced-trees.internal.scm
index 99be0bc..b220780 100644
--- a/mcgoron.weight-balanced-trees.internal.scm
+++ b/mcgoron.weight-balanced-trees.internal.scm
@@ -255,8 +255,8 @@
 
 (define search
   (case-lambda
-    ((tree cmp key) (search tree cmp key (lambda () #f)))
-    ((tree cmp key default)
+    ((cmp key tree) (search cmp key tree (lambda () #f)))
+    ((cmp key tree default)
      (let search ((tree tree))
        (if (null? tree)
            (default)
@@ -266,4 +266,71 @@
                data
                (search right))))))))
 
+(define the-sentinel-value
+  ;; A dummy list allocated so that it is not `eq?` to any other object.
+  ;; Used as the return value for internal set functions.
+  (cons #f '()))
+(define return-sentinel (lambda () the-sentinel-value))
+(define (sentinel? x) (eq? x the-sentinel-value))
+
+(: union (* node-type node-type --> node-type))
+(define (union cmp left right)
+  (let union ((left left)
+              (right right))
+    (cond
+      ((null? left) right)
+      ((null? right) left)
+      (else (with-node (right right-data
+                              ("<" left-of-right)
+                              (">" right-of-right))
+              (let-values (((new-left _ new-right)
+                            (split cmp left right-data return-sentinel)))
+                (join right-data
+                      (union new-left left-of-right)
+                      (union new-right right-of-right))))))))
+
+(: intersection (* node-type node-type --> node-type))
+(define (intersection cmp left right)
+  (let intersection ((left left)
+                     (right right))
+    (if (or (null? left) (null? right))
+        '()
+        (with-node (right right-data
+                          ("<" left-of-right)
+                          (">" right-of-right))
+          (let-values (((new-left new-key new-right)
+                        (split cmp left right-data return-sentinel)))
+            (let ((final-left (intersection new-left left-of-right))
+                  (final-right (intersection new-right right-of-right)))
+              (if (sentinel? new-key)
+                  (join2 final-left final-right) ; right-data not found
+                  (join new-key final-left final-right))))))))
+
+(: difference (* node-type node-type --> node-type))
+(define (difference cmp left right)
+  (let difference ((left left)
+                   (right right))
+    (cond
+      ((null? left) '())
+      ((null? right) left)
+      (else (with-node (right right-data
+                              ("<" left-of-right)
+                              (">" right-of-right))
+              (let-values (((new-left new-key new-right)
+                            (split cmp left right-data return-sentinel)))
+                (join2 (difference new-left left-of-right)
+                       (difference new-right right-of-right))))))))
+
+;;; ;;;;;;;;;;;;;;;;;
+;;; Generic tree functions
+;;; ;;;;;;;;;;;;;;;;;
+
+(: every (* node-type --> *))
+(define (every predicate? tree)
+  (if (null? tree)
+      #t
+      (with-node (tree data ("<" left) (">" right))
+        (and (predicate? data)
+             (every predicate? left)
+             (every predicate? right)))))