Fix call indirect (#96)

butterunderflow · web-flow · commit 099a71114c03 · 2026-03-10T08:59:43.000-04:00
diff --git a/benchmarks/wasm/call_indirect_test.wat b/benchmarks/wasm/call_indirect_test.wat
@@ -3,6 +3,9 @@
   (type (;1;) (func (param i32 i32) (result i32)))
   (type (;2;) (func (result i32)))
   (func (;0;) (type 0) (param i32) (result i32)
+    (local i32 i32 i32 i32 i32)
+    local.get 3
+    drop
     local.get 0
     i32.const 1
     i32.add)
diff --git a/headers/wasm/concrete_rt.hpp b/headers/wasm/concrete_rt.hpp
@@ -3,6 +3,7 @@
 
 #include "concrete_num.hpp"
 #include "controls.hpp"
+#include "immer/vector_transient.hpp"
 #include "wasm/profile.hpp"
 #include "wasm/utils.hpp"
 #include <cassert>
@@ -49,8 +50,8 @@ class Stack_t {
 
   Num pop() {
     Profile.step(StepProfileKind::POP);
-#ifdef DEBUG
     assert(count > 0 && "Stack underflow");
+#ifdef DEBUG
     printf("[Debug] popping a value %ld from stack, size of concrete stack is: "
            "%d\n",
            stack_ptr[count - 1].value, count);
@@ -123,45 +124,86 @@ class Stack_t {
   Num *stack_ptr;
 };
 static Stack_t Stack;
+class SymFrames_t;
 
 const int FRAME_SIZE = 1024 * 8;
 class Frames_t {
 public:
-  Frames_t() : count(0), stack_ptr(new Num[FRAME_SIZE]) {
+  Frames_t() : count(0), stack_ptr(new Num[FRAME_SIZE]), frame_ptrs() {
     size_t page_size = (size_t)sysconf(_SC_PAGESIZE);
     // pre touch the memory to avoid page faults during execution
     for (int i = 0; i < FRAME_SIZE; i += page_size) {
       stack_ptr[i] = Num(0);
     }
   }
 
-  std::monostate popFrame(std::int32_t size) {
+  std::monostate popFrameCaller(std::int32_t size) {
+    assert(size >= 0);
+    assert(size <= count);
+    assert(!frame_ptrs.empty());
+    auto frame_base = current_frame_base();
+    assert(frame_base + size == count);
+    count -= size;
+#ifdef USE_IMM
+    frame_ptrs.take(frame_ptrs.size() - 1);
+#else
+    frame_ptrs.pop_back();
+#endif
+    return std::monostate{};
+  }
+
+  std::monostate popFrameCallee(std::int32_t size) {
     assert(size >= 0);
+    assert(size <= count);
     count -= size;
     return std::monostate{};
   }
 
   Num get(std::int32_t index) {
+    assert(!frame_ptrs.empty() && "No active frame");
+    auto frame_base = current_frame_base();
+    assert(index >= 0 && frame_base + index < count && "Index out of bounds");
     Profile.step(StepProfileKind::GET);
-    auto ret = stack_ptr[count - 1 - index];
+    auto ret = stack_ptr[frame_base + index];
     return ret;
   }
 
   void set(std::int32_t index, Num num) {
+    assert(!frame_ptrs.empty() && "No active frame");
+    auto frame_base = current_frame_base();
+    assert(index >= 0 && frame_base + index < count && "Index out of bounds");
     Profile.step(StepProfileKind::SET);
-    stack_ptr[count - 1 - index] = num;
+    stack_ptr[frame_base + index] = num;
   }
 
-  void pushFrame(std::int32_t size) {
+  void pushFrameCaller(std::int32_t size) {
     assert(size >= 0);
+    frame_ptrs.push_back(count);
     count += size;
     // Zero-initialize the new stack frames.
     for (std::int32_t i = 0; i < size; ++i) {
-      stack_ptr[count - 1 - i] = Num(0);
+      stack_ptr[count - size + i] = Num(0);
     }
   }
 
-  void reset() { count = 0; }
+  void pushFrameCallee(std::int32_t size) {
+    assert(size >= 0);
+    assert(!frame_ptrs.empty() && "No active frame");
+    auto old_count = count;
+    count += size;
+    for (std::int32_t i = 0; i < size; ++i) {
+      stack_ptr[old_count + i] = Num(0);
+    }
+  }
+
+  void reset() {
+    count = 0;
+#ifdef USE_IMM
+    frame_ptrs = immer::vector_transient<size_t>();
+#else
+    frame_ptrs.clear();
+#endif
+  }
 
   size_t size() const { return count; }
 
@@ -176,8 +218,23 @@ class Frames_t {
   }
 
 private:
+  friend class SymFrames_t;
+
+  size_t current_frame_base() const {
+#ifdef USE_IMM
+    return *(frame_ptrs.end() - 1);
+#else
+    return frame_ptrs.back();
+#endif
+  }
+
   int32_t count;
   Num *stack_ptr;
+#ifdef USE_IMM
+  immer::vector_transient<size_t> frame_ptrs;
+#else
+  std::vector<size_t> frame_ptrs;
+#endif
 };
 
 static Frames_t Frames;
@@ -319,4 +376,4 @@ struct FuncTable_t {
 
 static FuncTable_t FuncTable;
 
-#endif // WASM_CONCRETE_RT_HPP
+#endif // WASM_CONCRETE_RT_HPP
diff --git a/headers/wasm/sym_rt.hpp b/headers/wasm/sym_rt.hpp
@@ -10,11 +10,11 @@
 #include "immer/vector.hpp"
 #include "immer/vector_transient.hpp"
 #include "profile.hpp"
+#include "symbolic_decl.hpp"
+#include "symbolic_impl.hpp"
 #include "symval_decl.hpp"
 #include "symval_factory.hpp"
 #include "symval_impl.hpp"
-#include "symbolic_decl.hpp"
-#include "symbolic_impl.hpp"
 #include "utils.hpp"
 #include "wasm/concrete_num.hpp"
 #include "wasm/z3_env.hpp"
@@ -35,7 +35,6 @@
 #include <variant>
 #include <vector>
 
-
 class Snapshot_t;
 
 class SymStack_t {
@@ -125,6 +124,16 @@ static SymStack_t SymStack;
 class SymFrames_t {
 
 public:
+  void restore_frame_ptr(Frames_t &frame) const;
+
+  void pushFramePtr() {
+#ifdef USE_IMM
+    frame_ptrs.push_back(stack.size());
+#else
+    frame_ptrs.push_back(stack.size());
+#endif
+  }
+
   void pushFrameSlot(int width) {
 #ifdef USE_IMM
     stack.push_back(SVFactory::make_concrete_bv(I64V(0), width));
@@ -133,7 +142,33 @@ class SymFrames_t {
 #endif
   }
 
-  std::monostate popFrame(int size) {
+  std::monostate popFrameCaller(int size) {
+    assert(size >= 0);
+    assert(static_cast<size_t>(size) <= stack.size());
+    assert(!frame_ptrs.empty());
+    auto frame_base = current_frame_base();
+    assert(frame_base + size == stack.size());
+
+    for (int i = 0; i < size; ++i) {
+      symbolic_size -= stack[stack.size() - 1 - i]->size();
+    }
+
+#ifdef USE_IMM
+    stack.take(stack.size() - size);
+#else
+    stack.erase(stack.end() - size, stack.end());
+#endif
+
+#ifdef USE_IMM
+    frame_ptrs.take(frame_ptrs.size() - 1);
+#else
+    frame_ptrs.pop_back();
+#endif
+
+    return std::monostate{};
+  }
+
+  std::monostate popFrameCallee(int size) {
     // Pop the frame of the given size
     assert(size >= 0);
     assert(static_cast<size_t>(size) <= stack.size());
@@ -153,18 +188,26 @@ class SymFrames_t {
 
   SymVal get(int index) {
     // Get the symbolic value at the given frame index
-    auto res = stack[stack.size() - 1 - index];
+    assert(!frame_ptrs.empty());
+    auto frame_base = current_frame_base();
+    assert(index >= 0 &&
+           static_cast<size_t>(frame_base + index) < stack.size());
+    auto res = stack[frame_base + index];
     return res;
   }
 
   void set(int index, SymVal val) {
     // Set the symbolic value at the given index
     assert(val.symptr != nullptr);
-    symbolic_size += val->size() - stack[stack.size() - 1 - index]->size();
+    assert(!frame_ptrs.empty());
+    auto frame_base = current_frame_base();
+    assert(index >= 0 &&
+           static_cast<size_t>(frame_base + index) < stack.size());
+    symbolic_size += val->size() - stack[frame_base + index]->size();
 #ifdef USE_IMM
-    stack.set(stack.size() - 1 - index, val);
+    stack.set(frame_base + index, val);
 #else
-    stack[stack.size() - 1 - index] = val;
+    stack[frame_base + index] = val;
 #endif
   }
 
@@ -173,8 +216,10 @@ class SymFrames_t {
 
 #ifdef USE_IMM
     stack = immer::vector_transient<SymVal>();
+    frame_ptrs = immer::vector_transient<size_t>();
 #else
     stack.clear();
+    frame_ptrs.clear();
 #endif
     symbolic_size = 0;
   }
@@ -193,10 +238,20 @@ class SymFrames_t {
   }
 
 private:
+  size_t current_frame_base() const {
+#ifdef USE_IMM
+    return *(frame_ptrs.end() - 1);
+#else
+    return frame_ptrs.back();
+#endif
+  }
+
   int symbolic_size = 0;
 #ifdef USE_IMM
+  immer::vector_transient<size_t> frame_ptrs;
   immer::vector_transient<SymVal> stack;
 #else
+  std::vector<size_t> frame_ptrs;
   std::vector<SymVal> stack;
 #endif
 };
@@ -324,6 +379,10 @@ class SymMemory_t {
   }
 };
 
+inline void SymFrames_t::restore_frame_ptr(Frames_t &frame) const {
+  frame.frame_ptrs = frame_ptrs;
+}
+
 static SymMemory_t SymMemory;
 
 static std::monostate memoryInitialize(int32_t offset,
@@ -1563,6 +1622,7 @@ static void resume_conc_frames(const SymFrames_t &sym_frame, Frames_t &frames,
     auto conc = res.value;
     frames.set_from_front(i, conc);
   }
+  sym_frame.restore_frame_ptr(frames);
 }
 
 static void resume_conc_frames_by_model(const SymFrames_t &sym_frame,
@@ -1576,6 +1636,7 @@ static void resume_conc_frames_by_model(const SymFrames_t &sym_frame,
     auto conc = res.value;
     frames.set_from_front(i, conc);
   }
+  sym_frame.restore_frame_ptr(frames);
 }
 
 static void resume_conc_memory(const SymMemory_t &sym_memory, Memory_t &memory,
diff --git a/src/main/scala/wasm/StagedConcolicMiniWasm.scala b/src/main/scala/wasm/StagedConcolicMiniWasm.scala
