diff options
-rw-r--r-- | source/blender/blenlib/BLI_set.hh | 33 | ||||
-rw-r--r-- | source/blender/blenlib/tests/BLI_set_test.cc | 14 |
2 files changed, 47 insertions, 0 deletions
diff --git a/source/blender/blenlib/BLI_set.hh b/source/blender/blenlib/BLI_set.hh index 9684f372db7..f1cf44018c9 100644 --- a/source/blender/blenlib/BLI_set.hh +++ b/source/blender/blenlib/BLI_set.hh @@ -351,6 +351,23 @@ class Set { } /** + * Returns the key in the set that compares equal to the given key. If it does not exist, the key + * is newly added. + */ + const Key &lookup_key_or_add(const Key &key) + { + return this->lookup_key_or_add_as(key); + } + const Key &lookup_key_or_add(Key &&key) + { + return this->lookup_key_or_add_as(std::move(key)); + } + template<typename ForwardKey> const Key &lookup_key_or_add_as(ForwardKey &&key) + { + return this->lookup_key_or_add__impl(std::forward<ForwardKey>(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. @@ -735,6 +752,22 @@ class Set { } template<typename ForwardKey> + const Key &lookup_key_or_add__impl(ForwardKey &&key, const uint64_t hash) + { + SET_SLOT_PROBING_BEGIN (hash, slot) { + if (slot.contains(key, is_equal_, hash)) { + return *slot.key(); + } + if (slot.is_empty()) { + slot.occupy(std::forward<ForwardKey>(key), hash); + occupied_and_removed_slots_++; + return *slot.key(); + } + } + SET_SLOT_PROBING_END(); + } + + template<typename ForwardKey> int64_t count_collisions__impl(const ForwardKey &key, const uint64_t hash) const { int64_t collisions = 0; diff --git a/source/blender/blenlib/tests/BLI_set_test.cc b/source/blender/blenlib/tests/BLI_set_test.cc index bf49bfb611b..dbe820c9d15 100644 --- a/source/blender/blenlib/tests/BLI_set_test.cc +++ b/source/blender/blenlib/tests/BLI_set_test.cc @@ -453,6 +453,20 @@ TEST(set, LookupKeyPtr) EXPECT_EQ(set.lookup_key_ptr({3, 50}), nullptr); } +TEST(set, LookupKeyOrAdd) +{ + Set<MyKeyType> set; + set.add({1, 10}); + set.add({2, 20}); + EXPECT_EQ(set.size(), 2); + EXPECT_EQ(set.lookup_key_or_add({2, 40}).attached_data, 20); + EXPECT_EQ(set.size(), 2); + EXPECT_EQ(set.lookup_key_or_add({3, 40}).attached_data, 40); + EXPECT_EQ(set.size(), 3); + EXPECT_EQ(set.lookup_key_or_add({3, 60}).attached_data, 40); + EXPECT_EQ(set.size(), 3); +} + TEST(set, StringViewKeys) { Set<std::string_view> set; |