From 7041568ff922318f462b06e1c833cc9287aaa04b Mon Sep 17 00:00:00 2001
From: Jacques Lucke <jacques@blender.org>
Date: Fri, 30 Apr 2021 11:09:19 +0200
Subject: 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.
---
 source/blender/blenlib/BLI_vector_set.hh           | 45 ++++++++++++++++++++++
 .../blender/blenlib/tests/BLI_vector_set_test.cc   | 26 +++++++++++++
 2 files changed, 71 insertions(+)

(limited to 'source')

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
@@ -397,6 +397,23 @@ class VectorSet {
     return this->index_of_try__impl(key, hash_(key));
   }
 
+  /**
+   * 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.
    */
@@ -483,6 +500,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);
diff --git a/source/blender/blenlib/tests/BLI_vector_set_test.cc b/source/blender/blenlib/tests/BLI_vector_set_test.cc
index bbbe96f9b7e..c535ac57774 100644
--- a/source/blender/blenlib/tests/BLI_vector_set_test.cc
+++ b/source/blender/blenlib/tests/BLI_vector_set_test.cc
@@ -232,4 +232,30 @@ TEST(vector_set, PopExceptions)
   EXPECT_EQ(set.size(), 4);
 }
 
+TEST(vector_set, IndexOfOrAdd)
+{
+  VectorSet<int> set;
+  EXPECT_EQ(set.index_of_or_add(3), 0);
+  EXPECT_EQ(set.index_of_or_add(3), 0);
+  EXPECT_EQ(set.index_of_or_add(2), 1);
+  EXPECT_EQ(set.index_of_or_add(0), 2);
+  EXPECT_EQ(set.index_of_or_add(2), 1);
+  EXPECT_EQ(set.index_of_or_add(3), 0);
+  EXPECT_EQ(set.index_of_or_add(5), 3);
+  EXPECT_EQ(set.index_of_or_add(8), 4);
+  EXPECT_EQ(set.index_of_or_add(5), 3);
+}
+
+TEST(vector_set, Clear)
+{
+  VectorSet<int> set = {4, 6, 2, 4};
+  EXPECT_EQ(set.size(), 3);
+  set.clear();
+  EXPECT_EQ(set.size(), 0);
+  set.add_multiple({4, 1, 6, 8, 3, 6, 9, 3});
+  EXPECT_EQ(set.size(), 6);
+  set.clear();
+  EXPECT_EQ(set.size(), 0);
+}
+
 }  // namespace blender::tests
-- 
cgit v1.2.3