diff --git a/AK/HashMap.h b/AK/HashMap.h index ec8d9464724..d3e1af9b630 100644 --- a/AK/HashMap.h +++ b/AK/HashMap.h @@ -199,6 +199,31 @@ public: m_table.remove(it); } + Optional take(K const& key) + { + if (auto it = find(key); it != end()) { + auto value = move(it->value); + m_table.remove(it); + + return value; + } + + return {}; + } + + template Key> + requires(IsSame>) Optional 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 bbb4876c5df..6e70540a9fd 100644 --- a/Tests/AK/TestHashMap.cpp +++ b/Tests/AK/TestHashMap.cpp @@ -9,6 +9,7 @@ #include #include #include +#include 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 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()); +}