AK: Define HashMap::take to find and remove a value from the map

2 files changed, 57 insertions, 0 deletions

diff --git a/AK/HashMap.h b/AK/HashMap.h
index ec8d946472..d3e1af9b63 100644
--- a/AK/HashMap.h
+++ b/AK/HashMap.h
@@ -199,6 +199,31 @@ public:
         m_table.remove(it);
     }
 
+    Optional<V> take(K const& key)
+    {
+        if (auto it = find(key); it != end()) {
+            auto value = move(it->value);
+            m_table.remove(it);
+
+            return value;
+        }
+
+        return {};
+    }
+
+    template<Concepts::HashCompatible<K> Key>
+    requires(IsSame<KeyTraits, Traits<K>>) Optional<V> take(Key const& key)
+    {
+        if (auto it = find(key); it != end()) {
+            auto value = move(it->value);
+            m_table.remove(it);
+
+            return value;
+        }
+
+        return {};
+    }
+
     V& ensure(K const& key)
     {
         auto it = find(key);
diff --git a/Tests/AK/TestHashMap.cpp b/Tests/AK/TestHashMap.cpp
index bbb4876c5d..6e70540a9f 100644
--- a/Tests/AK/TestHashMap.cpp
+++ b/Tests/AK/TestHashMap.cpp
@@ -9,6 +9,7 @@
 #include <AK/DeprecatedString.h>
 #include <AK/HashMap.h>
 #include <AK/OwnPtr.h>
+#include <AK/String.h>
 
 TEST_CASE(construct)
 {
@@ -219,3 +220,34 @@ TEST_CASE(in_place_rehashing_ordered_loop_bug)
     map.set("yt.innertube::nextId", "");
     VERIFY(map.keys().size() == 2);
 }
+
+TEST_CASE(take)
+{
+    HashMap<String, int> map;
+
+    EXPECT(!map.take("foo"sv).has_value());
+    EXPECT(!map.take("bar"sv).has_value());
+    EXPECT(!map.take(String::from_utf8_short_string("baz"sv)).has_value());
+
+    map.set(String::from_utf8_short_string("foo"sv), 1);
+    map.set(String::from_utf8_short_string("bar"sv), 2);
+    map.set(String::from_utf8_short_string("baz"sv), 3);
+
+    auto foo = map.take("foo"sv);
+    EXPECT_EQ(foo, 1);
+
+    foo = map.take("foo"sv);
+    EXPECT(!foo.has_value());
+
+    auto bar = map.take("bar"sv);
+    EXPECT_EQ(bar, 2);
+
+    bar = map.take("bar"sv);
+    EXPECT(!bar.has_value());
+
+    auto baz = map.take(String::from_utf8_short_string("baz"sv));
+    EXPECT_EQ(baz, 3);
+
+    baz = map.take(String::from_utf8_short_string("baz"sv));
+    EXPECT(!baz.has_value());
+}