chore(bptree): introduce ability comparing keys to non-key types (#4805)

Not yet used by sorted_map code, so no function changes.
This will allow to clean up a hackish comparator code in sorted_map.cc

Co-authored-by: mkaruza <[email protected]>

Author: romange

src/core/bptree_set.h

@@ -21,6 +21,9 @@ template <typename T> struct DefaultCompareTo {
 
 template <typename T> struct BPTreePolicy {
   using KeyT = T;
+
+  // The three way comparator that should accept a query ( or key) on the left, and the key
+  // on the right.
   using KeyCompareTo = DefaultCompareTo<T>;
 };
 
@@ -89,16 +92,15 @@ template <typename T, typename Policy = BPTreePolicy<T>> class BPTree {
   ///             Should return false to stop iteration.
   bool IterateReverse(uint32_t rank_start, uint32_t rank_end, std::function<bool(KeyT)> cb) const;
 
-  /// @brief Returns the path to the first item in the tree that is greater or equal to key.
+  /// @brief Returns the path to the first item in the tree for which comp(q, key) >= 0.
   /// @param item
   /// @return the path if such item exists, empty path otherwise.
-  /// @todo: to wrap the result into iterator to avoid the leakage of internal data structures.
-  detail::BPTreePath<T> GEQ(KeyT key) const;
+  template <typename Q> BPTreePath GEQ(Q&& query) const;
 
-  /// @brief Returns the path to the largest item in the tree that is less or equal to key.
+  /// @brief Returns the path to the largest item in the tree such that comp(q, key) <= 0.
   /// @param key
   /// @return the path if such item exists, empty path otherwise.
-  detail::BPTreePath<T> LEQ(KeyT key) const;
+  template <typename Q> BPTreePath LEQ(Q&& query) const;
 
   /// @brief Deletes the element pointed by path.
   /// @param path
@@ -122,10 +124,10 @@ template <typename T, typename Policy = BPTreePolicy<T>> class BPTree {
   void IncreaseSubtreeCounts(const BPTreePath& path, unsigned depth, int32_t delta);
 
   // Charts the path towards key. Returns true if key is found.
-  // In that case path->Last().first->Key(path->Last().second) == key.
+  // In that case comp(q, path->Last().first->Key(path->Last().second)) == 0.
   // Fills the tree path not including the key itself. In case key was not found,
   // returns the path to the item that is greater than the key.
-  bool Locate(KeyT key, BPTreePath* path) const;
+  template <typename Q> bool Locate(Q&& q, BPTreePath* path) const;
 
   // Sets the tree path to item at specified rank. Rank is 0-based and must be less than Size().
   // returns the index of the key in the last node of the path.
@@ -243,12 +245,15 @@ std::optional<uint32_t> BPTree<T, Policy>::GetRank(KeyT item, bool reverse) cons
 }
 
 template <typename T, typename Policy>
-bool BPTree<T, Policy>::Locate(KeyT key, BPTreePath* path) const {
+template <typename Q>
+bool BPTree<T, Policy>::Locate(Q&& q, BPTreePath* path) const {
   assert(root_);
   BPTreeNode* node = root_;
   typename Policy::KeyCompareTo cmp;
+  auto cmp_cb = [&](const KeyT& key) { return cmp(q, key); };
+
   while (true) {
-    typename BPTreeNode::SearchResult res = node->BSearch(key, cmp);
+    typename BPTreeNode::SearchResult res = node->BSearch(cmp_cb);
     path->Push(node, res.index);
     if (res.found) {
       return true;
@@ -486,10 +491,11 @@ void BPTree<T, Policy>::ToRank(uint32_t rank, BPTreePath* path) const {
 }
 
 template <typename T, typename Policy>
-detail::BPTreePath<T> BPTree<T, Policy>::GEQ(KeyT item) const {
+template <typename Q>
+auto BPTree<T, Policy>::GEQ(Q&& query) const -> BPTreePath {
   BPTreePath path;
 
-  bool res = Locate(item, &path);
+  bool res = Locate(query, &path);
 
   // if we did not find the item and the path does not lead to any key in the node,
   // adjust the path to point to the next key in the tree.
@@ -502,9 +508,10 @@ detail::BPTreePath<T> BPTree<T, Policy>::GEQ(KeyT item) const {
 }
 
 template <typename T, typename Policy>
-detail::BPTreePath<T> BPTree<T, Policy>::LEQ(KeyT item) const {
+template <typename Q>
+auto BPTree<T, Policy>::LEQ(Q&& query) const -> BPTreePath {
   BPTreePath path;
-  bool res = Locate(item, &path);
+  bool res = Locate(query, &path);
 
   if (!res) {  // fix the result in case the path leads to key greater than item.
     path.Prev();

src/core/detail/bptree_internal.h

@@ -143,8 +143,9 @@ template <typename T> class BPTreeNode {
   };
 
   // Searches for key in the node using binary search.
-  // Returns SearchResult with index of the key if found.
-  template <typename Comp> SearchResult BSearch(KeyT key, Comp&& comp) const;
+  // Returns SearchResult with index of the smallest key for which comp(key) >=0.
+  // comp: is a three way comparator.
+  template <typename Comp> SearchResult BSearch(Comp&& comp) const;
 
   void Split(BPTreeNode* right, KeyT* median);
 
@@ -367,13 +368,13 @@ template <typename T> class BPTreePath {
 // if all items are smaller than key, returns num_items_.
 template <typename T>
 template <typename Comp>
-auto BPTreeNode<T>::BSearch(KeyT key, Comp&& cmp_op) const -> SearchResult {
+auto BPTreeNode<T>::BSearch(Comp&& cmp_op) const -> SearchResult {
   uint16_t lo = 0;
   uint16_t hi = num_items_;
   assert(hi > 0);
 
   // optimization: check the last item first.
-  int cmp_res = cmp_op(key, Key(hi - 1));
+  int cmp_res = cmp_op(Key(hi - 1));
   if (cmp_res >= 0) {
     return cmp_res > 0 ? SearchResult{.index = hi, .found = false}
                        : SearchResult{.index = uint16_t(hi - 1), .found = true};
@@ -388,7 +389,7 @@ auto BPTreeNode<T>::BSearch(KeyT key, Comp&& cmp_op) const -> SearchResult {
 
     KeyT item = Key(mid);
 
-    int cmp_res = cmp_op(key, item);
+    int cmp_res = cmp_op(item);
     if (cmp_res == 0) {
       return SearchResult{.index = mid, .found = true};
     }

src/core/sorted_map.h

@@ -43,14 +43,6 @@ class SortedMap {
   SortedMap(const SortedMap&) = delete;
   SortedMap& operator=(const SortedMap&) = delete;
 
-  struct ScoreSdsPolicy {
-    using KeyT = ScoreSds;
-
-    struct KeyCompareTo {
-      int operator()(KeyT a, KeyT b) const;
-    };
-  };
-
   bool Reserve(size_t sz);
   int AddElem(double score, std::string_view ele, int in_flags, int* out_flags, double* newscore);
 
@@ -97,6 +89,14 @@ class SortedMap {
   bool DefragIfNeeded(float ratio);
 
  private:
+  struct ScoreSdsPolicy {
+    using KeyT = ScoreSds;
+
+    struct KeyCompareTo {
+      int operator()(KeyT a, KeyT b) const;
+    };
+  };
+
   using ScoreTree = BPTree<ScoreSds, ScoreSdsPolicy>;
 
   // hash map from fields to scores.