BLI: add methods to lookup a stored key in a set

2 files changed, 130 insertions, 0 deletions

diff --git a/source/blender/blenlib/BLI_set.hh b/source/blender/blenlib/BLI_set.hh
index 09b2d170eea..bf981154da4 100644
--- a/source/blender/blenlib/BLI_set.hh
+++ b/source/blender/blenlib/BLI_set.hh
@@ -313,6 +313,64 @@ class Set {
   }
 
   /**
+   * 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 set.
+   */
+  const Key &lookup_key(const Key &key) const
+  {
+    return this->lookup_key_as(key);
+  }
+
+  /**
+   * Same as `lookup_key`, but accepts other key types that are supported by the hash function.
+   */
+  template<typename ForwardKey> const Key &lookup_key_as(const ForwardKey &key) const
+  {
+    return this->lookup_key__impl(key, hash_(key));
+  }
+
+  /**
+   * Returns the key that is stored in the set that compares equal to the given key. If the key is
+   * not in the set, the given default value is returned instead.
+   */
+  const Key &lookup_key_default(const Key &key, const Key &default_value) const
+  {
+    return this->lookup_key_default_as(key, default_value);
+  }
+
+  /**
+   * Same as `lookup_key_default`, but accepts other key types that are supported by the hash
+   * function.
+   */
+  template<typename ForwardKey>
+  const Key &lookup_key_default_as(const ForwardKey &key, const Key &default_key) const
+  {
+    const Key *ptr = this->lookup_key_ptr__impl(key, hash_(key));
+    if (ptr == nullptr) {
+      return default_key;
+    }
+    return *ptr;
+  }
+
+  /**
+   * Returns a pointer to the key that is stored in the set that compares equal to the given key.
+   * If the key is not in the set, nullptr is returned instead.
+   */
+  const Key *lookup_key_ptr(const Key &key) const
+  {
+    return this->lookup_key_ptr_as(key);
+  }
+
+  /**
+   * Same as `lookup_key_ptr`, but accepts other key types that are supported by the hash
+   * function.
+   */
+  template<typename ForwardKey> const Key *lookup_key_ptr_as(const ForwardKey &key) const
+  {
+    return this->lookup_key_ptr__impl(key, hash_(key));
+  }
+
+  /**
    * Deletes the key from the set. Returns true when the key did exist beforehand, otherwise false.
    *
    * This is similar to std::unordered_set::erase.
@@ -596,6 +654,33 @@ class Set {
     SET_SLOT_PROBING_END();
   }
 
+  template<typename ForwardKey>
+  const Key &lookup_key__impl(const ForwardKey &key, const uint32_t hash) const
+  {
+    BLI_assert(this->contains_as(key));
+
+    SET_SLOT_PROBING_BEGIN (hash, slot) {
+      if (slot.contains(key, is_equal_, hash)) {
+        return *slot.key();
+      }
+    }
+    SET_SLOT_PROBING_END();
+  }
+
+  template<typename ForwardKey>
+  const Key *lookup_key_ptr__impl(const ForwardKey &key, const uint32_t hash) const
+  {
+    SET_SLOT_PROBING_BEGIN (hash, slot) {
+      if (slot.contains(key, is_equal_, hash)) {
+        return slot.key();
+      }
+      if (slot.is_empty()) {
+        return nullptr;
+      }
+    }
+    SET_SLOT_PROBING_END();
+  }
+
   template<typename ForwardKey> void add_new__impl(ForwardKey &&key, const uint32_t hash)
   {
     BLI_assert(!this->contains_as(key));
diff --git a/tests/gtests/blenlib/BLI_set_test.cc b/tests/gtests/blenlib/BLI_set_test.cc
index ac78eb786df..7d7ec401a68 100644
--- a/tests/gtests/blenlib/BLI_set_test.cc
+++ b/tests/gtests/blenlib/BLI_set_test.cc
@@ -403,6 +403,51 @@ TEST(set, IntrusiveIntKey)
   EXPECT_TRUE(set.remove(4));
 }
 
+struct MyKeyType {
+  uint32_t key;
+  uint32_t attached_data;
+
+  uint32_t hash() const
+  {
+    return key;
+  }
+
+  friend bool operator==(const MyKeyType &a, const MyKeyType &b)
+  {
+    return a.key == b.key;
+  }
+};
+
+TEST(set, LookupKey)
+{
+  Set<MyKeyType> set;
+  set.add({1, 10});
+  set.add({2, 20});
+  EXPECT_EQ(set.lookup_key({1, 30}).attached_data, 10);
+  EXPECT_EQ(set.lookup_key({2, 0}).attached_data, 20);
+}
+
+TEST(set, LookupKeyDefault)
+{
+  Set<MyKeyType> set;
+  set.add({1, 10});
+  set.add({2, 20});
+
+  MyKeyType fallback{5, 50};
+  EXPECT_EQ(set.lookup_key_default({1, 66}, fallback).attached_data, 10);
+  EXPECT_EQ(set.lookup_key_default({4, 40}, fallback).attached_data, 50);
+}
+
+TEST(set, LookupKeyPtr)
+{
+  Set<MyKeyType> set;
+  set.add({1, 10});
+  set.add({2, 20});
+  EXPECT_EQ(set.lookup_key_ptr({1, 50})->attached_data, 10);
+  EXPECT_EQ(set.lookup_key_ptr({2, 50})->attached_data, 20);
+  EXPECT_EQ(set.lookup_key_ptr({3, 50}), nullptr);
+}
+
 /**
  * Set this to 1 to activate the benchmark. It is disabled by default, because it prints a lot.
  */