From 77f4663bc8cd1de8315dc918a5a5bcc38a35e1fa Mon Sep 17 00:00:00 2001
From: Andrew Noyes <andrew@weaselab.dev>
Date: Wed, 1 May 2024 15:52:40 -0700
Subject: [PATCH] Add RootSet

---
 VersionedMap.cpp | 116 ++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 115 insertions(+), 1 deletion(-)

diff --git a/VersionedMap.cpp b/VersionedMap.cpp
index c9dae01..c5a84c9 100644
--- a/VersionedMap.cpp
+++ b/VersionedMap.cpp
@@ -216,6 +216,9 @@ struct MemManager {
       }
     };
     for (int i = 0; i < numRoots; ++i) {
+      if (roots[i] == 0) {
+        continue;
+      }
       tryPush(roots[i]);
       while (stackIndex > 0) {
         uint32_t p = stack[--stackIndex];
@@ -281,7 +284,100 @@ private:
   uint32_t freeList = 0;
 };
 
-struct VersionedMap::Impl {};
+struct RootSet {
+
+  /// Register the root node for version after adding mutations
+  void add(uint32_t node, int64_t version) {
+    if (end == 0) {
+      nodes[end] = node;
+      versions[end] = version;
+      ++end;
+      return;
+    }
+    if (nodes[end - 1] == node) {
+      return;
+    }
+    if (end == capacity) {
+      capacity *= 2;
+      nodes = (uint32_t *)realloc(nodes, capacity * sizeof(uint32_t));
+      versions = (int64_t *)realloc(versions, capacity * sizeof(int64_t));
+    }
+
+    nodes[end] = node;
+    versions[end] = version;
+    ++end;
+  }
+
+  /// Inform that there will be no calls to rootForVersion with a version less
+  /// than `oldestVersion`
+  void setOldestVersion(int64_t oldestVersion) {
+    const uint32_t firstToKeep = rootForVersion(oldestVersion);
+
+    if (firstToKeep != 0) {
+      memmove(nodes, nodes + firstToKeep,
+              (end - firstToKeep) * sizeof(uint32_t));
+      memmove(versions, versions + firstToKeep,
+              (end - firstToKeep) * sizeof(int64_t));
+      end -= firstToKeep;
+    }
+    assert(end > 0);
+    assert(nodes[0] <= oldestVersion);
+  }
+
+  /// Get a root node that can correctly be used for `version`
+  uint32_t rootForVersion(int64_t version) const {
+    assert(end > 0);
+    assert(nodes[0] <= version);
+
+    // Find the last version <= oldestVersion
+    int left = 0;
+    int right = end;
+    int result = 0;
+    while (left <= right) {
+      int mid = left + (right - left) / 2;
+      if (versions[mid] <= version) {
+        result = mid;
+        left = mid + 1;
+      } else {
+        right = mid - 1;
+      }
+    }
+    assert(result < end);
+    return result;
+  }
+
+  const uint32_t *roots() const { return nodes; }
+  int rootCount() const { return end; }
+
+  RootSet() {
+    nodes = (uint32_t *)malloc(kMinCapacity * sizeof(uint32_t));
+    versions = (int64_t *)malloc(kMinCapacity * sizeof(int64_t));
+    capacity = kMinCapacity;
+    nodes[0] = 0;
+    versions[0] = 0;
+    end = 1;
+  }
+
+  ~RootSet() {
+    free(versions);
+    free(nodes);
+  }
+
+private:
+  uint32_t *nodes;
+  // versions[i] is the version of nodes[i]
+  int64_t *versions;
+
+  constexpr static uint32_t kMinCapacity = 16;
+  uint32_t capacity;
+  uint32_t end;
+};
+
+struct VersionedMap::Impl {
+
+  MemManager mm;
+  RootSet roots;
+};
 } // namespace weaselab
 
 #ifdef ENABLE_MAIN
@@ -311,5 +407,23 @@ int main() {
   mm.base[root].pointers[1] = 0;
   mm.base[root].updated.store(false, std::memory_order_relaxed);
   bench.run("gc", [&]() { mm.gc(&root, 1, 0); });
+
+  {
+    int i = 0;
+    constexpr int kNumVersions = 1000;
+    weaselab::RootSet roots;
+    for (; i < kNumVersions; i += 2) {
+      roots.add(i, i);
+      roots.add(i, i + 1);
+    }
+    bench.run("roots - setOldestVersion", [&]() {
+      roots.add(i, i);
+      roots.setOldestVersion(i - kNumVersions);
+      ++i;
+    });
+    bench.run("roots - rootForVersion", [&]() {
+      bench.doNotOptimizeAway(roots.rootForVersion(i - kNumVersions / 2));
+    });
+  }
 }
 #endif
\ No newline at end of file