From 235938b5aaedd15dcc379360b3a89503510b1678 Mon Sep 17 00:00:00 2001
From: Andrew Noyes <andrew@weaselab.dev>
Date: Fri, 22 Nov 2024 18:27:45 -0800
Subject: [PATCH] Update eraseTree to account for leaves

---
 ConflictSet.cpp | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/ConflictSet.cpp b/ConflictSet.cpp
index 6c7891d..f2b2b9b 100644
--- a/ConflictSet.cpp
+++ b/ConflictSet.cpp
@@ -276,6 +276,8 @@ struct TaggedNodePointer {
     __builtin_prefetch((void *)withoutType());
   }
 
+  bool isLeaf() { return getType() > Type_Node256; }
+
 private:
   TaggedNodePointer(struct Node *p, Type t) : p((uintptr_t)p) {
     assert((this->p & 15) == 0);
@@ -3069,6 +3071,12 @@ void eraseTree(Node *root, WriteContext *writeContext) {
 
     removeKey(n);
 
+    auto pushIfNotLeaf = [&toFree](TaggedNodePointer p) {
+      if (!p.isLeaf()) {
+        toFree.push_back(p);
+      }
+    };
+
     switch (n->getType()) {
     case Type_Node0: {
       auto *n0 = static_cast<Node0 *>(n);
@@ -3077,31 +3085,33 @@ void eraseTree(Node *root, WriteContext *writeContext) {
     case Type_Node3: {
       auto *n3 = static_cast<Node3 *>(n);
       for (int i = 0; i < n3->numChildren; ++i) {
-        toFree.push_back(n3->children[i]);
+        pushIfNotLeaf(n3->children[i]);
       }
       writeContext->release(n3);
     } break;
     case Type_Node16: {
       auto *n16 = static_cast<Node16 *>(n);
       for (int i = 0; i < n16->numChildren; ++i) {
-        toFree.push_back(n16->children[i]);
+        pushIfNotLeaf(n16->children[i]);
       }
       writeContext->release(n16);
     } break;
     case Type_Node48: {
       auto *n48 = static_cast<Node48 *>(n);
       for (int i = 0; i < n48->numChildren; ++i) {
-        toFree.push_back(n48->children[i]);
+        pushIfNotLeaf(n48->children[i]);
       }
       writeContext->release(n48);
     } break;
     case Type_Node256: {
       auto *n256 = static_cast<Node256 *>(n);
-      auto *out = toFree.unsafePrepareAppend(n256->numChildren).data();
-      n256->bitSet.forEachSet([&](int i) { *out++ = n256->children[i]; });
-      assert(out == toFree.end());
+      n256->bitSet.forEachSet([&](int i) { pushIfNotLeaf(n256->children[i]); });
       writeContext->release(n256);
     } break;
+    case Type_Leaf3:           // GCOVR_EXCL_LINE
+    case Type_Leaf16:          // GCOVR_EXCL_LINE
+    case Type_Leaf48:          // GCOVR_EXCL_LINE
+    case Type_Leaf256:         // GCOVR_EXCL_LINE
     default:                   // GCOVR_EXCL_LINE
       __builtin_unreachable(); // GCOVR_EXCL_LINE
     }