Avoid some branches on node type while inserting

ConflictSet.cpp

@@ -902,6 +902,8 @@ Node *&getChildExists(Node *self, uint8_t index) {
 }
 
 InternalVersionT maxVersion(Node *n, ConflictSet::Impl *);
+InternalVersionT exchangeMaxVersion(Node *n, InternalVersionT newMax,
+                                    ConflictSet::Impl *);
 
 void setMaxVersion(Node *n, ConflictSet::Impl *, InternalVersionT maxVersion);
 
@@ -1123,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()) {
@@ -1134,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;
@@ -1141,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;
@@ -1148,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;
@@ -2881,85 +2902,93 @@ checkMaxBetweenExclusiveImpl<true>(Node *n, int begin, int end,
                                    InternalVersionT readVersion, ReadContext *);
 #endif
 
+// 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) {
+  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);
+
+    newSelf->parent = old->parent;
+    newSelf->parentsIndex = old->parentsIndex;
+    newSelf->partialKeyLen = partialKeyIndex;
+    newSelf->entryPresent = false;
+    newSelf->numChildren = 1;
+
+    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;
+
+    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);
+}
+
 // Returns a pointer to the newly inserted node.  Caller must set
-// `entryPresent`, and `entry` fields.  The search path of the result will have
-// `maxVersion` at least `writeVersion` as a postcondition. Nodes along the
-// search path to `key` may be invalidated.
-[[nodiscard]] Node *insert(Node **self, std::span<const uint8_t> key,
-                           InternalVersionT writeVersion, WriteContext *tls,
-                           ConflictSet::Impl *impl) {
+// `entryPresent`, and `entry` fields.  All nodes along the search path of the
+// result will have `maxVersion` set to `writeVersion` as a postcondition. Nodes
+// along the search path may be invalidated.
+[[nodiscard]]
+Node *insert(Node **self, std::span<const uint8_t> key,
+             InternalVersionT writeVersion, WriteContext *tls,
+             ConflictSet::Impl *impl) {
+
+  if ((*self)->partialKeyLen > 0) {
+    consumePartialKey(*self, key, writeVersion, tls, impl);
+  }
+  assert(maxVersion(*self, impl) <= writeVersion);
+  setMaxVersion(*self, impl, writeVersion);
 
   for (;; ++tls->accum.insert_iterations) {
 
-    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 = maxVersion(old, impl);
-
-        // *self will have one child
-        *self = tls->allocate<Node3>(partialKeyIndex);
-
-        memcpy((char *)*self + kNodeCopyBegin, (char *)old + kNodeCopyBegin,
-               kNodeCopySize);
-        (*self)->partialKeyLen = partialKeyIndex;
-
-        // Not necessary to call removeKey here, since this node is "synthetic"
-        (*self)->entryPresent = false;
-
-        (*self)->numChildren = 0;
-        memcpy((*self)->partialKey(), old->partialKey(),
-               (*self)->partialKeyLen);
-
-        getOrCreateChild(*self, old->partialKey()[partialKeyIndex], tls) = old;
-        old->parent = *self;
-        old->parentsIndex = old->partialKey()[partialKeyIndex];
-        setMaxVersion(old, impl, oldMaxVersion);
-
-        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);
-
-    } else {
-      // Consider adding a partial key
-      if ((*self)->numChildren == 0 && !(*self)->entryPresent) {
-        assert((*self)->getCapacity() >= int(key.size()));
-        (*self)->partialKeyLen = key.size();
-        memcpy((*self)->partialKey(), key.data(), (*self)->partialKeyLen);
-        key = key.subspan((*self)->partialKeyLen,
-                          key.size() - (*self)->partialKeyLen);
-      }
-    }
-
-    assert(maxVersion(*self, impl) <= writeVersion);
-    setMaxVersion(*self, impl, writeVersion);
-
     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;
       child->entryPresent = false;
-      child->partialKeyLen = 0;
+      child->partialKeyLen = key.size() - 1;
       child->parent = *self;
       child->parentsIndex = key.front();
-      setMaxVersion(child, impl, tls->zero);
+      setMaxVersion(child, impl, writeVersion);
+      memcpy(child->partialKey(), key.data() + 1, child->partialKeyLen);
+      return child;
     }
 
     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);
+    }
   }
 }
 
@@ -3556,6 +3585,41 @@ InternalVersionT maxVersion(Node *n, ConflictSet::Impl *impl) {
   }
 }
 
+InternalVersionT exchangeMaxVersion(Node *n, InternalVersionT newMax,
+                                    ConflictSet::Impl *impl) {
+  int index = n->parentsIndex;
+  n = n->parent;
+  if (n == nullptr) {
+    return std::exchange(impl->rootMaxVersion, newMax);
+  }
+  switch (n->getType()) {
+  case Type_Node0:           // GCOVR_EXCL_LINE
+    __builtin_unreachable(); // GCOVR_EXCL_LINE
+  case Type_Node3: {
+    auto *n3 = static_cast<Node3 *>(n);
+    int i = getNodeIndex(n3, index);
+    return std::exchange(n3->childMaxVersion[i], newMax);
+  }
+  case Type_Node16: {
+    auto *n16 = static_cast<Node16 *>(n);
+    int i = getNodeIndex(n16, index);
+    return std::exchange(n16->childMaxVersion[i], newMax);
+  }
+  case Type_Node48: {
+    auto *n48 = static_cast<Node48 *>(n);
+    assert(n48->bitSet.test(index));
+    return std::exchange(n48->childMaxVersion[n48->index[index]], newMax);
+  }
+  case Type_Node256: {
+    auto *n256 = static_cast<Node256 *>(n);
+    assert(n256->bitSet.test(index));
+    return std::exchange(n256->childMaxVersion[index], newMax);
+  }
+  default:                   // GCOVR_EXCL_LINE
+    __builtin_unreachable(); // GCOVR_EXCL_LINE
+  }
+}
+
 void setMaxVersion(Node *n, ConflictSet::Impl *impl, InternalVersionT newMax) {
   int index = n->parentsIndex;
   n = n->parent;