Fix SHOW_MEMORY crashes on valid input

`Node::getCapacity()` asserted `!releaseDeferred`, but releasing a deferred
node requires its capacity, so any node replacement aborted an assertion-
enabled `SHOW_MEMORY` build. `partialKeyCapacity` is not repurposed by
`deferRelease`, so there is no reason to reject deferred nodes. Drop the
assertion.

`eraseTree` released a node before visiting its children, so the `parent`
pointers walked by `removeKey` pointed to freed memory (and `deferRelease`
repurposes `parent` as a list link). Carry the search path length of each node
in the worklist, computed while its parent is still alive, and account for it
with the new `removeKey` overload.

Add a test covering the SHOW_MEMORY instrumentation, which the rest of the
test suite doesn't compile with.

CMakeLists.txt

@@ -390,6 +390,24 @@ if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR AND BUILD_TESTING)
                         PROPERTIES CXX_STANDARD_REQUIRED ON)
   add_test(NAME conflict_set_cxx_api_test COMMAND conflict_set_cxx_api_test)
 
+  # SHOW_MEMORY instrumentation, which the rest of the test suite doesn't
+  # compile with
+  add_executable(conflict_set_show_memory_test ConflictSet.cpp
+                                               ShowMemoryTest.cpp)
+  target_compile_options(conflict_set_show_memory_test PRIVATE ${TEST_FLAGS})
+  target_compile_definitions(conflict_set_show_memory_test
+                             PRIVATE SHOW_MEMORY=1)
+  target_include_directories(conflict_set_show_memory_test
+                             PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include)
+  if(NOT CMAKE_CROSSCOMPILING)
+    target_compile_options(conflict_set_show_memory_test
+                           PRIVATE -fsanitize=address,undefined)
+    target_link_options(conflict_set_show_memory_test PRIVATE
+                        -fsanitize=address,undefined)
+  endif()
+  add_test(NAME conflict_set_show_memory_test
+           COMMAND conflict_set_show_memory_test)
+
   # symbol visibility tests
   if(NOT CMAKE_BUILD_TYPE STREQUAL Debug)
     if(APPLE)

ConflictSet.cpp

@@ -85,11 +85,18 @@ void addNode(struct Node *);
 void removeNode(struct Node *);
 void addKey(struct Node *);
 void removeKey(struct Node *);
+// Like `removeKey`, but with a caller-provided search path length. Used when
+// the node's ancestors may already have been released, which makes walking
+// `parent` pointers invalid.
+void removeKey(struct Node *, int64_t searchPathLength);
+int64_t getSearchPathLength(struct Node *);
 #else
 constexpr void addNode(struct Node *) {}
 constexpr void removeNode(struct Node *) {}
 constexpr void addKey(struct Node *) {}
 constexpr void removeKey(struct Node *) {}
+constexpr void removeKey(struct Node *, int64_t) {}
+constexpr int64_t getSearchPathLength(struct Node *) { return 0; }
 #endif
 
 // ==================== BEGIN IMPLEMENTATION ====================
@@ -307,7 +314,9 @@ struct Node {
   uint8_t *partialKey();
   Type getType() const { return type; }
   int32_t getCapacity() const {
-    assert(!releaseDeferred);
+    // `partialKeyCapacity` is not repurposed by `deferRelease`, so this is also
+    // valid for nodes with `releaseDeferred` set. It's needed to account for
+    // the capacity of a deferred node when it's finally released.
     return partialKeyCapacity;
   }
 
@@ -2754,17 +2763,30 @@ checkMaxBetweenExclusiveImpl<true>(Node256 *n, int begin, int end,
 }
 
 void eraseTree(Node *root, WriteContext *writeContext) {
+  // A node is released before its children are visited, so the `parent`
+  // pointers of queued nodes can't be relied upon. Carry the search path length
+  // of each node along with it instead.
+  struct ToFree {
+    Node *n;
+    // Length of the key leading to `n`. See `getSearchPathLength`.
+    int64_t searchPathLength;
+  };
   Arena arena;
-  auto toFree = vector<Node *>(arena);
-  toFree.push_back(root);
+  auto toFree = vector<ToFree>(arena);
+  toFree.push_back({root, getSearchPathLength(root)});
 
   while (toFree.size() > 0) {
-    auto *n = toFree.back();
+    const auto item = toFree.back();
+    auto *n = item.n;
     toFree.pop_back();
     writeContext->accum.entries_erased += n->entryPresent;
     ++writeContext->accum.nodes_released;
 
-    removeKey(n);
+    removeKey(n, item.searchPathLength);
+
+    // Search path length of the children of `n`. Must be read before `n` is
+    // released.
+    const int64_t childSearchPathLengthBase = item.searchPathLength + 1;
 
     switch (n->getType()) {
     case Type_Node0: {
@@ -2774,28 +2796,37 @@ 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]);
+        Node *child = n3->children[i];
+        toFree.push_back(
+            {child, childSearchPathLengthBase + child->partialKeyLen});
       }
       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]);
+        Node *child = n16->children[i];
+        toFree.push_back(
+            {child, childSearchPathLengthBase + child->partialKeyLen});
       }
       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]);
+        Node *child = n48->children[i];
+        toFree.push_back(
+            {child, childSearchPathLengthBase + child->partialKeyLen});
       }
       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]; });
+      n256->bitSet.forEachSet([&](int i) {
+        Node *child = n256->children[i];
+        *out++ = {child, childSearchPathLengthBase + child->partialKeyLen};
+      });
       assert(out == toFree.end());
       writeContext->release(n256);
     } break;
@@ -6018,6 +6049,13 @@ void removeKey(Node *n) {
   }
 }
 
+void removeKey(Node *n, int64_t searchPathLength) {
+  if (n->entryPresent) {
+    --totalKeys;
+    keyBytes -= searchPathLength;
+  }
+}
+
 struct __attribute__((visibility("default"))) PeakPrinter {
   ~PeakPrinter() {
     printf("--- radix_tree ---\n");

Internal.h

@@ -23,7 +23,9 @@ using namespace weaselab;
 #include <callgrind.h>
 
 #define DEBUG_VERBOSE 0
+#ifndef SHOW_MEMORY
 #define SHOW_MEMORY 0
+#endif
 
 // std::span is not trivially constructible. We want a span that leaves its
 // members uninitialized for performance reasons.
@@ -96,6 +98,11 @@ static_assert(std::is_trivial_v<TrivialSpan>);
 #if SHOW_MEMORY
 inline int64_t mallocBytes = 0;
 inline int64_t peakMallocBytes = 0;
+// Defined in ConflictSet.cpp
+extern int64_t nodeBytes;
+extern int64_t partialCapacityBytes;
+extern int64_t totalKeys;
+extern int64_t keyBytes;
 #endif
 
 inline thread_local int64_t mallocBytesDelta = 0;

ShowMemoryTest.cpp

@@ -0,0 +1,89 @@
+// Exercises the SHOW_MEMORY instrumentation, which the other tests don't
+// compile with (this target passes -DSHOW_MEMORY=1). The instrumentation has
+// its own consistency checks, which abort on failure, and asserts are enabled
+// by TEST_FLAGS. Only one ConflictSet may exist at a time, since the
+// accounting is process-wide.
+#include "ConflictSet.h"
+#include "Internal.h"
+
+#include <cassert>
+#include <cstdint>
+#include <string>
+
+using namespace weaselab;
+
+namespace {
+
+void writePoint(ConflictSet &cs, const std::string &key, int64_t writeVersion) {
+  ConflictSet::WriteRange w{{(const uint8_t *)key.data(), int(key.size())},
+                            {nullptr, 0}};
+  cs.addWrites(&w, 1, writeVersion);
+}
+
+ConflictSet::Result checkPoint(const ConflictSet &cs, const std::string &key,
+                               int64_t readVersion) {
+  ConflictSet::ReadRange r{{(const uint8_t *)key.data(), int(key.size())},
+                           {nullptr, 0},
+                           readVersion};
+  ConflictSet::Result result;
+  cs.check(&r, &result, 1);
+  return result;
+}
+
+// Everything a ConflictSet allocated must have been accounted for and released
+// by the time it's destroyed.
+void checkFullyReleased() {
+  assert(nodeBytes == 0);
+  assert(partialCapacityBytes == 0);
+  assert(totalKeys == 0);
+  assert(keyBytes == 0);
+  assert(mallocBytes == 0);
+}
+
+} // namespace
+
+int main() {
+  // Splitting a node replaces it, which defers the release of the old node.
+  // Releasing a deferred node must account for its capacity, and the key it
+  // held must be accounted for without walking `parent` pointers of nodes that
+  // have already been released.
+  {
+    ConflictSet cs{0};
+    writePoint(cs, "a", 1);
+    writePoint(cs, "ab", 2);
+    assert(checkPoint(cs, "a", 0) == ConflictSet::Conflict);
+    assert(checkPoint(cs, "ab", 0) == ConflictSet::Conflict);
+    assert(checkPoint(cs, "b", 0) == ConflictSet::Commit);
+    assert(nodeBytes > 0);
+    assert(keyBytes > 0);
+  }
+  checkFullyReleased();
+
+  // Exceeding the version window erases the tree mid-workload.
+  {
+    ConflictSet cs{0};
+    writePoint(cs, "abc", 1);
+    const int64_t version = int64_t(1) << 40;
+    writePoint(cs, "abc", version);
+    assert(checkPoint(cs, "abc", 0) == ConflictSet::TooOld);
+    assert(checkPoint(cs, "abc", version) == ConflictSet::Commit);
+  }
+  checkFullyReleased();
+
+  // Erase a tree with many nodes, both mid-workload and at destruction.
+  {
+    ConflictSet cs{0};
+    for (int i = 0; i < 1000; ++i) {
+      writePoint(cs, "prefix/" + std::to_string(i), i + 1);
+    }
+    const int64_t version = int64_t(1) << 40;
+    writePoint(cs, "prefix/1", version);
+    assert(checkPoint(cs, "prefix/1", version) == ConflictSet::Commit);
+
+    for (int i = 0; i < 1000; ++i) {
+      writePoint(cs, "other/" + std::to_string(i), version + i + 1);
+    }
+    assert(checkPoint(cs, "other/999", version) == ConflictSet::Conflict);
+  }
+  checkFullyReleased();
+}