From 4594b99ef3871e4d4506348fb34ace98a4747589 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pawe=C5=82=20Bylica?= <chfast@gmail.com>
Date: Wed, 31 Mar 2021 19:33:30 +0200
Subject: [PATCH] Add segmented stack space to ExecutionContext

---
 lib/fizzy/execute.cpp           | 10 ++++--
 lib/fizzy/execution_context.hpp | 57 ++++++++++++++++++++++++++++++---
 2 files changed, 60 insertions(+), 7 deletions(-)

diff --git a/lib/fizzy/execute.cpp b/lib/fizzy/execute.cpp
index 447ff86505..d78b867fa7 100644
--- a/lib/fizzy/execute.cpp
+++ b/lib/fizzy/execute.cpp
@@ -565,6 +565,7 @@ ExecutionResult execute(
         return Trap;
 
     const auto& func_type = instance.module->get_function_type(func_idx);
+    const auto args_count = func_type.inputs.size();
 
     assert(instance.module->imported_function_types.size() == instance.imported_functions.size());
     if (func_idx < instance.imported_functions.size())
@@ -573,10 +574,13 @@ ExecutionResult execute(
     const auto& code = instance.module->get_code(func_idx);
     auto* const memory = instance.memory.get();
 
-    const auto local_ctx = ctx.create_local_context();
+    const auto required_stack_space =
+        args_count + code.local_count + static_cast<size_t>(code.max_stack_height);
 
-    OperandStack stack(args, func_type.inputs.size(), code.local_count,
-        static_cast<size_t>(code.max_stack_height));
+    const auto local_ctx = ctx.create_local_context(required_stack_space);
+
+    OperandStack stack(
+        args, args_count, code.local_count, static_cast<size_t>(code.max_stack_height));
 
     const uint8_t* pc = code.instructions.data();
 
diff --git a/lib/fizzy/execution_context.hpp b/lib/fizzy/execution_context.hpp
index 0b9be912d3..280316bf35 100644
--- a/lib/fizzy/execution_context.hpp
+++ b/lib/fizzy/execution_context.hpp
@@ -4,13 +4,19 @@
 
 #pragma once
 
+#include "value.hpp"
+#include <cassert>
+#include <cstddef>
+
 namespace fizzy
 {
 /// The storage for information shared by calls in the same execution "thread".
 /// Users may decide how to allocate the execution context, but some good defaults are available.
 class ExecutionContext
 {
-    /// Call local execution context.
+    static constexpr size_t DefaultStackSpaceSegmentSize = 100;
+
+    /// Call depth increment guard.
     /// It will automatically decrement the call depth to the original value
     /// when going out of scope.
     class [[nodiscard]] LocalContext
@@ -18,24 +24,67 @@ class ExecutionContext
         ExecutionContext& m_shared_ctx;  ///< Reference to the shared execution context.
 
     public:
+        Value* stack_space = nullptr;
+        Value* prev_stack_space_segment = nullptr;
+        size_t prev_free_stack_space = 0;
+
         LocalContext(const LocalContext&) = delete;
         LocalContext(LocalContext&&) = delete;
         LocalContext& operator=(const LocalContext&) = delete;
         LocalContext& operator=(LocalContext&&) = delete;
 
-        explicit LocalContext(ExecutionContext& ctx) noexcept : m_shared_ctx{ctx}
+        LocalContext(ExecutionContext& ctx, size_t required_stack_space) : m_shared_ctx{ctx}
         {
             ++m_shared_ctx.depth;
+
+            prev_free_stack_space = m_shared_ctx.free_stack_space;
+
+            if (required_stack_space <= m_shared_ctx.free_stack_space)
+            {
+                // Must be a segment of default size or required_stack_space is 0.
+                const auto offset =
+                    DefaultStackSpaceSegmentSize - m_shared_ctx.free_stack_space;
+                stack_space = m_shared_ctx.stack_space_segment + offset;
+                prev_free_stack_space = m_shared_ctx.free_stack_space;
+                m_shared_ctx.free_stack_space -= required_stack_space;
+            }
+            else
+            {
+                prev_stack_space_segment = m_shared_ctx.stack_space_segment;
+                const auto new_segment_size =
+                    std::max(DefaultStackSpaceSegmentSize, required_stack_space);
+                m_shared_ctx.stack_space_segment = new Value[new_segment_size];
+                stack_space = m_shared_ctx.stack_space_segment;
+                m_shared_ctx.free_stack_space = new_segment_size - required_stack_space;
+            }
         }
 
-        ~LocalContext() noexcept { --m_shared_ctx.depth; }
+        ~LocalContext() noexcept
+        {
+            --m_shared_ctx.depth;
+
+            m_shared_ctx.free_stack_space = prev_free_stack_space;
+            if (prev_stack_space_segment != nullptr)
+            {
+                assert(m_shared_ctx.stack_space_segment == stack_space);
+                delete[] stack_space;
+                m_shared_ctx.stack_space_segment = prev_stack_space_segment;
+            }
+        }
     };
 
 public:
+    Value first_stack_space_segment[DefaultStackSpaceSegmentSize];
+    Value* stack_space_segment = first_stack_space_segment;
+    size_t free_stack_space = DefaultStackSpaceSegmentSize;
+
     int depth = 0;  ///< Current call depth.
 
     /// Increments the call depth and returns the local call context which
     /// decrements the call depth back to the original value when going out of scope.
-    LocalContext create_local_context() noexcept { return LocalContext{*this}; }
+    LocalContext create_local_context(size_t required_stack_space = 0)
+    {
+        return LocalContext{*this, required_stack_space};
+    }
 };
 }  // namespace fizzy