Staged Miniwasm Interpreter

.github/workflows/scala.yml

@@ -78,3 +78,4 @@ jobs:
         sbt 'testOnly gensym.wasm.TestScriptRun'
         sbt 'testOnly gensym.wasm.TestConcolic'
         sbt 'testOnly gensym.wasm.TestDriver'
+        sbt 'testOnly gensym.wasm.TestStagedEval'

benchmarks/wasm/global.wat

@@ -0,0 +1,19 @@
+(module
+  (type (;0;) (func (result i32)))
+  (type (;1;) (func))
+
+  (func (;0;) (type 0) (result i32)
+    i32.const 42
+    global.set 0
+    global.get 0
+  )
+  (func (;1;) (type 1)
+    call 0
+    ;; should be 42
+    ;; drop
+    )
+  (start 1)
+  (memory (;0;) 2)
+  (export "main" (func 1))
+  (global (;0;) (mut i32) (i32.const 0))
+)

benchmarks/wasm/staged/brtable.wat

@@ -0,0 +1,11 @@
+(module $push-drop
+  (global (;0;) (mut i32) (i32.const 1048576))
+  (func (;0;) (type 1) (result i32)
+    i32.const 2
+    (block
+     (block
+      br_table 0 1 0
+     )
+    )
+  )
+  (start 0))

src/main/scala/wasm/AST.scala

@@ -270,7 +270,16 @@ case class RefType(kind: HeapType) extends ValueType
 
 case class GlobalType(ty: ValueType, mut: Boolean) extends WasmType
 
-abstract class BlockType extends WIR
+abstract class BlockType extends WIR {
+  def funcType: FuncType =
+    this match {
+      case VarBlockType(_, None) =>
+        ??? // TODO: fill this branch until we handle type index correctly
+      case VarBlockType(_, Some(tipe)) => tipe
+      case ValBlockType(Some(tipe))    => FuncType(List(), List(), List(tipe))
+      case ValBlockType(None)          => FuncType(List(), List(), List())
+    }
+}
 case class VarBlockType(index: Int, tipe: Option[FuncType]) extends BlockType
 case class ValBlockType(tipe: Option[ValueType]) extends BlockType;
 

src/main/scala/wasm/ConcolicMiniWasm.scala

@@ -229,15 +229,6 @@ object Primitives {
     case NumType(F32Type) => F32V(rng.nextFloat())
     case NumType(F64Type) => F64V(rng.nextDouble())
   }
-
-  def getFuncType(ty: BlockType): FuncType =
-    ty match {
-      case VarBlockType(_, None) =>
-        ??? // TODO: fill this branch until we handle type index correctly
-      case VarBlockType(_, Some(tipe)) => tipe
-      case ValBlockType(Some(tipe))    => FuncType(List(), List(), List(tipe))
-      case ValBlockType(None)          => FuncType(List(), List(), List())
-    }
 }
 
 case class Frame(module: ModuleInstance, locals: ArrayBuffer[Value], symLocals: ArrayBuffer[SymVal])
@@ -383,15 +374,15 @@ case class Evaluator(module: ModuleInstance) {
         eval(rest, concStack, symStack, frame, kont, trail)
       case Unreachable => throw new RuntimeException("Unreachable")
       case Block(ty, inner) =>
-        val funcTy = getFuncType(ty)
+        val funcTy = ty.funcType
         val (inputSize, outputSize) = (funcTy.inps.size, funcTy.out.size)
         val (inputs, restStack) = concStack.splitAt(inputSize)
         val (symInputs, restSymStack) = symStack.splitAt(inputSize)
         val restK: Cont = (retStack, retSymStack, tree) =>
           eval(rest, retStack.take(outputSize) ++ restStack, retSymStack.take(outputSize) ++ restSymStack, frame, kont, trail)(tree)
         eval(inner, inputs, symInputs, frame, restK, restK :: trail)
       case Loop(ty, inner) =>
-        val funcTy = getFuncType(ty)
+        val funcTy = ty.funcType
         val (inputSize, outputSize) = (funcTy.inps.size, funcTy.out.size)
         val (inputs, restStack) = concStack.splitAt(inputSize)
         val (symInputs, restSymStack) = symStack.splitAt(inputSize)

src/main/scala/wasm/MiniWasm.scala

@@ -229,15 +229,6 @@ object Primtives {
     case VecType(kind) => ???
     case RefType(kind) => RefNullV(kind)
   }
-
-  def getFuncType(ty: BlockType): FuncType =
-    ty match {
-      case VarBlockType(_, None) =>
-        ??? // TODO: fill this branch until we handle type index correctly
-      case VarBlockType(_, Some(tipe)) => tipe
-      case ValBlockType(Some(tipe))    => FuncType(List(), List(), List(tipe))
-      case ValBlockType(None)          => FuncType(List(), List(), List())
-    }
 }
 
 case class Frame(locals: ArrayBuffer[Value])
@@ -264,8 +255,8 @@ case class Evaluator(module: ModuleInstance) {
         val frameLocals = args ++ locals.map(zero(_))
         val newFrame = Frame(ArrayBuffer(frameLocals: _*))
         if (isTail)
-          // when tail call, share the continuation for returning with the callee
-          eval(body, List(), newFrame, kont, List(kont))
+          // when tail call, return to the caller's return continuation
+          eval(body, List(), newFrame, trail.last, List(trail.last))
         else {
           val restK: Cont[Ans] = (retStack) =>
             eval(rest, retStack.take(ty.out.size) ++ newStack, frame, kont, trail)
@@ -380,7 +371,7 @@ case class Evaluator(module: ModuleInstance) {
         eval(rest, stack, frame, kont, trail)
       case Unreachable => throw Trap()
       case Block(ty, inner) =>
-        val funcTy = getFuncType(ty)
+        val funcTy = ty.funcType
         val (inputs, restStack) = stack.splitAt(funcTy.inps.size)
         val restK: Cont[Ans] = (retStack) =>
           eval(rest, retStack.take(funcTy.out.size) ++ restStack, frame, kont, trail)
@@ -389,15 +380,15 @@ case class Evaluator(module: ModuleInstance) {
         // We construct two continuations, one for the break (to the begining of the loop),
         // and one for fall-through to the next instruction following the syntactic structure
         // of the program.
-        val funcTy = getFuncType(ty)
+        val funcTy = ty.funcType
         val (inputs, restStack) = stack.splitAt(funcTy.inps.size)
         val restK: Cont[Ans] = (retStack) =>
           eval(rest, retStack.take(funcTy.out.size) ++ restStack, frame, kont, trail)
         def loop(retStack: Stack): Ans =
           eval(inner, retStack.take(funcTy.inps.size), frame, restK, loop _ :: trail)
         loop(inputs)
       case If(ty, thn, els) =>
-        val funcTy = getFuncType(ty)
+        val funcTy = ty.funcType
         val I32V(cond) :: newStack = stack
         val inner = if (cond != 0) thn else els
         val (inputs, restStack) = newStack.splitAt(funcTy.inps.size)

src/main/scala/wasm/Parser.scala

@@ -314,7 +314,7 @@ class GSWasmVisitor extends WatParserBaseVisitor[WIR] {
     else if (ctx.LOCAL_GET() != null) LocalGet(getVar(ctx.idx(0)).toInt)
     else if (ctx.LOCAL_SET() != null) LocalSet(getVar(ctx.idx(0)).toInt)
     else if (ctx.LOCAL_TEE() != null) LocalTee(getVar(ctx.idx(0)).toInt)
-    else if (ctx.GLOBAL_SET() != null) GlobalGet(getVar(ctx.idx(0)).toInt)
+    else if (ctx.GLOBAL_SET() != null) GlobalSet(getVar(ctx.idx(0)).toInt)
     else if (ctx.GLOBAL_GET() != null) GlobalGet(getVar(ctx.idx(0)).toInt)
     else if (ctx.load() != null) {
       val ty = visitNumType(ctx.load.numType)