1 files changed, 36 insertions, 6 deletions

diff --git a/src/wasm-shell.cpp b/src/wasm-shell.cpp
index b3b249771..7eb4b847a 100644
--- a/src/wasm-shell.cpp
+++ b/src/wasm-shell.cpp
@@ -8,10 +8,15 @@
 using namespace cashew;
 using namespace wasm;
 
+IString ASSERT_RETURN("assert_return"),
+        ASSERT_TRAP("assert_trap"),
+        INVOKE("invoke");
+
 int main(int argc, char **argv) {
   debug = getenv("WASM_SHELL_DEBUG") ? getenv("WASM_SHELL_DEBUG")[0] - '0' : 0;
 
   char *infile = argv[1];
+  bool print_wasm = argc >= 3; // second arg means print it out
 
   if (debug) std::cerr << "loading '" << infile << "'...\n";
   FILE *f = fopen(argv[1], "r");
@@ -30,14 +35,39 @@ int main(int argc, char **argv) {
 
   if (debug) std::cerr << "parsing text to s-expressions...\n";
   SExpressionParser parser(input);
-  if (debug) std::cout << *parser.root << '\n';
+  Element& root = *parser.root;
+  if (debug) std::cout << root << '\n';
+
+  // A .wast may have multiple modules, with some asserts after them
+  size_t i = 0;
+  while (i < root.size()) {
+    if (debug) std::cerr << "parsing s-expressions to wasm...\n";
+    Module wasm;
+    SExpressionWasmBuilder builder(wasm, *root[i]);
+    i++;
 
-  if (debug) std::cerr << "parsing s-expressions to wasm...\n";
-  Module wasm;
-  SExpressionWasmBuilder builder(wasm, *(*parser.root)[0]);
+    if (print_wasm) {
+      if (debug) std::cerr << "printing...\n";
+      std::cout << wasm;
+    }
 
-  if (debug) std::cerr << "printing...\n";
-  std::cout << wasm;
+    // run asserts
+    while (i < root.size()) {
+      Element& curr = *root[i];
+      IString id = curr[0]->str();
+      if (id == MODULE) break;
+      Element& invoke = *curr[1];
+      assert(invoke[0]->str() == INVOKE);
+      IString name = invoke[1]->str();
+      LiteralList arguments;
+      for (size_t j = 2; j < invoke.size(); j++) {
+        Expression* argument = builder.parseExpression(*invoke[2]);
+        arguments.push_back(argument->dyn_cast<Const>()->value);
+      }
+      interpreter.callFunction(name, arguments);
+      i++;
+    }
+  }
 
   if (debug) std::cerr << "done.\n";
 }