BLI: support removing Map elements during iteration

While it was technically safe to call Map.remove while iterating over
a map, it wasn't really designed to work. Also it wasn't very efficient,
because to remove the element, the map would have to search it
again. Now it is possible to remove an element given an iterator
into the map. It is safe to remove the element while iterating over
the map. Obviously, the removed element must not be accessed
anymore after it has been removed.

2 files changed, 40 insertions, 0 deletions

diff --git a/source/blender/blenlib/BLI_map.hh b/source/blender/blenlib/BLI_map.hh
index c69d5986885..95afbfc2ec6 100644
--- a/source/blender/blenlib/BLI_map.hh
+++ b/source/blender/blenlib/BLI_map.hh
@@ -857,6 +857,19 @@ class Map {
   }
 
   /**
+   * Remove the key-value-pair that the iterator is currently pointing at.
+   * It is valid to call this method while iterating over the map. However, after this method has
+   * been called, the removed element must not be accessed anymore.
+   */
+  void remove(const BaseIterator &iterator)
+  {
+    Slot &slot = iterator.current_slot();
+    BLI_assert(slot.is_occupied());
+    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/tests/BLI_map_test.cc b/source/blender/blenlib/tests/BLI_map_test.cc
index bb15f7f0d8d..18be456bd20 100644
--- a/source/blender/blenlib/tests/BLI_map_test.cc
+++ b/source/blender/blenlib/tests/BLI_map_test.cc
@@ -613,6 +613,33 @@ TEST(map, AddAsVariadic)
   EXPECT_EQ(map.lookup(2), "t");
 }
 
+TEST(map, RemoveDuringIteration)
+{
+  Map<int, int> map;
+  map.add(2, 1);
+  map.add(5, 2);
+  map.add(1, 2);
+  map.add(6, 0);
+  map.add(3, 3);
+
+  EXPECT_EQ(map.size(), 5);
+
+  using Iter = Map<int, int>::MutableItemIterator;
+  Iter begin = map.items().begin();
+  Iter end = map.items().end();
+  for (Iter iter = begin; iter != end; ++iter) {
+    Map<int, int>::MutableItem item = *iter;
+    if (item.value == 2) {
+      map.remove(iter);
+    }
+  }
+
+  EXPECT_EQ(map.size(), 3);
+  EXPECT_EQ(map.lookup(2), 1);
+  EXPECT_EQ(map.lookup(6), 0);
+  EXPECT_EQ(map.lookup(3), 3);
+}
+
 /**
  * Set this to 1 to activate the benchmark. It is disabled by default, because it prints a lot.
  */