heterogeneous lookup for ConcurrentHashMap

Summary:
This diff adds heterogeneous lookup handling for `ConcurrentHashMap` (`find()`
and `at()`), which allows lookups on keys that the hasher and key equal
functors accept but are not exactly the same as the `KeyType` of the map.
Common usage of the feature include looking up with a `std::string_view` key in
a `map<string, X>`, thereby avoiding creating an unnecessary temporary string
object.

Reviewed By: yfeldblum

Differential Revision: D25250409

fbshipit-source-id: 78dad163d623794b179c1691180b54cf109784ee

folly/concurrency/ConcurrentHashMap.h

@@ -149,6 +149,10 @@ class ConcurrentHashMap {
       Atom,
       Mutex,
       Impl>;
+  template <typename K, typename T>
+  using EnableHeterogeneousFind = std::enable_if_t<
+      detail::EligibleForHeterogeneousFind<KeyType, HashFn, KeyEqual, K>::value,
+      T>;
 
   float load_factor_ = SegmentT::kDefaultLoadFactor;
 
@@ -239,14 +243,11 @@ class ConcurrentHashMap {
     return true;
   }
 
-  ConstIterator find(const KeyType& k) const {
-    auto segment = pickSegment(k);
-    ConstIterator res(this, segment);
-    auto seg = segments_[segment].load(std::memory_order_acquire);
-    if (!seg || !seg->find(res.it_, k)) {
-      res.segment_ = NumShards;
-    }
-    return res;
+  ConstIterator find(const KeyType& k) const { return findImpl(k); }
+
+  template <typename K, EnableHeterogeneousFind<K, int> = 0>
+  ConstIterator find(const K& k) const {
+    return findImpl(k);
   }
 
   ConstIterator cend() const noexcept { return ConstIterator(NumShards); }
@@ -374,12 +375,11 @@ class ConcurrentHashMap {
     return item.first->second;
   }
 
-  const ValueType at(const KeyType& key) const {
-    auto item = find(key);
-    if (item == cend()) {
-      throw_exception<std::out_of_range>("at(): value out of range");
-    }
-    return item->second;
+  const ValueType at(const KeyType& key) const { return atImpl(key); }
+
+  template <typename K, EnableHeterogeneousFind<K, int> = 0>
+  const ValueType at(const K& key) const {
+    return atImpl(key);
   }
 
   // TODO update assign interface, operator[], at
@@ -541,7 +541,28 @@ class ConcurrentHashMap {
   };
 
  private:
-  uint64_t pickSegment(const KeyType& k) const {
+  template <typename K>
+  ConstIterator findImpl(const K& k) const {
+    auto segment = pickSegment(k);
+    ConstIterator res(this, segment);
+    auto seg = segments_[segment].load(std::memory_order_acquire);
+    if (!seg || !seg->find(res.it_, k)) {
+      res.segment_ = NumShards;
+    }
+    return res;
+  }
+
+  template <typename K>
+  const ValueType atImpl(const K& k) const {
+    auto item = find(k);
+    if (item == cend()) {
+      throw_exception<std::out_of_range>("at(): key not in map");
+    }
+    return item->second;
+  }
+
+  template <typename K>
+  uint64_t pickSegment(const K& k) const {
     auto h = HashFn()(k);
     // Use the lowest bits for our shard bits.
     //

folly/concurrency/detail/ConcurrentHashMap-detail.h

@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <folly/container/HeterogeneousAccess.h>
 #include <folly/container/detail/F14Mask.h>
 #include <folly/lang/Exception.h>
 #include <folly/lang/Launder.h>
@@ -333,7 +334,8 @@ class alignas(64) BucketTable {
     oldbuckets->retire(concurrenthashmap::HazptrTableDeleter(oldcount));
   }
 
-  bool find(Iterator& res, const KeyType& k) {
+  template <typename K>
+  bool find(Iterator& res, const K& k) {
     auto& hazcurr = res.hazptrs_[1];
     auto& haznext = res.hazptrs_[2];
     auto h = HashFn()(k);
@@ -1252,7 +1254,8 @@ class alignas(64) SIMDTable {
     rehash_internal(folly::nextPowTwo(new_chunk_count), cohort);
   }
 
-  bool find(Iterator& res, const KeyType& k) {
+  template <typename K>
+  bool find(Iterator& res, const K& k) {
     auto& hazz = res.hazptrs_[1];
     auto h = HashFn()(k);
     auto hp = splitHash(h);
@@ -1770,7 +1773,10 @@ class alignas(64) ConcurrentHashMapSegment {
   // Must hold lock.
   void rehash(size_t bucket_count) { impl_.rehash(bucket_count, cohort_); }
 
-  bool find(Iterator& res, const KeyType& k) { return impl_.find(res, k); }
+  template <typename K>
+  bool find(Iterator& res, const K& k) {
+    return impl_.find(res, k);
+  }
 
   // Listed separately because we need a prev pointer.
   size_type erase(const key_type& key) {

folly/concurrency/test/ConcurrentHashMapTest.cpp

@@ -20,6 +20,7 @@
 #include <memory>
 #include <thread>
 
+#include <folly/Traits.h>
 #include <folly/hash/Hash.h>
 #include <folly/portability/GFlags.h>
 #include <folly/portability/GTest.h>
@@ -830,6 +831,43 @@ TYPED_TEST_P(ConcurrentHashMapTest, IteratorLoop) {
   EXPECT_EQ(count, kNum);
 }
 
+namespace {
+template <typename T, typename Arg>
+using detector_find = decltype(std::declval<T>().find(std::declval<Arg>()));
+}
+
+TYPED_TEST_P(ConcurrentHashMapTest, HeterogeneousLookup) {
+  using Hasher = folly::transparent<folly::hasher<folly::StringPiece>>;
+  using KeyEqual = folly::transparent<std::equal_to<folly::StringPiece>>;
+  using M = ConcurrentHashMap<std::string, bool, Hasher, KeyEqual>;
+
+  constexpr auto hello = "hello"_sp;
+  constexpr auto buddy = "buddy"_sp;
+  constexpr auto world = "world"_sp;
+
+  M map;
+  map.emplace(hello, true);
+  map.emplace(world, false);
+
+  auto checks = [hello, buddy](auto& ref) {
+    // find
+    EXPECT_TRUE(ref.end() == ref.find(buddy));
+    EXPECT_EQ(hello, ref.find(hello)->first);
+
+    // at
+    EXPECT_TRUE(ref.at(hello));
+    EXPECT_THROW(ref.at(buddy), std::out_of_range);
+
+    // invocability checks
+    static_assert(
+        !is_detected_v<detector_find, decltype(ref), int>,
+        "there shouldn't be a find() overload for this string map with an int param");
+  };
+
+  checks(map);
+  checks(folly::as_const(map));
+}
+
 REGISTER_TYPED_TEST_CASE_P(
     ConcurrentHashMapTest,
     MapTest,
@@ -864,7 +902,8 @@ REGISTER_TYPED_TEST_CASE_P(
     assignStressTest,
     insertStressTest,
     IteratorMove,
-    IteratorLoop);
+    IteratorLoop,
+    HeterogeneousLookup);
 
 using folly::detail::concurrenthashmap::bucket::BucketTable;
 

folly/container/F14Map.h

@@ -54,7 +54,7 @@ template <typename Policy>
 class F14BasicMap {
   template <typename K, typename T>
   using EnableHeterogeneousFind = std::enable_if_t<
-      EligibleForHeterogeneousFind<
+      ::folly::detail::EligibleForHeterogeneousFind<
           typename Policy::Key,
           typename Policy::Hasher,
           typename Policy::KeyEqual,
@@ -63,7 +63,7 @@ class F14BasicMap {
 
   template <typename K, typename T>
   using EnableHeterogeneousInsert = std::enable_if_t<
-      EligibleForHeterogeneousInsert<
+      ::folly::detail::EligibleForHeterogeneousInsert<
           typename Policy::Key,
           typename Policy::Hasher,
           typename Policy::KeyEqual,
@@ -77,7 +77,7 @@ class F14BasicMap {
 
   template <typename K, typename T>
   using EnableHeterogeneousErase = std::enable_if_t<
-      EligibleForHeterogeneousFind<
+      ::folly::detail::EligibleForHeterogeneousFind<
           typename Policy::Key,
           typename Policy::Hasher,
           typename Policy::KeyEqual,
@@ -395,7 +395,7 @@ class F14BasicMap {
 
  private:
   template <typename Arg>
-  using UsableAsKey =
+  using UsableAsKey = ::folly::detail::
       EligibleForHeterogeneousFind<key_type, hasher, key_equal, Arg>;
 
  public:
@@ -953,7 +953,7 @@ class F14VectorMapImpl : public F14BasicMap<MapPolicyWithDefaults<
 
   template <typename K, typename T>
   using EnableHeterogeneousVectorErase = std::enable_if_t<
-      EligibleForHeterogeneousFind<
+      ::folly::detail::EligibleForHeterogeneousFind<
           Key,
           Hasher,
           KeyEqual,

folly/container/F14Set.h

@@ -50,7 +50,7 @@ template <typename Policy>
 class F14BasicSet {
   template <typename K, typename T>
   using EnableHeterogeneousFind = std::enable_if_t<
-      EligibleForHeterogeneousFind<
+      ::folly::detail::EligibleForHeterogeneousFind<
           typename Policy::Value,
           typename Policy::Hasher,
           typename Policy::KeyEqual,
@@ -59,7 +59,7 @@ class F14BasicSet {
 
   template <typename K, typename T>
   using EnableHeterogeneousInsert = std::enable_if_t<
-      EligibleForHeterogeneousInsert<
+      ::folly::detail::EligibleForHeterogeneousInsert<
           typename Policy::Value,
           typename Policy::Hasher,
           typename Policy::KeyEqual,
@@ -71,7 +71,7 @@ class F14BasicSet {
 
   template <typename K, typename T>
   using EnableHeterogeneousErase = std::enable_if_t<
-      EligibleForHeterogeneousFind<
+      ::folly::detail::EligibleForHeterogeneousFind<
           typename Policy::Value,
           typename Policy::Hasher,
           typename Policy::KeyEqual,
@@ -311,7 +311,7 @@ class F14BasicSet {
 
  private:
   template <typename Arg>
-  using UsableAsKey =
+  using UsableAsKey = ::folly::detail::
       EligibleForHeterogeneousFind<key_type, hasher, key_equal, Arg>;
 
  public:
@@ -718,7 +718,7 @@ class F14VectorSetImpl : public F14BasicSet<SetPolicyWithDefaults<
 
   template <typename K, typename T>
   using EnableHeterogeneousVectorErase = std::enable_if_t<
-      EligibleForHeterogeneousFind<
+      ::folly::detail::EligibleForHeterogeneousFind<
           typename Policy::Value,
           typename Policy::Hasher,
           typename Policy::KeyEqual,

folly/container/HeterogeneousAccess.h

@@ -143,6 +143,27 @@ struct TransparentRangeHash {
   }
 };
 
+template <
+    typename TableKey,
+    typename Hasher,
+    typename KeyEqual,
+    typename ArgKey>
+struct EligibleForHeterogeneousFind
+    : Conjunction<
+          is_transparent<Hasher>,
+          is_transparent<KeyEqual>,
+          is_invocable<Hasher, ArgKey const&>,
+          is_invocable<KeyEqual, ArgKey const&, TableKey const&>> {};
+
+template <
+    typename TableKey,
+    typename Hasher,
+    typename KeyEqual,
+    typename ArgKey>
+using EligibleForHeterogeneousInsert = Conjunction<
+    EligibleForHeterogeneousFind<TableKey, Hasher, KeyEqual, ArgKey>,
+    std::is_constructible<TableKey, ArgKey>>;
+
 } // namespace detail
 
 template <typename T>

folly/container/detail/F14MapFallback.h

@@ -43,12 +43,14 @@ class F14BasicMap : public std::unordered_map<K, M, H, E, A> {
   using Super = std::unordered_map<K, M, H, E, A>;
 
   template <typename K2, typename T>
-  using EnableHeterogeneousFind =
-      std::enable_if_t<EligibleForHeterogeneousFind<K, H, E, K2>::value, T>;
+  using EnableHeterogeneousFind = std::enable_if_t<
+      ::folly::detail::EligibleForHeterogeneousFind<K, H, E, K2>::value,
+      T>;
 
   template <typename K2, typename T>
-  using EnableHeterogeneousInsert =
-      std::enable_if_t<EligibleForHeterogeneousInsert<K, H, E, K2>::value, T>;
+  using EnableHeterogeneousInsert = std::enable_if_t<
+      ::folly::detail::EligibleForHeterogeneousInsert<K, H, E, K2>::value,
+      T>;
 
   template <typename K2>
   using IsIter = Disjunction<
@@ -57,7 +59,7 @@ class F14BasicMap : public std::unordered_map<K, M, H, E, A> {
 
   template <typename K2, typename T>
   using EnableHeterogeneousErase = std::enable_if_t<
-      EligibleForHeterogeneousFind<
+      ::folly::detail::EligibleForHeterogeneousFind<
           K,
           H,
           E,
@@ -192,7 +194,7 @@ class F14BasicMap : public std::unordered_map<K, M, H, E, A> {
 
  private:
   template <typename Arg>
-  using UsableAsKey =
+  using UsableAsKey = ::folly::detail::
       EligibleForHeterogeneousFind<key_type, hasher, key_equal, Arg>;
 
  public:

folly/container/detail/F14SetFallback.h

@@ -54,12 +54,14 @@ class F14BasicSet
  private:
   template <typename K, typename T>
   using EnableHeterogeneousFind = std::enable_if_t<
-      EligibleForHeterogeneousFind<key_type, hasher, key_equal, K>::value,
+      ::folly::detail::
+          EligibleForHeterogeneousFind<key_type, hasher, key_equal, K>::value,
       T>;
 
   template <typename K, typename T>
   using EnableHeterogeneousInsert = std::enable_if_t<
-      EligibleForHeterogeneousInsert<key_type, hasher, key_equal, K>::value,
+      ::folly::detail::
+          EligibleForHeterogeneousInsert<key_type, hasher, key_equal, K>::value,
       T>;
 
   template <typename K>
@@ -69,7 +71,7 @@ class F14BasicSet
 
   template <typename K, typename T>
   using EnableHeterogeneousErase = std::enable_if_t<
-      EligibleForHeterogeneousFind<
+      ::folly::detail::EligibleForHeterogeneousFind<
           key_type,
           hasher,
           key_equal,
@@ -101,7 +103,7 @@ class F14BasicSet
 
  private:
   template <typename Arg>
-  using UsableAsKey =
+  using UsableAsKey = ::folly::detail::
       EligibleForHeterogeneousFind<key_type, hasher, key_equal, Arg>;
 
  public:

folly/container/detail/F14Table.h

@@ -206,27 +206,6 @@ typename Container::size_type erase_if_impl(
   return old_size - c.size();
 }
 
-template <
-    typename TableKey,
-    typename Hasher,
-    typename KeyEqual,
-    typename ArgKey>
-struct EligibleForHeterogeneousFind
-    : Conjunction<
-          is_transparent<Hasher>,
-          is_transparent<KeyEqual>,
-          is_invocable<Hasher, ArgKey const&>,
-          is_invocable<KeyEqual, ArgKey const&, TableKey const&>> {};
-
-template <
-    typename TableKey,
-    typename Hasher,
-    typename KeyEqual,
-    typename ArgKey>
-using EligibleForHeterogeneousInsert = Conjunction<
-    EligibleForHeterogeneousFind<TableKey, Hasher, KeyEqual, ArgKey>,
-    std::is_constructible<TableKey, ArgKey>>;
-
 } // namespace detail
 } // namespace f14