diff options
-rw-r--r-- | source/blender/blenlib/BLI_set.hh | 85 | ||||
-rw-r--r-- | tests/gtests/blenlib/BLI_set_test.cc | 45 |
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. */ |