LibWasm+LibWeb+test-wasm: Refcount Wasm::Module for function references

Prior to funcref, a partial chunk of an invalid module was never needed,
but funcref allows a partially instantiated module to modify imported
tables with references to its own functions, which means we need to keep
the second module alive while that function reference is present within
the imported table.
This was tested by the spectests, but very rarely caught as our GC does
not behave particularly predictably, making it so the offending module
remains in memory just long enough to let the tests pass.

This commit makes it so all function references keep their respective
modules alive.

Userland/Libraries/LibWasm/AbstractMachine/AbstractMachine.cpp

@@ -14,14 +14,14 @@
 
 namespace Wasm {
 
-Optional<FunctionAddress> Store::allocate(ModuleInstance& module, CodeSection::Code const& code, TypeIndex type_index)
+Optional<FunctionAddress> Store::allocate(ModuleInstance& instance, Module const& module, CodeSection::Code const& code, TypeIndex type_index)
 {
     FunctionAddress address { m_functions.size() };
-    if (type_index.value() > module.types().size())
+    if (type_index.value() > instance.types().size())
         return {};
 
-    auto& type = module.types()[type_index.value()];
-    m_functions.empend(WasmFunction { type, module, code });
+    auto& type = instance.types()[type_index.value()];
+    m_functions.empend(WasmFunction { type, instance, module, code });
     return address;
 }
 
@@ -84,6 +84,14 @@ FunctionInstance* Store::get(FunctionAddress address)
     return &m_functions[value];
 }
 
+Module const* Store::get_module_for(Wasm::FunctionAddress address)
+{
+    auto* function = get(address);
+    if (!function || function->has<HostFunction>())
+        return nullptr;
+    return function->get<WasmFunction>().module_ref().ptr();
+}
+
 TableInstance* Store::get(TableAddress address)
 {
     auto value = address.value();
@@ -223,7 +231,7 @@ InstantiationResult AbstractMachine::instantiate(Module const& module, Vector<Ex
     size_t i = 0;
     for (auto& code : module.code_section().functions()) {
         auto type_index = module.function_section().types()[i];
-        auto address = m_store.allocate(main_module_instance, code, type_index);
+        auto address = m_store.allocate(main_module_instance, module, code, type_index);
         VERIFY(address.has_value());
         auxiliary_instance.functions().append(*address);
         module_functions.append(*address);

Userland/Libraries/LibWasm/AbstractMachine/AbstractMachine.h

@@ -51,6 +51,7 @@ public:
     };
     struct Func {
         FunctionAddress address;
+        RefPtr<Module> source_module; // null if host function.
     };
     struct Extern {
         ExternAddress address;
@@ -139,7 +140,7 @@ public:
         // 2: null funcref
         // 3: null externref
         ref.ref().visit(
-            [&](Reference::Func const& func) { m_value = u128(bit_cast<u64>(func.address), 0); },
+            [&](Reference::Func const& func) { m_value = u128(bit_cast<u64>(func.address), bit_cast<u64>(func.source_module.ptr())); },
             [&](Reference::Extern const& func) { m_value = u128(bit_cast<u64>(func.address), 1); },
             [&](Reference::Null const& null) { m_value = u128(0, null.type.kind() == ValueType::Kind::FunctionReference ? 2 : 3); });
     }
@@ -177,17 +178,15 @@ public:
             return bit_cast<f64>(m_value.low());
         }
         if constexpr (IsSame<T, Reference>) {
-            switch (m_value.high()) {
+            switch (m_value.high() & 3) {
             case 0:
-                return Reference { Reference::Func { bit_cast<FunctionAddress>(m_value.low()) } };
+                return Reference { Reference::Func { bit_cast<FunctionAddress>(m_value.low()), bit_cast<Wasm::Module*>(m_value.high()) } };
             case 1:
                 return Reference { Reference::Extern { bit_cast<ExternAddress>(m_value.low()) } };
             case 2:
                 return Reference { Reference::Null { ValueType(ValueType::Kind::FunctionReference) } };
             case 3:
                 return Reference { Reference::Null { ValueType(ValueType::Kind::ExternReference) } };
-            default:
-                VERIFY_NOT_REACHED();
             }
         }
         VERIFY_NOT_REACHED();
@@ -341,20 +340,23 @@ private:
 
 class WasmFunction {
 public:
-    explicit WasmFunction(FunctionType const& type, ModuleInstance const& module, CodeSection::Code const& code)
+    explicit WasmFunction(FunctionType const& type, ModuleInstance const& instance, Module const& module, CodeSection::Code const& code)
         : m_type(type)
-        , m_module(module)
+        , m_module(module.make_weak_ptr())
+        , m_module_instance(instance)
         , m_code(code)
     {
     }
 
     auto& type() const { return m_type; }
-    auto& module() const { return m_module; }
+    auto& module() const { return m_module_instance; }
     auto& code() const { return m_code; }
+    RefPtr<Module const> module_ref() const { return m_module.strong_ref(); }
 
 private:
     FunctionType m_type;
-    ModuleInstance const& m_module;
+    WeakPtr<Module const> m_module;
+    ModuleInstance const& m_module_instance;
     CodeSection::Code const& m_code;
 };
 
@@ -554,7 +556,7 @@ class Store {
 public:
     Store() = default;
 
-    Optional<FunctionAddress> allocate(ModuleInstance&, CodeSection::Code const&, TypeIndex);
+    Optional<FunctionAddress> allocate(ModuleInstance&, Module const&, CodeSection::Code const&, TypeIndex);
     Optional<FunctionAddress> allocate(HostFunction&&);
     Optional<TableAddress> allocate(TableType const&);
     Optional<MemoryAddress> allocate(MemoryType const&);
@@ -562,6 +564,7 @@ public:
     Optional<GlobalAddress> allocate(GlobalType const&, Value);
     Optional<ElementAddress> allocate(ValueType const&, Vector<Reference>);
 
+    Module const* get_module_for(FunctionAddress);
     FunctionInstance* get(FunctionAddress);
     TableInstance* get(TableAddress);
     MemoryInstance* get(MemoryAddress);

Userland/Libraries/LibWasm/AbstractMachine/BytecodeInterpreter.cpp

@@ -864,7 +864,7 @@ ALWAYS_INLINE void BytecodeInterpreter::interpret_instruction(Configuration& con
         auto index = instruction.arguments().get<FunctionIndex>().value();
         auto& functions = configuration.frame().module().functions();
         auto address = functions[index];
-        configuration.value_stack().append(Value(address.value()));
+        configuration.value_stack().append(Value(Reference { Reference::Func { address, configuration.store().get_module_for(address) } }));
         return;
     }
     case Instructions::ref_is_null.value(): {