implement switch

6 files changed, 152 insertions, 15 deletions

diff --git a/src/asm2wasm.h b/src/asm2wasm.h
index 7d8d4d47a..04d4db21c 100644
--- a/src/asm2wasm.h
+++ b/src/asm2wasm.h
@@ -965,9 +965,6 @@ Function* Asm2WasmBuilder::processFunction(Ref ast) {
 #ifndef __EMSCRIPTEN__
       return allocator.alloc<Nop>(); // ignore in reference interpreter
 #else
-      return allocator.alloc<Host>(); // XXX abort in wasm.js
-#endif
-#if 0
       IString name = getNextId("switch");
       breakStack.push_back(name);
       auto ret = allocator.alloc<Switch>();
@@ -978,17 +975,34 @@ Function* Asm2WasmBuilder::processFunction(Ref ast) {
         Ref curr = cases[i];
         Ref condition = curr[0];
         Ref body = curr[1];
+        Switch::Case case_;
+        case_.body = processStatements(body, 0);
         if (condition->isNull()) {
-          ret->default_ = processStatements(body, 0);
+          case_.name = ret->default_ = getNextId("switch-default");
         } else {
           assert(condition[0] == NUM || condition[0] == UNARY_PREFIX);
-          Switch::Case case_;
-          case_.value = getLiteral(condition);
-          case_.body = processStatements(body, 0);
-          case_.fallthru = false; // XXX we assume no fallthru, ever
-          ret->cases.push_back(case_);
+          int32_t index = getLiteral(condition).geti32();
+          assert(index >= 0); // TODO: apply an offset, and have an option to abort switching at all (or do it in emscripten?)
+          case_.name = getNextId("switch-case");
+          if (ret->targets.size() <= index) {
+            ret->targets.resize(index+1);
+          }
+          ret->targets[index] = case_.name;
         }
+        ret->cases.push_back(case_);
+      }
+      // ensure a default
+      if (ret->default_.isNull()) {
+        Switch::Case defaultCase;
+        defaultCase.name = ret->default_ = getNextId("switch-default");
+        defaultCase.body = allocator.alloc<Nop>(); // ok if others fall through to this
+        ret->cases.push_back(defaultCase);
+      }
+      for (size_t i = 0; i < ret->targets.size(); i++) {
+        if (ret->targets[i].isNull()) ret->targets[i] = ret->default_;
       }
+      // finalize
+      ret->updateCaseMap();
       breakStack.pop_back();
       return ret;
 #endif
diff --git a/src/wasm-interpreter.h b/src/wasm-interpreter.h
index 022b274af..6e1c960bc 100644
--- a/src/wasm-interpreter.h
+++ b/src/wasm-interpreter.h
@@ -163,7 +163,33 @@ public:
       }
       Flow visitSwitch(Switch *curr) override {
         NOTE_ENTER("Switch");
-        abort();
+        Flow flow = visit(curr->value);
+        if (flow.breaking()) {
+          flow.clearIf(curr->name);
+          return flow;
+        }
+        int32_t index = flow.value.geti32();
+        Name target = curr->default_;
+        if (index >= 0 && index < curr->targets.size()) {
+          target = curr->targets[index];
+        }
+        auto iter = curr->caseMap.find(target);
+        if (iter == curr->caseMap.end()) {
+          // not in the cases, so this is a break outside
+          return Flow(target);
+        }
+        size_t caseIndex = iter->second;
+        assert(caseIndex < curr->cases.size());
+        while (caseIndex < curr->cases.size()) {
+          Switch::Case& c = curr->cases[caseIndex];
+          Flow flow = visit(c.body);
+          if (flow.breaking()) {
+            flow.clearIf(c.name);
+            return flow;
+          }
+          caseIndex++;
+        }
+        return Flow();
       }
 
       Flow generateArguments(const ExpressionList& operands, LiteralList& arguments) {
diff --git a/src/wasm.h b/src/wasm.h
index e1ffe76a1..065e2a2d7 100644
--- a/src/wasm.h
+++ b/src/wasm.h
@@ -402,21 +402,39 @@ public:
   Switch() : Expression(SwitchId) {}
 
   struct Case {
-    Literal value;
+    Name name;
     Expression *body;
-    bool fallthru;
   };
 
   Name name;
   Expression *value;
+  std::vector<Name> targets;
+  Name default_;
   std::vector<Case> cases;
-  Expression *default_;
+  std::map<Name, size_t> caseMap; // name => index in cases
+
+  void updateCaseMap() {
+    for (size_t i = 0; i < cases.size(); i++) {
+      caseMap[cases[i].name] = i;
+    }
+  }
 
   std::ostream& doPrint(std::ostream &o, unsigned indent) {
     printOpening(o, "switch ") << name;
     incIndent(o, indent);
     printFullLine(o, indent, value);
-    o << "TODO: cases/default\n";
+    doIndent(o, indent) << "[ ";
+    for (auto& t : targets) {
+      o << t.str << ' ';
+    }
+    o << "] (default " << default_.str << ")\n";
+    for (auto& c : cases) {
+      doIndent(o, indent);
+      printMinorOpening(o, "case ") << c.name.str;
+      incIndent(o, indent);
+      printFullLine(o, indent, c.body);
+      decIndent(o, indent) << '\n';
+    }
     return decIndent(o, indent);
   }
 
@@ -1051,7 +1069,6 @@ struct WasmWalker : public WasmVisitor<void> {
         for (auto& case_ : curr->cases) {
           parent.walk(case_.body);
         }
-        parent.walk(curr->default_);
       }
       void visitCall(Call *curr) override {
         ExpressionList& list = curr->operands;
diff --git a/test/control_flow.cpp b/test/control_flow.cpp
index 871e21281..504db2646 100644
--- a/test/control_flow.cpp
+++ b/test/control_flow.cpp
@@ -65,5 +65,53 @@ int EMSCRIPTEN_KEEPALIVE check_while_forever(int x) {
   return x;
 }
 
+int EMSCRIPTEN_KEEPALIVE check_switch(int x) {
+  switch (x) {
+    case 1: return 10;
+    case 3: return 20;
+    case 5: return 30;
+    case 10: return 40;
+    case 11: return 50;
+    default: return 60;
+  }
+  return 70;
+}
+
+int EMSCRIPTEN_KEEPALIVE check_switch_nodefault(int x) {
+  switch (x) {
+    case 1: return 10;
+    case 3: return 20;
+    case 5: return 30;
+    case 10: return 40;
+    case 11: return 50;
+  }
+  return 70;
+}
+
+int EMSCRIPTEN_KEEPALIVE check_switch_rdefault(int x) {
+  switch (x) {
+    default: return -60;
+    case 1: return 10;
+    case 3: return 20;
+    case 5: return 30;
+    case 10: return 40;
+    case 11: return 50;
+  }
+  return 70;
+}
+
+int EMSCRIPTEN_KEEPALIVE check_switch_fallthrough(int x) {
+  switch (x) {
+    case 1: return 10;
+    case 2:
+    case 3: x++;
+    case 5: return x;
+    case 10: return 40;
+    case 11: return 50;
+    default: return 60;
+  }
+  return 70;
+}
+
 }
 
diff --git a/test/control_flow.post.js b/test/control_flow.post.js
index 839806fe9..823580e7b 100644
--- a/test/control_flow.post.js
+++ b/test/control_flow.post.js
@@ -19,4 +19,8 @@ test('loop_continue');
 test('do_loop');
 test('do_once');
 test('while_forever');
+test('switch');
+test('switch_nodefault');
+test('switch_rdefault');
+test('switch_fallthrough');
 
diff --git a/test/control_flow.txt b/test/control_flow.txt
index c541546a6..68ee62bcd 100644
--- a/test/control_flow.txt
+++ b/test/control_flow.txt
@@ -47,3 +47,31 @@ while_forever
   4 ==> 1922
   11 ==> 1346
   90 ==> 1426
+switch
+  1 ==> 10
+  2 ==> 60
+  3 ==> 20
+  4 ==> 60
+  11 ==> 50
+  90 ==> 60
+switch_nodefault
+  1 ==> 10
+  2 ==> 70
+  3 ==> 20
+  4 ==> 70
+  11 ==> 50
+  90 ==> 70
+switch_rdefault
+  1 ==> 10
+  2 ==> -60
+  3 ==> 20
+  4 ==> -60
+  11 ==> 50
+  90 ==> -60
+switch_fallthrough
+  1 ==> 10
+  2 ==> 3
+  3 ==> 4
+  4 ==> 60
+  11 ==> 50
+  90 ==> 60