From b6219c06735d15431192fd54945693441457650e Mon Sep 17 00:00:00 2001
From: Luke Wilde <luke@ladybird.org>
Date: Thu, 16 Apr 2026 19:53:04 +0100
Subject: [PATCH] LibGC: Introduce WeakHashMap

Mirrors `WeakHashSet` for map shapes. Cell-typed key and/or value
slots are stored as `Weak<T>` so entries vanish when their referent
is collected; non-cell slots are stored directly.
---
 Libraries/LibGC/Forward.h        |   3 +
 Libraries/LibGC/WeakHashMap.h    | 213 +++++++++++++++++++++++++++++++
 Tests/LibGC/TestGCContainers.cpp | 152 ++++++++++++++++++++++
 3 files changed, 368 insertions(+)
 create mode 100644 Libraries/LibGC/WeakHashMap.h

diff --git a/Libraries/LibGC/Forward.h b/Libraries/LibGC/Forward.h
index e9504e8809..0741ea7220 100644
--- a/Libraries/LibGC/Forward.h
+++ b/Libraries/LibGC/Forward.h
@@ -52,4 +52,7 @@ class RootVector;
 template<typename T, typename TraitsForT, bool IsOrdered>
 class RootHashTable;
 
+template<typename K, typename V>
+class WeakHashMap;
+
 }
diff --git a/Libraries/LibGC/WeakHashMap.h b/Libraries/LibGC/WeakHashMap.h
new file mode 100644
index 0000000000..8dada60484
--- /dev/null
+++ b/Libraries/LibGC/WeakHashMap.h
@@ -0,0 +1,213 @@
+/*
+ * Copyright (c) 2026, Luke Wilde <luke@ladybird.org>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#pragma once
+
+#include <AK/HashMap.h>
+#include <AK/StdLibExtras.h>
+#include <LibGC/Cell.h>
+#include <LibGC/Weak.h>
+#include <LibGC/WeakInlines.h>
+
+namespace GC {
+
+// Cell-typed key and/or value slots are held as Weak<T> so entries vanish
+// when their referent is collected; non-cell slots are stored directly.
+template<typename K, typename V>
+class WeakHashMap {
+    static constexpr bool key_is_cell = IsBaseOf<Cell, K>;
+    static constexpr bool value_is_cell = IsBaseOf<Cell, V>;
+
+    using KeyStorage = Conditional<key_is_cell, Weak<K>, K>;
+    using ValueStorage = Conditional<value_is_cell, Weak<V>, V>;
+
+    struct KeyTraits : public DefaultTraits<KeyStorage> {
+        static unsigned hash(KeyStorage const& value)
+        {
+            if constexpr (key_is_cell)
+                return Traits<K const*>::hash(value.ptr());
+            else
+                return Traits<K>::hash(value);
+        }
+    };
+
+    using TableType = HashMap<KeyStorage, ValueStorage, KeyTraits>;
+
+public:
+    WeakHashMap() = default;
+
+    HashSetResult set(K const& key, V const& value)
+    {
+        maybe_prune();
+        return m_table.set(to_key_storage(key), to_value_storage(value));
+    }
+
+    HashSetResult set(K const& key, V&& value)
+    {
+        maybe_prune();
+        return m_table.set(to_key_storage(key), to_value_storage(move(value)));
+    }
+
+    bool remove(K const& key)
+    {
+        maybe_prune();
+        auto it = find_iterator(key);
+        if (it == m_table.end())
+            return false;
+        m_table.remove(it);
+        return true;
+    }
+
+    template<typename Callback>
+    V& ensure(K const& key, Callback initialization_callback)
+    {
+        maybe_prune();
+
+        auto it = find_iterator(key);
+        if (it != m_table.end()) {
+            if constexpr (value_is_cell) {
+                if (auto value = it->value.ptr())
+                    return *value;
+                m_table.remove(it);
+            } else {
+                return it->value;
+            }
+        }
+
+        if constexpr (value_is_cell) {
+            auto value = initialization_callback();
+            [[maybe_unused]] auto result = m_table.set(to_key_storage(key), to_value_storage(value));
+            VERIFY(result == HashSetResult::InsertedNewEntry);
+            return *value;
+        } else {
+            return m_table.ensure(to_key_storage(key), [&] {
+                return to_value_storage(initialization_callback());
+            });
+        }
+    }
+
+    V& ensure(K const& key)
+    requires(!value_is_cell)
+    {
+        return ensure(key, [] { return V(); });
+    }
+
+    bool contains(K const& key) const
+    {
+        return find_iterator(key) != m_table.end();
+    }
+
+    auto get(K const& key)
+    {
+        if constexpr (value_is_cell) {
+            auto it = find_iterator(key);
+            if (it == m_table.end())
+                return static_cast<V*>(nullptr);
+            return static_cast<V*>(it->value.ptr());
+        } else {
+            auto it = find_iterator(key);
+            if (it == m_table.end())
+                return Optional<V&> {};
+            return Optional<V&> { it->value };
+        }
+    }
+
+    auto get(K const& key) const
+    {
+        if constexpr (value_is_cell) {
+            auto it = find_iterator(key);
+            if (it == m_table.end())
+                return static_cast<V const*>(nullptr);
+            return static_cast<V const*>(it->value.ptr());
+        } else {
+            auto it = find_iterator(key);
+            if (it == m_table.end())
+                return Optional<V const&> {};
+            return Optional<V const&> { it->value };
+        }
+    }
+
+    bool is_empty() const
+    {
+        for (auto const& entry : m_table) {
+            if constexpr (key_is_cell) {
+                if (!entry.key.ptr())
+                    continue;
+            }
+            if constexpr (value_is_cell) {
+                if (!entry.value.ptr())
+                    continue;
+            }
+            return false;
+        }
+        return true;
+    }
+
+    void clear() { m_table.clear(); }
+
+private:
+    static KeyStorage to_key_storage(K const& key)
+    {
+        if constexpr (key_is_cell)
+            return Weak<K>(key);
+        else
+            return key;
+    }
+
+    template<typename U>
+    static ValueStorage to_value_storage(U&& value)
+    {
+        if constexpr (value_is_cell)
+            return Weak<V>(forward<U>(value));
+        else
+            return ValueStorage { forward<U>(value) };
+    }
+
+    auto find_iterator(K const& key) const -> typename TableType::ConstIteratorType
+    {
+        if constexpr (key_is_cell) {
+            return m_table.find(Traits<K const*>::hash(&key), [&](auto& entry) {
+                return entry.key.ptr() == &key;
+            });
+        } else {
+            return m_table.find(key);
+        }
+    }
+
+    auto find_iterator(K const& key) -> typename TableType::IteratorType
+    {
+        if constexpr (key_is_cell) {
+            return m_table.find(Traits<K const*>::hash(&key), [&](auto& entry) {
+                return entry.key.ptr() == &key;
+            });
+        } else {
+            return m_table.find(key);
+        }
+    }
+
+    void maybe_prune()
+    {
+        if (++m_mutations_since_last_prune < max(m_table.size(), static_cast<size_t>(64)))
+            return;
+        m_table.remove_all_matching([](auto const& key, auto const& value) {
+            if constexpr (key_is_cell) {
+                if (!key.ptr())
+                    return true;
+            }
+            if constexpr (value_is_cell) {
+                if (!value.ptr())
+                    return true;
+            }
+            return false;
+        });
+        m_mutations_since_last_prune = 0;
+    }
+
+    TableType m_table;
+    size_t m_mutations_since_last_prune { 0 };
+};
+
+}
diff --git a/Tests/LibGC/TestGCContainers.cpp b/Tests/LibGC/TestGCContainers.cpp
index 280d559285..5eae30610f 100644
--- a/Tests/LibGC/TestGCContainers.cpp
+++ b/Tests/LibGC/TestGCContainers.cpp
@@ -16,6 +16,7 @@
 #include <LibGC/RootHashMap.h>
 #include <LibGC/RootHashTable.h>
 #include <LibGC/RootVector.h>
+#include <LibGC/WeakHashMap.h>
 #include <LibTest/TestCase.h>
 
 class TestCell : public GC::Cell {
@@ -324,3 +325,154 @@ TEST_CASE(removed_conservative_hash_map_key_not_reported)
 
     EXPECT(!possible_values_contain(map, cell.ptr()));
 }
+
+TEST_CASE(weak_hash_map_non_cell_key_cell_value)
+{
+    auto& heap = test_heap();
+    GC::WeakHashMap<int, TestCell> map;
+
+    auto cell = heap.allocate<TestCell>();
+    map.set(42, *cell);
+
+    EXPECT(map.contains(42));
+    EXPECT_EQ(map.get(42), cell.ptr());
+
+    EXPECT(map.remove(42));
+    EXPECT(!map.contains(42));
+    EXPECT_EQ(map.get(42), static_cast<TestCell*>(nullptr));
+}
+
+TEST_CASE(weak_hash_map_cell_key_non_cell_value)
+{
+    auto& heap = test_heap();
+    GC::WeakHashMap<TestCell, int> map;
+
+    auto cell = heap.allocate<TestCell>();
+    map.set(*cell, 7);
+
+    EXPECT(map.contains(*cell));
+    EXPECT_EQ(map.get(*cell).value(), 7);
+
+    EXPECT(map.remove(*cell));
+    EXPECT(!map.contains(*cell));
+}
+
+TEST_CASE(weak_hash_map_cell_key_and_cell_value)
+{
+    auto& heap = test_heap();
+    GC::WeakHashMap<TestCell, TestCell> map;
+
+    auto key = heap.allocate<TestCell>();
+    auto value = heap.allocate<TestCell>();
+    map.set(*key, *value);
+
+    EXPECT(map.contains(*key));
+    EXPECT_EQ(map.get(*key), value.ptr());
+}
+
+TEST_CASE(weak_hash_map_value_collection_clears_entry)
+{
+    auto& heap = test_heap();
+    GC::WeakHashMap<int, TestCell> map;
+
+    {
+        auto cell = heap.allocate<TestCell>();
+        map.set(1, *cell);
+    }
+
+    heap.collect_garbage(GC::Heap::CollectionType::CollectEverything);
+
+    EXPECT_EQ(map.get(1), static_cast<TestCell*>(nullptr));
+}
+
+TEST_CASE(weak_hash_map_ensure_creates_missing_cell_value)
+{
+    auto& heap = test_heap();
+    GC::WeakHashMap<int, TestCell> map;
+
+    bool called = false;
+    auto& value = map.ensure(1, [&] {
+        called = true;
+        return heap.allocate<TestCell>();
+    });
+
+    EXPECT(called);
+    EXPECT_EQ(map.get(1), &value);
+}
+
+TEST_CASE(weak_hash_map_ensure_reuses_live_cell_value)
+{
+    auto& heap = test_heap();
+    GC::WeakHashMap<int, TestCell> map;
+
+    auto cell = heap.allocate<TestCell>();
+    map.set(1, *cell);
+
+    bool called = false;
+    auto& value = map.ensure(1, [&] {
+        called = true;
+        return heap.allocate<TestCell>();
+    });
+
+    EXPECT(!called);
+    EXPECT_EQ(&value, cell.ptr());
+}
+
+TEST_CASE(weak_hash_map_ensure_replaces_collected_cell_value)
+{
+    auto& heap = test_heap();
+    GC::WeakHashMap<int, TestCell> map;
+
+    {
+        auto cell = heap.allocate<TestCell>();
+        map.set(1, *cell);
+    }
+
+    heap.collect_garbage(GC::Heap::CollectionType::CollectEverything);
+    EXPECT_EQ(map.get(1), static_cast<TestCell*>(nullptr));
+
+    auto replacement = heap.allocate<TestCell>();
+    bool called = false;
+    auto& value = map.ensure(1, [&] {
+        called = true;
+        return replacement;
+    });
+
+    EXPECT(called);
+    EXPECT_EQ(&value, replacement.ptr());
+    EXPECT_EQ(map.get(1), replacement.ptr());
+}
+
+TEST_CASE(weak_hash_map_ensure_handles_non_cell_value)
+{
+    auto& heap = test_heap();
+    GC::WeakHashMap<TestCell, int> map;
+
+    auto key = heap.allocate<TestCell>();
+
+    bool called = false;
+    auto& value = map.ensure(*key, [&] {
+        called = true;
+        return 7;
+    });
+
+    EXPECT(called);
+    EXPECT_EQ(value, 7);
+
+    value = 9;
+    called = false;
+    auto& existing_value = map.ensure(*key, [&] {
+        called = true;
+        return 11;
+    });
+
+    EXPECT(!called);
+    EXPECT_EQ(existing_value, 9);
+
+    auto default_key = heap.allocate<TestCell>();
+    auto& default_value = map.ensure(*default_key);
+    EXPECT_EQ(default_value, 0);
+
+    default_value = 13;
+    EXPECT_EQ(map.get(*default_key).value(), 13);
+}