generator->node

4 files changed, 34 insertions, 17 deletions

diff --git a/doc/mcgoron.weight-balanced-trees.internal.scm b/doc/mcgoron.weight-balanced-trees.internal.scm
index 12d7e6e..ac3b921 100644
--- a/doc/mcgoron.weight-balanced-trees.internal.scm
+++ b/doc/mcgoron.weight-balanced-trees.internal.scm
@@ -155,9 +155,19 @@ Calls `predicate?` on each element of `tree` in an arbitrary order. If all
 calls return a truthy value, return a truthy value. Otherwise return `#f`."))
  ((name . "node->generator")
   (signature
-   lambda ((balanced? node)) procedure?)
+   lambda ((generator? node)) procedure?)
   (subsigs
    (return lambda () *))
   (desc "
 Returns a generator (see SRFI-158) that generates the elements of `node`
-in arbitrary order.")))
\ No newline at end of file
+in arbitrary order."))
+ ((name . "generator->node")
+  (signature
+   lambda ((comparator? cmp) (generator? generator)) balanced?)
+  (desc "
+* It is an error if `generator` is not finite.
+* It is an error if the values generated by `generator` are not comparable
+  by `cmp`.
+
+Returns a weight balanced tree whose elements are all of the elements
+of `generator`.")))
\ No newline at end of file
diff --git a/mcgoron/weight-balanced-trees/internal.scm b/mcgoron/weight-balanced-trees/internal.scm
index 9b9fdd1..ffbe6b4 100644
--- a/mcgoron/weight-balanced-trees/internal.scm
+++ b/mcgoron/weight-balanced-trees/internal.scm
@@ -435,17 +435,15 @@
           (loop (insert comparator node value))))))
 
 (define (node->in-order-generator node)
-  (if (null? node)
-      eof-object
-      (let ((queue (list-queue)))
-        (define (traverse-left! node)
-          (when (not (null? node))
-            (list-queue-add-front! node)
-            (traverse-left! (get-left node))))
-        (traverse-left! node)
-        (lambda ()
-          (if (list-queue-empty? queue)
-              (eof-object)
-              (let ((current (list-queue-remove-front! queue)))
-                (traverse-left! (get-right current))
-                (get-data current)))))))
+  (let ((queue (list-queue)))
+    (define (traverse-left! node)
+      (when (not (null? node))
+        (list-queue-add-front! node)
+        (traverse-left! (get-left node))))
+    (traverse-left! node)
+    (lambda ()
+      (if (list-queue-empty? queue)
+          (eof-object)
+          (let ((current (list-queue-remove-front! queue)))
+            (traverse-left! (get-right current))
+            (get-data current))))))
diff --git a/mcgoron/weight-balanced-trees/internal.sld b/mcgoron/weight-balanced-trees/internal.sld
index 3839126..0676e81 100644
--- a/mcgoron/weight-balanced-trees/internal.sld
+++ b/mcgoron/weight-balanced-trees/internal.sld
@@ -37,6 +37,6 @@
           union intersection difference xor
           update insert delete
           every
-          node->generator)
+          node->generator generator->node)
   (include "internal.scm"))
 
diff --git a/tests/run.scm b/tests/run.scm
index 8cfdabc..81c5a86 100644
--- a/tests/run.scm
+++ b/tests/run.scm
@@ -546,3 +546,12 @@
   (test-property node->generator-works
                  (list (list-generator-of (exact-integer-generator)))))
 
+(test-group "generator->node"
+  (define (generator->node-works lst)
+    (let* ((gen (list->generator lst))
+           (node (generator->node number-comparator gen)))
+      (and (node-every (cut member <> lst) node)
+           (every (cut search number-comparator <> node) lst))))
+  (test-property generator->node-works
+                 (list (list-generator-of (exact-integer-generator)))))
+