add node->generator

5 files changed, 102 insertions, 18 deletions

diff --git a/doc/mcgoron.weight-balanced-trees.internal.scm b/doc/mcgoron.weight-balanced-trees.internal.scm
index a48cb23..12d7e6e 100644
--- a/doc/mcgoron.weight-balanced-trees.internal.scm
+++ b/doc/mcgoron.weight-balanced-trees.internal.scm
@@ -126,7 +126,7 @@ resulting node is inserted into the tree.
 "))
  ((name . "delete")
   (signature
-   lambda ((comparator? cmp) element (balanced? set)) balanced?)
+   lambda ((comparator? cmp) (balanced? set) element) balanced?)
   (desc "
 * It is an error if `cmp` does not order the elements of `set` and the
   value `element`.
@@ -134,6 +134,17 @@ resulting node is inserted into the tree.
 Search `set` for an element `E` that compares equal to `element`. If `E`
 is found, the returned tree does not have `E`. Otherwise the returned
 tree has the same number of elements."))
+ ((name . "insert")
+  (signature
+   lambda ((comparator? cmp) (balanced? set) element) balanced?)
+  (desc "
+* It is an error if `cmp` does not order the elements of `set` and the
+  value `element`.
+
+Search `set` for an element `E` that compares equal to `element`. If `E`
+is found, the return tree element will have `E` replaced with `element`.
+Otherwise, the returned tree will have `element` in addition to the rest
+of the values in the tree."))
  ((name . "every")
   (signature
    lambda ((procedure? predicate?) (wb-tree-node? tree)) *)
@@ -141,4 +152,12 @@ tree has the same number of elements."))
    (predicate? lambda (*) *))
   (desc "
 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`.")))
\ No newline at end of file
+calls return a truthy value, return a truthy value. Otherwise return `#f`."))
+ ((name . "node->generator")
+  (signature
+   lambda ((balanced? 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
diff --git a/mcgoron/weight-balanced-trees/internal.scm b/mcgoron/weight-balanced-trees/internal.scm
index f0dc7c5..9b9fdd1 100644
--- a/mcgoron/weight-balanced-trees/internal.scm
+++ b/mcgoron/weight-balanced-trees/internal.scm
@@ -363,7 +363,7 @@
 ;;; Single value operations
 ;;; ;;;;;;;;;;;;;;;;;;;;
 
-(: update (* node-type * (-> *) (-> node-type) --> node-type))
+(: update (* node-type * (-> *) (-> node-type) -> node-type))
 (define (update cmp set to-search on-found on-not-found)
   (let update ((set set))
     (if (null? set)
@@ -374,13 +374,24 @@
             (wb-tree-node (on-found data) left right)
             (join data left (update right)))))))
 
-(: delete (* * node-type --> node-type))
-(define (delete cmp to-search set)
+(: insert (* node-type * --> node-type))
+(define (insert cmp set new-value)
+  (let insert ((set set))
+    (if (null? set)
+        (wb-tree-node new-value '() '())
+        (with-node (set data ("<" left) (">" right))
+          (comparator-if<=> cmp new-value data
+            (join data (insert left) right)
+            (wb-tree-node new-value left right)
+            (join data left (insert right)))))))
+
+(: delete (* node-type * --> node-type))
+(define (delete cmp set to-delete)
   (let delete ((set set))
     (if (null? set)
         '()
         (with-node (set data ("<" left) (">" right))
-          (comparator-if<=> cmp to-search data
+          (comparator-if<=> cmp to-delete data
             (join data (delete left) right)
             (join2 left right)
             (join data left (delete right)))))))
@@ -398,3 +409,43 @@
              (every predicate? left)
              (every predicate? right)))))
 
+;;; ;;;;;;;;;;;;;;;;;;;
+;;; Generators
+;;; ;;;;;;;;;;;;;;;;;;;
+
+(define (node->generator node)
+  (let ((queue (list-queue)))
+    (define (add-when-not-null! node)
+      (when (not (null? node))
+        (list-queue-add-front! queue node)))
+    (add-when-not-null! node)
+    (lambda ()
+      (if (list-queue-empty? queue)
+          (eof-object)
+          (let ((current-node (list-queue-remove-front! queue)))
+            (add-when-not-null! (get-left current-node))
+            (add-when-not-null! (get-right current-node))
+            (get-data current-node))))))
+
+(define (generator->node comparator gen)
+  (let loop ((node '()))
+    (let ((value (gen)))
+      (if (eof-object? value)
+          node
+          (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)))))))
diff --git a/mcgoron/weight-balanced-trees/internal.sld b/mcgoron/weight-balanced-trees/internal.sld
index 2b77a64..3839126 100644
--- a/mcgoron/weight-balanced-trees/internal.sld
+++ b/mcgoron/weight-balanced-trees/internal.sld
@@ -15,7 +15,7 @@
 
 (define-library (mcgoron weight-balanced-trees internal)
   (import (scheme base) (scheme case-lambda)
-          (srfi 128))
+          (srfi 117) (srfi 128))
   (cond-expand
     ;; Handle type declarations
     (chicken (import (chicken type)))
@@ -35,7 +35,8 @@
           join join2 split
           search
           union intersection difference xor
-          update delete
-          every)
+          update insert delete
+          every
+          node->generator)
   (include "internal.scm"))
 
diff --git a/tests/run.scm b/tests/run.scm
index 25cbc4c..8cfdabc 100644
--- a/tests/run.scm
+++ b/tests/run.scm
@@ -422,13 +422,11 @@
 ;;; Update
 ;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
+(define %insert insert)
 (define (insert elem set)
-  (update number-comparator
-          set
-          elem
-          (lambda (value) value)
-          (lambda ()
-            (wb-tree-node elem '() '()))))
+  (%insert number-comparator
+           set
+           elem))
 
 (test-group "update of unordered list is balanced"
   (define (insert-unordered lst)
@@ -471,7 +469,7 @@
                             (when (verbose)
                               (dynamic-property-set! 'elem elem)
                               (dynamic-property-set! 'set set))
-                            (delete number-comparator elem set))
+                            (delete number-comparator set elem))
                           set1
                           lst2)))
       (procedure lst1 lst2 diff deleted))))
@@ -480,7 +478,7 @@
   (define (delete-balanced? lst1 lst2)
     (let* ((set (fold insert '() lst1))
            (deleted (fold (lambda (elem set)
-                            (delete number-comparator elem set))
+                            (delete number-comparator set elem))
                           set
                           lst2)))
       (balanced? deleted)))
@@ -534,3 +532,17 @@
   (test-property xor-subset-of-lset
                  (list (list-generator-of (exact-integer-generator))
                        (list-generator-of (exact-integer-generator)))))
+
+;;; ;;;;;;;;;;
+;;; Generators
+;;; ;;;;;;;;;;
+
+(test-group "node->generator"
+  (define (node->generator-works lst)
+    (let* ((node (fold insert '() lst))
+           (result-lst
+            (generator->list (node->generator node))))
+      (lset= = lst result-lst)))
+  (test-property node->generator-works
+                 (list (list-generator-of (exact-integer-generator)))))
+
diff --git a/weight-balanced-trees.egg b/weight-balanced-trees.egg
index aa0f88c..470ca73 100644
--- a/weight-balanced-trees.egg
+++ b/weight-balanced-trees.egg
@@ -7,5 +7,6 @@
  (test-dependencies "srfi-133" "srfi-132" "srfi-194" "srfi-132" "srfi-128" "srfi-1" "srfi-64" "srfi-252" "srfi-158" "sexpr-srfi-64-runner" "srfi-197")
  (components (extension mcgoron.weight-balanced-trees.internal
                (source "mcgoron/weight-balanced-trees/internal.sld")
+               (types-file)
                (source-dependencies "mcgoron/weight-balanced-trees/internal.scm")
-               (csc-options "-R" "r7rs" "-X" "r7rs"))))
+               (csc-options "-O3" "-R" "r7rs" "-X" "r7rs"))))