LibJS: Notify WeakSets when heap cells are sweeped

This is an implementation of the following optional optimization:
https://tc39.es/ecma262/#sec-weakref-execution

Tests/LibJS/test-js.cpp

@@ -32,6 +32,19 @@ TESTJS_GLOBAL_FUNCTION(run_queued_promise_jobs, runQueuedPromiseJobs)
     return JS::js_undefined();
 }
 
+TESTJS_GLOBAL_FUNCTION(get_weak_set_size, getWeakSetSize)
+{
+    auto* object = vm.argument(0).to_object(global_object);
+    if (!object)
+        return {};
+    if (!is<JS::WeakSet>(object)) {
+        vm.throw_exception<JS::TypeError>(global_object, JS::ErrorType::NotA, "WeakSet");
+        return {};
+    }
+    auto* weak_set = static_cast<JS::WeakSet*>(object);
+    return JS::Value(weak_set->values().size());
+}
+
 TESTJS_RUN_FILE_FUNCTION(const String& test_file, JS::Interpreter&)
 {
     if (!test262_parser_tests)

Userland/Libraries/LibJS/Heap/Heap.cpp

@@ -16,6 +16,7 @@
 #include <LibJS/Heap/HeapBlock.h>
 #include <LibJS/Interpreter.h>
 #include <LibJS/Runtime/Object.h>
+#include <LibJS/Runtime/WeakSet.h>
 #include <setjmp.h>
 
 namespace JS {
@@ -182,18 +183,22 @@ void Heap::sweep_dead_cells(bool print_report, const Core::ElapsedTimer& measure
     dbgln_if(HEAP_DEBUG, "sweep_dead_cells:");
     Vector<HeapBlock*, 32> empty_blocks;
     Vector<HeapBlock*, 32> full_blocks_that_became_usable;
+    Vector<Cell*> sweeped_cells;
 
     size_t collected_cells = 0;
     size_t live_cells = 0;
     size_t collected_cell_bytes = 0;
     size_t live_cell_bytes = 0;
 
+    auto should_store_sweeped_cells = !m_weak_sets.is_empty();
     for_each_block([&](auto& block) {
         bool block_has_live_cells = false;
         bool block_was_full = block.is_full();
         block.template for_each_cell_in_state<Cell::State::Live>([&](Cell* cell) {
             if (!cell->is_marked()) {
                 dbgln_if(HEAP_DEBUG, "  ~ {}", cell);
+                if (should_store_sweeped_cells)
+                    sweeped_cells.append(cell);
                 block.deallocate(cell);
                 ++collected_cells;
                 collected_cell_bytes += block.cell_size();
@@ -221,6 +226,9 @@ void Heap::sweep_dead_cells(bool print_report, const Core::ElapsedTimer& measure
         allocator_for_size(block->cell_size()).block_did_become_usable({}, *block);
     }
 
+    for (auto* weak_set : m_weak_sets)
+        weak_set->remove_sweeped_cells({}, sweeped_cells);
+
     if constexpr (HEAP_DEBUG) {
         for_each_block([&](auto& block) {
             dbgln(" > Live HeapBlock @ {}: cell_size={}", &block, block.cell_size());
@@ -272,6 +280,18 @@ void Heap::did_destroy_marked_value_list(Badge<MarkedValueList>, MarkedValueList
     m_marked_value_lists.remove(&list);
 }
 
+void Heap::did_create_weak_set(Badge<WeakSet>, WeakSet& set)
+{
+    VERIFY(!m_weak_sets.contains(&set));
+    m_weak_sets.set(&set);
+}
+
+void Heap::did_destroy_weak_set(Badge<WeakSet>, WeakSet& set)
+{
+    VERIFY(m_weak_sets.contains(&set));
+    m_weak_sets.remove(&set);
+}
+
 void Heap::defer_gc(Badge<DeferGC>)
 {
     ++m_gc_deferrals;

Userland/Libraries/LibJS/Heap/Heap.h

@@ -70,6 +70,9 @@ public:
     void did_create_marked_value_list(Badge<MarkedValueList>, MarkedValueList&);
     void did_destroy_marked_value_list(Badge<MarkedValueList>, MarkedValueList&);
 
+    void did_create_weak_set(Badge<WeakSet>, WeakSet&);
+    void did_destroy_weak_set(Badge<WeakSet>, WeakSet&);
+
     void defer_gc(Badge<DeferGC>);
     void undefer_gc(Badge<DeferGC>);
 
@@ -106,6 +109,8 @@ private:
 
     HashTable<MarkedValueList*> m_marked_value_lists;
 
+    HashTable<WeakSet*> m_weak_sets;
+
     BlockAllocator m_block_allocator;
 
     size_t m_gc_deferrals { 0 };

Userland/Libraries/LibJS/Runtime/WeakSet.cpp

@@ -16,10 +16,18 @@ WeakSet* WeakSet::create(GlobalObject& global_object)
 WeakSet::WeakSet(Object& prototype)
     : Object(prototype)
 {
+    heap().did_create_weak_set({}, *this);
 }
 
 WeakSet::~WeakSet()
 {
+    heap().did_destroy_weak_set({}, *this);
+}
+
+void WeakSet::remove_sweeped_cells(Badge<Heap>, Vector<Cell*>& cells)
+{
+    for (auto* cell : cells)
+        m_values.remove(cell);
 }
 
 }

Userland/Libraries/LibJS/Runtime/WeakSet.h

@@ -21,11 +21,13 @@ public:
     explicit WeakSet(Object& prototype);
     virtual ~WeakSet() override;
 
-    HashTable<Object*> const& values() const { return m_values; };
+    HashTable<Cell*> const& values() const { return m_values; };
-    HashTable<Object*>& values() { return m_values; };
+    HashTable<Cell*>& values() { return m_values; };
+
+    void remove_sweeped_cells(Badge<Heap>, Vector<Cell*>&);
 
 private:
-    HashTable<Object*> m_values;
+    HashTable<Cell*> m_values; // This stores Cell pointers instead of Object pointers to aide with sweeping
 };
 
 }

Userland/Libraries/LibJS/Tests/builtins/WeakSet/WeakSet.prototype.add.js

@@ -14,3 +14,13 @@ test("invalid values", () => {
         }).toThrowWithMessage(TypeError, "is not an object");
     });
 });
+
+test("automatic removal of garbage-collected values", () => {
+    const weakSet = new WeakSet();
+    {
+        expect(weakSet.add({ a: 1 })).toBe(weakSet);
+        expect(getWeakSetSize(weakSet)).toBe(1);
+    }
+    gc();
+    expect(getWeakSetSize(weakSet)).toBe(0);
+});

Userland/Libraries/LibTest/JavaScriptTestRunner.h

@@ -26,6 +26,7 @@
 #include <LibJS/Runtime/GlobalObject.h>
 #include <LibJS/Runtime/JSONObject.h>
 #include <LibJS/Runtime/TypedArray.h>
+#include <LibJS/Runtime/WeakSet.h>
 #include <LibTest/Results.h>
 #include <fcntl.h>
 #include <sys/time.h>