BLI: support looking up a key stored in Map or VectorSet

Sometimes it is useful to find the key that compares equal
to a known key. Typically that happens when the key itself
has additional data attached that is not part of its hash.

Note that the returned key reference/pointer is const, because
the caller must not change the key in a way that changes its
hash or how it compares to other keys.

4 files changed, 89 insertions, 0 deletions

diff --git a/source/blender/blenlib/BLI_map.hh b/source/blender/blenlib/BLI_map.hh
index 95afbfc2ec6..4d254960f34 100644
--- a/source/blender/blenlib/BLI_map.hh
+++ b/source/blender/blenlib/BLI_map.hh
@@ -605,6 +605,37 @@ class Map {
   }
 
   /**
+   * Returns the key that is stored in the set that compares equal to the given key. This invokes
+   * undefined behavior when the key is not in the map.
+   */
+  const Key &lookup_key(const Key &key) const
+  {
+    return this->lookup_key_as(key);
+  }
+  template<typename ForwardKey> const Key &lookup_key_as(const ForwardKey &key) const
+  {
+    const Slot &slot = this->lookup_slot(key, hash_(key));
+    return *slot.key();
+  }
+
+  /**
+   * Returns a pointer to the key that is stored in the map that compares equal to the given key.
+   * If the key is not in the map, null is returned.
+   */
+  const Key *lookup_key_ptr(const Key &key) const
+  {
+    return this->lookup_key_ptr_as(key);
+  }
+  template<typename ForwardKey> const Key *lookup_key_ptr_as(const ForwardKey &key) const
+  {
+    const Slot *slot = this->lookup_slot_ptr(key, hash_(key));
+    if (slot == nullptr) {
+      return nullptr;
+    }
+    return slot->key();
+  }
+
+  /**
    * Calls the provided callback for every key-value-pair in the map. The callback is expected
    * to take a `const Key &` as first and a `const Value &` as second parameter.
    */
diff --git a/source/blender/blenlib/BLI_vector_set.hh b/source/blender/blenlib/BLI_vector_set.hh
index 2f5d04aefa2..567e4fd8128 100644
--- a/source/blender/blenlib/BLI_vector_set.hh
+++ b/source/blender/blenlib/BLI_vector_set.hh
@@ -415,6 +415,38 @@ class VectorSet {
   }
 
   /**
+   * Returns the key that is stored in the vector set that compares equal to the given key. This
+   * invokes undefined behavior when the key is not in the set.
+   */
+  const Key &lookup_key(const Key &key) const
+  {
+    return this->lookup_key_as(key);
+  }
+  template<typename ForwardKey> const Key &lookup_key_as(const ForwardKey &key) const
+  {
+    const Key *key_ptr = this->lookup_key_ptr_as(key);
+    BLI_assert(key_ptr != nullptr);
+    return *key_ptr;
+  }
+
+  /**
+   * Returns a pointer to the key that is stored in the vector set that compares equal to the given
+   * key. If the key is not in the set, null is returned.
+   */
+  const Key *lookup_key_ptr(const Key &key) const
+  {
+    return this->lookup_key_ptr_as(key);
+  }
+  template<typename ForwardKey> const Key *lookup_key_ptr_as(const ForwardKey &key) const
+  {
+    const int64_t index = this->index_of_try__impl(key, hash_(key));
+    if (index >= 0) {
+      return keys_ + index;
+    }
+    return nullptr;
+  }
+
+  /**
    * Get a pointer to the beginning of the array containing all keys.
    */
   const Key *data() const
diff --git a/source/blender/blenlib/tests/BLI_map_test.cc b/source/blender/blenlib/tests/BLI_map_test.cc
index 18be456bd20..679a10e9ce0 100644
--- a/source/blender/blenlib/tests/BLI_map_test.cc
+++ b/source/blender/blenlib/tests/BLI_map_test.cc
@@ -640,6 +640,19 @@ TEST(map, RemoveDuringIteration)
   EXPECT_EQ(map.lookup(3), 3);
 }
 
+TEST(map, LookupKey)
+{
+  Map<std::string, int> map;
+  map.add("a", 0);
+  map.add("b", 1);
+  map.add("c", 2);
+  EXPECT_EQ(map.lookup_key("a"), "a");
+  EXPECT_EQ(map.lookup_key_as("c"), "c");
+  EXPECT_EQ(map.lookup_key_ptr_as("d"), nullptr);
+  EXPECT_EQ(map.lookup_key_ptr_as("b")->size(), 1);
+  EXPECT_EQ(map.lookup_key_ptr("a"), map.lookup_key_ptr_as("a"));
+}
+
 /**
  * Set this to 1 to activate the benchmark. It is disabled by default, because it prints a lot.
  */
diff --git a/source/blender/blenlib/tests/BLI_vector_set_test.cc b/source/blender/blenlib/tests/BLI_vector_set_test.cc
index c535ac57774..c4016ca75e1 100644
--- a/source/blender/blenlib/tests/BLI_vector_set_test.cc
+++ b/source/blender/blenlib/tests/BLI_vector_set_test.cc
@@ -258,4 +258,17 @@ TEST(vector_set, Clear)
   EXPECT_EQ(set.size(), 0);
 }
 
+TEST(vector_set, LookupKey)
+{
+  VectorSet<std::string> set;
+  set.add("a");
+  set.add("b");
+  set.add("c");
+  EXPECT_EQ(set.lookup_key("a"), "a");
+  EXPECT_EQ(set.lookup_key_as("c"), "c");
+  EXPECT_EQ(set.lookup_key_ptr_as("d"), nullptr);
+  EXPECT_EQ(set.lookup_key_ptr_as("b")->size(), 1);
+  EXPECT_EQ(set.lookup_key_ptr("a"), set.lookup_key_ptr_as("a"));
+}
+
 }  // namespace blender::tests