2 files changed, 51 insertions, 0 deletions

diff --git a/source/blender/blenlib/BLI_map.hh b/source/blender/blenlib/BLI_map.hh
index 5819430fe97..82415dd5726 100644
--- a/source/blender/blenlib/BLI_map.hh
+++ b/source/blender/blenlib/BLI_map.hh
@@ -73,6 +73,7 @@
 #include "BLI_hash.hh"
 #include "BLI_hash_tables.hh"
 #include "BLI_map_slots.hh"
+#include "BLI_optional.hh"
 #include "BLI_probing_strategies.hh"
 
 namespace blender {
@@ -388,6 +389,23 @@ class Map {
   }
 
   /**
+   * Get the value that is stored for the given key and remove it from the map. If the key is not
+   * in the map, a value-less optional is returned.
+   */
+  Optional<Value> pop_try(const Key &key)
+  {
+    return this->pop_try_as(key);
+  }
+
+  /**
+   * Same as `pop_try`, but accepts other key types that are supported by the hash function.
+   */
+  template<typename ForwardKey> Optional<Value> pop_try_as(const ForwardKey &key)
+  {
+    return this->pop_try__impl(key, m_hash(key));
+  }
+
+  /**
    * This method can be used to implement more complex custom behavior without having to do
    * multiple lookups
    *
@@ -1029,6 +1047,22 @@ class Map {
     MAP_SLOT_PROBING_END();
   }
 
+  template<typename ForwardKey> Optional<Value> pop_try__impl(const ForwardKey &key, uint32_t hash)
+  {
+    MAP_SLOT_PROBING_BEGIN (hash, slot) {
+      if (slot.contains(key, m_is_equal, hash)) {
+        Optional<Value> value = *slot.value();
+        slot.remove();
+        m_removed_slots++;
+        return value;
+      }
+      if (slot.is_empty()) {
+        return {};
+      }
+    }
+    MAP_SLOT_PROBING_END();
+  }
+
   template<typename ForwardKey, typename CreateValueF, typename ModifyValueF>
   auto add_or_modify__impl(ForwardKey &&key,
                            const CreateValueF &create_value,
diff --git a/tests/gtests/blenlib/BLI_map_test.cc b/tests/gtests/blenlib/BLI_map_test.cc
index aa8903ed8ff..aa064e2f427 100644
--- a/tests/gtests/blenlib/BLI_map_test.cc
+++ b/tests/gtests/blenlib/BLI_map_test.cc
@@ -82,6 +82,23 @@ TEST(map, PopItem)
   EXPECT_FALSE(map.contains(1));
 }
 
+TEST(map, PopTry)
+{
+  Map<int, int> map;
+  map.add(1, 5);
+  map.add(2, 7);
+  EXPECT_EQ(map.size(), 2);
+  Optional<int> value = map.pop_try(4);
+  EXPECT_EQ(map.size(), 2);
+  EXPECT_FALSE(value.has_value());
+  value = map.pop_try(2);
+  EXPECT_EQ(map.size(), 1);
+  EXPECT_TRUE(value.has_value());
+  EXPECT_EQ(value.value(), 7);
+  EXPECT_EQ(*map.pop_try(1), 5);
+  EXPECT_EQ(map.size(), 0);
+}
+
 TEST(map, PopItemMany)
 {
   Map<int, int> map;