split and search

6 files changed, 263 insertions, 113 deletions

diff --git a/doc/mcgoron.weight-balanced-trees.internal.scm b/doc/mcgoron.weight-balanced-trees.internal.scm
index 5812d72..176d18e 100644
--- a/doc/mcgoron.weight-balanced-trees.internal.scm
+++ b/doc/mcgoron.weight-balanced-trees.internal.scm
@@ -18,10 +18,7 @@
   (desc "Returns the number of elements in this tree."))
  ((name . "balanced?")
   (signature lambda ((wb-tree-node? x)) => boolean?)
-  (tags internal)
-  (desc "Recursively traverses `x` and checks if it is weight balanced.
-This function is not called in normal code, but can be useful for
-debugging."))
+  (desc "Recursively traverses `x` and checks if it is weight balanced."))
  ((name . "in-order-vector->node")
   (signature lambda ((vector? x)) => wb-tree-node?)
   (desc "
@@ -45,4 +42,33 @@ elements of `x`."))
   `right`.
 
 Returns a balanced tree containing all elements of `left` and `right`
-with `data`.")))
\ No newline at end of file
+with `data`."))
+ ((name . "join2")
+  (signature lambda ((wb-tree-node? left) (wb-tree-node? right))
+             => wb-tree-node?)
+  (desc "
+* It is an error if any element in `left` is greater than or equal to any
+  element in `right`.
+
+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?))
+  "
+* 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.
+
+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.")
+ ((name . "search")
+  (signature case-lambda ((((wb-tree-node? tree) (comparator? cmp) key) *)
+                          (((wb-tree-node? tree) (comparator? cmp) key (procedure? default)) *)))
+  (desc "
+* It is an error if `cmp` does not order `key` and all elements in `tree`.
+* It is an error if `default` is not a thunk.
+
+Searches `tree` for `key` using `cmp`. If a node contains data that compares equal
+to `key`, return that data. Otherwise, tail call `default.` If `default` is not
+supplied, the function returns `#f`.")))
\ No newline at end of file
diff --git a/mcgoron.weight-balanced-trees.internal.scm b/mcgoron.weight-balanced-trees.internal.scm
index 8d9dc6b..c657178 100644
--- a/mcgoron.weight-balanced-trees.internal.scm
+++ b/mcgoron.weight-balanced-trees.internal.scm
@@ -236,3 +236,32 @@
       (let-values (((new-left new-data)
                     (split-last left)))
         (join new-data new-left right))))
+
+;;; 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)
+  (let split ((tree tree))
+    (if (null? tree)
+        (values '() #f '())
+        (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)
+            (let-values (((new-left bool new-right) (split right)))
+              (values (join data left new-left) bool new-right)))))))
+
+(define search
+  (case-lambda
+    ((tree cmp key) (search tree cmp key (lambda () #f)))
+    ((tree cmp key default)
+     (let search ((tree tree))
+       (if (null? tree)
+           (default)
+           (with-node (tree data ("<" left) (">" right))
+             (comparator-if<=> cmp key data
+               (search left)
+               data
+               (search right))))))))
+
diff --git a/mcgoron.weight-balanced-trees.internal.sld b/mcgoron.weight-balanced-trees.internal.sld
index 804a8bf..cdb47ee 100644
--- a/mcgoron.weight-balanced-trees.internal.sld
+++ b/mcgoron.weight-balanced-trees.internal.sld
@@ -14,7 +14,8 @@
  |#
 
 (define-library (mcgoron weight-balanced-trees internal)
-  (import (scheme base))
+  (import (scheme base) (scheme case-lambda)
+          (srfi 128))
   (cond-expand
     ;; Handle type declarations
     (chicken (import (chicken type)))
@@ -30,6 +31,7 @@
           balanced?
           in-order-vector->node
           node->in-order-list
-          join join2)
+          join join2 split
+          search)
   (include "mcgoron.weight-balanced-trees.internal.scm"))
 
diff --git a/tests/random-number-vector.sld b/tests/random-number-vector.sld
new file mode 100644
index 0000000..2e40cd0
--- /dev/null
+++ b/tests/random-number-vector.sld
@@ -0,0 +1,49 @@
+#| Copyright 2024 Peter McGoron
+ |
+ | Licensed under the Apache License, Version 2.0 (the "License");
+ | you may not use this file except in compliance with the License.
+ | You may obtain a copy of the License at
+ |
+ |     http://www.apache.org/licenses/LICENSE-2.0
+ |
+ | Unless required by applicable law or agreed to in writing, software
+ | distributed under the License is distributed on an "AS IS" BASIS,
+ | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ | See the License for the specific language governing permissions and
+ | limitations under the License.
+ |#
+
+(define-library (mcgoron weight-balanced-trees test random-number-vector)
+  (import (scheme base)
+          (srfi 158) (srfi 197) (srfi 252) (srfi 133) (srfi 26) (srfi 132))
+  (export in-order-vector-generator)
+  (begin
+    (define (random-number-vector of)
+      ;; Generate a random vector of exact integers with at most `of` number
+      ;; of elements.
+      (vector-generator-of (exact-integer-generator) of))
+    (define (filter-non-empty-vectors gen)
+      ;; Filter out the empty vector.
+      (gremove (lambda (vec) (zero? (vector-length vec))) gen))
+    (define (remove-duplicates gen)
+      ;; Filter out vectors with duplicate entries.
+      (define (filter vec)
+        (call/cc
+         (lambda (return)
+           (vector-fold (lambda (prev elem)
+                          (cond
+                            ((null? prev) elem)
+                            ((= elem prev) (return #f))
+                            (else elem)))
+                        '()
+                        vec)
+           #t)))
+      (gfilter filter gen))
+    (define (make-in-order gen)
+      ;; Sort all vectors to be in-order.
+      (gmap (cut vector-sort < <>) gen))
+    (define (in-order-vector-generator of)
+      (chain (random-number-vector of)
+             (filter-non-empty-vectors _)
+             (make-in-order _)
+             (remove-duplicates _)))))
diff --git a/tests/run.scm b/tests/run.scm
index 1aea3b0..c2e06ec 100644
--- a/tests/run.scm
+++ b/tests/run.scm
@@ -13,7 +13,14 @@
  | limitations under the License.
  |#
 
+(import r7rs)
+
+(load "random-number-vector.sld")
+(load "util.sld")
+
 (import (mcgoron weight-balanced-trees internal)
+        (mcgoron weight-balanced-trees test util)
+        (mcgoron weight-balanced-trees test random-number-vector)
         (prefix (only (mcgoron srfi 64) factory) mcgoron-test-)
         (except (mcgoron srfi 64) factory)
         (srfi 1) (srfi 26) (srfi 64) (srfi 128) (srfi 132) (srfi 133)
@@ -32,57 +39,13 @@
 (test-runner-factory mcgoron-test-factory)
 (test-runner-current (test-runner-create))
 
-;;; ;;;;;;;;;;;
-;;; Utils
-;;; ;;;;;;;;;;;
-
-(define (vector-copy-exclusive vec start end)
-  ;; Copy `[start,end)` from `vec`.
-  (if (= start end)
-      #()
-      (vector-copy vec start (- end 1))))
-
 ;;; ;;;;;;;;;;;;;;;;;;;
 ;;; Random in-order vector generator
 ;;; ;;;;;;;;;;;;;;;;;;;
 
-(define (random-number-vector of)
-  ;; Generate a random vector of exact integers with at most `of` number
-  ;; of elements.
-  (vector-generator-of (exact-integer-generator) of))
-
-(define (filter-non-empty-vectors gen)
-  ;; Filter out the empty vector.
-  (gremove (lambda (vec) (zero? (vector-length vec))) gen))
-
-(define (remove-duplicates gen)
-  ;; Filter out vectors with duplicate entries.
-  (define (filter vec)
-    (call/cc
-     (lambda (return)
-       (vector-fold (lambda (prev elem)
-                      (cond
-                        ((null? prev) elem)
-                        ((= elem prev) (return #f))
-                        (else elem)))
-                    '()
-                    vec)
-       #t)))
-  (gfilter filter gen))
-
-(define (make-in-order gen)
-  ;; Sort all vectors to be in-order.
-  (gmap (cut vector-sort < <>) gen))
-
-(define (in-order-vector-generator of)
-  (chain (random-number-vector of)
-         (filter-non-empty-vectors _)
-         (make-in-order _)
-         (remove-duplicates _)))
-
-(define (vector->node-generator of)
-  (chain (in-order-vector-generator of)
-         (gmap in-order-vector->node _)))
+(define (vector->node-generator)
+  (gmap in-order-vector->node (in-order-vector-generator
+                               (max-vector-length))))
 
 (define number-comparator
   (make-comparator
@@ -108,8 +71,7 @@
 
 (test-group "vector->node balanced?"
   (test-assert (balanced? '()))
-  (test-property balanced? (list (vector->node-generator
-                                  (max-vector-length)))))
+  (test-property balanced? (list (vector->node-generator))))
 
 ;;; ;;;;;;;;;
 ;;; Make a "split" vector. A "split" vector is two sorted vectors and a 
@@ -129,82 +91,138 @@
             (vector-ref vec pivot)
             (vector-copy-exclusive vec (+ pivot 1) len))))
 
-(define (split-vector-generator of)
-  (chain (random-number-vector of)
-         (filter-non-empty-vectors _)
-         (remove-duplicates _)
-         (make-in-order _)
+(define (split-vector-generator)
+  (chain (in-order-vector-generator (max-vector-length))
          (gmap split-in-order-vector _)))
 
+(define (call/split-vecs procedure)
+  (lambda (vec)
+    (let ((left (vector-ref vec 0))
+          (middle (vector-ref vec 1))
+          (right (vector-ref vec 2)))
+      (procedure left middle right))))
+
+(define (join-vectors left middle right)
+  (join middle (in-order-vector->node left) (in-order-vector->node right)))
+
 ;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;
 ;;; join
 ;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
-(define (join-preserves-order-of vec)
-  ;; Test if joining two in-order vectors with a node in between will
-  ;; preserve order in the tree.
-  (let ((left (vector-ref vec 0))
-        (middle (vector-ref vec 1))
-        (right (vector-ref vec 2)))
-    ;; Log the left, middle, and right vectors.
-    (when (verbose)
-      (dynamic-property-set! 'right right)
-      (dynamic-property-set! 'middle middle)
-      (dynamic-property-set! 'left left))
-    (let ((ret (node->in-order-list
-                (join middle
-                      (in-order-vector->node left)
-                      (in-order-vector->node right))))
-          (orig (append (vector->list left)
-                        (list middle)
-                        (vector->list right))))
-      (when (verbose)
-        (dynamic-property-set! 'ret ret))
-      (equal? ret orig))))
-
-(test-group "join-preserves-order-of"
+(define join-preserves-order-of
+  (call/split-vecs
+   (lambda (left middle right)
+     ;; Test if joining two in-order vectors with a node in between will
+     ;; preserve order in the tree.
+     ;; Log the left, middle, and right vectors.
+     (when (verbose)
+       (dynamic-property-set! 'right right)
+       (dynamic-property-set! 'middle middle)
+       (dynamic-property-set! 'left left))
+     (let ((ret (node->in-order-list
+                 (join-vectors left middle right)))
+           (orig (append (vector->list left)
+                         (list middle)
+                         (vector->list right))))
+       (when (verbose)
+         (dynamic-property-set! 'ret ret))
+       (equal? ret orig)))))
+
+(test-group "join of in-order vectors preserves order"
   (parameterize ((verbose #f))
-    (test-property join-preserves-order-of (list (split-vector-generator
-                                                  (max-vector-length))))))
+    (test-property join-preserves-order-of
+                   (list (split-vector-generator)))))
 
-(test-group "join balanced?"
-  (define (join-split-vectors vec)
-    (let ((left (vector-ref vec 0))
-          (middle (vector-ref vec 1))
-          (right (vector-ref vec 2)))
-      (join middle
-            (in-order-vector->node left)
-            (in-order-vector->node right))))
-  (define (joined-vector-generator of)
-    (gmap join-split-vectors (split-vector-generator of)))
-  (test-property balanced? (list (joined-vector-generator (max-vector-length)))))
+(test-group "join of in-order vectors is balanced"
+  (define (joined-vector-generator)
+    (gmap (call/split-vecs join-vectors)
+          (split-vector-generator)))
+  (test-property balanced? (list (joined-vector-generator))))
 
 ;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;
 ;;; join2
 ;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;
 
-(define (join2-preserves-order-of vec)
-  (let* ((left (vector-ref vec 0))
-         (right (vector-ref vec 2))
-         (ret (node->in-order-list
-               (join2 (in-order-vector->node left)
-                      (in-order-vector->node right))))
-         (orig (append (vector->list left)
-                       (vector->list right))))
-    (equal? ret orig)))
-
-(test-group "join2 preserves order"
+(define join2-preserves-order-of
+  (call/split-vecs
+   (lambda (left _ right)
+     (let ((ret (node->in-order-list
+                 (join2 (in-order-vector->node left)
+                        (in-order-vector->node right))))
+           (orig (append (vector->list left)
+                         (vector->list right))))
+       (equal? ret orig)))))
+
+(test-group "join2 of in-order vectors preserves order"
   (test-property join2-preserves-order-of
-                 (list (split-vector-generator
-                        (max-vector-length)))))
+                 (list (split-vector-generator))))
 
-(test-group "join2 is balanced"
+(test-group "join2 of in-order vectors is balanced"
   (define (join2-split-vector-generator)
     (define (join2-of-split vec)
       (join2 (in-order-vector->node (vector-ref vec 0))
              (in-order-vector->node (vector-ref vec 2))))
-    (gmap join2-of-split (split-vector-generator
-                          (max-vector-length))))
+    (gmap join2-of-split (split-vector-generator)))
   (test-property balanced?
                  (list (join2-split-vector-generator))))
 
+;;; ;;;;;;
+;;; search
+;;; ;;;;;;
+
+(define (call/inserted procedure)
+  (call/split-vecs
+   (lambda (left middle right)
+     (let ((node (join-vectors left middle right)))
+       (procedure left middle right node)))))
+
+(define (call-w/o-inserted procedure)
+  (call/split-vecs
+   (lambda (left middle right)
+     (let ((node (join2 (in-order-vector->node left)
+                        (in-order-vector->node right))))
+       (procedure left middle right node)))))
+
+(test-group "search with inserted node"
+  (define search-finds
+    (call/inserted
+     (lambda (left middle right node)
+       (when (verbose)
+         (dynamic-property-set! 'expected middle)
+         (dynamic-property-set! 'right right)
+         (dynamic-property-set! 'left left)
+         (dynamic-property-set! 'node (node->in-order-list node)))
+       (search node number-comparator middle))))
+  (parameterize ((verbose #f))
+    (test-property search-finds (list (split-vector-generator)))))
+
+(test-group "search without inserted node"
+  (define search-does-not-find
+    (call-w/o-inserted
+     (lambda (left middle right node)
+       (not (search node number-comparator middle)))))
+  (test-property search-does-not-find (list (split-vector-generator))))
+
+;;; ;;;;;;;;;;;;;;;;;;;;;;;;
+;;; split
+;;; ;;;;;;;;;;;;;;;;;;;;;;;;
+
+(define split-finds
+  (call/inserted
+   (lambda (left middle right node)
+     (let-values (((new-left found? new-right)
+                   (split number-comparator node middle)))
+       found?))))
+
+(test-group "split finds"
+  (test-property split-finds (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)))
+       (not found?)))))
+
+(test-group "split does not find"
+  (test-property split-does-not-find (list (split-vector-generator))))
diff --git a/tests/util.sld b/tests/util.sld
new file mode 100644
index 0000000..daaa9db
--- /dev/null
+++ b/tests/util.sld
@@ -0,0 +1,26 @@
+#| Copyright 2024 Peter McGoron
+ |
+ | Licensed under the Apache License, Version 2.0 (the "License");
+ | you may not use this file except in compliance with the License.
+ | You may obtain a copy of the License at
+ |
+ |     http://www.apache.org/licenses/LICENSE-2.0
+ |
+ | Unless required by applicable law or agreed to in writing, software
+ | distributed under the License is distributed on an "AS IS" BASIS,
+ | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ | See the License for the specific language governing permissions and
+ | limitations under the License.
+ |#
+
+(define-library (mcgoron weight-balanced-trees test util)
+  (import (scheme base) (srfi 158) (srfi 26))
+  (export vector-copy-exclusive)
+  (begin
+    (define (vector-copy-exclusive vec start end)
+      ;; Copy `start,end` (start inclusive end exclusive) from `vec`.
+      (if (= start end)
+          #()
+          (vector-copy vec start (- end 1))))
+    (define (boolean-equivalent? x y)
+      (boolean=? (not (not x)) (not (not y))))))
\ No newline at end of file