chore: add ForceUpdate to BPTree (#3993)

ReallocIfNeeded deleted and reinserted the same key in the BPTree stored in SortedMap. This is not really needed since the key is actually the same and we can just update the pointer within BPTree instead of doing the deletion + insertion chore.

---------

Signed-off-by: kostas <kostas@dragonflydb.io>

src/core/bptree_set.h

@@ -104,6 +104,10 @@ template <typename T, typename Policy = BPTreePolicy<T>> class BPTree {
   /// @param path
   void Delete(BPTreePath path);
 
+  /// @brief Forces an update to the key. Assumes key has the same value.
+  /// Replaces old with new_obj.
+  void ForceUpdate(KeyT old, KeyT new_obj);
+
  private:
   BPTreeNode* CreateNode(bool leaf);
 
@@ -589,4 +593,15 @@ template <typename T, typename Policy> void BPTree<T, Policy>::DestroyNode(BPTre
   num_nodes_--;
 }
 
+template <typename T, typename Policy> void BPTree<T, Policy>::ForceUpdate(KeyT old, KeyT new_obj) {
+  BPTreePath path;
+  [[maybe_unused]] bool found = Locate(old, &path);
+
+  assert(path.Depth() > 0u);
+  assert(found);
+
+  BPTreeNode* node = path.Last().first;
+  node->SetKey(path.Last().second, new_obj);
+}
+
 }  // namespace dfly

src/core/bptree_set_test.cc

@@ -470,4 +470,49 @@ void RegisterBPTreeBench() {
 
 REGISTER_MODULE_INITIALIZER(Bptree, RegisterBPTreeBench());
 
+TEST_F(BPTreeSetTest, ForceUpdate) {
+  struct Policy {
+    // Similar to how it's used in SortedMap just a little simpler.
+    using KeyT = int*;
+
+    struct KeyCompareTo {
+      int operator()(KeyT a, KeyT b) const {
+        if (*a < *b)
+          return -1;
+        if (*a > *b)
+          return 1;
+        return 0;
+      }
+    };
+  };
+
+  auto gen_vector = []() {
+    std::vector<std::unique_ptr<int>> tmp;
+    for (size_t i = 0; i < 1000; ++i) {
+      tmp.push_back(std::make_unique<int>(i));
+    }
+    return tmp;
+  };
+
+  std::vector<std::unique_ptr<int>> original = gen_vector();
+  std::vector<std::unique_ptr<int>> modified = gen_vector();
+
+  BPTree<int*, Policy> bptree;
+  for (auto& item : original) {
+    bptree.Insert(item.get());
+  }
+
+  for (auto& item : modified) {
+    bptree.ForceUpdate(item.get(), item.get());
+  }
+
+  original.clear();
+  size_t index = 0;
+  bptree.Iterate(0, 1000, [&](int* ptr) {
+    EXPECT_EQ(modified[index].get(), ptr);
+    ++index;
+    return true;
+  });
+}
+
 }  // namespace dfly

src/core/score_map.cc

@@ -148,7 +148,7 @@ detail::SdsScorePair ScoreMap::iterator::BreakToPair(void* obj) {
 
 namespace {
 // Does not Release obj. Callers must do so explicitly if a `Reallocation` happened
-pair<sds, bool> ReallocIfNeededGeneric(void* obj, float ratio) {
+pair<sds, bool> DuplicateEntryIfFragmented(void* obj, float ratio) {
   sds key = (sds)obj;
   size_t key_len = sdslen(key);
 
@@ -167,15 +167,15 @@ bool ScoreMap::iterator::ReallocIfNeeded(float ratio, std::function<void(sds, sd
   bool reallocated = false;
   auto body = [ratio, &cb, &reallocated](auto* ptr) {
     auto* obj = ptr->GetObject();
-    auto [new_obj, realloc] = ReallocIfNeededGeneric(obj, ratio);
-    if (realloc) {
+    auto [new_obj, duplicate] = DuplicateEntryIfFragmented(obj, ratio);
+    if (duplicate) {
       if (cb) {
         cb((sds)obj, (sds)new_obj);
       }
       sdsfree((sds)obj);
       ptr->SetObject(new_obj);
     }
-    reallocated |= realloc;
+    reallocated |= duplicate;
   };
 
   TraverseApply(curr_entry_, body);

src/core/sorted_map.cc

@@ -770,10 +770,7 @@ SortedMap* SortedMap::FromListPack(PMR_NS::memory_resource* res, const uint8_t*
 }
 
 bool SortedMap::DefragIfNeeded(float ratio) {
-  auto cb = [this](sds old_obj, sds new_obj) {
-    score_tree->Delete(old_obj);
-    score_tree->Insert(new_obj);
-  };
+  auto cb = [this](sds old_obj, sds new_obj) { score_tree->ForceUpdate(old_obj, new_obj); };
   bool reallocated = false;
 
   for (auto it = score_map->begin(); it != score_map->end(); ++it) {