chore: several improvements around sorted map

1. Pass memory_resource to sorted_map.
2. Get rid of GetDict leaky accessor in SortedMap and introduce a proper
   Scan method.
3. Introduce correct BPTree type inside SortedMap::DFImpl.
4. Added a test for bptree_test that covers sds comparison
   (apparently, sdscmp can return values outside of [-1, 1] range).
   Fixed bptree code to support a proper spec for three-way comparison.
5. Expose pointers to internal objects allocated by score_map so we could insert them
   into bptree.

Signed-off-by: Roman Gershman <[email protected]>

Author: romange

src/core/bptree_set.h

@@ -65,7 +65,7 @@ template <typename T, typename Policy = BPTreePolicy<T>> class BPTree {
 
   void Clear();
 
-  BPTreeNode* DEBUG_root() {
+  const BPTreeNode* DEBUG_root() const {
     return root_;
   }
 
@@ -299,10 +299,10 @@ void BPTree<T, Policy>::InsertToFullLeaf(KeyT item, const BPTreePath& path) {
   assert(node->NumItems() < Layout::kMaxLeafKeys);
 
   if (insert_pos <= node->NumItems()) {
-    assert(Comp()(item, median) == -1);
+    assert(Comp()(item, median) < 0);
     node->LeafInsert(insert_pos, item);
   } else {
-    assert(Comp()(item, median) == 1);
+    assert(Comp()(item, median) > 0);
     right->LeafInsert(insert_pos - node->NumItems() - 1, item);
   }
 
@@ -359,12 +359,12 @@ void BPTree<T, Policy>::InsertToFullLeaf(KeyT item, const BPTreePath& path) {
     assert(node->NumItems() < Layout::kMaxInnerKeys);
 
     if (pos <= node->NumItems()) {
-      assert(Comp()(median, next_median) == -1);
+      assert(Comp()(median, next_median) < 0);
 
       node->InnerInsert(pos, median, right);
       node->IncreaseTreeCount(1);
     } else {
-      assert(Comp()(median, next_median) == 1);
+      assert(Comp()(median, next_median) > 0);
 
       next_right->InnerInsert(pos - node->NumItems() - 1, median, right);
 
@@ -566,7 +566,7 @@ template <typename T, typename Policy> void BPTree<T, Policy>::Delete(BPTreePath
     path.DigRight();
 
     BPTreeNode* leaf = path.Last().first;
-    assert(Comp()(leaf->Key(leaf->NumItems() - 1), node->Key(key_pos)) == -1);
+    assert(Comp()(leaf->Key(leaf->NumItems() - 1), node->Key(key_pos)) < 0);
 
     // set a new separator.
     node->SetKey(key_pos, leaf->Key(leaf->NumItems() - 1));

src/core/bptree_set_test.cc

@@ -22,35 +22,14 @@ using namespace std;
 
 namespace dfly {
 
-class BPTreeSetTest : public ::testing::Test {
-  using Node = detail::BPTreeNode<uint64_t>;
-
- protected:
-  static constexpr size_t kNumElems = 7000;
-
-  BPTreeSetTest() : mi_alloc_(mi_heap_get_backing()), bptree_(&mi_alloc_) {
-  }
-  static void SetUpTestSuite() {
-  }
+namespace {
 
-  void FillTree(unsigned factor = 1) {
-    for (unsigned i = 0; i < kNumElems; ++i) {
-      bptree_.Insert(i * factor);
-    }
-  }
-
-  bool Validate();
-
-  static bool Validate(const Node* node, uint64_t ubound);
-
-  MiMemoryResource mi_alloc_;
-  BPTree<uint64_t> bptree_;
-  mt19937 generator_{1};
-};
+template <typename Node, typename Policy>
+bool ValidateNode(const Node* node, typename Node::KeyT ubound) {
+  typename Policy::KeyCompareTo cmp;
 
-bool BPTreeSetTest::Validate(const Node* node, uint64_t ubound) {
   for (unsigned i = 1; i < node->NumItems(); ++i) {
-    if (node->Key(i - 1) >= node->Key(i))
+    if (cmp(node->Key(i - 1), node->Key(i)) > -1)
       return false;
   }
 
@@ -71,25 +50,72 @@ bool BPTreeSetTest::Validate(const Node* node, uint64_t ubound) {
     }
   }
 
-  return node->Key(node->NumItems() - 1) < ubound;
+  return cmp(node->Key(node->NumItems() - 1), ubound) == -1;
 }
 
+struct ZsetPolicy {
+  struct KeyT {
+    double d;
+    sds s;
+  };
+
+  struct KeyCompareTo {
+    int operator()(const KeyT& left, const KeyT& right) {
+      if (left.d < right.d)
+        return -1;
+      if (left.d > right.d)
+        return 1;
+
+      // Note that sdscmp can return values outside of [-1, 1] range.
+      return sdscmp(left.s, right.s);
+    }
+  };
+};
+
+using SDSTree = BPTree<ZsetPolicy::KeyT, ZsetPolicy>;
+
+}  // namespace
+
+class BPTreeSetTest : public ::testing::Test {
+  using Node = detail::BPTreeNode<uint64_t>;
+
+ protected:
+  static constexpr size_t kNumElems = 7000;
+
+  BPTreeSetTest() : mi_alloc_(mi_heap_get_backing()), bptree_(&mi_alloc_) {
+  }
+  static void SetUpTestSuite() {
+  }
+
+  void FillTree(unsigned factor = 1) {
+    for (unsigned i = 0; i < kNumElems; ++i) {
+      bptree_.Insert(i * factor);
+    }
+  }
+
+  bool Validate();
+
+  MiMemoryResource mi_alloc_;
+  BPTree<uint64_t> bptree_;
+  mt19937 generator_{1};
+};
+
 bool BPTreeSetTest::Validate() {
   auto* root = bptree_.DEBUG_root();
   if (!root)
     return true;
 
   // node, upper bound
-  std::vector<pair<Node*, uint64_t>> stack;
+  vector<pair<const Node*, uint64_t>> stack;
 
   stack.emplace_back(root, UINT64_MAX);
 
   while (!stack.empty()) {
-    Node* node = stack.back().first;
+    const Node* node = stack.back().first;
     uint64_t ubound = stack.back().second;
     stack.pop_back();
 
-    if (!Validate(node, ubound))
+    if (!ValidateNode<Node, BPTreePolicy<uint64_t>>(node, ubound))
       return false;
 
     if (!node->IsLeaf()) {
@@ -353,25 +379,24 @@ TEST_F(BPTreeSetTest, MemoryUsage) {
   LOG(INFO) << "btree after: " << mi_alloc.used() << " bytes";
 }
 
-struct ZsetPolicy {
-  struct KeyT {
-    double d;
-    sds s;
-  };
+TEST_F(BPTreeSetTest, InsertSDS) {
+  vector<ZsetPolicy::KeyT> vals;
+  for (unsigned i = 0; i < 256; ++i) {
+    sds s = sdsempty();
 
-  struct KeyCompareTo {
-    int operator()(const KeyT& left, const KeyT& right) {
-      if (left.d < right.d)
-        return -1;
-      if (left.d > right.d)
-        return 1;
+    s = sdscatfmt(s, "a%u", i);
+    vals.emplace_back(ZsetPolicy::KeyT{.d = 1000, .s = s});
+  }
 
-      return sdscmp(left.s, right.s);
-    }
-  };
-};
+  SDSTree tree(&mi_alloc_);
+  for (size_t i = 0; i < vals.size(); ++i) {
+    ASSERT_TRUE(tree.Insert(vals[i]));
+  }
 
-using SDSTree = BPTree<ZsetPolicy::KeyT, ZsetPolicy>;
+  for (auto v : vals) {
+    sdsfree(v.s);
+  }
+}
 
 static string RandomString(mt19937& rand, unsigned len) {
   const string_view alpanum = "1234567890abcdefghijklmnopqrstuvwxyz";

src/core/compact_object.cc

@@ -430,7 +430,7 @@ int RobjWrapper::ZsetAdd(double score, sds ele, int in_flags, int* out_flags, do
        * becomes too long *before* executing zzlInsert. */
       if (zl_len + 1 > server.zset_max_listpack_entries ||
           sdslen(ele) > server.zset_max_listpack_value || !lpSafeToAdd(lp, sdslen(ele))) {
-        unique_ptr<SortedMap> ss = SortedMap::FromListPack(lp);
+        unique_ptr<SortedMap> ss = SortedMap::FromListPack(tl.local_mr, lp);
         lpFree(lp);
         inner_obj_ = ss.release();
         encoding_ = OBJ_ENCODING_SKIPLIST;

src/core/score_map.cc

@@ -31,13 +31,7 @@ inline double GetValue(sds key) {
   return u.d;
 }
 
-}  // namespace
-
-ScoreMap::~ScoreMap() {
-  Clear();
-}
-
-bool ScoreMap::AddOrUpdate(string_view field, double value) {
+void* AllocateScored(string_view field, double value) {
   size_t meta_offset = field.size() + 1;
   DoubleUnion u;
   u.d = value;
@@ -52,21 +46,33 @@ bool ScoreMap::AddOrUpdate(string_view field, double value) {
 
   absl::little_endian::Store64(newkey + meta_offset, u.u);
 
+  return newkey;
+}
+
+}  // namespace
+
+ScoreMap::~ScoreMap() {
+  Clear();
+}
+
+pair<void*, bool> ScoreMap::AddOrUpdate(string_view field, double value) {
+  void* newkey = AllocateScored(field, value);
+
   // Replace the whole entry.
   sds prev_entry = (sds)AddOrReplaceObj(newkey, false);
   if (prev_entry) {
     ObjDelete(prev_entry, false);
-    return false;
+    return {newkey, false};
   }
 
-  return true;
+  return {newkey, true};
 }
 
-bool ScoreMap::AddOrSkip(std::string_view field, double value) {
+std::pair<void*, bool> ScoreMap::AddOrSkip(std::string_view field, double value) {
   void* obj = FindInternal(&field, 1);  // 1 - string_view
 
   if (obj)
-    return false;
+    return {obj, false};
 
   return AddOrUpdate(field, value);
 }

src/core/score_map.h

@@ -79,11 +79,11 @@ class ScoreMap : public DenseSet {
 
   // Returns true if field was added
   // otherwise updates its value and returns false.
-  bool AddOrUpdate(std::string_view field, double value);
+  std::pair<void*, bool> AddOrUpdate(std::string_view field, double value);
 
   // Returns true if field was added
   // false, if already exists. In that case no update is done.
-  bool AddOrSkip(std::string_view field, double value);
+  std::pair<void*, bool> AddOrSkip(std::string_view field, double value);
 
   bool Erase(std::string_view s1);
 
@@ -92,6 +92,11 @@ class ScoreMap : public DenseSet {
   /// @return sds
   std::optional<double> Find(std::string_view key);
 
+  // returns the internal object if found, otherwise nullptr.
+  void* FindObj(sds ele) {
+    return FindInternal(ele, 0);
+  }
+
   void Clear();
 
   iterator begin() {

src/core/score_map_test.cc

@@ -54,7 +54,7 @@ class ScoreMapTest : public ::testing::Test {
 };
 
 TEST_F(ScoreMapTest, Basic) {
-  EXPECT_TRUE(sm_->AddOrUpdate("foo", 5));
+  EXPECT_TRUE(sm_->AddOrUpdate("foo", 5).second);
   EXPECT_EQ(5, sm_->Find("foo"));
 
   auto it = sm_->begin();
@@ -70,13 +70,13 @@ TEST_F(ScoreMapTest, Basic) {
   }
 
   size_t sz = sm_->ObjMallocUsed();
-  EXPECT_FALSE(sm_->AddOrUpdate("foo", 17));
+  EXPECT_FALSE(sm_->AddOrUpdate("foo", 17).second);
   EXPECT_EQ(sm_->ObjMallocUsed(), sz);
 
   it = sm_->begin();
   EXPECT_EQ(17, it->second);
 
-  EXPECT_FALSE(sm_->AddOrSkip("foo", 31));
+  EXPECT_FALSE(sm_->AddOrSkip("foo", 31).second);
   EXPECT_EQ(17, it->second);
 }