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.
This commit is contained in:
Luke Wilde 2026-04-16 19:53:04 +01:00 committed by Alexander Kalenik
parent ec793f7c32
commit b6219c0673
3 changed files with 368 additions and 0 deletions

View file

@ -52,4 +52,7 @@ class RootVector;
template<typename T, typename TraitsForT, bool IsOrdered>
class RootHashTable;
template<typename K, typename V>
class WeakHashMap;
}

View file

@ -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 };
};
}

View file

@ -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);
}