From dc67f0ad4e97df15941f92d0f18e888a45976c27 Mon Sep 17 00:00:00 2001
From: Ali Mohammad Pur <ali.mpfard@gmail.com>
Date: Fri, 6 Jun 2025 14:12:28 +0200
Subject: [PATCH] LibWasm: Hold on to the stack depth for expressions in the
 validator

This allows preallocating the value stack when pushing frames, avoiding
repeated reallocs and copies.
---
 Libraries/LibWasm/AbstractMachine/Configuration.h |  2 ++
 Libraries/LibWasm/AbstractMachine/Validator.cpp   |  2 ++
 Libraries/LibWasm/AbstractMachine/Validator.h     | 10 +++++++++-
 Libraries/LibWasm/Types.h                         |  4 ++++
 4 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/Libraries/LibWasm/AbstractMachine/Configuration.h b/Libraries/LibWasm/AbstractMachine/Configuration.h
index eff1b922798..0012bb6c156 100644
--- a/Libraries/LibWasm/AbstractMachine/Configuration.h
+++ b/Libraries/LibWasm/AbstractMachine/Configuration.h
@@ -21,6 +21,8 @@ public:
     {
         Label label(frame.arity(), frame.expression().instructions().size(), m_value_stack.size());
         frame.label_index() = m_label_stack.size();
+        if (auto hint = frame.expression().stack_usage_hint(); hint.has_value())
+            m_value_stack.ensure_capacity(*hint);
         m_frame_stack.append(move(frame));
         m_label_stack.append(label);
     }
diff --git a/Libraries/LibWasm/AbstractMachine/Validator.cpp b/Libraries/LibWasm/AbstractMachine/Validator.cpp
index 950377039ba..b9bdf365e3d 100644
--- a/Libraries/LibWasm/AbstractMachine/Validator.cpp
+++ b/Libraries/LibWasm/AbstractMachine/Validator.cpp
@@ -3736,6 +3736,8 @@ ErrorOr<Validator::ExpressionTypeResult, ValidationError> Validator::validate(Ex
     m_frames.take_last();
     VERIFY(m_frames.is_empty());
 
+    expression.set_stack_usage_hint(stack.max_known_size());
+
     return ExpressionTypeResult { stack.release_vector(), is_constant_expression };
 }
 
diff --git a/Libraries/LibWasm/AbstractMachine/Validator.h b/Libraries/LibWasm/AbstractMachine/Validator.h
index 13863401196..f06db4d950d 100644
--- a/Libraries/LibWasm/AbstractMachine/Validator.h
+++ b/Libraries/LibWasm/AbstractMachine/Validator.h
@@ -228,6 +228,7 @@ public:
         void append(StackEntry entry)
         {
             Vector<StackEntry>::append(entry);
+            m_max_known_size = max(m_max_known_size, size());
         }
 
         ErrorOr<StackEntry, ValidationError> take(ValueType type, SourceLocation location = SourceLocation::current())
@@ -254,10 +255,17 @@ public:
             return {};
         }
 
-        Vector<StackEntry> release_vector() { return exchange(static_cast<Vector<StackEntry>&>(*this), Vector<StackEntry> {}); }
+        Vector<StackEntry> release_vector()
+        {
+            m_max_known_size = 0;
+            return exchange(static_cast<Vector<StackEntry>&>(*this), Vector<StackEntry> {});
+        }
+
+        size_t max_known_size() const { return m_max_known_size; }
 
     private:
         Vector<Frame> const& m_frames;
+        size_t m_max_known_size { 0 };
     };
 
     struct ExpressionTypeResult {
diff --git a/Libraries/LibWasm/Types.h b/Libraries/LibWasm/Types.h
index d95275b256e..9d14c29e4ba 100644
--- a/Libraries/LibWasm/Types.h
+++ b/Libraries/LibWasm/Types.h
@@ -708,8 +708,12 @@ public:
 
     static ParseResult<Expression> parse(ConstrainedStream& stream, Optional<size_t> size_hint = {});
 
+    void set_stack_usage_hint(size_t value) const { m_stack_usage_hint = value; }
+    auto stack_usage_hint() const { return m_stack_usage_hint; }
+
 private:
     Vector<Instruction> m_instructions;
+    mutable Optional<size_t> m_stack_usage_hint;
 };
 
 class GlobalSection {