From 882f5397f78211fe785f4c9ad7db1842116ba0ff Mon Sep 17 00:00:00 2001
From: Kelvin Nishikawa <kelnishi@gmail.com>
Date: Mon, 18 Nov 2024 00:58:25 -0800
Subject: [PATCH] Stack memory optimization

---
 Wacs.Core/Instructions/Control.cs             |  4 +-
 Wacs.Core/Instructions/Variable.cs            | 16 +++---
 Wacs.Core/Runtime/ExecContext.cs              | 55 ++++++++++---------
 Wacs.Core/Runtime/Frame.cs                    | 29 ++++++++--
 Wacs.Core/Runtime/Label.cs                    |  2 +-
 Wacs.Core/Runtime/RuntimeAttributes.cs        |  4 ++
 Wacs.Core/Runtime/WasmRuntime.cs              |  5 +-
 Wacs.Core/Types/IndexSpace.cs                 | 45 ++++++++-------
 .../Utilities/{IReusable.cs => IPoolable.cs}  |  4 +-
 Wacs.Core/Utilities/ObjectPoolExtension.cs    | 38 +++++++++++++
 Wacs.Core/Utilities/ReusableStack.cs          |  4 +-
 ...oledObjectPolicy.cs => StackPoolPolicy.cs} | 12 ++--
 Wacs.Core/Validation/WasmValidationContext.cs |  5 +-
 13 files changed, 148 insertions(+), 75 deletions(-)
 rename Wacs.Core/Utilities/{IReusable.cs => IPoolable.cs} (91%)
 create mode 100644 Wacs.Core/Utilities/ObjectPoolExtension.cs
 rename Wacs.Core/Utilities/{LocalsSpacePooledObjectPolicy.cs => StackPoolPolicy.cs} (77%)

diff --git a/Wacs.Core/Instructions/Control.cs b/Wacs.Core/Instructions/Control.cs
index 9d7d9c5..418c486 100644
--- a/Wacs.Core/Instructions/Control.cs
+++ b/Wacs.Core/Instructions/Control.cs
@@ -723,8 +723,8 @@ public override void Execute(ExecContext context)
                 $"Instruction return failed. Operand stack underflow");
             //We're managing separate stacks, so we won't need to shift the operands
             // var vals = context.OpStack.PopResults(context.Frame.Type.ResultType);
-            var frame = context.PopFrame();
-            context.ResumeSequence(frame.ContinuationAddress);
+            var address = context.PopFrame();
+            context.ResumeSequence(address);
         }
     }
 
diff --git a/Wacs.Core/Instructions/Variable.cs b/Wacs.Core/Instructions/Variable.cs
index 3c356c4..8255aa7 100644
--- a/Wacs.Core/Instructions/Variable.cs
+++ b/Wacs.Core/Instructions/Variable.cs
@@ -54,7 +54,7 @@ public override string RenderText(ExecContext? context)
             if (!context.Frame.Locals.Contains(Index))
                 return $"{base.RenderText(context)} {Index.Value}";
             
-            var value = context.Frame.Locals[Index];
+            var value = context.Frame.Locals.Get(Index);
             string valStr = $" (;>{value}<;)";
             return $"{base.RenderText(context)} {Index.Value}{valStr}";
         }
@@ -70,7 +70,7 @@ private static void ValidateLocalGet(IWasmValidationContext context, LocalIdx lo
         {
             context.Assert(context.Locals.Contains(localIndex),
                 $"Instruction local.get was invalid. Context Locals did not contain {localIndex}");
-            var value = context.Locals[localIndex];
+            var value = context.Locals.Get(localIndex);
             context.OpStack.PushType(value.Type);
         }
 
@@ -81,7 +81,7 @@ private static void ExecuteLocalGet(ExecContext context, LocalIdx localIndex)
             context.Assert( context.Frame.Locals.Contains(localIndex),
                 $"Instruction local.get could not get Local {localIndex}");
             //3.
-            var value = context.Frame.Locals[localIndex];
+            var value = context.Frame.Locals.Get(localIndex);
             //4.
             context.OpStack.PushValue(value);
         }
@@ -91,7 +91,7 @@ private static void ValidateLocalSet(IWasmValidationContext context, LocalIdx lo
         {
             context.Assert(context.Locals.Contains(localIndex),
                 $"Instruction local.set was invalid. Context Locals did not contain {localIndex}");
-            var value = context.Locals[localIndex];
+            var value = context.Locals.Get(localIndex);
             context.OpStack.PopType(value.Type);
         }
 
@@ -104,12 +104,12 @@ private static void ExecuteLocalSet(ExecContext context, LocalIdx localIndex)
             //3.
             context.Assert( context.OpStack.HasValue,
                 $"Operand Stack underflow in instruction local.set");
-            var localValue = context.Frame.Locals[localIndex];
+            var localValue = context.Frame.Locals.Get(localIndex);
             var type = localValue.Type;
             //4.
             var value = context.OpStack.PopType(type);
             //5.
-            context.Frame.Locals[localIndex] = value;
+            context.Frame.Locals.Set(localIndex, value);
         }
 
         //0x22
@@ -118,7 +118,7 @@ private static void ValidateLocalTee(IWasmValidationContext context, LocalIdx lo
         {
             context.Assert(context.Locals.Contains(localIndex),
                 $"Instruction local.tee was invalid. Context Locals did not contain {localIndex}");
-            var value = context.Locals[localIndex];
+            var value = context.Locals.Get(localIndex);
             context.OpStack.PopType(value.Type);
             context.OpStack.PushType(value.Type);
             context.OpStack.PushType(value.Type);
@@ -131,7 +131,7 @@ private static void ExecuteLocalTee(ExecContext context, LocalIdx localIndex)
             //1.
             context.Assert( context.OpStack.HasValue,
                 $"Operand Stack underflow in instruction local.tee");
-            var localValue = context.Frame.Locals[localIndex];
+            var localValue = context.Frame.Locals.Get(localIndex);
             //2.
             var value = context.OpStack.PopType(localValue.Type);
             //3.
diff --git a/Wacs.Core/Runtime/ExecContext.cs b/Wacs.Core/Runtime/ExecContext.cs
index 51a1f79..8757e68 100644
--- a/Wacs.Core/Runtime/ExecContext.cs
+++ b/Wacs.Core/Runtime/ExecContext.cs
@@ -15,6 +15,7 @@
 //  */
 
 using System;
+using System.Buffers;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
@@ -51,13 +52,15 @@ public ExecContext(Store store, RuntimeAttributes? attributes = default)
         {
             Store = store;
             Attributes = attributes ?? new RuntimeAttributes();
-            FrameStack = new(Attributes.InitialCallStack, Attributes.GrowCallStack);
-            LabelStack = new(Attributes.InitialCallStack * 2, Attributes.GrowCallStack);
-            CallStack = FrameStack.GetSubStack();
 
-            var poolPolicy = new LocalsSpacePooledObjectPolicy();
-            LocalsPool = new DefaultObjectPool<LocalsSpace>(poolPolicy, Attributes.InitialCallStack);
+            FramePool = new DefaultObjectPool<Frame>(new StackPoolPolicy<Frame>(), Attributes.MaxCallStack)
+                .Prime(Attributes.InitialCallStack);
             
+            LocalsDataPool = ArrayPool<Value>.Create(Attributes.MaxFunctionLocals, Attributes.LocalPoolSize);
+            
+            CallStack = new (Attributes.InitialCallStack);
+            LabelStack = new(Attributes.InitialLabelsStack, Attributes.GrowLabelsStack);
+
             _hostReturnSequence = new InstructionSequence(
                 InstructionFactory.CreateInstruction<InstFuncReturn>(OpCode.Func)
             );
@@ -75,11 +78,10 @@ public ExecContext(Store store, RuntimeAttributes? attributes = default)
         public Store Store { get; }
         public OpStack OpStack { get; }
 
-        private ReusableStack<Frame> FrameStack { get; }
-        private SubStack<Frame> CallStack { get; }
         private ReusableStack<Label> LabelStack { get; }
-
-        private DefaultObjectPool<LocalsSpace> LocalsPool { get; }
+        private ObjectPool<Frame> FramePool { get; }
+        private ArrayPool<Value> LocalsDataPool { get; }
+        private Stack<Frame> CallStack { get; }
 
         public Frame Frame => CallStack.Peek();
 
@@ -107,14 +109,16 @@ public Frame ReserveFrame(
         {
             locals ??= Array.Empty<ValType>();
             
-            var frame = CallStack.Reserve();
+            var frame = FramePool.Get();
             frame.Module = module;
             frame.Type = type;
             frame.Labels = LabelStack.GetSubStack();
             frame.Index = index;
             frame.ContinuationAddress = GetPointer();
-            frame.Locals = LocalsPool.Get();
-            frame.Locals.Enscribe(type.ParameterTypes.Types, locals);
+            
+            int capacity = type.ParameterTypes.Types.Length + locals.Length;
+            var localData = LocalsDataPool.Rent(capacity);
+            frame.Locals = new(localData, type.ParameterTypes.Types, locals);
             
             return frame;
         }
@@ -125,18 +129,18 @@ public void PushFrame(Frame frame)
                 throw new WasmRuntimeException($"Runtime call stack exhausted {CallStack.Count}");
             
             CallStack.Push(frame);
-            frame.Labels = LabelStack.GetSubStack();
         }
 
-        public Frame PopFrame()
+        public InstructionPointer PopFrame()
         {
             var frame = CallStack.Pop();
             frame.Labels.Drop();
-            
-            LocalsPool.Return(frame.Locals);
-            frame.Locals = null!;
-            
-            return frame;
+            if (frame.Locals.Data != null)
+                frame.ReturnLocals(LocalsDataPool);
+
+            var address = frame.ContinuationAddress;
+            FramePool.Return(frame);
+            return address;
         }
 
         public void ResetStack(Label label)
@@ -147,7 +151,6 @@ public void ResetStack(Label label)
             }
         }
 
-
         public void FlushCallStack()
         {
             while (CallStack.Count > 0)
@@ -193,10 +196,10 @@ public void EnterBlock(Block block, ResultType resultType, ByteCode inst, Stack<
         // @Spec 4.4.9.2. Exit Block
         public void ExitBlock()
         {
-            var label = Frame.Labels.Pop();
+            var addr = Frame.PopLabel();
             // We manage separate stacks, so we don't need to relocate the operands
             // var vals = OpStack.PopResults(label.Type);
-            ResumeSequence(label.ContinuationAddress);
+            ResumeSequence(addr);
         }
 
         // @Spec 4.4.10.1 Function Invocation
@@ -243,13 +246,13 @@ private void Invoke(FunctionInstance wasmFunc, FuncIdx idx)
             //Load parameters
             while (_asideVals.Count > 0)
             {
-                frame.Locals[(LocalIdx)li] = _asideVals.Pop();
+                frame.Locals.Set((LocalIdx)li, _asideVals.Pop());
                 li += 1;
             }
             //Set the Locals to default
             for (int ti = 0; li < localCount; ++li, ++ti)
             {
-                frame.Locals[(LocalIdx)li] = new Value(t[ti]);
+                frame.Locals.Set((LocalIdx)li, new Value(t[ti]));
             }
 
             //9.
@@ -286,10 +289,10 @@ public void FunctionReturn()
             // var vals = OpStack.PopResults(Frame.Type.ResultType);
             //5.
             //6.
-            var frame = PopFrame();
+            var address = PopFrame();
             //7. split stack, values left in place 
             //8.
-            ResumeSequence(frame.ContinuationAddress);
+            ResumeSequence(address);
         }
 
         public IInstruction? Next()
diff --git a/Wacs.Core/Runtime/Frame.cs b/Wacs.Core/Runtime/Frame.cs
index f355ea8..b4d1e06 100644
--- a/Wacs.Core/Runtime/Frame.cs
+++ b/Wacs.Core/Runtime/Frame.cs
@@ -14,6 +14,9 @@
 //  * limitations under the License.
 //  */
 
+using System.Buffers;
+using System.Collections.Generic;
+using Microsoft.Extensions.ObjectPool;
 using Wacs.Core.OpCodes;
 using Wacs.Core.Runtime.Types;
 using Wacs.Core.Types;
@@ -21,11 +24,15 @@
 
 namespace Wacs.Core.Runtime
 {
-    public class Frame : IReusable<Frame>
+    public class Frame : IPoolable
     {
+        // public ObjectPool<Label>? LabelPool = null;
+        
         public SubStack<Label> Labels;
+
+        // public readonly Stack<Label> Labels = new();
         public ModuleInstance Module { get; set; } = null!;
-        public LocalsSpace Locals { get; set; } = null!;
+        public LocalsSpace Locals { get; set; }
         public InstructionPointer ContinuationAddress { get; set; } = InstructionPointer.Nil;
         public FunctionType Type { get; set; } = null!;
         public FuncIdx Index { get; set; }
@@ -38,13 +45,19 @@ public void Clear()
         {
             Labels = default;
             Module = default!;
-            Locals = default!;
+            Locals = default;
             ContinuationAddress = default;
             Type = default!;
             Index = default!;
             FuncId = string.Empty;
         }
 
+        public void ReturnLocals(ArrayPool<Value> dataPool)
+        {
+            dataPool.Return(Locals.Data);
+            Locals = default!;
+        }
+
         public bool Equals(Frame other)
         {
             return ReferenceEquals(this, other);
@@ -66,8 +79,16 @@ public void ForceLabels(int depth)
 
             while (Labels.Count > depth)
             {
-                Labels.Pop();
+                PopLabel();
             }
         }
+
+        public InstructionPointer PopLabel()
+        {
+            var label = Labels.Pop();
+            var addr = label.ContinuationAddress;
+            // LabelPool?.Return(label);
+            return addr;
+        }
     }
 }
\ No newline at end of file
diff --git a/Wacs.Core/Runtime/Label.cs b/Wacs.Core/Runtime/Label.cs
index 096055a..83b352e 100644
--- a/Wacs.Core/Runtime/Label.cs
+++ b/Wacs.Core/Runtime/Label.cs
@@ -20,7 +20,7 @@
 
 namespace Wacs.Core.Runtime
 {
-    public class Label : IReusable<Label>
+    public class Label : IPoolable
     {
         public int Arity;
 
diff --git a/Wacs.Core/Runtime/RuntimeAttributes.cs b/Wacs.Core/Runtime/RuntimeAttributes.cs
index 43dc93c..68f394b 100644
--- a/Wacs.Core/Runtime/RuntimeAttributes.cs
+++ b/Wacs.Core/Runtime/RuntimeAttributes.cs
@@ -27,7 +27,11 @@ public class RuntimeAttributes
         public bool Live = true;
         public int MaxCallStack = 2048;
 
+        public int GrowLabelsStack = 512;
+        public int InitialLabelsStack = 2048;
+
         public int MaxFunctionLocals = 2048;
+        public int LocalPoolSize = 64;
 
         public int MaxOpStack = 1024;
         public IInstructionFactory InstructionFactory { get; set; } = SpecFactory.Factory;
diff --git a/Wacs.Core/Runtime/WasmRuntime.cs b/Wacs.Core/Runtime/WasmRuntime.cs
index c27de5b..5cdbab5 100644
--- a/Wacs.Core/Runtime/WasmRuntime.cs
+++ b/Wacs.Core/Runtime/WasmRuntime.cs
@@ -675,9 +675,8 @@ private ModuleInstance AllocateModule(Module module)
             }
 
             Context.PopFrame();
-            initFrame.Index = FuncIdx.ElementInitializers;
-            initFrame.ForceLabels(0);
-            Context.PushFrame(initFrame);
+            Frame initFrame2 = Context.ReserveFrame(moduleInstance, FunctionType.Empty, FuncIdx.ElementInitializers);
+            Context.PushFrame(initFrame2);
 
             //6. Allocate Elements
             //12. index ordered element addresses
diff --git a/Wacs.Core/Types/IndexSpace.cs b/Wacs.Core/Types/IndexSpace.cs
index 27f8fff..cce1fed 100644
--- a/Wacs.Core/Types/IndexSpace.cs
+++ b/Wacs.Core/Types/IndexSpace.cs
@@ -15,6 +15,7 @@
 //  */
 
 using System;
+using System.Buffers;
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
 using Wacs.Core.Runtime;
@@ -179,40 +180,46 @@ public override bool Contains(GlobalIdx idx) =>
             idx.Value < _imports.Count + _globals.Count;
     }
 
-    public class LocalsSpace : AbstractIndexSpace<LocalIdx, Value>
+    public struct LocalsSpace
     {
-        private readonly List<Value> _data = new();
+        public Value[]? Data;
+        private int _capacity;
 
-        public LocalsSpace() { }
+        public int Capacity => _capacity;
+        
+        public Value Get(LocalIdx idx)
+        {
+            if (Data == null)
+                throw new InvalidOperationException("LocalSpace was used uninitialized.");
+                
+            return Data[(Index)idx];
+        }
 
-        public override Value this[LocalIdx idx]
+        public void Set(LocalIdx idx, Value value)
         {
-            get => _data[(Index)idx];
-            set => _data[(Index)idx] = value;
+            if (Data == null)
+                throw new InvalidOperationException("LocalSpace was used uninitialized.");
+                
+            Data[(Index)idx] = value;
         }
 
-        public void Enscribe(ValType[] parameters, ValType[] locals)
+        public LocalsSpace(Value[] data, ValType[] parameters, ValType[] locals)
         {
-            _data.Clear();
-            // _data.Capacity = parameters.Length + locals.Length;
-            
+            _capacity = parameters.Length + locals.Length;
+            Data = data;
+            int idx = 0;
             foreach (var t in parameters)
             {
-                _data.Add(new Value(t));
+                Data[idx++] = new Value(t);
             }
             foreach (var t in locals)
             {
-                _data.Add(new Value(t));
+                Data[idx++] = new Value(t);
             }
         }
 
-        public void Reset()
-        {
-            _data.Clear();
-        }
-
-        public override bool Contains(LocalIdx idx) =>
-            idx.Value < _data.Count;
+        public bool Contains(LocalIdx idx) =>
+            idx.Value < _capacity;
     }
 
     public class ElementsSpace : AbstractIndexSpace<ElemIdx, Module.ElementSegment>
diff --git a/Wacs.Core/Utilities/IReusable.cs b/Wacs.Core/Utilities/IPoolable.cs
similarity index 91%
rename from Wacs.Core/Utilities/IReusable.cs
rename to Wacs.Core/Utilities/IPoolable.cs
index ce553d8..cffbfe5 100644
--- a/Wacs.Core/Utilities/IReusable.cs
+++ b/Wacs.Core/Utilities/IPoolable.cs
@@ -14,11 +14,9 @@
 //  * limitations under the License.
 //  */
 
-using System;
-
 namespace Wacs.Core.Utilities
 {
-    public interface IReusable<T> : IEquatable<T>
+    public interface IPoolable
     {
         public void Clear();
     }
diff --git a/Wacs.Core/Utilities/ObjectPoolExtension.cs b/Wacs.Core/Utilities/ObjectPoolExtension.cs
new file mode 100644
index 0000000..55e3a1c
--- /dev/null
+++ b/Wacs.Core/Utilities/ObjectPoolExtension.cs
@@ -0,0 +1,38 @@
+// /*
+//  * Copyright 2024 Kelvin Nishikawa
+//  *
+//  * Licensed under the Apache License, Version 2.0 (the "License");
+//  * you may not use this file except in compliance with the License.
+//  * You may obtain a copy of the License at
+//  *
+//  *     http://www.apache.org/licenses/LICENSE-2.0
+//  *
+//  * Unless required by applicable law or agreed to in writing, software
+//  * distributed under the License is distributed on an "AS IS" BASIS,
+//  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  * See the License for the specific language governing permissions and
+//  * limitations under the License.
+//  */
+
+using System;
+using System.Collections.Generic;
+using Microsoft.Extensions.ObjectPool;
+
+namespace Wacs.Core.Utilities
+{
+    public static class ObjectPoolExtension
+    {
+        public static ObjectPool<T> Prime<T>(this ObjectPool<T> pool, int count) where T : class, new()
+        {
+            var list = new List<T>();
+            for (int i = 0; i < count; i++)
+            {
+                list.Add(pool.Get());
+            }
+            foreach (var obj in list)
+                pool.Return(obj);
+            
+            return pool;
+        }
+    }
+}
\ No newline at end of file
diff --git a/Wacs.Core/Utilities/ReusableStack.cs b/Wacs.Core/Utilities/ReusableStack.cs
index 52a6487..889b40c 100644
--- a/Wacs.Core/Utilities/ReusableStack.cs
+++ b/Wacs.Core/Utilities/ReusableStack.cs
@@ -21,7 +21,7 @@
 namespace Wacs.Core.Utilities
 {
     public class ReusableStack<T>
-        where T : IReusable<T>, new()
+        where T : IPoolable, new()
     {
         private readonly int _growSize;
         internal T[] _backingArray;
@@ -92,7 +92,7 @@ public SubStack<T> GetSubStack()
         }
     }
 
-    public struct SubStack<T> : IEnumerable<T> where T : IReusable<T>, new()
+    public struct SubStack<T> : IEnumerable<T> where T : IPoolable, new()
     {
         private readonly int _zeroElement;
         private readonly ReusableStack<T> _stack;
diff --git a/Wacs.Core/Utilities/LocalsSpacePooledObjectPolicy.cs b/Wacs.Core/Utilities/StackPoolPolicy.cs
similarity index 77%
rename from Wacs.Core/Utilities/LocalsSpacePooledObjectPolicy.cs
rename to Wacs.Core/Utilities/StackPoolPolicy.cs
index 0f5bc4a..5d41394 100644
--- a/Wacs.Core/Utilities/LocalsSpacePooledObjectPolicy.cs
+++ b/Wacs.Core/Utilities/StackPoolPolicy.cs
@@ -15,19 +15,21 @@
 //  */
 
 using Microsoft.Extensions.ObjectPool;
+using Wacs.Core.Runtime;
 using Wacs.Core.Types;
 
 namespace Wacs.Core.Utilities
 {
-    public class LocalsSpacePooledObjectPolicy : PooledObjectPolicy<LocalsSpace>
+    public class StackPoolPolicy<T> : PooledObjectPolicy<T>
+    where T : class, IPoolable, new()
     {
-        public override LocalsSpace Create() => new LocalsSpace();
-
-        public override bool Return(LocalsSpace? obj)
+        public override T Create() => new T();
+        public override bool Return(T? obj)
         {
             if (obj == null)
                 return false;
-            obj.Reset();
+            
+            obj.Clear();
             return true; 
         }
     }
diff --git a/Wacs.Core/Validation/WasmValidationContext.cs b/Wacs.Core/Validation/WasmValidationContext.cs
index 9d7a65e..a5fc287 100644
--- a/Wacs.Core/Validation/WasmValidationContext.cs
+++ b/Wacs.Core/Validation/WasmValidationContext.cs
@@ -195,8 +195,9 @@ public ValidationContext<T> PushSubContext<T>(T child, int index = -1)
 
         public void SetExecFrame(FunctionType funcType, ValType[] localTypes)
         {
-            var locals = new LocalsSpace();
-            locals.Enscribe(funcType.ParameterTypes.Types, localTypes);
+            int capacity = funcType.ParameterTypes.Types.Length + localTypes.Length;
+            var data = new Value[capacity];
+            var locals = new LocalsSpace(data, funcType.ParameterTypes.Types, localTypes);
             ExecFrame = new Frame
             {
                 Module = ValidationModule,