From db8a856a747042b6a544f4dc778219dcaaf6540f Mon Sep 17 00:00:00 2001
From: Jacques Lucke <jacques@blender.org>
Date: Mon, 12 Oct 2020 15:10:01 +0200
Subject: [PATCH] BLI: support looking up a key from a set or adding it when
 not existant

---
 source/blender/blenlib/BLI_set.hh            | 33 ++++++++++++++++++++
 source/blender/blenlib/tests/BLI_set_test.cc | 14 +++++++++
 2 files changed, 47 insertions(+)

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
@@ -350,6 +350,23 @@ class Set {
     return this->lookup_key_ptr__impl(key, hash_(key));
   }
 
+  /**
+   * 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.
    *
@@ -734,6 +751,22 @@ class Set {
     SET_SLOT_PROBING_END();
   }
 
+  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
   {
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;