AK+LibJS: Introduce JS::HeapFunction

This change introduces HeapFunction, which is intended to be used as a
replacement for SafeFunction. The new type behaves like a regular
GC-allocated object, which means it needs to be visited from
visit_edges, and unlike SafeFunction, it does not create new roots for
captured parameters.

Co-Authored-By: Andreas Kling <kling@serenityos.org>

AK/Function.h

@@ -1,6 +1,7 @@
 /*
  * Copyright (C) 2016 Apple Inc. All rights reserved.
  * Copyright (c) 2021, Gunnar Beutner <gbeutner@serenityos.org>
+ * Copyright (c) 2018-2023, Andreas Kling <kling@serenityos.org>
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions
@@ -31,6 +32,7 @@
 #include <AK/BitCast.h>
 #include <AK/Noncopyable.h>
 #include <AK/ScopeGuard.h>
+#include <AK/Span.h>
 #include <AK/StdLibExtras.h>
 #include <AK/Types.h>
 
@@ -51,6 +53,7 @@ class Function<Out(In...)> {
     AK_MAKE_NONCOPYABLE(Function);
 
 public:
+    using FunctionType = Out(In...);
     using ReturnType = Out;
 
     Function() = default;
@@ -63,18 +66,27 @@ public:
         clear(false);
     }
 
+    [[nodiscard]] ReadonlyBytes raw_capture_range() const
+    {
+        if (!m_size)
+            return {};
+        if (auto* wrapper = callable_wrapper())
+            return ReadonlyBytes { wrapper, m_size };
+        return {};
+    }
+
     template<typename CallableType>
     Function(CallableType&& callable)
     requires((IsFunctionObject<CallableType> && IsCallableWithArguments<CallableType, Out, In...> && !IsSame<RemoveCVReference<CallableType>, Function>))
     {
-        init_with_callable(forward<CallableType>(callable));
+        init_with_callable(forward<CallableType>(callable), CallableKind::FunctionObject);
     }
 
     template<typename FunctionType>
     Function(FunctionType f)
     requires((IsFunctionPointer<FunctionType> && IsCallableWithArguments<RemovePointer<FunctionType>, Out, In...> && !IsSame<RemoveCVReference<FunctionType>, Function>))
     {
-        init_with_callable(move(f));
+        init_with_callable(move(f), CallableKind::FunctionPointer);
     }
 
     Function(Function&& other)
@@ -102,7 +114,7 @@ public:
     requires((IsFunctionObject<CallableType> && IsCallableWithArguments<CallableType, Out, In...>))
     {
         clear();
-        init_with_callable(forward<CallableType>(callable));
+        init_with_callable(forward<CallableType>(callable), CallableKind::FunctionObject);
         return *this;
     }
 
@@ -112,7 +124,7 @@ public:
     {
         clear();
         if (f)
-            init_with_callable(move(f));
+            init_with_callable(move(f), CallableKind::FunctionPointer);
         return *this;
     }
 
@@ -132,6 +144,11 @@ public:
     }
 
 private:
+    enum class CallableKind {
+        FunctionPointer,
+        FunctionObject,
+    };
+
     class CallableWrapperBase {
     public:
         virtual ~CallableWrapperBase() = default;
@@ -215,7 +232,7 @@ private:
     }
 
     template<typename Callable>
-    void init_with_callable(Callable&& callable)
+    void init_with_callable(Callable&& callable, CallableKind callable_kind)
     {
         VERIFY(m_call_nesting_level == 0);
         using WrapperType = CallableWrapper<Callable>;
@@ -231,12 +248,17 @@ private:
 #ifndef KERNEL
         }
 #endif
+        if (callable_kind == CallableKind::FunctionObject)
+            m_size = sizeof(WrapperType);
+        else
+            m_size = 0;
     }
 
     void move_from(Function&& other)
     {
         VERIFY(m_call_nesting_level == 0 && other.m_call_nesting_level == 0);
         auto* other_wrapper = other.callable_wrapper();
+        m_size = other.m_size;
         switch (other.m_kind) {
         case FunctionKind::NullPointer:
             break;
@@ -254,9 +276,11 @@ private:
         other.m_kind = FunctionKind::NullPointer;
     }
 
+    size_t m_size { 0 };
     FunctionKind m_kind { FunctionKind::NullPointer };
     bool m_deferred_clear { false };
     mutable Atomic<u16> m_call_nesting_level { 0 };
+
 #ifndef KERNEL
     // Empirically determined to fit most lambdas and functions.
     static constexpr size_t inline_capacity = 4 * sizeof(void*);

Userland/Libraries/LibJS/Heap/Cell.h

@@ -82,6 +82,8 @@ public:
         {
         }
 
+        virtual void visit_possible_values(ReadonlyBytes) = 0;
+
     protected:
         virtual void visit_impl(Cell&) = 0;
         virtual ~Visitor() = default;

Userland/Libraries/LibJS/Heap/Heap.cpp

@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2020-2022, Andreas Kling <kling@serenityos.org>
+ * Copyright (c) 2023, Aliaksandr Kalenik <kalenik.aliaksandr@gmail.com>
  *
  * SPDX-License-Identifier: BSD-2-Clause
  */
@@ -94,10 +95,49 @@ Cell* Heap::allocate_cell(size_t size)
     return allocator.allocate_cell(*this);
 }
 
+static void add_possible_value(HashMap<FlatPtr, HeapRootTypeOrLocation>& possible_pointers, FlatPtr data, HeapRootTypeOrLocation origin)
+{
+    if constexpr (sizeof(FlatPtr*) == sizeof(Value)) {
+        // Because Value stores pointers in non-canonical form we have to check if the top bytes
+        // match any pointer-backed tag, in that case we have to extract the pointer to its
+        // canonical form and add that as a possible pointer.
+        if ((data & SHIFTED_IS_CELL_PATTERN) == SHIFTED_IS_CELL_PATTERN)
+            possible_pointers.set(Value::extract_pointer_bits(data), move(origin));
+        else
+            possible_pointers.set(data, move(origin));
+    } else {
+        static_assert((sizeof(Value) % sizeof(FlatPtr*)) == 0);
+        // In the 32-bit case we will look at the top and bottom part of Value separately we just
+        // add both the upper and lower bytes as possible pointers.
+        possible_pointers.set(data, move(origin));
+    }
+}
+
+template<typename Callback>
+static void for_each_cell_among_possible_pointers(HashTable<HeapBlock*> all_live_heap_blocks, HashMap<FlatPtr, HeapRootTypeOrLocation>& possible_pointers, Callback callback)
+{
+    for (auto possible_pointer : possible_pointers.keys()) {
+        if (!possible_pointer)
+            continue;
+        auto* possible_heap_block = HeapBlock::from_cell(reinterpret_cast<Cell const*>(possible_pointer));
+        if (!all_live_heap_blocks.contains(possible_heap_block))
+            continue;
+        if (auto* cell = possible_heap_block->cell_from_possible_pointer(possible_pointer)) {
+            callback(cell, possible_pointer);
+        }
+    }
+}
+
 class GraphConstructorVisitor final : public Cell::Visitor {
 public:
-    explicit GraphConstructorVisitor(HashMap<Cell*, HeapRootTypeOrLocation> const& roots)
+    explicit GraphConstructorVisitor(Heap& heap, HashMap<Cell*, HeapRootTypeOrLocation> const& roots)
+        : m_heap(heap)
     {
+        m_heap.for_each_block([&](auto& block) {
+            m_all_live_heap_blocks.set(&block);
+            return IterationDecision::Continue;
+        });
+
         for (auto* root : roots.keys()) {
             visit(root);
             auto& graph_node = m_graph.ensure(reinterpret_cast<FlatPtr>(root));
@@ -117,6 +157,24 @@ public:
         m_work_queue.append(cell);
     }
 
+    virtual void visit_possible_values(ReadonlyBytes bytes) override
+    {
+        HashMap<FlatPtr, HeapRootTypeOrLocation> possible_pointers;
+
+        auto* raw_pointer_sized_values = reinterpret_cast<FlatPtr const*>(bytes.data());
+        for (size_t i = 0; i < (bytes.size() / sizeof(FlatPtr)); ++i)
+            add_possible_value(possible_pointers, raw_pointer_sized_values[i], HeapRootType::HeapFunctionCapturedPointer);
+
+        for_each_cell_among_possible_pointers(m_all_live_heap_blocks, possible_pointers, [&](Cell* cell, FlatPtr) {
+            if (m_node_being_visited)
+                m_node_being_visited->edges.set(reinterpret_cast<FlatPtr>(&cell));
+
+            if (m_graph.get(reinterpret_cast<FlatPtr>(&cell)).has_value())
+                return;
+            m_work_queue.append(*cell);
+        });
+    }
+
     void visit_all_cells()
     {
         while (!m_work_queue.is_empty()) {
@@ -183,13 +241,16 @@ private:
     GraphNode* m_node_being_visited { nullptr };
     Vector<Cell&> m_work_queue;
     HashMap<FlatPtr, GraphNode> m_graph;
+
+    Heap& m_heap;
+    HashTable<HeapBlock*> m_all_live_heap_blocks;
 };
 
 void Heap::dump_graph()
 {
     HashMap<Cell*, HeapRootTypeOrLocation> roots;
     gather_roots(roots);
-    GraphConstructorVisitor visitor(roots);
+    GraphConstructorVisitor visitor(*this, roots);
     vm().bytecode_interpreter().visit_edges(visitor);
     visitor.visit_all_cells();
     visitor.dump();
@@ -240,24 +301,6 @@ void Heap::gather_roots(HashMap<Cell*, HeapRootTypeOrLocation>& roots)
     }
 }
 
-static void add_possible_value(HashMap<FlatPtr, HeapRootTypeOrLocation>& possible_pointers, FlatPtr data, HeapRootTypeOrLocation origin)
-{
-    if constexpr (sizeof(FlatPtr*) == sizeof(Value)) {
-        // Because Value stores pointers in non-canonical form we have to check if the top bytes
-        // match any pointer-backed tag, in that case we have to extract the pointer to its
-        // canonical form and add that as a possible pointer.
-        if ((data & SHIFTED_IS_CELL_PATTERN) == SHIFTED_IS_CELL_PATTERN)
-            possible_pointers.set(Value::extract_pointer_bits(data), move(origin));
-        else
-            possible_pointers.set(data, move(origin));
-    } else {
-        static_assert((sizeof(Value) % sizeof(FlatPtr*)) == 0);
-        // In the 32-bit case we will look at the top and bottom part of Value separately we just
-        // add both the upper and lower bytes as possible pointers.
-        possible_pointers.set(data, move(origin));
-    }
-}
-
 #ifdef HAS_ADDRESS_SANITIZER
 __attribute__((no_sanitize("address"))) void Heap::gather_asan_fake_stack_roots(HashMap<FlatPtr, HeapRootTypeOrLocation>& possible_pointers, FlatPtr addr)
 {
@@ -322,28 +365,26 @@ __attribute__((no_sanitize("address"))) void Heap::gather_conservative_roots(Has
         return IterationDecision::Continue;
     });
 
-    for (auto possible_pointer : possible_pointers.keys()) {
-        if (!possible_pointer)
-            continue;
-        dbgln_if(HEAP_DEBUG, "  ? {}", (void const*)possible_pointer);
-        auto* possible_heap_block = HeapBlock::from_cell(reinterpret_cast<Cell const*>(possible_pointer));
-        if (all_live_heap_blocks.contains(possible_heap_block)) {
-            if (auto* cell = possible_heap_block->cell_from_possible_pointer(possible_pointer)) {
-                if (cell->state() == Cell::State::Live) {
-                    dbgln_if(HEAP_DEBUG, "  ?-> {}", (void const*)cell);
-                    roots.set(cell, *possible_pointers.get(possible_pointer));
-                } else {
-                    dbgln_if(HEAP_DEBUG, "  #-> {}", (void const*)cell);
-                }
-            }
+    for_each_cell_among_possible_pointers(all_live_heap_blocks, possible_pointers, [&](Cell* cell, FlatPtr possible_pointer) {
+        if (cell->state() == Cell::State::Live) {
+            dbgln_if(HEAP_DEBUG, "  ?-> {}", (void const*)cell);
+            roots.set(cell, *possible_pointers.get(possible_pointer));
+        } else {
+            dbgln_if(HEAP_DEBUG, "  #-> {}", (void const*)cell);
         }
-    }
+    });
 }
 
 class MarkingVisitor final : public Cell::Visitor {
 public:
-    explicit MarkingVisitor(HashMap<Cell*, HeapRootTypeOrLocation> const& roots)
+    explicit MarkingVisitor(Heap& heap, HashMap<Cell*, HeapRootTypeOrLocation> const& roots)
+        : m_heap(heap)
     {
+        m_heap.for_each_block([&](auto& block) {
+            m_all_live_heap_blocks.set(&block);
+            return IterationDecision::Continue;
+        });
+
         for (auto* root : roots.keys()) {
             visit(root);
         }
@@ -359,6 +400,24 @@ public:
         m_work_queue.append(cell);
     }
 
+    virtual void visit_possible_values(ReadonlyBytes bytes) override
+    {
+        HashMap<FlatPtr, HeapRootTypeOrLocation> possible_pointers;
+
+        auto* raw_pointer_sized_values = reinterpret_cast<FlatPtr const*>(bytes.data());
+        for (size_t i = 0; i < (bytes.size() / sizeof(FlatPtr)); ++i)
+            add_possible_value(possible_pointers, raw_pointer_sized_values[i], HeapRootType::HeapFunctionCapturedPointer);
+
+        for_each_cell_among_possible_pointers(m_all_live_heap_blocks, possible_pointers, [&](Cell* cell, FlatPtr) {
+            if (cell->is_marked())
+                return;
+            if (cell->state() != Cell::State::Live)
+                return;
+            cell->set_marked(true);
+            m_work_queue.append(*cell);
+        });
+    }
+
     void mark_all_live_cells()
     {
         while (!m_work_queue.is_empty()) {
@@ -367,14 +426,16 @@ public:
     }
 
 private:
+    Heap& m_heap;
     Vector<Cell&> m_work_queue;
+    HashTable<HeapBlock*> m_all_live_heap_blocks;
 };
 
 void Heap::mark_live_cells(HashMap<Cell*, HeapRootTypeOrLocation> const& roots)
 {
     dbgln_if(HEAP_DEBUG, "mark_live_cells:");
 
-    MarkingVisitor visitor(roots);
+    MarkingVisitor visitor(*this, roots);
 
     vm().bytecode_interpreter().visit_edges(visitor);
 

Userland/Libraries/LibJS/Heap/Heap.h

@@ -81,6 +81,9 @@ public:
     void uproot_cell(Cell* cell);
 
 private:
+    friend class MarkingVisitor;
+    friend class GraphConstructorVisitor;
+
     static bool cell_must_survive_garbage_collection(Cell const&);
 
     Cell* allocate_cell(size_t);

Userland/Libraries/LibJS/Heap/HeapFunction.h

@@ -0,0 +1,50 @@
+/*
+ * Copyright (c) 2023, Andreas Kling <kling@serenityos.org>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#pragma once
+
+#include <AK/Function.h>
+#include <LibJS/Heap/Cell.h>
+#include <LibJS/Heap/Heap.h>
+
+namespace JS {
+
+template<typename T>
+class HeapFunction final : public JS::Cell {
+    JS_CELL(HeapFunction, Cell);
+
+public:
+    static NonnullGCPtr<HeapFunction> create(Heap& heap, Function<T> function)
+    {
+        return heap.allocate_without_realm<HeapFunction>(move(function));
+    }
+
+    virtual ~HeapFunction() override = default;
+
+    Function<T> const& function() const { return m_function; }
+
+private:
+    HeapFunction(Function<T> function)
+        : m_function(move(function))
+    {
+    }
+
+    virtual void visit_edges(Visitor& visitor) override
+    {
+        Base::visit_edges(visitor);
+        visitor.visit_possible_values(m_function.raw_capture_range());
+    }
+
+    Function<T> m_function;
+};
+
+template<typename T>
+static NonnullGCPtr<HeapFunction<T>> create_heap_function(Heap& heap, Function<T> function)
+{
+    return HeapFunction<T>::create(heap, move(function));
+}
+
+}

Userland/Libraries/LibJS/Heap/HeapRootTypeOrLocation.h

@@ -11,6 +11,7 @@
 namespace JS {
 
 enum class HeapRootType {
+    HeapFunctionCapturedPointer,
     Handle,
     MarkedVector,
     RegisterPointer,