BLI: improve VectorSet data structure

This adds two new methods:
* `clear` just removes all keys from the vector set.
* `index_of_or_add` returns the index of a key and adds it if has not
  been added before.

1 files changed, 45 insertions, 0 deletions

diff --git a/source/blender/blenlib/BLI_vector_set.hh b/source/blender/blenlib/BLI_vector_set.hh
index 14b38d564cb..2f5d04aefa2 100644
--- a/source/blender/blenlib/BLI_vector_set.hh
+++ b/source/blender/blenlib/BLI_vector_set.hh
@@ -398,6 +398,23 @@ class VectorSet {
   }
 
   /**
+   * Return the index of the key in the vector. If the key is not in the set, add it and return its
+   * index.
+   */
+  int64_t index_of_or_add(const Key &key)
+  {
+    return this->index_of_or_add_as(key);
+  }
+  int64_t index_of_or_add(Key &&key)
+  {
+    return this->index_of_or_add_as(std::move(key));
+  }
+  template<typename ForwardKey> int64_t index_of_or_add_as(ForwardKey &&key)
+  {
+    return this->index_of_or_add__impl(std::forward<ForwardKey>(key), hash_(key));
+  }
+
+  /**
    * Get a pointer to the beginning of the array containing all keys.
    */
   const Key *data() const
@@ -484,6 +501,14 @@ class VectorSet {
   }
 
   /**
+   * Remove all keys from the vector set.
+   */
+  void clear()
+  {
+    this->noexcept_reset();
+  }
+
+  /**
    * Get the number of collisions that the probing strategy has to go through to find the key or
    * determine that it is not in the set.
    */
@@ -652,6 +677,26 @@ class VectorSet {
     VECTOR_SET_SLOT_PROBING_END();
   }
 
+  template<typename ForwardKey>
+  int64_t index_of_or_add__impl(ForwardKey &&key, const uint64_t hash)
+  {
+    this->ensure_can_add();
+
+    VECTOR_SET_SLOT_PROBING_BEGIN (hash, slot) {
+      if (slot.contains(key, is_equal_, hash, keys_)) {
+        return slot.index();
+      }
+      if (slot.is_empty()) {
+        const int64_t index = this->size();
+        new (keys_ + index) Key(std::forward<ForwardKey>(key));
+        slot.occupy(index, hash);
+        occupied_and_removed_slots_++;
+        return index;
+      }
+    }
+    VECTOR_SET_SLOT_PROBING_END();
+  }
+
   Key pop__impl()
   {
     BLI_assert(this->size() > 0);