1 files changed, 33 insertions, 4 deletions

diff --git a/src/wasm-shell.cpp b/src/binaryen-shell.cpp
index 701b6b6c2..775252485 100644
--- a/src/wasm-shell.cpp
+++ b/src/binaryen-shell.cpp
@@ -146,12 +146,41 @@ struct Invocation {
 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
+  char *infile = nullptr;
+  bool print_before = false;
+
+  for (size_t i = 1; i < argc; i++) {
+    char* curr = argv[i];
+    if (curr[0] == '-') {
+      std::string arg = curr;
+      if (arg == "--print-before") {
+        print_before = true;
+      } else {
+        if (infile) {
+          printf("error: unrecognized argument: %s\n", curr);
+          exit(1);
+        }
+      }
+    } else {
+      if (infile) {
+        printf("error: too many input files provided.\n");
+        exit(1);
+      }
+      infile = curr;
+    }
+  }
+
+  if (!infile) {
+    printf("error: no input file provided.\n");
+    exit(1);
+  }
 
   if (debug) std::cerr << "loading '" << infile << "'...\n";
   FILE *f = fopen(argv[1], "r");
-  assert(f);
+  if (!f) {
+    printf("error: could not open input file: %s\n", infile);
+    exit(1);
+  }
   fseek(f, 0, SEEK_END);
   int size = ftell(f);
   char *input = new char[size+1];
@@ -180,7 +209,7 @@ int main(int argc, char **argv) {
     auto interface = new ShellExternalInterface();
     auto instance = new ModuleInstance(wasm, interface);
 
-    if (print_wasm) {
+    if (print_before) {
       if (debug) std::cerr << "printing...\n";
       std::cout << wasm;
     }