Avoid some branches on node type while inserting

ConflictSet.cpp

@@ -1125,8 +1125,10 @@ Node *getFirstChildExists(Node *self) {
 }
 
 // Caller is responsible for assigning a non-null pointer to the returned
-// reference if null
-Node *&getOrCreateChild(Node *&self, uint8_t index, WriteContext *tls) {
+// reference if null. Updates child's max version to `newMaxVersion` if child
+// exists but does not have a partial key.
+Node *&getOrCreateChild(Node *&self, uint8_t index,
+                        InternalVersionT newMaxVersion, WriteContext *tls) {
 
   // Fast path for if it exists already
   switch (self->getType()) {
@@ -1136,6 +1138,9 @@ Node *&getOrCreateChild(Node *&self, uint8_t index, WriteContext *tls) {
     auto *self3 = static_cast<Node3 *>(self);
     int i = getNodeIndex(self3, index);
     if (i >= 0) {
+      if (self3->children[i]->partialKeyLen == 0) {
+        self3->childMaxVersion[i] = newMaxVersion;
+      }
       return self3->children[i];
     }
   } break;
@@ -1143,6 +1148,9 @@ Node *&getOrCreateChild(Node *&self, uint8_t index, WriteContext *tls) {
     auto *self16 = static_cast<Node16 *>(self);
     int i = getNodeIndex(self16, index);
     if (i >= 0) {
+      if (self16->children[i]->partialKeyLen == 0) {
+        self16->childMaxVersion[i] = newMaxVersion;
+      }
       return self16->children[i];
     }
   } break;
@@ -1150,12 +1158,23 @@ Node *&getOrCreateChild(Node *&self, uint8_t index, WriteContext *tls) {
     auto *self48 = static_cast<Node48 *>(self);
     int secondIndex = self48->index[index];
     if (secondIndex >= 0) {
+      if (self48->children[secondIndex]->partialKeyLen == 0) {
+        self48->childMaxVersion[secondIndex] = newMaxVersion;
+        self48->maxOfMax[secondIndex >> Node48::kMaxOfMaxShift] =
+            std::max(self48->maxOfMax[secondIndex >> Node48::kMaxOfMaxShift],
+                     newMaxVersion);
+      }
       return self48->children[secondIndex];
     }
   } break;
   case Type_Node256: {
     auto *self256 = static_cast<Node256 *>(self);
     if (auto &result = self256->children[index]; result != nullptr) {
+      if (self256->children[index]->partialKeyLen == 0) {
+        self256->childMaxVersion[index] = newMaxVersion;
+        self256->maxOfMax[index >> Node256::kMaxOfMaxShift] = std::max(
+            self256->maxOfMax[index >> Node256::kMaxOfMaxShift], newMaxVersion);
+      }
       return result;
     }
   } break;
@@ -2883,51 +2902,49 @@ checkMaxBetweenExclusiveImpl<true>(Node *n, int begin, int end,
                                    InternalVersionT readVersion, ReadContext *);
 #endif
 
-// Consume the partial key of `self` if it exists, and update `self` and `key`
-// such that `self` is along the search path of `key`, and has maxVersion
-// `writeVersion`
+// Consume the partial key of `self` (which must exist), and update `self` and
+// `key` such that `self` is along the search path of `key`
 void consumePartialKey(Node *&self, std::span<const uint8_t> &key,
                        InternalVersionT writeVersion, WriteContext *tls,
                        ConflictSet::Impl *impl) {
-  if (self->partialKeyLen > 0) {
-    // Handle an existing partial key
-    int commonLen = std::min<int>(self->partialKeyLen, key.size());
-    int partialKeyIndex =
-        longestCommonPrefix(self->partialKey(), key.data(), commonLen);
-    if (partialKeyIndex < self->partialKeyLen) {
-      auto *old = self;
-      InternalVersionT oldMaxVersion =
-          exchangeMaxVersion(old, writeVersion, impl);
+  assert(self->partialKeyLen > 0);
+  // Handle an existing partial key
+  int commonLen = std::min<int>(self->partialKeyLen, key.size());
+  int partialKeyIndex =
+      longestCommonPrefix(self->partialKey(), key.data(), commonLen);
+  if (partialKeyIndex < self->partialKeyLen) {
+    auto *old = self;
+    InternalVersionT oldMaxVersion =
+        exchangeMaxVersion(old, writeVersion, impl);
 
-      // *self will have one child (old)
-      auto *newSelf = tls->allocate<Node3>(partialKeyIndex);
+    // *self will have one child (old)
+    auto *newSelf = tls->allocate<Node3>(partialKeyIndex);
 
-      newSelf->parent = old->parent;
-      newSelf->parentsIndex = old->parentsIndex;
-      newSelf->partialKeyLen = partialKeyIndex;
-      newSelf->entryPresent = false;
-      newSelf->numChildren = 1;
+    newSelf->parent = old->parent;
+    newSelf->parentsIndex = old->parentsIndex;
+    newSelf->partialKeyLen = partialKeyIndex;
+    newSelf->entryPresent = false;
+    newSelf->numChildren = 1;
 
-      memcpy(newSelf->partialKey(), old->partialKey(), newSelf->partialKeyLen);
+    memcpy(newSelf->partialKey(), old->partialKey(), newSelf->partialKeyLen);
 
-      uint8_t oldDistinguishingByte = old->partialKey()[partialKeyIndex];
-      old->parent = newSelf;
-      old->parentsIndex = oldDistinguishingByte;
-      newSelf->index[0] = oldDistinguishingByte;
-      newSelf->children[0] = old;
-      newSelf->childMaxVersion[0] = oldMaxVersion;
-      self = newSelf;
+    uint8_t oldDistinguishingByte = old->partialKey()[partialKeyIndex];
+    old->parent = newSelf;
+    old->parentsIndex = oldDistinguishingByte;
+    newSelf->index[0] = oldDistinguishingByte;
+    newSelf->children[0] = old;
+    newSelf->childMaxVersion[0] = oldMaxVersion;
+    self = newSelf;
 
-      memmove(old->partialKey(), old->partialKey() + partialKeyIndex + 1,
-              old->partialKeyLen - (partialKeyIndex + 1));
-      old->partialKeyLen -= partialKeyIndex + 1;
+    memmove(old->partialKey(), old->partialKey() + partialKeyIndex + 1,
+            old->partialKeyLen - (partialKeyIndex + 1));
+    old->partialKeyLen -= partialKeyIndex + 1;
 
-      // We would consider decreasing capacity here, but we can't invalidate
-      // old since it's not on the search path. setOldestVersion will clean it
-      // up.
-    }
-    key = key.subspan(partialKeyIndex, key.size() - partialKeyIndex);
+    // We would consider decreasing capacity here, but we can't invalidate
+    // old since it's not on the search path. setOldestVersion will clean it
+    // up.
   }
+  key = key.subspan(partialKeyIndex, key.size() - partialKeyIndex);
 }
 
 // Returns a pointer to the newly inserted node.  Caller must set
@@ -2939,18 +2956,19 @@ Node *insert(Node **self, std::span<const uint8_t> key,
              InternalVersionT writeVersion, WriteContext *tls,
              ConflictSet::Impl *impl) {
 
-  for (;; ++tls->accum.insert_iterations) {
-
+  if ((*self)->partialKeyLen > 0) {
     consumePartialKey(*self, key, writeVersion, tls, impl);
+  }
+  assert(maxVersion(*self, impl) <= writeVersion);
+  setMaxVersion(*self, impl, writeVersion);
 
-    assert(maxVersion(*self, impl) <= writeVersion);
-    setMaxVersion(*self, impl, writeVersion);
+  for (;; ++tls->accum.insert_iterations) {
 
     if (key.size() == 0) {
       return *self;
     }
 
-    auto &child = getOrCreateChild(*self, key.front(), tls);
+    auto &child = getOrCreateChild(*self, key.front(), writeVersion, tls);
     if (!child) {
       child = tls->allocate<Node0>(key.size() - 1);
       child->numChildren = 0;
@@ -2965,6 +2983,12 @@ Node *insert(Node **self, std::span<const uint8_t> key,
 
     self = &child;
     key = key.subspan(1, key.size() - 1);
+
+    if ((*self)->partialKeyLen > 0) {
+      consumePartialKey(*self, key, writeVersion, tls, impl);
+      assert(maxVersion(*self, impl) <= writeVersion);
+      setMaxVersion(*self, impl, writeVersion);
+    }
   }
 }