Relooper switch support (#617)

* support switches in relooper and c api

* update relooper fuzzer for switches

4 files changed, 162 insertions, 83 deletions

diff --git a/src/binaryen-c.cpp b/src/binaryen-c.cpp
index ff8fae06e..39a6d0149 100644
--- a/src/binaryen-c.cpp
+++ b/src/binaryen-c.cpp
@@ -538,10 +538,14 @@ RelooperBlockRef RelooperAddBlockWithSwitch(RelooperRef relooper, BinaryenExpres
   return RelooperRef(ret);
 }
 
-void RelooperAddBranchForSwitch(RelooperBlockRef from, RelooperBlockRef to, BinaryenIndex index, BinaryenExpressionRef code) {
+void RelooperAddBranchForSwitch(RelooperBlockRef from, RelooperBlockRef to, BinaryenIndex* indexes, BinaryenIndex numIndexes, BinaryenExpressionRef code) {
   auto* fromBlock = (CFG::Block*)from;
   auto* toBlock = (CFG::Block*)to;
-  fromBlock->AddBranchTo(toBlock, (wasm::Index)index, (Expression*)code);
+  std::vector<Index> values;
+  for (Index i = 0; i < numIndexes; i++) {
+    values.push_back(indexes[i]);
+  }
+  fromBlock->AddSwitchBranchTo(toBlock, std::move(values), (Expression*)code);
 }
 
 BinaryenExpressionRef RelooperRenderAndDispose(RelooperRef relooper, RelooperBlockRef entry, BinaryenIndex labelHelper, BinaryenModuleRef module) {
diff --git a/src/binaryen-c.h b/src/binaryen-c.h
index 1421325d9..7f9a775d5 100644
--- a/src/binaryen-c.h
+++ b/src/binaryen-c.h
@@ -386,10 +386,10 @@ RelooperBlockRef RelooperAddBlock(RelooperRef relooper, BinaryenExpressionRef co
 void RelooperAddBranch(RelooperBlockRef from, RelooperBlockRef to, BinaryenExpressionRef condition, BinaryenExpressionRef code);
 
 // Create a basic block that ends a switch on a condition
-// TODO RelooperBlockRef RelooperAddBlockWithSwitch(RelooperRef relooper, BinaryenExpressionRef code, BinaryenExpressionRef condition);
+RelooperBlockRef RelooperAddBlockWithSwitch(RelooperRef relooper, BinaryenExpressionRef code, BinaryenExpressionRef condition);
 
-// Create a switch-style branch to another basic block. The block's switch table will have an index for this branch
-// TODO void RelooperAddBranchForSwitch(RelooperBlockRef from, RelooperBlockRef to, BinaryenIndex index, BinaryenExpressionRef code);
+// Create a switch-style branch to another basic block. The block's switch table will have these indexes going to that target
+void RelooperAddBranchForSwitch(RelooperBlockRef from, RelooperBlockRef to, BinaryenIndex* indexes, BinaryenIndex numIndexes, BinaryenExpressionRef code);
 
 // Generate structed wasm control flow from the CFG of blocks and branches that were created
 // on this relooper instance. This returns the rendered output, and also disposes of the
diff --git a/src/cfg/Relooper.cpp b/src/cfg/Relooper.cpp
index 78b3f3e1c..4151f7885 100644
--- a/src/cfg/Relooper.cpp
+++ b/src/cfg/Relooper.cpp
@@ -42,7 +42,12 @@ static void PrintDebug(const char *Format, ...);
 
 Branch::Branch(wasm::Expression* ConditionInit, wasm::Expression* CodeInit) : Ancestor(nullptr), Condition(ConditionInit), Code(CodeInit) {}
 
-Branch::Branch(wasm::Index IndexInit, wasm::Expression* CodeInit) : Ancestor(nullptr), Index(IndexInit), Code(CodeInit) {}
+Branch::Branch(std::vector<wasm::Index>&& ValuesInit, wasm::Expression* CodeInit) : Ancestor(nullptr), Code(CodeInit) {
+  if (ValuesInit.size() > 0) {
+    SwitchValues = std::unique_ptr<std::vector<wasm::Index>>(new std::vector<wasm::Index>(ValuesInit));
+  }
+  // otherwise, it is the default
+}
 
 wasm::Expression* Branch::Render(RelooperBuilder& Builder, Block *Target, bool SetLabel) {
   auto* Ret = Builder.makeBlock();
@@ -50,9 +55,9 @@ wasm::Expression* Branch::Render(RelooperBuilder& Builder, Block *Target, bool S
   if (SetLabel) Ret->list.push_back(Builder.makeSetLabel(Target->Id));
   if (Ancestor) {
     if (Type == Break) {
-      Ret->list.push_back(Builder.makeBreak(Ancestor->Id));
+      Ret->list.push_back(Builder.makeShapeBreak(Ancestor->Id));
     } else if (Type == Continue) {
-      Ret->list.push_back(Builder.makeContinue(Ancestor->Id));
+      Ret->list.push_back(Builder.makeShapeContinue(Ancestor->Id));
     }
   }
   Ret->finalize();
@@ -74,9 +79,9 @@ void Block::AddBranchTo(Block *Target, wasm::Expression* Condition, wasm::Expres
   BranchesOut[Target] = new Branch(Condition, Code);
 }
 
-void Block::AddBranchTo(Block *Target, wasm::Index Index, wasm::Expression* Code) {
+void Block::AddSwitchBranchTo(Block *Target, std::vector<wasm::Index>&& Values, wasm::Expression* Code) {
   assert(!contains(BranchesOut, Target)); // cannot add more than one branch to the same target
-  BranchesOut[Target] = new Branch(Index, Code);
+  BranchesOut[Target] = new Branch(std::move(Values), Code);
 }
 
 wasm::Expression* Block::Render(RelooperBuilder& Builder, bool InLoop) {
@@ -123,8 +128,10 @@ wasm::Expression* Block::Render(RelooperBuilder& Builder, bool InLoop) {
     Parent->Next = Parent->Next->Next;
 
     // When the Multiple has the same number of groups as we have branches,
-    // they will all be fused, so it is safe to not set the label at all
-    if (SetLabel && Fused->InnerMap.size() == ProcessedBranchesOut.size()) {
+    // they will all be fused, so it is safe to not set the label at all.
+    // If a switch, then we can have multiple branches to the same target
+    // (in different table indexes), and so this check is not sufficient TODO: optimize
+    if (SetLabel && Fused->InnerMap.size() == ProcessedBranchesOut.size() && !SwitchCondition) {
       SetLabel = false;
     }
   }
@@ -133,84 +140,146 @@ wasm::Expression* Block::Render(RelooperBuilder& Builder, bool InLoop) {
 
   // Find the default target, the one without a condition
   for (BlockBranchMap::iterator iter = ProcessedBranchesOut.begin(); iter != ProcessedBranchesOut.end(); iter++) {
-    if (!iter->second->Condition) {
+    if ((!SwitchCondition && !iter->second->Condition) || (SwitchCondition && !iter->second->SwitchValues)) {
       assert(!DefaultTarget && "block has branches without a default (nullptr for the condition)"); // Must be exactly one default // nullptr
       DefaultTarget = iter->first;
     }
   }
   assert(DefaultTarget); // Since each block *must* branch somewhere, this must be set
 
-  if (SwitchCondition) {
-    abort(); // TODO: switch
-  }
+  wasm::Expression* Root = nullptr; // root of the main part, that we are about to emit
 
-  // We'll emit a chain of if-elses
-  wasm::Expression* Root = nullptr;
-  wasm::If* CurrIf = nullptr;
+  if (!SwitchCondition) {
+    // We'll emit a chain of if-elses
+    wasm::If* CurrIf = nullptr;
 
-  wasm::Expression* RemainingConditions = nullptr;
+    wasm::Expression* RemainingConditions = nullptr;
 
-  for (BlockBranchMap::iterator iter = ProcessedBranchesOut.begin();; iter++) {
-    Block *Target;
-    Branch *Details;
-    if (iter != ProcessedBranchesOut.end()) {
-      Target = iter->first;
-      if (Target == DefaultTarget) continue; // done at the end
-      Details = iter->second;
-      assert(Details->Condition); // must have a condition if this is not the default target
-    } else {
-      Target = DefaultTarget;
-      Details = ProcessedBranchesOut[DefaultTarget];
-    }
-    bool SetCurrLabel = (SetLabel && Target->IsCheckedMultipleEntry) || ForceSetLabel;
-    bool HasFusedContent = Fused && contains(Fused->InnerMap, Target->Id);
-    wasm::Expression* CurrContent = nullptr;
-    if (SetCurrLabel || Details->Type != Branch::Direct || HasFusedContent || Details->Code) {
-      CurrContent = Details->Render(Builder, Target, SetCurrLabel);
-      if (HasFusedContent) {
-        CurrContent = Builder.blockify(CurrContent, Fused->InnerMap.find(Target->Id)->second->Render(Builder, InLoop));
-      } else if (Details->Type == Branch::Nested) {
-        // Nest the parent content here, and remove it from showing up afterwards as Next
-        assert(Parent->Next);
-        CurrContent = Builder.blockify(CurrContent, Parent->Next->Render(Builder, InLoop));
-        Parent->Next = nullptr;
+    for (BlockBranchMap::iterator iter = ProcessedBranchesOut.begin();; iter++) {
+      Block *Target;
+      Branch *Details;
+      if (iter != ProcessedBranchesOut.end()) {
+        Target = iter->first;
+        if (Target == DefaultTarget) continue; // done at the end
+        Details = iter->second;
+        assert(Details->Condition); // must have a condition if this is not the default target
+      } else {
+        Target = DefaultTarget;
+        Details = ProcessedBranchesOut[DefaultTarget];
       }
-    }
-    bool isDefault = iter == ProcessedBranchesOut.end();
-    // If there is nothing to show in this branch, omit the condition
-    if (CurrContent) {
-      if (isDefault) {
-        wasm::Expression* Now;
-        if (RemainingConditions) {
-          Now = Builder.makeIf(RemainingConditions, CurrContent);
-        } else {
-          Now = CurrContent;
+      bool SetCurrLabel = (SetLabel && Target->IsCheckedMultipleEntry) || ForceSetLabel;
+      bool HasFusedContent = Fused && contains(Fused->InnerMap, Target->Id);
+      wasm::Expression* CurrContent = nullptr;
+      if (SetCurrLabel || Details->Type != Branch::Direct || HasFusedContent || Details->Code) {
+        CurrContent = Details->Render(Builder, Target, SetCurrLabel);
+        if (HasFusedContent) {
+          CurrContent = Builder.blockify(CurrContent, Fused->InnerMap.find(Target->Id)->second->Render(Builder, InLoop));
+        } else if (Details->Type == Branch::Nested) {
+          // Nest the parent content here, and remove it from showing up afterwards as Next
+          assert(Parent->Next);
+          CurrContent = Builder.blockify(CurrContent, Parent->Next->Render(Builder, InLoop));
+          Parent->Next = nullptr;
         }
-        if (!CurrIf) {
-          assert(!Root);
-          Root = Now;
+      }
+      bool isDefault = iter == ProcessedBranchesOut.end();
+      // If there is nothing to show in this branch, omit the condition
+      if (CurrContent) {
+        if (isDefault) {
+          wasm::Expression* Now;
+          if (RemainingConditions) {
+            Now = Builder.makeIf(RemainingConditions, CurrContent);
+          } else {
+            Now = CurrContent;
+          }
+          if (!CurrIf) {
+            assert(!Root);
+            Root = Now;
+          } else {
+            CurrIf->ifFalse = Now;
+          }
         } else {
-          CurrIf->ifFalse = Now;
+          auto* Now = Builder.makeIf(Details->Condition, CurrContent);
+          if (!CurrIf) {
+            assert(!Root);
+            Root = CurrIf = Now;
+          } else {
+            CurrIf->ifFalse = Now;
+            CurrIf = Now;
+          }
         }
       } else {
-        auto* Now = Builder.makeIf(Details->Condition, CurrContent);
-        if (!CurrIf) {
-          assert(!Root);
-          Root = CurrIf = Now;
+        auto* Now = Builder.makeUnary(wasm::EqZInt32, Details->Condition);
+        if (RemainingConditions) {
+          RemainingConditions = Builder.makeBinary(wasm::AndInt32, RemainingConditions, Now);
         } else {
-          CurrIf->ifFalse = Now;
-          CurrIf = Now;
+          RemainingConditions = Now;
         }
       }
-    } else {
-      auto* Now = Builder.makeUnary(wasm::EqZInt32, Details->Condition);
-      if (RemainingConditions) {
-        RemainingConditions = Builder.makeBinary(wasm::AndInt32, RemainingConditions, Now);
+      if (isDefault) break;
+    }
+  } else {
+    // Emit a switch
+    auto Base = std::string("switch$") + std::to_string(Id);
+    auto SwitchDefault = wasm::Name(Base + "$default");
+    auto SwitchLeave = wasm::Name(Base + "$leave");
+    std::map<Block*, wasm::Name> BlockNameMap;
+    auto* Outer = Builder.makeBlock();
+    auto* Inner = Outer;
+    std::vector<wasm::Name> Table;
+    for (auto& iter : ProcessedBranchesOut) {
+      Block *Target = iter.first;
+      Branch *Details = iter.second;
+      wasm::Name CurrName;
+      if (Details->SwitchValues) {
+        CurrName = wasm::Name(Base + "$case$" + std::to_string(Target->Id));
+      } else {
+        CurrName = SwitchDefault;
+      }
+      // generate the content for this block
+      bool SetCurrLabel = (SetLabel && Target->IsCheckedMultipleEntry) || ForceSetLabel;
+      bool HasFusedContent = Fused && contains(Fused->InnerMap, Target->Id);
+      wasm::Expression* CurrContent = nullptr;
+      if (SetCurrLabel || Details->Type != Branch::Direct || HasFusedContent || Details->Code) {
+        CurrContent = Details->Render(Builder, Target, SetCurrLabel);
+        if (HasFusedContent) {
+          CurrContent = Builder.blockify(CurrContent, Fused->InnerMap.find(Target->Id)->second->Render(Builder, InLoop));
+        } else if (Details->Type == Branch::Nested) {
+          // Nest the parent content here, and remove it from showing up afterwards as Next
+          assert(Parent->Next);
+          CurrContent = Builder.blockify(CurrContent, Parent->Next->Render(Builder, InLoop));
+          Parent->Next = nullptr;
+        }
+      }
+      // generate a block to branch to, if we have content
+      if (CurrContent) {
+        auto* NextOuter = Builder.makeBlock();
+        NextOuter->list.push_back(Outer);
+        Outer->name = CurrName; // breaking on Outer leads to the content in NextOuter
+        NextOuter->list.push_back(CurrContent);
+        NextOuter->list.push_back(Builder.makeBreak(SwitchLeave));
+        // prepare for more nesting
+        Outer = NextOuter;
       } else {
-        RemainingConditions = Now;
+        CurrName = SwitchLeave; // just go out straight from the table
+        if (!Details->SwitchValues) {
+          // this is the default, and it has no content. So make the default be the leave
+          for (auto& Value : Table) {
+            if (Value == SwitchDefault) Value = SwitchLeave;
+          }
+          SwitchDefault = SwitchLeave;
+        }
+      }
+      if (Details->SwitchValues) {
+        for (auto Value : *Details->SwitchValues) {
+          while (Table.size() <= Value) Table.push_back(SwitchDefault);
+          Table[Value] = CurrName;
+        }
       }
     }
-    if (isDefault) break;
+    // finish up the whole pattern
+    Outer->name = SwitchLeave;
+    Inner->list.push_back(Builder.makeSwitch(Table, SwitchDefault, SwitchCondition));
+    Root = Outer;
   }
 
   if (Root) {
diff --git a/src/cfg/Relooper.h b/src/cfg/Relooper.h
index 3040b0fc2..74a95c8bf 100644
--- a/src/cfg/Relooper.h
+++ b/src/cfg/Relooper.h
@@ -27,10 +27,11 @@ added since the original academic paper [1] was published about it.
 #include <stdarg.h>
 #include <stdlib.h>
 
-#include <map>
 #include <deque>
-#include <set>
 #include <list>
+#include <map>
+#include <memory>
+#include <set>
 
 #include "wasm.h"
 #include "wasm-builder.h"
@@ -52,10 +53,10 @@ public:
   wasm::Binary* makeCheckLabel(wasm::Index value) {
     return makeBinary(wasm::EqInt32, makeGetLabel(), makeConst(wasm::Literal(int32_t(value))));
   }
-  wasm::Break* makeBreak(int id) {
+  wasm::Break* makeShapeBreak(int id) {
     return wasm::Builder::makeBreak(getBreakName(id));
   }
-  wasm::Break* makeContinue(int id) {
+  wasm::Break* makeShapeContinue(int id) {
     return wasm::Builder::makeBreak(getContinueName(id));
   }
 
@@ -82,20 +83,19 @@ struct Branch {
   Shape *Ancestor; // If not NULL, this shape is the relevant one for purposes of getting to the target block. We break or continue on it
   Branch::FlowType Type; // If Ancestor is not NULL, this says whether to break or continue
 
-  // A branch either has a condition expression if the block ends in ifs, or if the block ends in a switch, then an index, which
-  // becomes the index in the table of the switch.
-  union {
-    wasm::Expression* Condition; // The condition for which we branch. For example, "my_var == 1". Conditions are checked one by one. One of the conditions should have NULL as the condition, in which case it is the default
-    wasm::Index Index; // The index in the table of the switch that we correspond do
-  };
+  // A branch either has a condition expression if the block ends in ifs, or if the block ends in a switch, then a list of indexes, which
+  // becomes the indexes in the table of the switch. If not a switch, the condition can be any expression.
+  wasm::Expression* Condition;
+  std::unique_ptr<std::vector<wasm::Index>> SwitchValues; // switches are rare, so have just a pointer here
 
   wasm::Expression* Code; // If provided, code that is run right before the branch is taken. This is useful for phis
 
   Branch(wasm::Expression* ConditionInit, wasm::Expression* CodeInit = nullptr);
-  Branch(wasm::Index IndexInit, wasm::Expression* CodeInit = nullptr);
+
+  Branch(std::vector<wasm::Index>&& ValuesInit, wasm::Expression* CodeInit = nullptr);
 
   // Emits code for branch
-   wasm::Expression* Render(RelooperBuilder& Builder, Block *Target, bool SetLabel);
+  wasm::Expression* Render(RelooperBuilder& Builder, Block *Target, bool SetLabel);
 };
 
 // like std::set, except that begin() -> end() iterates in the
@@ -226,8 +226,14 @@ struct Block {
   Block(wasm::Expression* CodeInit, wasm::Expression* SwitchConditionInit = nullptr);
   ~Block();
 
+  // Add a branch: if the condition holds we branch (or if null, we branch if all others failed)
+  // Note that there can be only one branch from A to B (if you need multiple conditions for the branch,
+  // create a more interesting expression in the Condition).
   void AddBranchTo(Block *Target, wasm::Expression* Condition, wasm::Expression* Code = nullptr);
-  void AddBranchTo(Block *Target, wasm::Index Index, wasm::Expression* Code = nullptr);
+
+  // Add a switch branch: if the switch condition is one of these values, we branch (or if the list is empty, we are the default)
+  // Note that there can be only one branch from A to B (if you need multiple values for the branch, that's what the array and default are for).
+  void AddSwitchBranchTo(Block *Target, std::vector<wasm::Index>&& Values, wasm::Expression* Code = nullptr);
 
   // Emit code for the block, including its contents and branchings out
   wasm::Expression* Render(RelooperBuilder& Builder, bool InLoop);