Optimize PrefixMap find by walking through small prefixes

Summary:
When the bucket's SmallPrefix differs from the query's (or query < 8 bytes), skip the binary search entirely and walk backwards from roughTo using pre-built SmallPrefix values and precomputed masks. This avoids touching fullPrefixes_ string data entirely — we only do uint64 comparisons.

* Add SmallPrefix::starts_with_mask() that uses a precomputed mask for the fast-path walk, avoiding runtime bit-counting.
* Add bucketInfo_ vector of (bucketIdx, prefixMask) pairs: the bucket index locates the pre-built SmallPrefix in smallPrefixes_, and the mask has 1-bits for the significant prefix bytes (derived from min(prefix.size(), 8 ) at construction). Both are read in the same cache line during the walk.
*  Populate bucketInfo_ in a separate loop after smallPrefixes_ is fully built, so we don't depend on iterator stability during insertions.
* Template findPrefix into findPrefixImpl<enableSmallPrefixLookup> so the constructor can use the slow path (findPrefixImpl<false>) while bucketInfo_ is not yet populated.

Benchmark command:
```
buck run fbcode/mode/opt fbcode//mcrouter/lib/fbi/cpp/test/facebook:string_prefix_map_benchmark -- --bm_regex _Lb --bm_mode=adaptive --bm_min_secs=5 --bm_max_secs=30
```

# x86

Before:
```
============================================================================
[...]facebook/StringPrefixMapBenchmark.cpp     relative  time/iter   iters/s
============================================================================
PrefixMapInCache_Lb                                        68.87ns    14.52M
PrefixMapNotInCache_Lb                                    305.95ns     3.27M
```

After:
```
============================================================================
[...]facebook/StringPrefixMapBenchmark.cpp     relative  time/iter   iters/s
============================================================================
PrefixMapInCache_Lb                                        54.56ns    18.33M
PrefixMapNotInCache_Lb                                    233.65ns     4.28M
```

- PrefixMapInCache_Lb:    68.87ns → 54.56ns (21% less time/iter)
- PrefixMapNotInCache_Lb:   305.95ns → 233.65ns (24% less time/iter)

# ARM

Before
```
============================================================================
[...]facebook/StringPrefixMapBenchmark.cpp     relative  time/iter   iters/s
============================================================================
PrefixMapInCache_Lb                                        93.37ns    10.71M
PrefixMapNotInCache_Lb                                    458.93ns     2.18M
```

After
```
============================================================================
[...]facebook/StringPrefixMapBenchmark.cpp     relative  time/iter   iters/s
============================================================================
PrefixMapInCache_Lb                                        76.96ns    12.99M
PrefixMapNotInCache_Lb                                    376.88ns     2.65M
```

- PrefixMapInCache_Lb:    93.37ns → 76.96ns  (18% less time/iter)
- PrefixMapNotInCache_Lb:   458.93ns → 376.88ns (18% less time/iter)

Reviewed By: DenisYaroshevskiy

Differential Revision: D97313182

fbshipit-source-id: ba0e9ddd59e4eb9a8df367860260e30cb06d5d5a

Author: Alfredo Altamirano

mcrouter/lib/fbi/cpp/LowerBoundPrefixMap.cpp

@@ -7,9 +7,12 @@
 
 #include "mcrouter/lib/fbi/cpp/LowerBoundPrefixMap.h"
 
+#include <cassert>
 #include <cstring>
 #include <ostream>
 
+#include <folly/lang/Bits.h>
+
 namespace facebook::memcache::detail {
 namespace {
 
@@ -46,6 +49,14 @@ I branchlessUpperBound(I f, I l, const T& x, Compare comp) {
 
 } // namespace
 
+FOLLY_ALWAYS_INLINE
+bool checkSmallPrefixStartsWith(
+    SmallPrefix full,
+    SmallPrefix start,
+    std::uint64_t lengthMask) {
+  return (full.data_ & lengthMask) == start.data_;
+}
+
 std::ostream& operator<<(std::ostream& os, const SmallPrefix& self) {
   std::uint64_t correctOrder = folly::Endian::swap(self.data_);
   std::string_view s(reinterpret_cast<const char*>(correctOrder), 8u);
@@ -57,6 +68,7 @@ LowerBoundPrefixMapCommon::LowerBoundPrefixMapCommon(
     : fullPrefixes_(std::move(sortedUniquePrefixes)) {
   smallPrefixes_.reserve(fullPrefixes_.size() + 1);
   previousPrefix_.reserve(fullPrefixes_.size());
+  smallPrefixExtraInfo_.reserve(fullPrefixes_.size());
 
   // Adding an empty string with no matches.
   // This acts as a sentinel so we are always guranteed to find an
@@ -68,11 +80,18 @@ LowerBoundPrefixMapCommon::LowerBoundPrefixMapCommon(
 
   for (std::size_t i = 0; i != fullPrefixes_.size(); ++i) {
     const auto& prefix = fullPrefixes_[i];
+    // Null bytes are not expected. If present it would break the
+    // findPrefix fast SmallPrefix path (e.g. SmallPrefix "ab\0", query "ab").
+    assert(prefix.find('\0') == std::string_view::npos);
 
     // Prefixes always come before lexicographically, so if there is one
     // we will find it.
     previousPrefix_.push_back(findPrefix(prefix));
 
+    auto len = std::min<size_t>(prefix.size(), (size_t)8);
+    auto mask = folly::n_most_significant_bits<std::uint64_t>(len * 8);
+    smallPrefixExtraInfo_.push_back({SmallPrefix{prefix}, mask});
+
     // Add small prefix ---
     // if it is there already - update existing, otherwise insert [i, i+1]
     //
@@ -90,26 +109,38 @@ LowerBoundPrefixMapCommon::LowerBoundPrefixMapCommon(
 
 std::uint32_t LowerBoundPrefixMapCommon::findPrefix(
     std::string_view query) const noexcept {
+  SmallPrefix qSmall{query};
 #ifdef __aarch64__
   auto afterPrefix = branchlessUpperBound(
       smallPrefixes_.begin(),
       smallPrefixes_.end(),
-      SmallPrefix{query},
+      qSmall,
       [](const SmallPrefix& v, const auto& elem) { return v < elem.first; });
 #else
   // Due to a sentinel - guaranteed to not be .begin()
-  auto afterPrefix = smallPrefixes_.upper_bound(SmallPrefix{query});
+  auto afterPrefix = smallPrefixes_.upper_bound(qSmall);
 #endif
-  auto [roughFrom, roughTo] = std::prev(afterPrefix)->second;
-
-  // Binary search complete strings between rough boundaries.
-  // NOTE: which array we search - doesn't matter -
-  //       we just want indexes.
-  auto cur = std::upper_bound(
-                 fullPrefixes_.begin() + roughFrom,
-                 fullPrefixes_.begin() + roughTo,
-                 query) -
-      fullPrefixes_.begin();
+  const auto bucketIt = std::prev(afterPrefix);
+  auto [roughFrom, roughTo] = bucketIt->second;
+
+  if (bucketIt->first < qSmall || query.size() < 8) {
+    std::uint32_t cur = roughTo;
+    while (cur != 0 &&
+           !checkSmallPrefixStartsWith(
+               qSmall,
+               smallPrefixExtraInfo_[cur - 1].smallPrefix,
+               smallPrefixExtraInfo_[cur - 1].lengthMask)) {
+      cur = previousPrefix_[cur - 1];
+    }
+    return cur;
+  }
+
+  std::uint32_t cur = static_cast<std::uint32_t>(
+      std::upper_bound(
+          fullPrefixes_.begin() + roughFrom,
+          fullPrefixes_.begin() + roughTo,
+          query) -
+      fullPrefixes_.begin());
 
   while (cur != 0 &&
          !std_string_view_starts_with(query, fullPrefixes_[cur - 1])) {

mcrouter/lib/fbi/cpp/LowerBoundPrefixMap.h

@@ -81,11 +81,20 @@ class SmallPrefix {
   }
 
   friend std::ostream& operator<<(std::ostream& os, const SmallPrefix& self);
+  friend bool checkSmallPrefixStartsWith(
+      SmallPrefix full,
+      SmallPrefix start,
+      std::uint64_t lengthMask);
 
  private:
   std::uint64_t data_ = 0;
 };
 
+struct SmallPrefixExtraInfo {
+  SmallPrefix smallPrefix;
+  std::uint64_t lengthMask = 0;
+};
+
 // What is the idea
 //
 // If a string (query) has a prefix in a sorted lexicographically array
@@ -125,6 +134,14 @@ struct LowerBoundPrefixMapCommon {
   //     ^________|
   // This is the index of that prefix, base 1 (0 means absence).
   std::vector<std::uint32_t> previousPrefix_;
+
+  // Sometimes we get to check SmallPrefixes for starts with instead of full
+  // prefixes. This is faster but we need some extra information - which we
+  // store for each full prefix.
+  // NOTE: in theory this is suboptimal, because we iterate through a list of
+  // full prefixes and there are less SmallPrefixes but our experiments with
+  // that so far were unfruitful.
+  std::vector<SmallPrefixExtraInfo> smallPrefixExtraInfo_;
 };
 
 template <typename Storage>

mcrouter/lib/fbi/cpp/test/LowerBoundPrefixMapTest.cpp

@@ -210,5 +210,53 @@ TEST(LowerBoundPrefixMapTest, NullKey) {
   ASSERT_EQ(lbMap.findPrefix(std::string_view{}), lbMap.end());
 }
 
+TEST(LowerBoundPrefixMapTest, LongPrefixShortQuery) {
+  // Entries sharing the same 8-byte prefix but with different lengths.
+  // Queries of exactly 8 bytes must not falsely match longer entries.
+  std::vector<KeyValue> kv = {
+      KeyValue{"abcdefgh", 1},
+      KeyValue{"abcdefghijk", 2},
+  };
+  LowerBoundPrefixMap<int> lbMap{kv};
+
+  // "abcdefgh" matches exactly
+  ASSERT_NE(lbMap.end(), lbMap.findPrefix("abcdefgh"));
+  ASSERT_EQ("abcdefgh", lbMap.findPrefix("abcdefgh")->key());
+  ASSERT_EQ(1, lbMap.findPrefix("abcdefgh")->value());
+
+  // "abcdefghij" — "abcdefgh" is a prefix, not "abcdefghijk"
+  ASSERT_NE(lbMap.end(), lbMap.findPrefix("abcdefghij"));
+  ASSERT_EQ("abcdefgh", lbMap.findPrefix("abcdefghij")->key());
+
+  // "abcdefghijk" matches the longer entry exactly
+  ASSERT_NE(lbMap.end(), lbMap.findPrefix("abcdefghijk"));
+  ASSERT_EQ("abcdefghijk", lbMap.findPrefix("abcdefghijk")->key());
+  ASSERT_EQ(2, lbMap.findPrefix("abcdefghijk")->value());
+
+  // "abcdefghijkz" — "abcdefghijk" is the longest prefix
+  ASSERT_NE(lbMap.end(), lbMap.findPrefix("abcdefghijkz"));
+  ASSERT_EQ("abcdefghijk", lbMap.findPrefix("abcdefghijkz")->key());
+
+  // "abcdefg" (7 bytes) — no entry is a prefix of this
+  ASSERT_EQ(lbMap.end(), lbMap.findPrefix("abcdefg"));
+}
+
+TEST(LowerBoundPrefixMapTest, SameBucketWalkSkipsSuffix) {
+  // Entries share the same 8-byte prefix. "abcdefghXXX" is NOT a prefix of
+  // "abcdefghYYY", but "abcdefgh" IS. The walk must check the suffix (bytes
+  // 8+) before accepting a match, not stop at the first SmallPrefix match.
+  std::vector<KeyValue> kv = {
+      KeyValue{"abcdefgh", 1},
+      KeyValue{"abcdefghXXX", 2},
+  };
+  LowerBoundPrefixMap<int> lbMap{kv};
+
+  // Query "abcdefghYYY": "abcdefghXXX" shares the 8-byte prefix but is not
+  // a prefix. Must walk past it to find "abcdefgh".
+  ASSERT_NE(lbMap.end(), lbMap.findPrefix("abcdefghYYY"));
+  ASSERT_EQ("abcdefgh", lbMap.findPrefix("abcdefghYYY")->key());
+  ASSERT_EQ(1, lbMap.findPrefix("abcdefghYYY")->value());
+}
+
 } // namespace
 } // namespace facebook::memcache