set operations

3 files changed, 271 insertions, 18 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)))))
 
diff --git a/mcgoron.weight-balanced-trees.internal.sld b/mcgoron.weight-balanced-trees.internal.sld
index cdb47ee..ac85dee 100644
--- a/mcgoron.weight-balanced-trees.internal.sld
+++ b/mcgoron.weight-balanced-trees.internal.sld
@@ -32,6 +32,8 @@
           in-order-vector->node
           node->in-order-list
           join join2 split
-          search)
+          search
+          union intersection difference
+          every)
   (include "mcgoron.weight-balanced-trees.internal.scm"))
 
diff --git a/tests/run.scm b/tests/run.scm
index 14800ba..f9218bc 100644
--- a/tests/run.scm
+++ b/tests/run.scm
@@ -18,7 +18,11 @@
 (load "random-number-vector.sld")
 (load "util.sld")
 
-(import (mcgoron weight-balanced-trees internal)
+(import (except (mcgoron weight-balanced-trees internal)
+                every)
+        (prefix (only (mcgoron weight-balanced-trees internal)
+                      every)
+                node-)
         (mcgoron weight-balanced-trees test util)
         (mcgoron weight-balanced-trees test random-number-vector)
         (prefix (only (mcgoron srfi 64) factory) mcgoron-test-)
@@ -38,6 +42,7 @@
 
 (test-runner-factory mcgoron-test-factory)
 (test-runner-current (test-runner-create))
+(set-verbosity! 'fails)
 
 ;;; ;;;;;;;;;;;;;;;;;;;
 ;;; Random in-order vector generator
@@ -192,7 +197,7 @@
          (dynamic-property-set! 'right right)
          (dynamic-property-set! 'left left)
          (dynamic-property-set! 'node (node->in-order-list node)))
-       (search node number-comparator middle))))
+       (search number-comparator middle node))))
   (parameterize ((verbose #f))
     (test-property search-finds (list (split-vector-generator)))))
 
@@ -200,31 +205,210 @@
   (define search-does-not-find
     (call-w/o-inserted
      (lambda (left middle right node)
-       (not (search node number-comparator middle)))))
+       (not (search number-comparator middle node)))))
   (test-property search-does-not-find (list (split-vector-generator))))
 
 ;;; ;;;;;;;;;;;;;;;;;;;;;;;;
 ;;; split
 ;;; ;;;;;;;;;;;;;;;;;;;;;;;;
 
-(define split-finds
-  (call/inserted
+(define (call/split wrapper procedure)
+  (wrapper
    (lambda (left middle right node)
-     (let-values (((new-left found? new-right)
-                   (split number-comparator node middle
-                          (lambda () #f))))
-       found?))))
+     (call-with-values (lambda ()
+                         (split number-comparator node middle
+                                (lambda () #f)))
+       procedure))))
+
+(define split-finds
+  (call/split
+   call/inserted
+   (lambda (new-left found? new-right)
+     found?)))
 
 (test-group "split finds"
   (test-property split-finds (list (split-vector-generator))))
 
+(define split-left-balanced
+  (call/split
+   call/inserted
+   (lambda (new-found found? new-right)
+     (balanced? new-found))))
+
+(test-group "split left balanced"
+  (test-property split-left-balanced (list (split-vector-generator))))
+
+(define split-right-balanced
+  (call/split
+   call/inserted
+   (lambda (new-found found? new-right)
+     (balanced? new-right))))
+
+(test-group "split right balanced"
+  (test-property split-right-balanced (list (split-vector-generator))))
+
 (define split-does-not-find
-  (call-w/o-inserted
-   (lambda (left middle right node)
-     (let-values (((new-left found? new-right)
-                   (split number-comparator node middle
-                          (lambda () #f))))
-       (not found?)))))
+  (call/split
+   call-w/o-inserted
+   (lambda (new-left found? new-right)
+     (not found?))))
 
 (test-group "split does not find"
   (test-property split-does-not-find (list (split-vector-generator))))
+
+(define split-left-balanced-w/o-inserted
+  (call/split
+   call-w/o-inserted
+   (lambda (new-left found? new-right)
+     (balanced? new-left))))
+
+(test-group "split left balanced without insertion"
+  (test-property split-left-balanced-w/o-inserted (list (split-vector-generator))))
+
+(define split-right-balanced-w/o-inserted
+  (call/split
+   call-w/o-inserted
+   (lambda (new-left found? new-right)
+     (balanced? new-right))))
+
+(test-group "split right balanced without insertion"
+  (test-property split-right-balanced-w/o-inserted (list (split-vector-generator))))
+
+;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;;; Union
+;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+(test-group "union balanced"
+  (define (union-balanced node1 node2)
+    (balanced? (union number-comparator node1 node2)))
+  (test-property union-balanced
+                 (list (vector->node-generator)
+                       (vector->node-generator))))
+
+(define (union-subset-of-lset node1 node2)
+    (let ((node (union number-comparator node1 node2))
+          (union (lset-union =
+                             (node->in-order-list node1)
+                             (node->in-order-list node2))))
+      (if (verbose)
+          (dynamic-property-set! 'union
+                                 (node->in-order-list node)))
+      (every (cut search number-comparator <> node) union)))
+
+(test-group "union subset of lset"
+  (parameterize ((verbose #t))
+    (test-property union-subset-of-lset
+                   (list (vector->node-generator)
+                         (vector->node-generator)))))
+
+(test-group "union subset of lset, right null"
+  (test-property union-subset-of-lset
+                 (list (vector->node-generator)
+                       (circular-generator '()))))
+
+(test-group "union subset of lset, left null"
+  (test-property union-subset-of-lset
+                 (list (circular-generator '())
+                       (vector->node-generator))))
+
+(define (lset-subset-of-union node1 node2)
+  (let ((node (union number-comparator node1 node2))
+        (union (lset-union =
+                           (node->in-order-list node1)
+                           (node->in-order-list node2))))
+    (node-every (cut memq <> union) node)))
+
+(test-group "lset subset of union"
+  (test-property lset-subset-of-union
+                 (list (vector->node-generator)
+                       (vector->node-generator))))
+
+(test-group "lset subset of union, right null"
+  (test-property lset-subset-of-union
+                 (list (vector->node-generator)
+                       (circular-generator '()))))
+
+(test-group "lset subset of union, left null"
+  (test-property lset-subset-of-union
+                 (list (circular-generator '())
+                       (vector->node-generator))))
+
+;;; ;;;;;;;;;;;;;;;;;;;;
+;;; Intersection
+;;; ;;;;;;;;;;;;;;;;;;;;
+
+(test-group "intersection balanced?"
+  (define (intersection-balanced? node1 node2)
+    (balanced? (intersection number-comparator node1 node2)))
+  (test-property intersection-balanced?
+                 (list (vector->node-generator)
+                       (vector->node-generator))))
+
+(test-group "intersection of itself"
+  (define (intersection-of-itself node1)
+    (let ((node (intersection number-comparator node1 node1))
+          (lst (node->in-order-list node1)))
+      (every (cut search number-comparator <> node) lst)))
+  (test-property intersection-of-itself
+                 (list (vector->node-generator))))
+
+(test-group "intersection of null, left"
+  (define (intersection-of-null-left node1)
+    (null? (intersection number-comparator node1 '())))
+  (test-property intersection-of-null-left
+                 (list (vector->node-generator))))
+
+(test-group "intersection of null, right"
+  (define (intersection-of-null-right node1)
+    (null? (intersection number-comparator '() node1)))
+  (test-property intersection-of-null-right 
+                 (list (vector->node-generator))))
+
+(test-group "intersection subset of lset"
+  (define (intersection-lset node1 node2)
+    (let ((node (intersection number-comparator node1 node2))
+          (intersection (lset-intersection =
+                                           (node->in-order-list node1)
+                                           (node->in-order-list node2))))
+      (every (cut search number-comparator <> node) intersection)))
+  (test-property intersection-lset
+                 (list (vector->node-generator)
+                       (vector->node-generator))))
+
+(test-group "lset subset of intersection"
+  (define (lset-subset-of-intersection node1 node2)
+    (let ((node (intersection number-comparator node1 node2))
+          (intersection (lset-intersection =
+                                           (node->in-order-list node1)
+                                           (node->in-order-list node2))))
+      (node-every (cut memv <> intersection) node)))
+  (test-property lset-subset-of-intersection
+                 (list (vector->node-generator)
+                       (vector->node-generator))))
+
+;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;
+;;; Difference
+;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+(test-group "difference subset of lset"
+  (define (difference-subset-of-lset node1 node2)
+    (let ((node (difference number-comparator node1 node2))
+          (intersection (lset-difference =
+                                         (node->in-order-list node1)
+                                         (node->in-order-list node2))))
+      (every (cut search number-comparator <> node) intersection)))
+  (test-property difference-subset-of-lset
+                 (list (vector->node-generator)
+                       (vector->node-generator))))
+
+(test-group "lset subset of difference"
+  (define (lset-subset-of-difference node1 node2)
+    (let ((node (difference number-comparator node1 node2))
+          (intersection (lset-difference =
+                                         (node->in-order-list node1)
+                                         (node->in-order-list node2))))
+      (node-every (cut memq <> intersection) node)))
+  (test-property lset-subset-of-difference
+                 (list (vector->node-generator)
+                       (vector->node-generator))))
+