LibGC: Prune weak containers in stop-the-world phase of GC

Move weak container cleanup (remove_dead_cells) out of both
sweep_dead_cells() and start_incremental_sweep() to the place
where it is actually safe to inspect cell state: collect_garbage().

Previously, remove_dead_cells could access cells that had already
been swept and poisoned by ASAN, causing use-after-poison crashes
when a new GC triggered while an incremental sweep was in progress.

Libraries/LibGC/Heap.cpp

@@ -538,6 +538,15 @@ void Heap::collect_garbage(CollectionType collection_type, bool print_report)
             sweep_weak_blocks();
         }
 
+        // Prune weak containers while we're still stop-the-world; doing this
+        // during incremental sweep risks reading cells that have already been
+        // freed and ASAN-poisoned.
+        for (auto& weak_container : m_weak_containers) {
+            if (!weak_container.owner_cell({}).is_marked())
+                continue;
+            weak_container.remove_dead_cells({});
+        }
+
         // Run sweep callbacks at STW so they fire for every collection,
         // not just CollectEverything. Static caches like
         // StaticPropertyLookupCache prune by mark state and must see valid
@@ -1051,12 +1060,6 @@ void Heap::sweep_dead_cells(bool print_report, Core::ElapsedTimer const& measure
         });
     }
 
-    {
-        ScopedPhaseTimer timer { g_recording_phase_timings, g_phase_timings.sweep_weak_containers_us };
-        for (auto& weak_container : m_weak_containers)
-            weak_container.remove_dead_cells({});
-    }
-
     {
         ScopedPhaseTimer timer { g_recording_phase_timings, g_phase_timings.sweep_block_reclassify_us };
         for (auto* block : empty_blocks) {
@@ -1175,9 +1178,6 @@ void Heap::start_incremental_sweep()
     m_sweep_live_cell_bytes = 0;
     m_sweep_live_external_bytes = 0;
 
-    for (auto& weak_container : m_weak_containers)
-        weak_container.remove_dead_cells({});
-
     // Populate each allocator's pending sweep list with its current blocks.
     // Blocks allocated during incremental sweep won't be on these lists
     // and don't need sweeping.

Libraries/LibGC/WeakContainer.h

@@ -16,6 +16,7 @@ public:
     explicit WeakContainer(Heap&);
     virtual ~WeakContainer();
 
+    virtual Cell const& owner_cell(Badge<Heap>) const = 0;
     virtual void remove_dead_cells(Badge<Heap>) = 0;
 
 protected:

Libraries/LibJS/Bytecode/Executable.cpp

@@ -341,13 +341,16 @@ StaticPropertyLookupCache::StaticPropertyLookupCache()
 
 static void clear_cache_entry_if_dead(PropertyLookupCache::Entry& entry)
 {
-    if (entry.from_shape && entry.from_shape->state() != Cell::State::Live)
+    auto cell_is_dead = [](Cell const* cell) {
+        return cell->state() != Cell::State::Live || !cell->is_marked();
+    };
+    if (entry.from_shape && cell_is_dead(entry.from_shape))
         entry.from_shape = nullptr;
-    if (entry.shape && entry.shape->state() != Cell::State::Live)
+    if (entry.shape && cell_is_dead(entry.shape))
         entry.shape = nullptr;
-    if (entry.prototype && entry.prototype->state() != Cell::State::Live)
+    if (entry.prototype && cell_is_dead(entry.prototype))
         entry.prototype = nullptr;
-    if (entry.prototype_chain_validity && entry.prototype_chain_validity->state() != Cell::State::Live)
+    if (entry.prototype_chain_validity && cell_is_dead(entry.prototype_chain_validity))
         entry.prototype_chain_validity = nullptr;
 }
 
@@ -370,7 +373,7 @@ void Executable::remove_dead_cells(Badge<GC::Heap>)
             clear_cache_entry_if_dead(entry);
     }
     for (auto& cache : object_shape_caches) {
-        if (cache.shape && cache.shape->state() != Cell::State::Live)
+        if (cache.shape && (cache.shape->state() != Cell::State::Live || !cache.shape->is_marked()))
             cache.shape = nullptr;
     }
 }

Libraries/LibJS/Bytecode/Executable.h

@@ -239,6 +239,7 @@ public:
 
     [[nodiscard]] Operand original_operand_from_raw(u32) const;
 
+    virtual Cell const& owner_cell(Badge<GC::Heap>) const override { return *this; }
     virtual void remove_dead_cells(Badge<GC::Heap>) override;
 
 private:

Libraries/LibJS/Runtime/FinalizationRegistry.cpp

@@ -54,7 +54,7 @@ void FinalizationRegistry::remove_dead_cells(Badge<GC::Heap>)
 {
     auto any_cells_were_removed = false;
     for (auto& record : m_records) {
-        if (!record.target || record.target->state() == Cell::State::Live)
+        if (!record.target || (record.target->state() == Cell::State::Live && record.target->is_marked()))
             continue;
         record.target = nullptr;
         any_cells_were_removed = true;

Libraries/LibJS/Runtime/FinalizationRegistry.h

@@ -31,6 +31,7 @@ public:
     bool remove_by_token(Cell& unregister_token);
     ThrowCompletionOr<void> cleanup(GC::Ptr<JobCallback> = {});
 
+    virtual Cell const& owner_cell(Badge<GC::Heap>) const override { return *this; }
     virtual void remove_dead_cells(Badge<GC::Heap>) override;
 
     Realm& realm() { return *m_realm; }

Libraries/LibJS/Runtime/WeakMap.cpp

@@ -53,7 +53,7 @@ bool WeakMap::weak_map_remove(GC::Ptr<Cell> key)
 void WeakMap::remove_dead_cells(Badge<GC::Heap>)
 {
     m_values.remove_all_matching([](Cell* key, Value) {
-        return key->state() != Cell::State::Live;
+        return key->state() != Cell::State::Live || !key->is_marked();
     });
 }
 

Libraries/LibJS/Runtime/WeakMap.h

@@ -31,6 +31,7 @@ public:
     bool weak_map_remove(GC::Ptr<Cell>);
     size_t weak_map_size() const { return m_values.size(); }
 
+    virtual Cell const& owner_cell(Badge<GC::Heap>) const override { return *this; }
     virtual void remove_dead_cells(Badge<GC::Heap>) override;
     virtual size_t external_memory_size() const override;
 

Libraries/LibJS/Runtime/WeakRef.cpp

@@ -38,7 +38,7 @@ WeakRef::WeakRef(Symbol& value, Object& prototype)
 
 void WeakRef::remove_dead_cells(Badge<GC::Heap>)
 {
-    if (m_value.visit([](Cell* cell) -> bool { return cell->state() == Cell::State::Live; }, [](Empty) -> bool { return true; }))
+    if (m_value.visit([](Cell* cell) -> bool { return cell->state() == Cell::State::Live && cell->is_marked(); }, [](Empty) -> bool { return true; }))
         return;
 
     m_value = Empty {};

Libraries/LibJS/Runtime/WeakRef.h

@@ -28,6 +28,7 @@ public:
 
     void update_execution_generation() { m_last_execution_generation = vm().execution_generation(); }
 
+    virtual Cell const& owner_cell(Badge<GC::Heap>) const override { return *this; }
     virtual void remove_dead_cells(Badge<GC::Heap>) override;
 
 private:

Libraries/LibJS/Runtime/WeakSet.cpp

@@ -46,7 +46,7 @@ bool WeakSet::weak_set_remove(GC::Ptr<Cell> value)
 void WeakSet::remove_dead_cells(Badge<GC::Heap>)
 {
     m_values.remove_all_matching([](Cell* cell) {
-        return cell->state() != Cell::State::Live;
+        return cell->state() != Cell::State::Live || !cell->is_marked();
     });
 }
 

Libraries/LibJS/Runtime/WeakSet.h

@@ -30,6 +30,7 @@ public:
     bool weak_set_remove(GC::Ptr<Cell>);
     size_t weak_set_size() const { return m_values.size(); }
 
+    virtual Cell const& owner_cell(Badge<GC::Heap>) const override { return *this; }
     virtual void remove_dead_cells(Badge<GC::Heap>) override;
     virtual size_t external_memory_size() const override;