tint/utils: Support hetrogeneous hashmap key lookups

Allows map with std::string keys to be looked up, using a c-string or
stringview without incuring a temporary heap allocation.

Change-Id: Id5b7fd5ac1ab7febf545472f9767273f8637a0de
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/131623
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Ben Clayton <bclayton@google.com>
Reviewed-by: Antonio Maiorano <amaiorano@google.com>

src/tint/utils/hash.h

@@ -18,6 +18,7 @@
 #include <stdint.h>
 #include <cstdio>
 #include <functional>
+#include <string>
 #include <tuple>
 #include <utility>
 #include <variant>
@@ -152,6 +153,29 @@ struct Hasher<std::variant<TYPES...>> {
     }
 };
 
+/// Hasher specialization for std::string, which also supports hashing of const char* and
+/// std::string_view without first constructing a std::string.
+template <>
+struct Hasher<std::string> {
+    /// @param str the string to hash
+    /// @returns a hash of the string
+    size_t operator()(const std::string& str) const {
+        return std::hash<std::string_view>()(std::string_view(str));
+    }
+
+    /// @param str the string to hash
+    /// @returns a hash of the string
+    size_t operator()(const char* str) const {
+        return std::hash<std::string_view>()(std::string_view(str));
+    }
+
+    /// @param str the string to hash
+    /// @returns a hash of the string
+    size_t operator()(const std::string_view& str) const {
+        return std::hash<std::string_view>()(str);
+    }
+};
+
 /// @returns a hash of the variadic list of arguments.
 ///          The returned hash is dependent on the order of the arguments.
 template <typename... ARGS>
@@ -176,6 +200,47 @@ size_t HashCombine(size_t hash, const ARGS&... values) {
     return hash;
 }
 
+/// A STL-compatible equal_to implementation that specializes for types.
+template <typename T>
+struct EqualTo {
+    /// @param lhs the left hand side value
+    /// @param rhs the right hand side value
+    /// @returns true if the two values are equal
+    constexpr bool operator()(const T& lhs, const T& rhs) const {
+        return std::equal_to<T>()(lhs, rhs);
+    }
+};
+
+/// A specialization for EqualTo for std::string, which supports additional comparision with
+/// std::string_view and const char*.
+template <>
+struct EqualTo<std::string> {
+    /// @param lhs the left hand side value
+    /// @param rhs the right hand side value
+    /// @returns true if the two values are equal
+    bool operator()(const std::string& lhs, const std::string& rhs) const { return lhs == rhs; }
+
+    /// @param lhs the left hand side value
+    /// @param rhs the right hand side value
+    /// @returns true if the two values are equal
+    bool operator()(const std::string& lhs, const char* rhs) const { return lhs == rhs; }
+
+    /// @param lhs the left hand side value
+    /// @param rhs the right hand side value
+    /// @returns true if the two values are equal
+    bool operator()(const std::string& lhs, std::string_view rhs) const { return lhs == rhs; }
+
+    /// @param lhs the left hand side value
+    /// @param rhs the right hand side value
+    /// @returns true if the two values are equal
+    bool operator()(const char* lhs, const std::string& rhs) const { return lhs == rhs; }
+
+    /// @param lhs the left hand side value
+    /// @param rhs the right hand side value
+    /// @returns true if the two values are equal
+    bool operator()(std::string_view lhs, const std::string& rhs) const { return lhs == rhs; }
+};
+
 /// Wrapper for a hashable type enabling the wrapped value to be used as a key
 /// for an unordered_map or unordered_set.
 template <typename T>

src/tint/utils/hash_test.cc

@@ -73,5 +73,41 @@ TEST(HashTests, UnorderedKeyWrapper) {
     EXPECT_EQ(m[W({2, 1})], 0);
 }
 
+TEST(EqualTo, String) {
+    std::string str_a = "hello";
+    std::string str_b = "world";
+    const char* cstr_a = "hello";
+    const char* cstr_b = "world";
+    std::string_view sv_a = "hello";
+    std::string_view sv_b = "world";
+    EXPECT_TRUE(EqualTo<std::string>()(str_a, str_a));
+    EXPECT_TRUE(EqualTo<std::string>()(str_a, cstr_a));
+    EXPECT_TRUE(EqualTo<std::string>()(str_a, sv_a));
+    EXPECT_TRUE(EqualTo<std::string>()(str_a, str_a));
+    EXPECT_TRUE(EqualTo<std::string>()(cstr_a, str_a));
+    EXPECT_TRUE(EqualTo<std::string>()(sv_a, str_a));
+
+    EXPECT_FALSE(EqualTo<std::string>()(str_a, str_b));
+    EXPECT_FALSE(EqualTo<std::string>()(str_a, cstr_b));
+    EXPECT_FALSE(EqualTo<std::string>()(str_a, sv_b));
+    EXPECT_FALSE(EqualTo<std::string>()(str_a, str_b));
+    EXPECT_FALSE(EqualTo<std::string>()(cstr_a, str_b));
+    EXPECT_FALSE(EqualTo<std::string>()(sv_a, str_b));
+
+    EXPECT_FALSE(EqualTo<std::string>()(str_b, str_a));
+    EXPECT_FALSE(EqualTo<std::string>()(str_b, cstr_a));
+    EXPECT_FALSE(EqualTo<std::string>()(str_b, sv_a));
+    EXPECT_FALSE(EqualTo<std::string>()(str_b, str_a));
+    EXPECT_FALSE(EqualTo<std::string>()(cstr_b, str_a));
+    EXPECT_FALSE(EqualTo<std::string>()(sv_b, str_a));
+
+    EXPECT_TRUE(EqualTo<std::string>()(str_b, str_b));
+    EXPECT_TRUE(EqualTo<std::string>()(str_b, cstr_b));
+    EXPECT_TRUE(EqualTo<std::string>()(str_b, sv_b));
+    EXPECT_TRUE(EqualTo<std::string>()(str_b, str_b));
+    EXPECT_TRUE(EqualTo<std::string>()(cstr_b, str_b));
+    EXPECT_TRUE(EqualTo<std::string>()(sv_b, str_b));
+}
+
 }  // namespace
 }  // namespace tint::utils

src/tint/utils/hashmap.h

@@ -31,11 +31,14 @@ template <typename KEY,
           typename VALUE,
           size_t N,
           typename HASH = Hasher<KEY>,
-          typename EQUAL = std::equal_to<KEY>>
+          typename EQUAL = EqualTo<KEY>>
 class Hashmap : public HashmapBase<KEY, VALUE, N, HASH, EQUAL> {
     using Base = HashmapBase<KEY, VALUE, N, HASH, EQUAL>;
     using PutMode = typename Base::PutMode;
 
+    template <typename T>
+    using ReferenceKeyType = traits::CharArrayToCharPtr<std::remove_reference_t<T>>;
+
   public:
     /// The key type
     using Key = KEY;
@@ -51,7 +54,7 @@ class Hashmap : public HashmapBase<KEY, VALUE, N, HASH, EQUAL> {
     /// The value returned by the Reference reflects the current state of the Hashmap, and so the
     /// referenced value may change, or transition between valid or invalid based on the current
     /// state of the Hashmap.
-    template <bool IS_CONST>
+    template <bool IS_CONST, typename K>
     class ReferenceT {
         /// `const Value` if IS_CONST, or `Value` if !IS_CONST
         using T = std::conditional_t<IS_CONST, const Value, Value>;
@@ -89,24 +92,34 @@ class Hashmap : public HashmapBase<KEY, VALUE, N, HASH, EQUAL> {
         friend Hashmap;
 
         /// Constructor
-        ReferenceT(Map& map, const Key& key)
+        template <typename K_ARG>
-            : map_(map), key_(key), cached_(nullptr), generation_(map.Generation() - 1) {}
+        ReferenceT(Map& map, K_ARG&& key)
+            : map_(map),
+              key_(std::forward<K_ARG>(key)),
+              cached_(nullptr),
+              generation_(map.Generation() - 1) {}
 
         /// Constructor
-        ReferenceT(Map& map, const Key& key, T* value)
+        template <typename K_ARG>
-            : map_(map), key_(key), cached_(value), generation_(map.Generation()) {}
+        ReferenceT(Map& map, K_ARG&& key, T* value)
+            : map_(map),
+              key_(std::forward<K_ARG>(key)),
+              cached_(value),
+              generation_(map.Generation()) {}
 
         Map& map_;
-        const Key key_;
+        const K key_;
         mutable T* cached_ = nullptr;
         mutable size_t generation_ = 0;
     };
 
     /// A mutable reference returned by Find()
-    using Reference = ReferenceT</*IS_CONST*/ false>;
+    template <typename K>
+    using Reference = ReferenceT</*IS_CONST*/ false, K>;
 
     /// An immutable reference returned by Find()
-    using ConstReference = ReferenceT</*IS_CONST*/ true>;
+    template <typename K>
+    using ConstReference = ReferenceT</*IS_CONST*/ true, K>;
 
     /// Adds a value to the map, if the map does not already contain an entry with the key @p key.
     /// @param key the entry key.
@@ -129,7 +142,8 @@ class Hashmap : public HashmapBase<KEY, VALUE, N, HASH, EQUAL> {
     /// @param key the key to search for.
     /// @returns the value of the entry that is equal to `value`, or no value if the entry was not
     ///          found.
-    std::optional<Value> Get(const Key& key) const {
+    template <typename K>
+    std::optional<Value> Get(K&& key) const {
         if (auto [found, index] = this->IndexOf(key); found) {
             return this->slots_[index].entry->value;
         }
@@ -169,18 +183,24 @@ class Hashmap : public HashmapBase<KEY, VALUE, N, HASH, EQUAL> {
     /// @param key the entry's key value to search for.
     /// @returns the value of the entry.
     template <typename K>
-    Reference GetOrZero(K&& key) {
+    auto GetOrZero(K&& key) {
         auto res = Add(std::forward<K>(key), Value{});
-        return Reference(*this, key, res.value);
+        return Reference<ReferenceKeyType<K>>(*this, key, res.value);
     }
 
     /// @param key the key to search for.
     /// @returns a reference to the entry that is equal to the given value.
-    Reference Find(const Key& key) { return Reference(*this, key); }
+    template <typename K>
+    auto Find(K&& key) {
+        return Reference<ReferenceKeyType<K>>(*this, std::forward<K>(key));
+    }
 
     /// @param key the key to search for.
     /// @returns a reference to the entry that is equal to the given value.
-    ConstReference Find(const Key& key) const { return ConstReference(*this, key); }
+    template <typename K>
+    auto Find(K&& key) const {
+        return ConstReference<ReferenceKeyType<K>>(*this, std::forward<K>(key));
+    }
 
     /// @returns the keys of the map as a vector.
     /// @note the order of the returned vector is non-deterministic between compilers.
@@ -232,14 +252,16 @@ class Hashmap : public HashmapBase<KEY, VALUE, N, HASH, EQUAL> {
     }
 
   private:
-    Value* Lookup(const Key& key) {
+    template <typename K>
+    Value* Lookup(K&& key) {
         if (auto [found, index] = this->IndexOf(key); found) {
             return &this->slots_[index].entry->value;
         }
         return nullptr;
     }
 
-    const Value* Lookup(const Key& key) const {
+    template <typename K>
+    const Value* Lookup(K&& key) const {
         if (auto [found, index] = this->IndexOf(key); found) {
             return &this->slots_[index].entry->value;
         }

src/tint/utils/hashmap_base.h

@@ -106,7 +106,7 @@ template <typename KEY,
           typename VALUE,
           size_t N,
           typename HASH = Hasher<KEY>,
-          typename EQUAL = std::equal_to<KEY>>
+          typename EQUAL = EqualTo<KEY>>
 class HashmapBase {
     static constexpr bool ValueIsVoid = std::is_same_v<VALUE, void>;
 
@@ -157,8 +157,9 @@ class HashmapBase {
     /// A slot can either be empty or filled with a value. If the slot is empty, #hash and #distance
     /// will be zero.
     struct Slot {
-        bool Equals(size_t key_hash, const Key& key) const {
+        template <typename K>
-            return key_hash == hash && EQUAL()(key, KeyOf(*entry));
+        bool Equals(size_t key_hash, K&& key) const {
+            return key_hash == hash && EQUAL()(std::forward<K>(key), KeyOf(*entry));
         }
 
         /// The slot value. If this does not contain a value, then the slot is vacant.
@@ -502,8 +503,9 @@ class HashmapBase {
     /// @param key the key to hash
     /// @returns a tuple holding the target slot index for the given value, and the hash of the
     /// value, respectively.
-    HashResult Hash(const Key& key) const {
+    template <typename K>
-        size_t hash = HASH()(key);
+    HashResult Hash(K&& key) const {
+        size_t hash = HASH()(std::forward<K>(key));
         size_t index = Wrap(hash);
         return {index, hash};
     }
@@ -512,7 +514,8 @@ class HashmapBase {
     /// @param key the key to search for.
     /// @returns a tuple holding a boolean representing whether the key was found in the map, and
     /// if found, the index of the slot that holds the key.
-    std::tuple<bool, size_t> IndexOf(const Key& key) const {
+    template <typename K>
+    std::tuple<bool, size_t> IndexOf(K&& key) const {
         const auto hash = Hash(key);
         const auto count = slots_.Length();
         for (size_t distance = 0, index = hash.scan_start; distance < count; distance++) {

src/tint/utils/hashmap_test.cc

@@ -155,6 +155,134 @@ TEST(Hashmap, Index) {
     EXPECT_FALSE(one);
 }
 
+TEST(Hashmap, StringKeys) {
+    Hashmap<std::string, int, 4> map;
+    EXPECT_FALSE(map.Find("zero"));
+    EXPECT_FALSE(map.Find(std::string("zero")));
+    EXPECT_FALSE(map.Find(std::string_view("zero")));
+
+    map.Add("three", 3);
+    auto three_cstr = map.Find("three");
+    auto three_str = map.Find(std::string("three"));
+    auto three_sv = map.Find(std::string_view("three"));
+    map.Add(std::string("two"), 2);
+    auto two_cstr = map.Find("two");
+    auto two_str = map.Find(std::string("two"));
+    auto two_sv = map.Find(std::string_view("two"));
+    map.Add("four", 4);
+    auto four_cstr = map.Find("four");
+    auto four_str = map.Find(std::string("four"));
+    auto four_sv = map.Find(std::string_view("four"));
+    map.Add(std::string("eight"), 8);
+    auto eight_cstr = map.Find("eight");
+    auto eight_str = map.Find(std::string("eight"));
+    auto eight_sv = map.Find(std::string_view("eight"));
+
+    ASSERT_TRUE(three_cstr);
+    ASSERT_TRUE(three_str);
+    ASSERT_TRUE(three_sv);
+    ASSERT_TRUE(two_cstr);
+    ASSERT_TRUE(two_str);
+    ASSERT_TRUE(two_sv);
+    ASSERT_TRUE(four_cstr);
+    ASSERT_TRUE(four_str);
+    ASSERT_TRUE(four_sv);
+    ASSERT_TRUE(eight_cstr);
+    ASSERT_TRUE(eight_str);
+    ASSERT_TRUE(eight_sv);
+
+    EXPECT_EQ(*three_cstr, 3);
+    EXPECT_EQ(*three_str, 3);
+    EXPECT_EQ(*three_sv, 3);
+    EXPECT_EQ(*two_cstr, 2);
+    EXPECT_EQ(*two_str, 2);
+    EXPECT_EQ(*two_sv, 2);
+    EXPECT_EQ(*four_cstr, 4);
+    EXPECT_EQ(*four_str, 4);
+    EXPECT_EQ(*four_sv, 4);
+    EXPECT_EQ(*eight_cstr, 8);
+    EXPECT_EQ(*eight_str, 8);
+    EXPECT_EQ(*eight_sv, 8);
+
+    map.Add("zero", 0);  // Note: Find called before Add() is okay!
+    auto zero_cstr = map.Find("zero");
+    auto zero_str = map.Find(std::string("zero"));
+    auto zero_sv = map.Find(std::string_view("zero"));
+
+    map.Add(std::string("five"), 5);
+    auto five_cstr = map.Find("five");
+    auto five_str = map.Find(std::string("five"));
+    auto five_sv = map.Find(std::string_view("five"));
+    map.Add("six", 6);
+    auto six_cstr = map.Find("six");
+    auto six_str = map.Find(std::string("six"));
+    auto six_sv = map.Find(std::string_view("six"));
+    map.Add("one", 1);
+    auto one_cstr = map.Find("one");
+    auto one_str = map.Find(std::string("one"));
+    auto one_sv = map.Find(std::string_view("one"));
+    map.Add(std::string("seven"), 7);
+    auto seven_cstr = map.Find("seven");
+    auto seven_str = map.Find(std::string("seven"));
+    auto seven_sv = map.Find(std::string_view("seven"));
+
+    ASSERT_TRUE(zero_cstr);
+    ASSERT_TRUE(zero_str);
+    ASSERT_TRUE(zero_sv);
+    ASSERT_TRUE(three_cstr);
+    ASSERT_TRUE(three_str);
+    ASSERT_TRUE(three_sv);
+    ASSERT_TRUE(two_cstr);
+    ASSERT_TRUE(two_str);
+    ASSERT_TRUE(two_sv);
+    ASSERT_TRUE(four_cstr);
+    ASSERT_TRUE(four_str);
+    ASSERT_TRUE(four_sv);
+    ASSERT_TRUE(eight_cstr);
+    ASSERT_TRUE(eight_str);
+    ASSERT_TRUE(eight_sv);
+    ASSERT_TRUE(five_cstr);
+    ASSERT_TRUE(five_str);
+    ASSERT_TRUE(five_sv);
+    ASSERT_TRUE(six_cstr);
+    ASSERT_TRUE(six_str);
+    ASSERT_TRUE(six_sv);
+    ASSERT_TRUE(one_cstr);
+    ASSERT_TRUE(one_str);
+    ASSERT_TRUE(one_sv);
+    ASSERT_TRUE(seven_cstr);
+    ASSERT_TRUE(seven_str);
+    ASSERT_TRUE(seven_sv);
+
+    EXPECT_EQ(*zero_cstr, 0);
+    EXPECT_EQ(*zero_str, 0);
+    EXPECT_EQ(*zero_sv, 0);
+    EXPECT_EQ(*three_cstr, 3);
+    EXPECT_EQ(*three_str, 3);
+    EXPECT_EQ(*three_sv, 3);
+    EXPECT_EQ(*two_cstr, 2);
+    EXPECT_EQ(*two_str, 2);
+    EXPECT_EQ(*two_sv, 2);
+    EXPECT_EQ(*four_cstr, 4);
+    EXPECT_EQ(*four_str, 4);
+    EXPECT_EQ(*four_sv, 4);
+    EXPECT_EQ(*eight_cstr, 8);
+    EXPECT_EQ(*eight_str, 8);
+    EXPECT_EQ(*eight_sv, 8);
+    EXPECT_EQ(*five_cstr, 5);
+    EXPECT_EQ(*five_str, 5);
+    EXPECT_EQ(*five_sv, 5);
+    EXPECT_EQ(*six_cstr, 6);
+    EXPECT_EQ(*six_str, 6);
+    EXPECT_EQ(*six_sv, 6);
+    EXPECT_EQ(*one_cstr, 1);
+    EXPECT_EQ(*one_str, 1);
+    EXPECT_EQ(*one_sv, 1);
+    EXPECT_EQ(*seven_cstr, 7);
+    EXPECT_EQ(*seven_str, 7);
+    EXPECT_EQ(*seven_sv, 7);
+}
+
 TEST(Hashmap, Iterator) {
     using Map = Hashmap<int, std::string, 8>;
     using Entry = typename Map::Entry;

src/tint/utils/traits.h

@@ -159,16 +159,16 @@ using SliceTuple =
 
 namespace detail {
 /// Base template for IsTypeIn
-template <class T, class TypeList>
+template <typename T, typename TypeList>
 struct IsTypeIn;
 
 /// Specialization for IsTypeIn
-template <class T, template <class...> class TypeContainer, class... Ts>
+template <typename T, template <typename...> typename TypeContainer, typename... Ts>
 struct IsTypeIn<T, TypeContainer<Ts...>> : std::disjunction<std::is_same<T, Ts>...> {};
 }  // namespace detail
 
 /// Evaluates to true if T is one of the types in the TypeContainer's template arguments.
-/// Works for std::variant, std::tuple, std::pair, or any class template where all parameters are
+/// Works for std::variant, std::tuple, std::pair, or any typename template where all parameters are
 /// types.
 template <typename T, typename TypeContainer>
 static constexpr bool IsTypeIn = detail::IsTypeIn<T, TypeContainer>::value;
@@ -183,6 +183,32 @@ static constexpr bool IsStringLike =
     std::is_same_v<Decay<T>, std::string> || std::is_same_v<Decay<T>, std::string_view> ||
     std::is_same_v<Decay<T>, const char*>;
 
+namespace detail {
+/// Helper for CharArrayToCharPtr
+template <typename T>
+struct CharArrayToCharPtrImpl {
+    /// Evaluates to T
+    using type = T;
+};
+/// Specialization of CharArrayToCharPtrImpl for `char[N]`
+template <size_t N>
+struct CharArrayToCharPtrImpl<char[N]> {
+    /// Evaluates to `char*`
+    using type = char*;
+};
+/// Specialization of CharArrayToCharPtrImpl for `const char[N]`
+template <size_t N>
+struct CharArrayToCharPtrImpl<const char[N]> {
+    /// Evaluates to `const char*`
+    using type = const char*;
+};
+}  // namespace detail
+
+/// Evaluates to `char*` or `const char*` if `T` is `char[N]` or `const char[N]`, respectively,
+/// otherwise T.
+template <typename T>
+using CharArrayToCharPtr = typename detail::CharArrayToCharPtrImpl<T>::type;
+
 }  // namespace tint::utils::traits
 
 #endif  // SRC_TINT_UTILS_TRAITS_H_

src/tint/utils/traits_test.cc

@@ -241,4 +241,9 @@ TEST(SliceTuple, MixedTupleSliceHighPart) {
     static_assert(std::is_same_v<std::tuple_element_t<1, sliced>, float>);
 }
 
+static_assert(std::is_same_v<char*, CharArrayToCharPtr<char[2]>>);
+static_assert(std::is_same_v<const char*, CharArrayToCharPtr<const char[2]>>);
+static_assert(std::is_same_v<int, CharArrayToCharPtr<int>>);
+static_assert(std::is_same_v<int[2], CharArrayToCharPtr<int[2]>>);
+
 }  // namespace tint::utils::traits