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

1 files changed, 85 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));