BLI: new add_overwrite methods for Set and VectorSet

These methods can be useful when storing keys that contain more data than just
what affects their hash. This came up in #134000. `Map` already has a method
with the same name.

Pull Request: https://projects.blender.org/blender/blender/pulls/135456

source/blender/blenlib/BLI_set.hh

@@ -258,6 +258,25 @@ class Set {
     return this->add__impl(std::forward<ForwardKey>(key), hash_(key));
   }
 
+  /**
+   * Similar to #add but reinserts the key if it already exists. Using this only makes sense if the
+   * key contains additional data besides what affects the hash.
+   *
+   * \return True if the key was newly added, false if it was already present and was overwritten.
+   */
+  bool add_overwrite(const Key &key)
+  {
+    return this->add_overwrite_as(key);
+  }
+  bool add_overwrite(Key &&key)
+  {
+    return this->add_overwrite_as(std::move(key));
+  }
+  template<typename ForwardKey> bool add_overwrite_as(ForwardKey &&key)
+  {
+    return this->add_overwrite__impl(std::forward<ForwardKey>(key), hash_(key));
+  }
+
   /**
    * Convenience function to add many keys to the set at once. Duplicates are removed
    * automatically.
@@ -817,6 +836,27 @@ class Set {
     SET_SLOT_PROBING_END();
   }
 
+  template<typename ForwardKey> bool add_overwrite__impl(ForwardKey &&key, const uint64_t hash)
+  {
+    this->ensure_can_add();
+
+    SET_SLOT_PROBING_BEGIN (hash, slot) {
+      if (slot.is_empty()) {
+        slot.occupy(std::forward<ForwardKey>(key), hash);
+        BLI_assert(hash_(*slot.key()) == hash);
+        occupied_and_removed_slots_++;
+        return true;
+      }
+      if (slot.contains(key, is_equal_, hash)) {
+        Key &stored_key = *slot.key();
+        stored_key = std::forward<ForwardKey>(key);
+        BLI_assert(hash_(stored_key) == hash);
+        return false;
+      }
+    }
+    SET_SLOT_PROBING_END();
+  }
+
   template<typename ForwardKey> bool remove__impl(const ForwardKey &key, const uint64_t hash)
   {
     SET_SLOT_PROBING_BEGIN (hash, slot) {

source/blender/blenlib/BLI_vector_set.hh

@@ -293,6 +293,29 @@ class VectorSet {
     return this->add__impl(std::forward<ForwardKey>(key), hash_(key));
   }
 
+  /**
+   * Similar to #add but reinserts the key if it already exists. Using this only makes sense if the
+   * key contains additional data besides what affects the hash.
+   *
+   * \note This is different from first removing and then adding the key again, because
+   * #add_overwrite does not change the index where the value is stored. Removing an element can
+   * change the order of elements.
+   *
+   * \return True if the key was newly added, false if it was already present and was overwritten.
+   */
+  bool add_overwrite(const Key &key)
+  {
+    return this->add_overwrite_as(key);
+  }
+  bool add_overwrite(Key &&key)
+  {
+    return this->add_overwrite_as(std::move(key));
+  }
+  template<typename ForwardKey> bool add_overwrite_as(ForwardKey &&key)
+  {
+    return this->add_overwrite__impl(std::forward<ForwardKey>(key), hash_(key));
+  }
+
   /**
    * Convenience function to add many keys to the vector set at once. Duplicates are removed
    * automatically.
@@ -737,7 +760,7 @@ class VectorSet {
 
     VECTOR_SET_SLOT_PROBING_BEGIN (hash, slot) {
       if (slot.is_empty()) {
-        int64_t index = this->size();
+        const int64_t index = this->size();
         Key *dst = keys_ + index;
         new (dst) Key(std::forward<ForwardKey>(key));
         BLI_assert(hash_(*dst) == hash);
@@ -752,6 +775,31 @@ class VectorSet {
     VECTOR_SET_SLOT_PROBING_END();
   }
 
+  template<typename ForwardKey> bool add_overwrite__impl(ForwardKey &&key, const uint64_t hash)
+  {
+    this->ensure_can_add();
+
+    VECTOR_SET_SLOT_PROBING_BEGIN (hash, slot) {
+      if (slot.is_empty()) {
+        const int64_t index = this->size();
+        Key *dst = keys_ + index;
+        new (dst) Key(std::forward<ForwardKey>(key));
+        BLI_assert(hash_(*dst) == hash);
+        slot.occupy(index, hash);
+        occupied_and_removed_slots_++;
+        return true;
+      }
+      if (slot.contains(key, is_equal_, hash, keys_)) {
+        const int64_t index = slot.index();
+        Key &stored_key = keys_[index];
+        stored_key = std::forward<ForwardKey>(key);
+        BLI_assert(hash_(stored_key) == hash);
+        return false;
+      }
+    }
+    VECTOR_SET_SLOT_PROBING_END();
+  }
+
   template<typename ForwardKey>
   int64_t index_of__impl(const ForwardKey &key, const uint64_t hash) const
   {

source/blender/blenlib/tests/BLI_set_test.cc

@@ -626,6 +626,41 @@ TEST(set, Equality)
   EXPECT_NE(f, a);
 }
 
+namespace {
+struct KeyWithData {
+  int key;
+  std::string data;
+
+  uint64_t hash() const
+  {
+    return uint64_t(this->key);
+  }
+
+  friend bool operator==(const KeyWithData &a, const KeyWithData &b)
+  {
+    return a.key == b.key;
+  }
+};
+}  // namespace
+
+TEST(set, AddOverwrite)
+{
+  Set<KeyWithData> set;
+  EXPECT_TRUE(set.add_overwrite(KeyWithData{1, "a"}));
+  EXPECT_EQ(set.size(), 1);
+  EXPECT_FALSE(set.add(KeyWithData{1, "b"}));
+  EXPECT_EQ(set.size(), 1);
+  EXPECT_EQ(set.lookup_key(KeyWithData{1, "_"}).data, "a");
+  EXPECT_FALSE(set.add_overwrite(KeyWithData{1, "c"}));
+  EXPECT_EQ(set.size(), 1);
+  EXPECT_EQ(set.lookup_key(KeyWithData{1, "_"}).data, "c");
+
+  const KeyWithData key{2, "d"};
+  EXPECT_TRUE(set.add_overwrite(key));
+  EXPECT_EQ(set.size(), 2);
+  EXPECT_EQ(set.lookup_key(key).data, "d");
+}
+
 /**
  * Set this to 1 to activate the benchmark. It is disabled by default, because it prints a lot.
  */

source/blender/blenlib/tests/BLI_vector_set_test.cc

@@ -339,4 +339,44 @@ TEST(vector_set, CustomIDVectorSet)
   EXPECT_EQ(set.size(), 2);
 }
 
+namespace {
+struct KeyWithData {
+  int key;
+  std::string data;
+
+  uint64_t hash() const
+  {
+    return uint64_t(this->key);
+  }
+
+  friend bool operator==(const KeyWithData &a, const KeyWithData &b)
+  {
+    return a.key == b.key;
+  }
+};
+}  // namespace
+
+TEST(vector_set, AddOverwrite)
+{
+  VectorSet<KeyWithData> set;
+  EXPECT_TRUE(set.add_overwrite(KeyWithData{1, "a"}));
+  EXPECT_EQ(set.size(), 1);
+  EXPECT_EQ(set[0].data, "a");
+  EXPECT_FALSE(set.add(KeyWithData{1, "b"}));
+  EXPECT_EQ(set.size(), 1);
+  EXPECT_EQ(set[0].data, "a");
+  EXPECT_EQ(set.lookup_key(KeyWithData{1, "_"}).data, "a");
+  EXPECT_FALSE(set.add_overwrite(KeyWithData{1, "c"}));
+  EXPECT_EQ(set.size(), 1);
+  EXPECT_EQ(set[0].data, "c");
+  EXPECT_EQ(set.lookup_key(KeyWithData{1, "_"}).data, "c");
+
+  const KeyWithData key{2, "d"};
+  EXPECT_TRUE(set.add_overwrite(key));
+  EXPECT_EQ(set.size(), 2);
+  EXPECT_EQ(set[0].data, "c");
+  EXPECT_EQ(set[1].data, "d");
+  EXPECT_EQ(set.lookup_key(key).data, "d");
+}
+
 }  // namespace blender::tests