add node->generator

2 files changed, 59 insertions, 7 deletions

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"))