BLI: simplify removing elements from containers with predicate

Previously removing elements based on a predicate was a bit cumbersome,
especially for hash tables. Now there is a new `remove_if` method in some
data structures which is similar to `std::erase_if`. We could consider adding
`blender::erase_if` in the future to more closely mimic the standard library,
but for now this is using the api design of the surrounding code is used.

8 files changed, 120 insertions, 0 deletions

diff --git a/source/blender/blenlib/BLI_map.hh b/source/blender/blenlib/BLI_map.hh
index 4322e3c0f2d..d5bc53c4d6c 100644
--- a/source/blender/blenlib/BLI_map.hh
+++ b/source/blender/blenlib/BLI_map.hh
@@ -887,6 +887,25 @@ class Map {
   }
 
   /**
+   * Remove all key-value-pairs for that the given predicate is true.
+   *
+   * This is similar to std::erase_if.
+   */
+  template<typename Predicate> void remove_if(Predicate &&predicate)
+  {
+    for (Slot &slot : slots_) {
+      if (slot.is_occupied()) {
+        const Key &key = *slot.key();
+        Value &value = *slot.value();
+        if (predicate(MutableItem{key, value})) {
+          slot.remove();
+          removed_slots_++;
+        }
+      }
+    }
+  }
+
+  /**
    * Print common statistics like size and collision count. This is useful for debugging purposes.
    */
   void print_stats(StringRef name = "") const
diff --git a/source/blender/blenlib/BLI_set.hh b/source/blender/blenlib/BLI_set.hh
index a5f62001798..afe18529762 100644
--- a/source/blender/blenlib/BLI_set.hh
+++ b/source/blender/blenlib/BLI_set.hh
@@ -493,6 +493,24 @@ class Set {
   }
 
   /**
+   * Remove all values for which the given predicate is true.
+   *
+   * This is similar to std::erase_if.
+   */
+  template<typename Predicate> void remove_if(Predicate &&predicate)
+  {
+    for (Slot &slot : slots_) {
+      if (slot.is_occupied()) {
+        const Key &key = *slot.key();
+        if (predicate(key)) {
+          slot.remove();
+          removed_slots_++;
+        }
+      }
+    }
+  }
+
+  /**
    * Print common statistics like size and collision count. This is useful for debugging purposes.
    */
   void print_stats(StringRef name = "") const
diff --git a/source/blender/blenlib/BLI_vector.hh b/source/blender/blenlib/BLI_vector.hh
index d07c9aeeb3d..cba58945546 100644
--- a/source/blender/blenlib/BLI_vector.hh
+++ b/source/blender/blenlib/BLI_vector.hh
@@ -805,6 +805,17 @@ class Vector {
   }
 
   /**
+   * Remove all values for which the given predicate is true.
+   *
+   * This is similar to std::erase_if.
+   */
+  template<typename Predicate> void remove_if(Predicate &&predicate)
+  {
+    end_ = std::remove_if(this->begin(), this->end(), predicate);
+    UPDATE_VECTOR_SIZE(this);
+  }
+
+  /**
    * Do a linear search to find the value in the vector.
    * When found, return the first index, otherwise return -1.
    */
diff --git a/source/blender/blenlib/BLI_vector_set.hh b/source/blender/blenlib/BLI_vector_set.hh
index 9e8cf5af59f..d182e1f1678 100644
--- a/source/blender/blenlib/BLI_vector_set.hh
+++ b/source/blender/blenlib/BLI_vector_set.hh
@@ -349,6 +349,25 @@ class VectorSet {
   }
 
   /**
+   * Remove all values for which the given predicate is true. This may change the order of elements
+   * in the vector.
+   *
+   * This is similar to std::erase_if.
+   */
+  template<typename Predicate> void remove_if(Predicate &&predicate)
+  {
+    for (Slot &slot : slots_) {
+      if (slot.is_occupied()) {
+        const int64_t index = slot.index();
+        const Key &key = keys_[index];
+        if (predicate(key)) {
+          this->remove_key_internal(slot);
+        }
+      }
+    }
+  }
+
+  /**
    * Delete and return a key from the set. This will remove the last element in the vector. The
    * order of the remaining elements in the set is not changed.
    */
diff --git a/source/blender/blenlib/tests/BLI_map_test.cc b/source/blender/blenlib/tests/BLI_map_test.cc
index dbbd4843abc..0ca07309836 100644
--- a/source/blender/blenlib/tests/BLI_map_test.cc
+++ b/source/blender/blenlib/tests/BLI_map_test.cc
@@ -640,6 +640,24 @@ TEST(map, RemoveDuringIteration)
   EXPECT_EQ(map.lookup(3), 3);
 }
 
+TEST(map, RemoveIf)
+{
+  Map<int64_t, int64_t> map;
+  for (const int64_t i : IndexRange(100)) {
+    map.add(i * i, i);
+  }
+  map.remove_if([](auto item) { return item.key > 100; });
+  EXPECT_EQ(map.size(), 11);
+  for (const int64_t i : IndexRange(100)) {
+    if (i <= 10) {
+      EXPECT_EQ(map.lookup(i * i), i);
+    }
+    else {
+      EXPECT_FALSE(map.contains(i * i));
+    }
+  }
+}
+
 TEST(map, LookupKey)
 {
   Map<std::string, int> map;
diff --git a/source/blender/blenlib/tests/BLI_set_test.cc b/source/blender/blenlib/tests/BLI_set_test.cc
index 9dfa48b5822..79439eb9079 100644
--- a/source/blender/blenlib/tests/BLI_set_test.cc
+++ b/source/blender/blenlib/tests/BLI_set_test.cc
@@ -576,6 +576,19 @@ TEST(set, RemoveDuringIteration)
   EXPECT_TRUE(set.contains(3));
 }
 
+TEST(set, RemoveIf)
+{
+  Set<int64_t> set;
+  for (const int64_t i : IndexRange(100)) {
+    set.add(i * i);
+  }
+  set.remove_if([](const int64_t key) { return key > 100; });
+  EXPECT_EQ(set.size(), 11);
+  for (const int64_t i : IndexRange(100)) {
+    EXPECT_EQ(set.contains(i * i), i <= 10);
+  }
+}
+
 /**
  * Set this to 1 to activate the benchmark. It is disabled by default, because it prints a lot.
  */
diff --git a/source/blender/blenlib/tests/BLI_vector_set_test.cc b/source/blender/blenlib/tests/BLI_vector_set_test.cc
index 34685c76e00..c3fbac1200a 100644
--- a/source/blender/blenlib/tests/BLI_vector_set_test.cc
+++ b/source/blender/blenlib/tests/BLI_vector_set_test.cc
@@ -128,6 +128,19 @@ TEST(vector_set, RemoveContained)
   EXPECT_EQ(set.size(), 0);
 }
 
+TEST(vector_set, RemoveIf)
+{
+  VectorSet<int64_t> set;
+  for (const int64_t i : IndexRange(100)) {
+    set.add(i * i);
+  }
+  set.remove_if([](const int64_t key) { return key % 2 == 0; });
+  EXPECT_EQ(set.size(), 50);
+  for (const int64_t i : IndexRange(100)) {
+    EXPECT_EQ(set.contains(i * i), i % 2 == 1);
+  }
+}
+
 TEST(vector_set, AddMultipleTimes)
 {
   VectorSet<int> set;
diff --git a/source/blender/blenlib/tests/BLI_vector_test.cc b/source/blender/blenlib/tests/BLI_vector_test.cc
index bbc4cecee4a..53e8cd1047b 100644
--- a/source/blender/blenlib/tests/BLI_vector_test.cc
+++ b/source/blender/blenlib/tests/BLI_vector_test.cc
@@ -419,6 +419,15 @@ TEST(vector, Remove)
   EXPECT_EQ(vec.begin(), vec.end());
 }
 
+TEST(vector, RemoveIf)
+{
+  Vector<int> vec = {1, 2, 3, 4, 5, 6, 7, 8};
+  vec.remove_if([](const int x) { return x % 2 == 0; });
+  const Vector<int> expected_vec = {1, 3, 5, 7};
+  EXPECT_EQ(vec.size(), expected_vec.size());
+  EXPECT_EQ_ARRAY(vec.data(), expected_vec.data(), size_t(vec.size()));
+}
+
 TEST(vector, ExtendSmallVector)
 {
   Vector<int> a = {2, 3, 4};