Execution thread shared stack space

lib/fizzy/execute.cpp

@@ -54,14 +54,14 @@ void branch(const Code& code, OperandStack& stack, const Instr*& pc, const uint8
 
 template <class F>
 bool invoke_function(const FuncType& func_type, const F& func, Instance& instance,
-    OperandStack& stack, int depth) noexcept
+    OperandStack& stack, ThreadContext& thread_context) noexcept
 {
     const auto num_args = func_type.inputs.size();
     assert(stack.size() >= num_args);
     span<const Value> call_args{stack.rend() - num_args, num_args};
     stack.drop(num_args);
 
-    const auto ret = func(instance, call_args, depth + 1);
+    const auto ret = func(instance, call_args, thread_context);
     // Bubble up traps
     if (ret.trapped)
         return false;
@@ -78,12 +78,13 @@ bool invoke_function(const FuncType& func_type, const F& func, Instance& instanc
 }
 
 inline bool invoke_function(const FuncType& func_type, uint32_t func_idx, Instance& instance,
-    OperandStack& stack, int depth) noexcept
+    OperandStack& stack, ThreadContext& thread_context) noexcept
 {
-    const auto func = [func_idx](Instance& _instance, span<const Value> args, int _depth) noexcept {
-        return execute(_instance, func_idx, args, _depth);
+    const auto func = [func_idx](Instance& _instance, span<const Value> args,
+                          ThreadContext& _thread_context) noexcept {
+        return execute(_instance, func_idx, args, _thread_context);
     };
-    return invoke_function(func_type, func, instance, stack, depth);
+    return invoke_function(func_type, func, instance, stack, thread_context);
 }
 
 template <typename T>
@@ -459,25 +460,52 @@ __attribute__((no_sanitize("float-cast-overflow"))) inline constexpr float demot
     return static_cast<float>(value);
 }
 
+class ThreadContextGuard
+{
+    ThreadContext& m_thread_context;
+    Value* const m_free_space;
+
+public:
+    explicit ThreadContextGuard(ThreadContext& thread_context) noexcept
+      : m_thread_context{thread_context}, m_free_space{m_thread_context.free_space}
+    {
+        m_thread_context.lock();
+    }
+
+    ~ThreadContextGuard() noexcept
+    {
+        m_thread_context.unlock();
+        m_thread_context.free_space = m_free_space;
+    }
+};
+
 }  // namespace
 
-ExecutionResult execute(
-    Instance& instance, FuncIdx func_idx, span<const Value> args, int depth) noexcept
+ExecutionResult execute(Instance& instance, FuncIdx func_idx, span<const Value> args,
+    ThreadContext& thread_context) noexcept
 {
-    assert(depth >= 0);
-    if (depth > CallStackLimit)
+    if (thread_context.depth > CallStackLimit)
         return Trap;
 
+    ThreadContextGuard guard{thread_context};
+
     assert(args.size() == instance.module.get_function_type(func_idx).inputs.size());
 
     assert(instance.module.imported_function_types.size() == instance.imported_functions.size());
     if (func_idx < instance.imported_functions.size())
-        return instance.imported_functions[func_idx].function(instance, args, depth);
+        return instance.imported_functions[func_idx].function(instance, args, thread_context);
 
     const auto& code = instance.module.get_code(func_idx);
     auto* const memory = instance.memory.get();
 
-    OperandStack stack(args, code.local_count, static_cast<size_t>(code.max_stack_height));
+    auto* current_thread_context = &thread_context;
+
+    OperandStack stack(args, code.local_count, static_cast<size_t>(code.max_stack_height),
+        current_thread_context->free_space, current_thread_context->space_left());
+    const auto storage_size_required =
+        args.size() + code.local_count + static_cast<size_t>(code.max_stack_height);
+    if (storage_size_required <= current_thread_context->space_left())
+        current_thread_context->free_space += storage_size_required;
 
     const Instr* pc = code.instructions.data();
     const uint8_t* immediates = code.immediates.data();
@@ -562,7 +590,8 @@ ExecutionResult execute(
             const auto called_func_idx = read<uint32_t>(immediates);
             const auto& called_func_type = instance.module.get_function_type(called_func_idx);
 
-            if (!invoke_function(called_func_type, called_func_idx, instance, stack, depth))
+            if (!invoke_function(
+                    called_func_type, called_func_idx, instance, stack, *current_thread_context))
                 goto trap;
             break;
         }
@@ -587,7 +616,8 @@ ExecutionResult execute(
             if (expected_type != actual_type)
                 goto trap;
 
-            if (!invoke_function(actual_type, called_func->function, instance, stack, depth))
+            if (!invoke_function(
+                    actual_type, called_func->function, instance, stack, *current_thread_context))
                 goto trap;
             break;
         }

lib/fizzy/execute.hpp

@@ -36,9 +36,17 @@ struct ExecutionResult
 constexpr ExecutionResult Void{true};
 constexpr ExecutionResult Trap{false};
 
-// Execute a function on an instance.
-ExecutionResult execute(
-    Instance& instance, FuncIdx func_idx, span<const Value> args, int depth = 0) noexcept;
+/// Execute a function on an instance.
+ExecutionResult execute(Instance& instance, FuncIdx func_idx, span<const Value> args,
+    ThreadContext& thread_context) noexcept;
+
+/// Execute a function on an instance with implicit depth 0.
+inline ExecutionResult execute(
+    Instance& instance, FuncIdx func_idx, span<const Value> args) noexcept
+{
+    ThreadContext thread_context;
+    return execute(instance, func_idx, args, thread_context);
+}
 
 inline ExecutionResult execute(
     Instance& instance, FuncIdx func_idx, std::initializer_list<Value> args) noexcept

lib/fizzy/instantiate.cpp

@@ -328,7 +328,9 @@ std::unique_ptr<Instance> instantiate(Module module,
         for (const auto idx : instance->module.elementsec[i].init)
         {
             auto func = [idx, &instance_ref = *instance](fizzy::Instance&, span<const Value> args,
-                            int depth) { return execute(instance_ref, idx, args, depth); };
+                            ThreadContext& thread_context) {
+                return execute(instance_ref, idx, args, thread_context);
+            };
 
             *it_table++ =
                 ExternalFunction{std::move(func), instance->module.get_function_type(idx)};
@@ -361,7 +363,9 @@ std::unique_ptr<Instance> instantiate(Module module,
                         auto& table_function = (*it_table)->function;
                         table_function = [shared_instance, func = std::move(table_function)](
                                              fizzy::Instance& _instance, span<const Value> args,
-                                             int depth) { return func(_instance, args, depth); };
+                                             ThreadContext& thread_context) {
+                            return func(_instance, args, thread_context);
+                        };
                         ++it_table;
                     }
                 }
@@ -432,8 +436,9 @@ std::optional<ExternalFunction> find_exported_function(Instance& instance, std::
         return std::nullopt;
 
     const auto idx = *opt_index;
-    auto func = [idx, &instance](fizzy::Instance&, span<const Value> args, int depth) {
-        return execute(instance, idx, args, depth);
+    auto func = [idx, &instance](
+                    fizzy::Instance&, span<const Value> args, ThreadContext& thread_context) {
+        return execute(instance, idx, args, thread_context);
     };
 
     return ExternalFunction{std::move(func), instance.module.get_function_type(idx)};

lib/fizzy/instantiate.hpp

@@ -22,9 +22,30 @@ namespace fizzy
 struct ExecutionResult;
 struct Instance;
 
+class ThreadContext
+{
+public:
+    static constexpr size_t N = 128;
+
+    Value stack_space[N];
+
+    Value* free_space = stack_space;
+
+    size_t space_left() const noexcept
+    {
+        return static_cast<size_t>((stack_space + N) - free_space);
+    }
+
+    int depth = 0;
+
+    void lock() noexcept { ++depth; }
+
+    void unlock() noexcept { --depth; }
+};
+
 struct ExternalFunction
 {
-    std::function<ExecutionResult(Instance&, span<const Value>, int depth)> function;
+    std::function<ExecutionResult(Instance&, span<const Value>, ThreadContext&)> function;
     FuncType type;
 };
 
@@ -102,7 +123,7 @@ struct ImportedFunction
     std::string name;
     std::vector<ValType> inputs;
     std::optional<ValType> output;
-    std::function<ExecutionResult(Instance&, span<const Value>, int depth)> function;
+    std::function<ExecutionResult(Instance&, span<const Value>, ThreadContext&)> function;
 };
 
 // Create vector of ExternalFunctions ready to be passed to instantiate.

lib/fizzy/stack.hpp

@@ -59,9 +59,6 @@ class Stack
 /// from the stack itself.
 class OperandStack
 {
-    /// The size of the pre-allocated internal storage: 128 bytes.
-    static constexpr auto small_storage_size = 128 / sizeof(Value);
-
     /// The pointer to the top item of the operand stack,
     /// or below the stack bottom if stack is empty.
     ///
@@ -77,9 +74,6 @@ class OperandStack
     /// The pointer to the bottom of the operand stack.
     Value* m_bottom;
 
-    /// The pre-allocated internal storage.
-    Value m_small_storage[small_storage_size];
-
     /// The unbounded storage for items.
     std::unique_ptr<Value[]> m_large_storage;
 
@@ -96,14 +90,14 @@ class OperandStack
     ///                             after the arguments.
     /// @param max_stack_height     The maximum operand stack height in the function. This excludes
     ///                             args and num_local_variables.
-    OperandStack(span<const Value> args, size_t num_local_variables, size_t max_stack_height)
+    OperandStack(span<const Value> args, size_t num_local_variables, size_t max_stack_height, Value* external_storage, size_t external_storage_size)
     {
         const auto num_args = args.size();
         const auto storage_size_required = num_args + num_local_variables + max_stack_height;
 
-        if (storage_size_required <= small_storage_size)
+        if (storage_size_required <= external_storage_size)
         {
-            m_locals = &m_small_storage[0];
+            m_locals = external_storage;
         }
         else
         {
@@ -112,7 +106,10 @@ class OperandStack
         }
 
         m_bottom = m_locals + num_args + num_local_variables;
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Warray-bounds"
         m_top = m_bottom - 1;
+#pragma GCC diagnostic pop
 
         std::copy(std::begin(args), std::end(args), m_locals);
         std::fill_n(m_locals + num_args, num_local_variables, 0);

test/unittests/api_test.cpp

@@ -16,10 +16,10 @@ namespace
 {
 auto function_returning_value(Value value) noexcept
 {
-    return [value](Instance&, span<const Value>, int) { return value; };
+    return [value](Instance&, span<const Value>, ThreadContext&) { return value; };
 }
 
-ExecutionResult function_returning_void(Instance&, span<const Value>, int) noexcept
+ExecutionResult function_returning_void(Instance&, span<const Value>, ThreadContext&) noexcept
 {
     return Void;
 }
@@ -243,7 +243,8 @@ TEST(api, find_exported_function)
 
     auto opt_function = find_exported_function(*instance, "foo");
     ASSERT_TRUE(opt_function);
-    EXPECT_THAT(opt_function->function(*instance, {}, 0), Result(42));
+    ThreadContext thread_context;
+    EXPECT_THAT(opt_function->function(*instance, {}, thread_context), Result(42));
     EXPECT_TRUE(opt_function->type.inputs.empty());
     ASSERT_EQ(opt_function->type.outputs.size(), 1);
     EXPECT_EQ(opt_function->type.outputs[0], ValType::i32);
@@ -262,15 +263,15 @@ TEST(api, find_exported_function)
         "0061736d010000000105016000017f021001087370656374657374036261720000040401700000050401010102"
         "0606017f0041000b07170403666f6f000001670300037461620100036d656d0200");
 
-    auto bar = [](Instance&, span<const Value>, int) { return Value{42}; };
+    auto bar = [](Instance&, span<const Value>, ThreadContext&) { return Value{42}; };
     const auto bar_type = FuncType{{}, {ValType::i32}};
 
     auto instance_reexported_function =
         instantiate(parse(wasm_reexported_function), {{bar, bar_type}});
 
     opt_function = find_exported_function(*instance_reexported_function, "foo");
     ASSERT_TRUE(opt_function);
-    EXPECT_THAT(opt_function->function(*instance, {}, 0), Result(42));
+    EXPECT_THAT(opt_function->function(*instance, {}, thread_context), Result(42));
     EXPECT_TRUE(opt_function->type.inputs.empty());
     ASSERT_EQ(opt_function->type.outputs.size(), 1);
     EXPECT_EQ(opt_function->type.outputs[0], ValType::i32);
@@ -411,8 +412,10 @@ TEST(api, find_exported_table)
     ASSERT_TRUE(opt_table);
     EXPECT_EQ(opt_table->table, instance->table.get());
     EXPECT_EQ(opt_table->table->size(), 2);
-    EXPECT_THAT((*opt_table->table)[0]->function(*instance, {}, 0), Result(2));
-    EXPECT_THAT((*opt_table->table)[1]->function(*instance, {}, 0), Result(1));
+
+    ThreadContext thread_context;
+    EXPECT_THAT((*opt_table->table)[0]->function(*instance, {}, thread_context), Result(2));
+    EXPECT_THAT((*opt_table->table)[1]->function(*instance, {}, thread_context), Result(1));
     EXPECT_EQ(opt_table->limits.min, 2);
     ASSERT_TRUE(opt_table->limits.max.has_value());
     EXPECT_EQ(opt_table->limits.max, 20);