From 313a1c072d6d730e4665350befb6cfae5e03b1aa Mon Sep 17 00:00:00 2001
From: Jacques Lucke <jacques@blender.org>
Date: Fri, 30 Apr 2021 16:22:24 +0200
Subject: 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.
---
 source/blender/blenlib/BLI_map.hh            | 13 +++++++++++++
 source/blender/blenlib/tests/BLI_map_test.cc | 27 +++++++++++++++++++++++++++
 2 files changed, 40 insertions(+)

(limited to 'source/blender/blenlib')

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
@@ -856,6 +856,19 @@ class Map {
     return MutableItemIterator(slots_.data(), slots_.size(), 0);
   }
 
+  /**
+   * 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.
    */
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.
  */
-- 
cgit v1.2.3