From 9fcfc44dc33e8b25fad6d84a9878f1607c0f8461 Mon Sep 17 00:00:00 2001
From: Andrew Noyes <andrew@weaselab.dev>
Date: Sun, 25 Feb 2024 21:04:39 -0800
Subject: [PATCH] Add forEachInRange

---
 ConflictSet.cpp | 55 +++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 55 insertions(+)

diff --git a/ConflictSet.cpp b/ConflictSet.cpp
index 6cba44e..2a8609d 100644
--- a/ConflictSet.cpp
+++ b/ConflictSet.cpp
@@ -98,6 +98,61 @@ struct BitSet {
   void reset(int i);
   int firstSetGeq(int i) const;
 
+  // Calls `f` with the index of each bit set in [begin, end)
+  template <class F> void forEachInRange(F f, int begin, int end) {
+    // See section 3.1 in https://arxiv.org/pdf/1709.07821.pdf for details about
+    // this approach
+
+    if ((begin >> 6) == (end >> 6)) {
+      {
+        uint64_t word = words[begin >> 6] & (uint64_t(-1) << (begin & 63)) &
+                        ~(uint64_t(-1) << (end & 63));
+        while (word) {
+          uint64_t temp = word & -word;
+          int index = (begin & ~63) + std::countr_zero(word);
+          f(index);
+          word ^= temp;
+        }
+      }
+      return;
+    }
+
+    // Check begin partial word
+    {
+      uint64_t word = words[begin >> 6] & (uint64_t(-1) << (begin & 63));
+      while (word) {
+        uint64_t temp = word & -word;
+        int index = (begin & ~63) + std::countr_zero(word);
+        f(index);
+        word ^= temp;
+      }
+    }
+
+    // Check inner, full words
+    begin += 64;
+    while ((begin >> 6) != (end >> 6)) {
+      uint64_t word = words[begin >> 6];
+      while (word) {
+        uint64_t temp = word & -word;
+        int index = (begin & ~63) + std::countr_zero(word);
+        f(index);
+        word ^= temp;
+      }
+      begin += 64;
+    }
+
+    if (end < 256) {
+      // Check end partial word
+      uint64_t word = words[end >> 6] & ~(uint64_t(-1) << (end & 63));
+      while (word) {
+        uint64_t temp = word & -word;
+        int index = (begin & ~63) + std::countr_zero(word);
+        f(index);
+        word ^= temp;
+      }
+    }
+  }
+
 private:
   uint64_t words[4] = {};
 };