Use AddConstRefIfExpensiveToCopy for non-transparent key arguments and internal raw_hash_set lookup helpers. Previously, KeyArg always resolved non-transparent keys to key_type and public lookup methods took `const key_arg<K>&`, while internal raw_hash_set lookup and preparation functions (e.g. find_or_prepare_insert()) took `const K&`. When these functions were not inlined, callers with cheap-to-copy keys had to spill the key to stack memory, pass a stack address in `%rsi`, and reload the key inside the callee before hashing, comparing, or tail-calling (e.g. for find_or_prepare_insert_large()). Add KeyArg::const_ref_if_expensive<K, key_type>, which aliases `const K&` for transparent containers and type_traits_internal::AddConstRefIfExpensiveToCopy<std::conditional_t<false, K, key_type>> for non-transparent containers (keeping key_type dependent on `K` so incomplete recursive key types remain supported), and use it in read-only key_arg<K> member functions (e.g. find()) while preserving `const ...&` and `...&&` forwarding overloads (e.g. insert()). Also introduce raw_hash_set::KeyParam<K> (AddConstRefIfExpensiveToCopy<absl::remove_cvref_t<K>>) and use it across internal lookup and insert-preparation helpers so cheap-to-copy keys are passed directly in registers. PiperOrigin-RevId: 1009900685
diff --git a/absl/container/internal/common.h b/absl/container/internal/common.h index 344c9fe..4c9cd9d 100644 --- a/absl/container/internal/common.h +++ b/absl/container/internal/common.h
@@ -84,16 +84,25 @@ // Transparent. Forward `K`. template <typename K, typename key_type> using type = K; + + // In the transparent case, the argument type is deduced so we can't use + // AddConstRefIfExpensiveToCopy. + template <typename K, typename key_type> + using const_ref_if_expensive = const K&; }; template <> struct KeyArg<false> { // Not transparent. Always use `key_type`. - // TODO(ezb): use AddConstRefIfExpensiveToCopy for non-transparent key - // arguments. Also experiment with using it more widely within non-user-facing - // functions in raw_hash_set. template <typename K, typename key_type> using type = key_type; + + // The std::conditional_t<false, K, key_type> is to allow key_type to be an + // incomplete type by making it a dependent type. + template <typename K, typename key_type> + using const_ref_if_expensive = + type_traits_internal::AddConstRefIfExpensiveToCopy< + std::conditional_t<false, K, key_type>>; }; // The node_handle concept from C++17.
diff --git a/absl/container/internal/raw_hash_map.h b/absl/container/internal/raw_hash_map.h index 746cb6e..e4ced0a 100644 --- a/absl/container/internal/raw_hash_map.h +++ b/absl/container/internal/raw_hash_map.h
@@ -63,11 +63,15 @@ using Hash = typename raw_hash_map::raw_hash_set::hasher; using Eq = typename raw_hash_map::raw_hash_set::key_equal; using Alloc = typename raw_hash_map::raw_hash_set::allocator_type; + using KeyArgImpl = + KeyArg<IsTransparent<Eq>::value && IsTransparent<Hash>::value>; template <class K> using key_arg = - typename KeyArg<IsTransparent<Eq>::value && IsTransparent<Hash>::value>:: - template type<K, typename Policy::key_type>; + typename KeyArgImpl::template type<K, typename Policy::key_type>; + template <class K> + using const_key_arg = typename KeyArgImpl::template const_ref_if_expensive< + K, typename Policy::key_type>; // NOTE: The mess here is to shorten the code for the (very repetitive) // function overloads, and to allow the lifetime-bound overloads to dispatch @@ -301,7 +305,7 @@ } template <class K = key_type, class P = Policy> - MappedReference<P> at(const key_arg<K>& key) ABSL_ATTRIBUTE_LIFETIME_BOUND { + MappedReference<P> at(const_key_arg<K> key) ABSL_ATTRIBUTE_LIFETIME_BOUND { auto it = this->find(key); if (this->unchecked_equals(it, this->end())) { ThrowStdOutOfRange("absl::container_internal::raw_hash_map<>::at"); @@ -310,7 +314,7 @@ } template <class K = key_type, class P = Policy> - MappedConstReference<P> at(const key_arg<K>& key) const + MappedConstReference<P> at(const_key_arg<K> key) const ABSL_ATTRIBUTE_LIFETIME_BOUND { auto it = this->find(key); if (this->unchecked_equals(it, this->end())) { @@ -365,7 +369,8 @@ template <class K, class V> std::pair<iterator, bool> insert_or_assign_impl(K&& k, V&& v) ABSL_ATTRIBUTE_LIFETIME_BOUND { - auto res = this->find_or_prepare_insert(k); + auto res = this->template find_or_prepare_insert< + typename raw_hash_map::template KeyParam<K>>(k); if (res.second) { this->emplace_at(res.first, std::forward<K>(k), std::forward<V>(v)); } else { @@ -377,7 +382,8 @@ template <class K = key_type, class... Args> std::pair<iterator, bool> try_emplace_impl(K&& k, Args&&... args) ABSL_ATTRIBUTE_LIFETIME_BOUND { - auto res = this->find_or_prepare_insert(k); + auto res = this->template find_or_prepare_insert< + typename raw_hash_map::template KeyParam<K>>(k); if (res.second) { // Try to avoid using std::tuple since that slows compilation. // (example: https://godbolt.org/z/nefW93T6K)
diff --git a/absl/container/internal/raw_hash_set.h b/absl/container/internal/raw_hash_set.h index 8ca006f..e1771fb 100644 --- a/absl/container/internal/raw_hash_set.h +++ b/absl/container/internal/raw_hash_set.h
@@ -2500,11 +2500,12 @@ private: // Alias used for heterogeneous lookup functions. - // `key_arg<K>` evaluates to `K` when the functors are transparent and to - // `key_type` otherwise. It permits template argument deduction on `K` for the - // transparent case. + // `const_key_arg<K>` evaluates to `const K&` when the functors are + // transparent and to `AddConstRefIfExpensiveToCopy<key_type>` otherwise. It + // permits template argument deduction on `K` for the transparent case. template <class K> - using key_arg = typename KeyArgImpl::template type<K, key_type>; + using const_key_arg = + typename KeyArgImpl::template const_ref_if_expensive<K, key_type>; using slot_type = typename PolicyTraits::slot_type; @@ -3294,9 +3295,9 @@ }; template <class K = key_type, class F> - iterator lazy_emplace(const key_arg<K>& key, + iterator lazy_emplace(const_key_arg<K> key, F&& f) ABSL_ATTRIBUTE_LIFETIME_BOUND { - auto res = find_or_prepare_insert(key); + auto res = find_or_prepare_insert<KeyParam<const_key_arg<K>>>(key); if (res.second) { slot_type* slot = res.first; allocator_type alloc(char_alloc_ref()); @@ -3316,7 +3317,7 @@ // // Uses "abc" directly without copying it into std::string. // s.erase("abc"); template <class K = key_type> - size_type erase(const key_arg<K>& key) { + size_type erase(const_key_arg<K> key) { auto it = find(key); if (it.unchecked_equals(end())) return 0; erase(it); @@ -3416,7 +3417,7 @@ template <class K = key_type, std::enable_if_t<!std::is_same_v<K, iterator>, int> = 0> - node_type extract(const key_arg<K>& key) { + node_type extract(const_key_arg<K> key) { auto it = find(key); return it.unchecked_equals(end()) ? node_type() : extract(const_iterator{it}); @@ -3456,7 +3457,7 @@ // // Uses "abc" directly without copying it into std::string. // s.count("abc"); template <class K = key_type> - size_t count(const key_arg<K>& key) const { + size_t count(const_key_arg<K> key) const { return contains(key) ? 1 : 0; } @@ -3466,7 +3467,7 @@ // NOTE: This is a very low level operation and should not be used without // specific benchmarks indicating its importance. template <class K = key_type> - void prefetch([[maybe_unused]] const key_arg<K>& key) const { + void prefetch([[maybe_unused]] const_key_arg<K> key) const { if (capacity() == DefaultCapacity()) return; // Avoid probing if we won't be able to prefetch the addresses received. #ifdef ABSL_HAVE_PREFETCH @@ -3480,33 +3481,32 @@ template <class K = key_type> ABSL_DEPRECATE_AND_INLINE() - iterator find(const key_arg<K>& key, - size_t) ABSL_ATTRIBUTE_LIFETIME_BOUND { + iterator find(const_key_arg<K> key, size_t) ABSL_ATTRIBUTE_LIFETIME_BOUND { return find(key); } // The API of find() has one extension: the type of the key argument doesn't // have to be key_type. This is so called heterogeneous key support. template <class K = key_type> - iterator find(const key_arg<K>& key) ABSL_ATTRIBUTE_LIFETIME_BOUND { - AssertOnFind(key); - if (is_small()) return find_small(key); - return find_large<FindLargeKey<key_arg<K>>>(key); + iterator find(const_key_arg<K> key) ABSL_ATTRIBUTE_LIFETIME_BOUND { + AssertOnFind<KeyParam<const_key_arg<K>>>(key); + if (is_small()) return find_small<KeyParam<const_key_arg<K>>>(key); + return find_large<FindLargeKey<const_key_arg<K>>>(key); } template <class K = key_type> ABSL_DEPRECATE_AND_INLINE() - const_iterator find(const key_arg<K>& key, - size_t) const ABSL_ATTRIBUTE_LIFETIME_BOUND { + const_iterator + find(const_key_arg<K> key, size_t) const ABSL_ATTRIBUTE_LIFETIME_BOUND { return find(key); } template <class K = key_type> - const_iterator find(const key_arg<K>& key) const + const_iterator find(const_key_arg<K> key) const ABSL_ATTRIBUTE_LIFETIME_BOUND { return const_cast<raw_hash_set*>(this)->find(key); } template <class K = key_type> - bool contains(const key_arg<K>& key) const { + bool contains(const_key_arg<K> key) const { // Here neither the iterator returned by `find()` nor `end()` can be invalid // outside of potential thread-safety issues. // `find()`'s return value is constructed, used, and then destructed @@ -3515,7 +3515,7 @@ } template <class K = key_type> - std::pair<iterator, iterator> equal_range(const key_arg<K>& key) + std::pair<iterator, iterator> equal_range(const_key_arg<K> key) ABSL_ATTRIBUTE_LIFETIME_BOUND { auto it = find(key); if (!it.unchecked_equals(end())) return {it, std::next(it)}; @@ -3523,7 +3523,7 @@ } template <class K = key_type> std::pair<const_iterator, const_iterator> equal_range( - const key_arg<K>& key) const ABSL_ATTRIBUTE_LIFETIME_BOUND { + const_key_arg<K> key) const ABSL_ATTRIBUTE_LIFETIME_BOUND { auto it = find(key); if (!it.unchecked_equals(end())) return {it, std::next(it)}; return {it, it}; @@ -3601,7 +3601,7 @@ struct EmplaceDecomposable { template <class K, class... Args> std::pair<iterator, bool> operator()(const K& key, Args&&... args) const { - auto res = s.find_or_prepare_insert(key); + auto res = s.find_or_prepare_insert<KeyParam<K>>(key); if (res.second) { s.emplace_at(res.first, std::forward<Args>(args)...); } @@ -3614,7 +3614,7 @@ struct InsertSlot { template <class K, class... Args> std::pair<iterator, bool> operator()(const K& key, Args&&...) && { - auto res = s.find_or_prepare_insert(key); + auto res = s.find_or_prepare_insert<KeyParam<K>>(key); if (res.second) { s.transfer(res.first, &slot); } else if (do_destroy) { @@ -3647,8 +3647,10 @@ // TODO(b/289225379): consider having a helper class that has the impls for // SOO functionality. template <class K> - ABSL_ATTRIBUTE_ALWAYS_INLINE iterator find_small(const K& key) { - return empty_when_small() || !equal_to(key, single_slot()) + ABSL_ATTRIBUTE_ALWAYS_INLINE iterator + find_small(absl::type_identity_t<K> key) { + static_assert(std::is_same_v<K, KeyParam<K>>); + return empty_when_small() || !equal_to<K>(key, single_slot()) ? end() : single_iterator(); } @@ -3674,7 +3676,7 @@ if (auto match = g.Match(h2); ABSL_PREDICT_TRUE(match)) { do { const size_t offset = seq.offset(*match); - if (ABSL_PREDICT_TRUE(equal_to(key, slot_array + offset))) + if (ABSL_PREDICT_TRUE(equal_to<K>(key, slot_array + offset))) return iterator_at_ptr(ctrl + offset, slot_array + offset); } while (ABSL_PREDICT_FALSE(++match)); } @@ -3769,10 +3771,12 @@ } template <class K> - ABSL_ATTRIBUTE_ALWAYS_INLINE bool equal_to(const K& key, + ABSL_ATTRIBUTE_ALWAYS_INLINE bool equal_to(absl::type_identity_t<K> key, slot_type* slot) const { - return PolicyTraits::apply(EqualElement<K, key_equal>{key, eq_ref()}, - PolicyTraits::element(slot)); + static_assert(std::is_same_v<K, KeyParam<K>>); + return PolicyTraits::apply( + EqualElement<absl::remove_cvref_t<K>, key_equal>{key, eq_ref()}, + PolicyTraits::element(slot)); } template <class K> ABSL_ATTRIBUTE_ALWAYS_INLINE size_t hash_of(const K& key) const { @@ -3893,7 +3897,8 @@ template <class K> ABSL_ATTRIBUTE_ALWAYS_INLINE std::pair<slot_type*, bool> - find_or_prepare_insert_soo(const K& key) { + find_or_prepare_insert_soo(absl::type_identity_t<K> key) { + static_assert(std::is_same_v<K, KeyParam<K>>); bool force_sampling; slot_type* slot = single_slot(); if (empty_when_small()) { @@ -3902,14 +3907,15 @@ return {slot, true}; } force_sampling = true; - } else if (equal_to(key, slot)) { + } else if (equal_to<K>(key, slot)) { return {slot, false}; } else { force_sampling = false; } constexpr bool kUseMemcpy = PolicyTraits::transfer_uses_memcpy() && SooEnabled(); - using Hasher = HashKey<hasher, K, kIsAbslHash, kSeedShift>; + using Hasher = + HashKey<hasher, absl::remove_cvref_t<K>, kIsAbslHash, kSeedShift>; Hasher hash_key{hash_ref(), key}; slot = to_slot( GrowSooTableToNextCapacityAndPrepareInsert< @@ -3922,16 +3928,18 @@ template <class K> ABSL_ATTRIBUTE_ALWAYS_INLINE std::pair<slot_type*, bool> - find_or_prepare_insert_small(const K& key) { + find_or_prepare_insert_small(absl::type_identity_t<K> key) { + static_assert(std::is_same_v<K, KeyParam<K>>); if constexpr (SooEnabled()) { - return find_or_prepare_insert_soo(key); + return find_or_prepare_insert_soo<K>(key); } if (!empty_when_small()) { - if (equal_to(key, single_slot())) { + if (equal_to<K>(key, single_slot())) { return {single_slot(), false}; } } - using Hasher = HashKey<hasher, K, kIsAbslHash, kSeedShift>; + using Hasher = + HashKey<hasher, absl::remove_cvref_t<K>, kIsAbslHash, kSeedShift>; Hasher hash_key{hash_ref(), key}; return {to_slot(PrepareInsertSmallNonSoo(common(), GetPolicyFunctions(), &hash_key, &Hasher::GetHash)), @@ -3959,7 +3967,7 @@ if (auto match = g.Match(h2); ABSL_PREDICT_TRUE(match)) { do { slot_type* slot = slot_array + seq.offset(*match); - if (ABSL_PREDICT_TRUE(equal_to(key, slot))) return {slot, false}; + if (ABSL_PREDICT_TRUE(equal_to<K>(key, slot))) return {slot, false}; } while (ABSL_PREDICT_FALSE(++match)); } auto mask_empty = g.MaskEmpty(); @@ -3986,10 +3994,15 @@ } protected: + template <class K> + using KeyParam = type_traits_internal::AddConstRefIfExpensiveToCopy< + absl::remove_cvref_t<K>>; + // Asserts for correctness that we run on find/find_or_prepare_insert. template <class K> - void AssertOnFind([[maybe_unused]] const K& key) { - AssertHashEqConsistent(key); + void AssertOnFind([[maybe_unused]] absl::type_identity_t<K> key) { + static_assert(std::is_same_v<K, KeyParam<K>>); + AssertHashEqConsistent<K>(key); AssertNotDebugCapacity(); } @@ -3998,14 +4011,16 @@ template <class K> static bool equal_to_key_fn(void* set, const void* key, void* slot) { - return static_cast<const raw_hash_set*>(set)->equal_to( - *static_cast<const K*>(key), to_slot(slot)); + static_assert(std::is_same_v<K, KeyParam<K>>); + return static_cast<const raw_hash_set*>(set)->equal_to<K>( + *static_cast<const absl::remove_cvref_t<K>*>(key), to_slot(slot)); } // Asserts that hash and equal functors provided by the user are consistent, // meaning that `eq(k1, k2)` implies `hash(k1)==hash(k2)`. template <class K> - void AssertHashEqConsistent(const K& key) { + void AssertHashEqConsistent(absl::type_identity_t<K> key) { + static_assert(std::is_same_v<K, KeyParam<K>>); // If the hash/eq functors are known to be consistent, then skip // validation. constexpr bool kIsStringAbslHashEq = @@ -4027,9 +4042,11 @@ // where the value can be inserted into, with the control byte already set to // `key`'s H2. Returns a bool indicating whether an insertion can take place. template <class K> - std::pair<slot_type*, bool> find_or_prepare_insert(const K& key) { - AssertOnFind(key); - return is_small() ? find_or_prepare_insert_small(key) + std::pair<slot_type*, bool> find_or_prepare_insert( + absl::type_identity_t<K> key) { + static_assert(std::is_same_v<K, KeyParam<K>>); + AssertOnFind<K>(key); + return is_small() ? find_or_prepare_insert_small<K>(key) : find_or_prepare_insert_large<FindLargeKey<K>>(key); } @@ -4109,13 +4126,12 @@ // (type_identity_t ensures that the template argument is not deduced and must // be specified explicitly.) template <class K> - using FindLargeKey = - type_traits_internal::AddConstRefIfExpensiveToCopy<std::conditional_t< - std::is_same_v<hasher, absl::container_internal::StringHash> && - std::is_same_v<key_equal, absl::container_internal::StringEq> && - std::is_convertible_v<const absl::remove_cvref_t<K>&, - absl::string_view>, - absl::string_view, absl::remove_cvref_t<K>>>; + using FindLargeKey = KeyParam<std::conditional_t< + std::is_same_v<hasher, absl::container_internal::StringHash> && + std::is_same_v<key_equal, absl::container_internal::StringEq> && + std::is_convertible_v<const absl::remove_cvref_t<K>&, + absl::string_view>, + absl::string_view, absl::remove_cvref_t<K>>>; // Prefetch the heap-allocated memory region to resolve potential TLB and // cache misses. This is intended to overlap with execution of calculating the @@ -4478,8 +4494,9 @@ constexpr static bool kIsAbslHash = Set::kIsAbslHash; - static size_t GetNumProbes(const Set& set, - const typename Set::key_type& key) { + static size_t GetNumProbes( + const Set& set, + typename Set::template KeyParam<typename Set::key_type> key) { if (set.is_small()) return 0; size_t num_probes = 0; const size_t hash = set.hash_of(key); @@ -4489,7 +4506,9 @@ while (true) { container_internal::Group g{ctrl + seq.offset()}; for (uint32_t i : g.Match(h2)) { - if (set.equal_to(key, set.slot_array(set.capacity()) + seq.offset(i))) + if (set.template equal_to< + typename Set::template KeyParam<typename Set::key_type>>( + key, set.slot_array(set.capacity()) + seq.offset(i))) return num_probes; ++num_probes; }