add a generating thunk for split

3 files changed, 17 insertions, 10 deletions

diff --git a/doc/mcgoron.weight-balanced-trees.internal.scm b/doc/mcgoron.weight-balanced-trees.internal.scm
index 176d18e..8baf5e7 100644
--- a/doc/mcgoron.weight-balanced-trees.internal.scm
+++ b/doc/mcgoron.weight-balanced-trees.internal.scm
@@ -52,16 +52,19 @@ with `data`."))
 
 Returns a new tree with the elements of `left` and `right`."))
  ((name . "split")
-  (signature lambda ((comparator? cmp) (wb-tree-node? tree) key)
-             => (values wb-tree-node? boolean? wb-tree-node?))
+  (signature lambda ((comparator? cmp) (wb-tree-node? tree) key (procedure? default))
+             => (values wb-tree-node? * wb-tree-node?))
   "
 * It is an error if any element of `tree` is not comparable by `cmp`.
 * It is an error if `key` is not comparable by `cmp`.
 * It is an error if `cmp` does not have an order.
+* It is an error if `default` is not a thunk.
 
 Returns two trees, one of all values to the left of `key`, and one with
-all values to the right of `key`. The boolean is true if the key was
-found and false otherwise.")
+all values to the right of `key`.
+
+If a value that compares equal to `key` is found, that value is returned.
+Otherwise the result of calling `default` with no arguments is returned.")
  ((name . "search")
   (signature case-lambda ((((wb-tree-node? tree) (comparator? cmp) key) *)
                           (((wb-tree-node? tree) (comparator? cmp) key (procedure? default)) *)))
diff --git a/mcgoron.weight-balanced-trees.internal.scm b/mcgoron.weight-balanced-trees.internal.scm
index c657178..99be0bc 100644
--- a/mcgoron.weight-balanced-trees.internal.scm
+++ b/mcgoron.weight-balanced-trees.internal.scm
@@ -18,6 +18,7 @@
   (left get-left)
   (right get-right))
 
+(: wb-tree-node? (* -> boolean : node-type))
 (define (wb-tree-node? x)
   (or (null? x) (non-null-wb-tree-node? x)))
 
@@ -239,16 +240,16 @@
 
 ;;; XXX: The comparator library does not export the struct type for
 ;;; the comparator.
-(: split (* node-type * --> node-type boolean node-type))
-(define (split cmp tree key)
+(: split (* node-type * (-> *) --> node-type * node-type))
+(define (split cmp tree key default)
   (let split ((tree tree))
     (if (null? tree)
-        (values '() #f '())
+        (values '() (default) '())
         (with-node (tree data ("<" left) (">" right))
           (comparator-if<=> cmp key data
             (let-values (((new-left bool new-right) (split left)))
               (values new-left bool (join data new-right right)))
-            (values left #t right)
+            (values left data right)
             (let-values (((new-left bool new-right) (split right)))
               (values (join data left new-left) bool new-right)))))))
 
@@ -265,3 +266,4 @@
                data
                (search right))))))))
 
+
diff --git a/tests/run.scm b/tests/run.scm
index c2e06ec..14800ba 100644
--- a/tests/run.scm
+++ b/tests/run.scm
@@ -211,7 +211,8 @@
   (call/inserted
    (lambda (left middle right node)
      (let-values (((new-left found? new-right)
-                   (split number-comparator node middle)))
+                   (split number-comparator node middle
+                          (lambda () #f))))
        found?))))
 
 (test-group "split finds"
@@ -221,7 +222,8 @@
   (call-w/o-inserted
    (lambda (left middle right node)
      (let-values (((new-left found? new-right)
-                   (split number-comparator node middle)))
+                   (split number-comparator node middle
+                          (lambda () #f))))
        (not found?)))))
 
 (test-group "split does not find"