Allow reification of slsets to get fast adjoin and membership testing

2 files changed, 84 insertions, 2 deletions

diff --git a/slsets.scm b/slsets.scm
index 9b93676..d5f4020 100644
--- a/slsets.scm
+++ b/slsets.scm
@@ -14,7 +14,10 @@
      slset-difference+intersection
      slset-union slset-xor
      slset-deduplicate
-     slset-delete slset-contains?)
+     slset-delete slset-contains?
+
+     with-reified-slset reified-slset? reified-slset->slset
+     reified-slset-adjoin! reified-slset-delete! reified-slset-contains?)
 
 (import scheme (chicken base) (chicken plist))
 
@@ -171,4 +174,52 @@
                     (lambda (m)
                       (filter (lambda (x) (not (marked? x m))) slset))))
 
+(define-record-type :reified-slset
+  (reified-slset marking slset)
+  reified-slset?
+  (marking reified-slset-marking)
+  (slset reified-slset-list reified-slset-list-set!))
+
+(define reified-slset->slset reified-slset-list)
+
+(define (with-reified-slset lst fun)
+  ;; Note: we can't use with-marked-list here because the reified
+  ;; slset will be mutated in-place, so the enter/leave thunks must
+  ;; extract the *current* list from the reified slset every time.
+  (let* ((marking (gensym 'm))
+         (rs (reified-slset marking lst)))
+    (dynamic-wind
+        (lambda () (mark-list! (reified-slset-list rs) marking))
+        (lambda ()
+          (fun rs)
+          (reified-slset-list rs))
+        (lambda () (unmark-list! (reified-slset-list rs) marking)))))
+
+(define (reified-slset-adjoin! rs . els)
+  (let ((m (reified-slset-marking rs)))
+    (let lp ((lst (reified-slset-list rs))
+             (els els))
+      (cond ((null? els)
+             (reified-slset-list-set! rs lst)
+             rs)
+            ((marked? (car els) m) (lp lst (cdr els)))
+            (else (mark! (car els) m)
+                  (lp (cons (car els) lst) (cdr els)))))))
+
+;; NOTE: Deleting elements not already occurring in the list is O(1)
+(define (reified-slset-delete! rs . els)
+  (let* ((m1 (reified-slset-marking rs))
+         (els (filter (lambda (e) (not (marked? m1 e))) els)))
+    (unless (null? els)
+      (let ((lst (with-marked-list els
+                                   (lambda (m2)
+                                     (filter (lambda (el) (not (marked? el m2)))
+                                             (reified-slset-list rs))))))
+        (for-each (lambda (e) (unmark! e m1)) els)
+        (reified-slset-list-set! rs lst)))
+    rs))
+
+(define (reified-slset-contains? rs el)
+  (marked? el (reified-slset-marking rs)))
+
 )
diff --git a/tests/run.scm b/tests/run.scm
index 4d15ea7..571a7fe 100644
--- a/tests/run.scm
+++ b/tests/run.scm
@@ -154,6 +154,37 @@
       (test-assert (clean? '(a b c z)))))
 
   (test-group "finally, all symbols used in any test are still clean"
-    (test-assert (clean? all-symbols))))
+    (test-assert (clean? all-symbols)))
+
+  (test-group "reified slsets"
+    (let ((result (with-reified-slset '(a b c d a c e)
+                                      (lambda (r)
+                                        (test #t (reified-slset-contains? r 'a))
+                                        (test #t (reified-slset-contains? r 'b))
+                                        (test #t (reified-slset-contains? r 'e))
+                                        (test #f (reified-slset-contains? r 'i))
+                                        (test #f (reified-slset-contains? r 'x))
+
+                                        (reified-slset-delete! r 'b)
+                                        (test '(a c d a c e)
+                                              (reified-slset->slset r))
+
+                                        (reified-slset-adjoin! r 'x)
+                                        (test #t (reified-slset-contains? r 'x))
+
+                                        (reified-slset-delete! r 'a 'e 'i)
+                                        (test '(x c d c)
+                                              (reified-slset->slset r))
+
+                                        ;; Extra check to ensure we unmark these
+                                        (test #f (reified-slset-contains? r 'a))
+
+                                        ;; Basic sanity check, remaining contents still in there
+                                        (test #t (reified-slset-contains? r 'c))))))
+      ;; Put it in an extra list to avoid test form complaining
+      (test '((x c d c)) (list result)))
+
+    (test-group "all symbols used here still clean"
+      (test-assert (clean? '(a b c d e i x))))))
 
 (test-exit)