From 5e1fb1dac571a08ae355a8253b851df7ebfcf059 Mon Sep 17 00:00:00 2001
From: Andrew Noyes <andrew@weaselab.dev>
Date: Thu, 7 Mar 2024 13:33:39 -0800
Subject: [PATCH] Use entry bytes in partial key if entry not present

Closes #8
---
 ConflictSet.cpp | 37 ++++++++++++++++++++++++-------------
 Internal.h      |  2 +-
 2 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/ConflictSet.cpp b/ConflictSet.cpp
index 0c4d81b..ff2a87b 100644
--- a/ConflictSet.cpp
+++ b/ConflictSet.cpp
@@ -227,23 +227,28 @@ enum class Type : int8_t {
   Node256,
   Invalid,
 };
+
+constexpr static int kPartialKeyMaxLenEntryPresent = 24;
+
 struct Node {
   /* begin section that's copied to the next node */
   Node *parent = nullptr;
-  // The max write version over all keys that start with the search path up to
-  // this point
+  uint8_t partialKey[kPartialKeyMaxLenEntryPresent];
+  // If not entryPresent, then the partial key might spill over into entry
   Entry entry;
-  int16_t numChildren = 0;
+  int32_t numChildren = 0;
   bool entryPresent = false;
   uint8_t parentsIndex = 0;
-  constexpr static auto kPartialKeyMaxLen = 26;
-  uint8_t partialKey[kPartialKeyMaxLen];
   int8_t partialKeyLen = 0;
   /* end section that's copied to the next node */
 
   Type type = Type::Invalid;
 };
 
+static_assert(offsetof(Node, entry) ==
+              offsetof(Node, partialKey) + kPartialKeyMaxLenEntryPresent);
+static_assert(std::is_pod_v<Entry>);
+
 struct Child {
   int64_t childMaxVersion;
   Node *child;
@@ -888,7 +893,7 @@ bytes:
 
 int longestCommonPrefixPartialKey(const uint8_t *ap, const uint8_t *bp,
                                   int cl) {
-  assert(cl <= Node::kPartialKeyMaxLen);
+  assert(cl <= kPartialKeyMaxLenEntryPresent + int(sizeof(Entry)));
   int i = 0;
   for (; i < cl; ++i) {
     if (*ap++ != *bp++) {
@@ -1515,10 +1520,10 @@ bool checkRangeRead(Node *n, std::span<const uint8_t> begin,
   return checkRangeLeftSide.ok & checkRangeRightSide.ok;
 }
 
-// Returns a pointer to the newly inserted node.  caller is reponsible for
-// setting 'entry' fields and `maxVersion` on the result, which may have
-// !entryPresent. The search path of the result's parent will have
-// `maxVersion` at least `writeVersion` as a postcondition.
+// Returns a pointer to the newly inserted node.  Caller must set
+// `entryPresent`, `entry` fields and `maxVersion` on the result.  The search
+// path of the result's parent will have `maxVersion` at least `writeVersion` as
+// a postcondition.
 template <bool kBegin>
 [[nodiscard]] Node *insert(Node **self, std::span<const uint8_t> key,
                            int64_t writeVersion, NodeAllocators *allocators,
@@ -1557,8 +1562,12 @@ template <bool kBegin>
     } else {
       // Consider adding a partial key
       if ((*self)->numChildren == 0 && !(*self)->entryPresent) {
-        (*self)->partialKeyLen =
-            std::min<int>(key.size(), (*self)->kPartialKeyMaxLen);
+        const bool willNotBePresent =
+            key.size() > kPartialKeyMaxLenEntryPresent + int(sizeof(Entry));
+        (*self)->partialKeyLen = std::min<int>(
+            key.size(), willNotBePresent
+                            ? kPartialKeyMaxLenEntryPresent + int(sizeof(Entry))
+                            : kPartialKeyMaxLenEntryPresent);
         memcpy((*self)->partialKey, key.data(), (*self)->partialKeyLen);
         key = key.subspan((*self)->partialKeyLen,
                           key.size() - (*self)->partialKeyLen);
@@ -1713,6 +1722,7 @@ void addWriteRange(Node *&root, int64_t oldestVersion,
   if (insertedEnd) {
     // beginNode may have been invalidated
     beginNode = insert<true>(useAsRoot, begin, writeVersion, allocators, impl);
+    assert(beginNode->entryPresent);
   }
 
   for (beginNode = nextLogical(beginNode); beginNode != endNode;) {
@@ -2280,7 +2290,8 @@ int main(void) {
 #ifdef ENABLE_FUZZ
 extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
   TestDriver<ConflictSet::Impl> driver{data, size};
-  static_assert(driver.kMaxKeyLen > Node::kPartialKeyMaxLen);
+  static_assert(driver.kMaxKeyLen >
+                kPartialKeyMaxLenEntryPresent + sizeof(Entry));
 
   for (;;) {
     bool done = driver.next();
diff --git a/Internal.h b/Internal.h
index 0cedb91..88635c3 100644
--- a/Internal.h
+++ b/Internal.h
@@ -482,7 +482,7 @@ template <class ConflictSetImpl> struct TestDriver {
   ConflictSetImpl cs{oldestVersion};
   ReferenceImpl refImpl{oldestVersion};
 
-  constexpr static auto kMaxKeyLen = 32;
+  constexpr static auto kMaxKeyLen = 64;
 
   bool ok = true;