From f2c9267b48660bc7ed80aa63f44f95cb3d32b87f Mon Sep 17 00:00:00 2001
From: Andrew Noyes <andrew@weaselab.dev>
Date: Tue, 30 Jan 2024 14:21:30 -0800
Subject: [PATCH] Use a bitset to speed up Node256 scans

---
 ConflictSet.cpp | 81 ++++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 73 insertions(+), 8 deletions(-)

diff --git a/ConflictSet.cpp b/ConflictSet.cpp
index b535d8e..c1e1f19 100644
--- a/ConflictSet.cpp
+++ b/ConflictSet.cpp
@@ -71,7 +71,66 @@ struct Node48 : Node {
   }
 };
 
+struct PointerSet {
+  bool test(int i) const {
+    assert(0 <= i);
+    assert(i < 256);
+    if (i < 128) {
+      return (lo >> i) & 1;
+    } else {
+      return (hi >> (i - 128)) & 1;
+    }
+  }
+
+  void set(int i) {
+    assert(0 <= i);
+    assert(i < 256);
+    if (i < 128) {
+      lo |= __uint128_t(1) << i;
+    } else {
+      hi |= __uint128_t(1) << (i - 128);
+    }
+  }
+
+  int firstSetGeq(int i) const {
+    if (i < 128) {
+      int a = std::countr_zero(lo >> i);
+      if (a < 128) {
+        assert(i + a < 128);
+        return i + a;
+      }
+      i = 128;
+    }
+    int b = std::countr_zero(hi >> (i - 128));
+    if (b < 128) {
+      assert(i + b < 256);
+      return i + b;
+    }
+    return -1;
+  }
+
+  int lastSetLeq(int i) const {
+    if (i >= 128) {
+      int a = std::countl_zero(hi << (255 - i));
+      if (a < 128) {
+        return i - a;
+      }
+      i = 127;
+    }
+    int b = std::countl_zero(lo << (127 - i));
+    if (b < 128) {
+      return i - b;
+    }
+    return -1;
+  }
+
+private:
+  __uint128_t lo = 0;
+  __uint128_t hi = 0;
+};
+
 struct Node256 : Node {
+  PointerSet pointerSet;
   Node *children[256] = {};
   Node256() { this->type = Type::Node256; }
 };
@@ -334,11 +393,12 @@ int getChildGeq(Node *self, int child) {
 #endif
   } else {
     auto *self256 = static_cast<Node256 *>(self);
-    for (int i = child; i < 256; ++i) {
-      if (self256->children[i]) {
-        return i;
-      }
+#ifndef NDEBUG
+    for (int i = 0; i < 256; ++i) {
+      assert(self256->pointerSet.test(i) == (self256->children[i] != nullptr));
     }
+#endif
+    return self256->pointerSet.firstSetGeq(child);
   }
   return -1;
 }
@@ -444,11 +504,12 @@ int getChildLeq(Node *self, int child) {
 #endif
   } else {
     auto *self256 = static_cast<Node256 *>(self);
-    for (int i = child; i >= 0; --i) {
-      if (self256->children[i]) {
-        return i;
-      }
+#ifndef NDEBUG
+    for (int i = 0; i < 256; ++i) {
+      assert(self256->pointerSet.test(i) == (self256->children[i] != nullptr));
     }
+#endif
+    return self256->pointerSet.lastSetLeq(child);
   }
   return -1;
 }
@@ -459,6 +520,8 @@ void setChildrenParents(Node *node) {
   }
 }
 
+// Caller is responsible for assigning a non-null pointer to the returned
+// reference if null
 Node *&getOrCreateChild(Node *&self, uint8_t index) {
   if (self->type == Type::Node4) {
     auto *self4 = static_cast<Node4 *>(self);
@@ -545,6 +608,7 @@ Node *&getOrCreateChild(Node *&self, uint8_t index) {
       memcpy((void *)newSelf, self, offsetof(Node, type));
       for (int i = 0; i < 256; ++i) {
         if (self48->index[i] >= 0) {
+          newSelf->pointerSet.set(i);
           newSelf->children[i] = self48->children[self48->index[i]];
         }
       }
@@ -565,6 +629,7 @@ Node *&getOrCreateChild(Node *&self, uint8_t index) {
     if (!self256->children[index]) {
       ++self->numChildren;
     }
+    self256->pointerSet.set(index);
     return self256->children[index];
   }
 }