LibJS: Make WeakContainer pruning do less work

Instead of iterating *all* swept cells when pruning weak containers,
only iterate the cells actually *in* the container.

Also, instead of compiling a list of all swept cells, we can simply
check the Cell::state() flag to know if something should be pruned.

Userland/Libraries/LibJS/Heap/Heap.cpp

@@ -216,22 +216,18 @@ 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*> swept_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_swept_cells = !m_weak_containers.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_swept_cells)
-                    swept_cells.append(cell);
 #ifdef JS_TRACK_ZOMBIE_CELLS
                 if (m_zombify_dead_cells) {
                     cell->set_state(Cell::State::Zombie);
@@ -269,7 +265,7 @@ void Heap::sweep_dead_cells(bool print_report, const Core::ElapsedTimer& measure
     }
 
     for (auto& weak_container : m_weak_containers)
-        weak_container.remove_swept_cells({}, swept_cells.span());
+        weak_container.remove_dead_cells({});
 
     if constexpr (HEAP_DEBUG) {
         for_each_block([&](auto& block) {

Userland/Libraries/LibJS/Runtime/FinalizationRegistry.cpp

@@ -42,19 +42,17 @@ bool FinalizationRegistry::remove_by_token(Object& unregister_token)
     return removed;
 }
 
-void FinalizationRegistry::remove_swept_cells(Badge<Heap>, Span<Cell*> cells)
+void FinalizationRegistry::remove_dead_cells(Badge<Heap>)
 {
-    auto any_cells_were_swept = false;
-    for (auto* cell : cells) {
-        for (auto& record : m_records) {
-            if (record.target != cell)
-                continue;
-            record.target = nullptr;
-            any_cells_were_swept = true;
-            break;
-        }
+    auto any_cells_were_removed = false;
+    for (auto& record : m_records) {
+        if (!record.target || record.target->state() == Cell::State::Live)
+            continue;
+        record.target = nullptr;
+        any_cells_were_removed = true;
+        break;
     }
-    if (any_cells_were_swept)
+    if (any_cells_were_removed)
         vm().enqueue_finalization_registry_cleanup_job(*this);
 }
 

Userland/Libraries/LibJS/Runtime/FinalizationRegistry.h

@@ -30,7 +30,7 @@ public:
     bool remove_by_token(Object& unregister_token);
     void cleanup(FunctionObject* callback = nullptr);
 
-    virtual void remove_swept_cells(Badge<Heap>, Span<Cell*>) override;
+    virtual void remove_dead_cells(Badge<Heap>) override;
 
 private:
     virtual void visit_edges(Visitor& visitor) override;

Userland/Libraries/LibJS/Runtime/WeakContainer.h

@@ -15,7 +15,7 @@ public:
     explicit WeakContainer(Heap&);
     virtual ~WeakContainer();
 
-    virtual void remove_swept_cells(Badge<Heap>, Span<Cell*>) = 0;
+    virtual void remove_dead_cells(Badge<Heap>) = 0;
 
 protected:
     void deregister();

Userland/Libraries/LibJS/Runtime/WeakMap.cpp

@@ -23,9 +23,15 @@ WeakMap::~WeakMap()
 {
 }
 
-void WeakMap::remove_swept_cells(Badge<Heap>, Span<Cell*> cells)
+void WeakMap::remove_dead_cells(Badge<Heap>)
 {
-    for (auto* cell : cells)
+    // FIXME: Do this in a single pass.
+    Vector<Cell*> to_remove;
+    for (auto& it : m_values) {
+        if (it.key->state() != Cell::State::Live)
+            to_remove.append(it.key);
+    }
+    for (auto* cell : to_remove)
         m_values.remove(cell);
 }
 

Userland/Libraries/LibJS/Runtime/WeakMap.h

@@ -27,7 +27,7 @@ public:
     HashMap<Cell*, Value> const& values() const { return m_values; };
     HashMap<Cell*, Value>& values() { return m_values; };
 
-    virtual void remove_swept_cells(Badge<Heap>, Span<Cell*>) override;
+    virtual void remove_dead_cells(Badge<Heap>) override;
 
 private:
 #ifdef JS_TRACK_ZOMBIE_CELLS

Userland/Libraries/LibJS/Runtime/WeakRef.cpp

@@ -25,18 +25,16 @@ WeakRef::~WeakRef()
 {
 }
 
-void WeakRef::remove_swept_cells(Badge<Heap>, Span<Cell*> cells)
+void WeakRef::remove_dead_cells(Badge<Heap>)
 {
     VERIFY(m_value);
-    for (auto* cell : cells) {
-        if (m_value != cell)
-            continue;
-        m_value = nullptr;
-        // This is an optimization, we deregister from the garbage collector early (even if we were not garbage collected ourself yet)
-        // to reduce the garbage collection overhead, which we can do because a cleared weak ref cannot be reused.
-        WeakContainer::deregister();
-        break;
-    }
+    if (m_value->state() == Cell::State::Live)
+        return;
+
+    m_value = nullptr;
+    // This is an optimization, we deregister from the garbage collector early (even if we were not garbage collected ourself yet)
+    // to reduce the garbage collection overhead, which we can do because a cleared weak ref cannot be reused.
+    WeakContainer::deregister();
 }
 
 void WeakRef::visit_edges(Visitor& visitor)

Userland/Libraries/LibJS/Runtime/WeakRef.h

@@ -27,7 +27,7 @@ public:
 
     void update_execution_generation() { m_last_execution_generation = vm().execution_generation(); };
 
-    virtual void remove_swept_cells(Badge<Heap>, Span<Cell*>) override;
+    virtual void remove_dead_cells(Badge<Heap>) override;
 
 private:
     virtual void visit_edges(Visitor&) override;

Userland/Libraries/LibJS/Runtime/WeakSet.cpp

@@ -23,9 +23,15 @@ WeakSet::~WeakSet()
 {
 }
 
-void WeakSet::remove_swept_cells(Badge<Heap>, Span<Cell*> cells)
+void WeakSet::remove_dead_cells(Badge<Heap>)
 {
-    for (auto* cell : cells)
+    // FIXME: Do this in a single pass.
+    Vector<Cell*> to_remove;
+    for (auto* cell : m_values) {
+        if (cell->state() != Cell::State::Live)
+            to_remove.append(cell);
+    }
+    for (auto* cell : to_remove)
         m_values.remove(cell);
 }
 

Userland/Libraries/LibJS/Runtime/WeakSet.h

@@ -27,7 +27,7 @@ public:
     HashTable<Cell*> const& values() const { return m_values; };
     HashTable<Cell*>& values() { return m_values; };
 
-    virtual void remove_swept_cells(Badge<Heap>, Span<Cell*>) override;
+    virtual void remove_dead_cells(Badge<Heap>) override;
 
 private:
 #ifdef JS_TRACK_ZOMBIE_CELLS