more tests for set-delete, more efficient implementation

1 files changed, 16 insertions, 12 deletions

diff --git a/mcgoron/weight-balanced-trees/srfi/113/sets.scm b/mcgoron/weight-balanced-trees/srfi/113/sets.scm
index 03f5045..964ff87 100644
--- a/mcgoron/weight-balanced-trees/srfi/113/sets.scm
+++ b/mcgoron/weight-balanced-trees/srfi/113/sets.scm
@@ -104,13 +104,15 @@
 (define set-replace! set-replace)
 
 (define (set-delete-all set elements)
-  (let ((cmp (set-element-comparator set)))
-    (raw-set
-     cmp
-     (fold (lambda (element node)
-             (delete cmp node element))
-           (get-node set)
-           elements))))
+  (define cmp (set-element-comparator set))
+  (define (delete node)
+    (if (null? node)
+        '()
+        (with-node (node data ("<" left) (">" right))
+          (if (member data elements (cute =? cmp <> <>))
+              (join2 (delete left) (delete right))
+              (join data (delete left) (delete right))))))
+  (raw-set cmp (delete (get-node set))))
 (define set-delete-all! set-delete-all)
 
 (define (set-delete set . elements)
@@ -357,19 +359,21 @@
 ;;; ;;;;;;;;;;;;
 
 (define (set-adjoin-all set elements)
+  ;; TODO: replace with version that only crawls the tree once.
   (let ((cmp (set-element-comparator set)))
     (raw-set
      cmp
      (fold (lambda (new set)
-        (update cmp
-                set
-                new
-                (lambda (old) old)
-                (lambda () (wb-tree-node new '() '()))))
+             (update cmp
+                     set
+                     new
+                     (lambda (old) old)
+                     (lambda () (wb-tree-node new '() '()))))
            (get-node set)
            elements))))
 
 (define (set-replace-all set elements)
+  ;; TODO: replace with version that only crawls the tree once.
   (let ((cmp (set-element-comparator set)))
     (fold (lambda (new set)
             (update cmp