--simplify-locals-nonesting (#1525)

Add a version of simplify-locals which does not create nesting. This keeps the IR flat (in the sense of --flatten).

Also refactor simpify-locals to be a template, so the various modes are all template parameters.

3 files changed, 99 insertions, 58 deletions

diff --git a/src/passes/SimplifyLocals.cpp b/src/passes/SimplifyLocals.cpp
index f6997d14e..4db85eb3e 100644
--- a/src/passes/SimplifyLocals.cpp
+++ b/src/passes/SimplifyLocals.cpp
@@ -32,7 +32,7 @@
 // can get rid of those (the operation is trivial there after it sorts by use
 // frequency).
 //
-// This pass has two main options:
+// This pass has options:
 //
 //   * Tee: allow teeing, i.e., sinking a local with more than one use,
 //          and so after sinking we have a tee for the first use.
@@ -40,6 +40,11 @@
 //                internal set_locals into one on the outside,
 //                that can itself then be sunk further.
 //
+// There is also an option to disallow nesting entirely, which disallows
+// Tee and Structure from those 2 options, and also disallows any sinking
+// operation that would create nesting. This keeps the IR flat while
+// removing redundant locals.
+//
 
 #include <wasm.h>
 #include <wasm-builder.h>
@@ -73,14 +78,11 @@ struct SetLocalRemover : public PostWalker<SetLocalRemover> {
 
 // Main class
 
-struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals>> {
+template<bool allowTee = true, bool allowStructure = true, bool allowNesting = true>
+struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals<allowTee, allowStructure, allowNesting>>> {
   bool isFunctionParallel() override { return true; }
 
-  Pass* create() override { return new SimplifyLocals(allowTee, allowStructure); }
-
-  bool allowTee, allowStructure;
-
-  SimplifyLocals(bool allowTee, bool allowStructure) : allowTee(allowTee), allowStructure(allowStructure) {}
+  Pass* create() override { return new SimplifyLocals<allowTee, allowStructure, allowNesting>(); }
 
   // information for a set_local we can sink
   struct SinkableInfo {
@@ -127,7 +129,7 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals>>
   // local => # of get_locals for it
   GetLocalCounter getCounter;
 
-  static void doNoteNonLinear(SimplifyLocals* self, Expression** currp) {
+  static void doNoteNonLinear(SimplifyLocals<allowTee, allowStructure, allowNesting>* self, Expression** currp) {
     auto* curr = *currp;
     if (curr->is<Break>()) {
       auto* br = curr->cast<Break>();
@@ -152,25 +154,25 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals>>
     self->sinkables.clear();
   }
 
-  static void doNoteIfElseCondition(SimplifyLocals* self, Expression** currp) {
+  static void doNoteIfElseCondition(SimplifyLocals<allowTee, allowStructure, allowNesting>* self, Expression** currp) {
     // we processed the condition of this if-else, and now control flow branches
     // into either the true or the false sides
     assert((*currp)->cast<If>()->ifFalse);
     self->sinkables.clear();
   }
 
-  static void doNoteIfElseTrue(SimplifyLocals* self, Expression** currp) {
+  static void doNoteIfElseTrue(SimplifyLocals<allowTee, allowStructure, allowNesting>* self, Expression** currp) {
     // we processed the ifTrue side of this if-else, save it on the stack
     assert((*currp)->cast<If>()->ifFalse);
     self->ifStack.push_back(std::move(self->sinkables));
   }
 
-  static void doNoteIfElseFalse(SimplifyLocals* self, Expression** currp) {
+  static void doNoteIfElseFalse(SimplifyLocals<allowTee, allowStructure, allowNesting>* self, Expression** currp) {
     // we processed the ifFalse side of this if-else, we can now try to
     // mere with the ifTrue side and optimize a return value, if possible
     auto* iff = (*currp)->cast<If>();
     assert(iff->ifFalse);
-    if (self->allowStructure) {
+    if (allowStructure) {
       self->optimizeIfReturn(iff, currp, self->ifStack.back());
     }
     self->ifStack.pop_back();
@@ -202,12 +204,39 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals>>
   void visitGetLocal(GetLocal *curr) {
     auto found = sinkables.find(curr->index);
     if (found != sinkables.end()) {
+      auto* set = (*found->second.item)->template cast<SetLocal>(); // the set we may be sinking
+      bool oneUse = firstCycle || getCounter.num[curr->index] == 1;
+      auto* get = set->value->template dynCast<GetLocal>(); // the set's value may be a get (i.e., the set is a copy)
+      // if nesting is not allowed, and this might cause nesting, check if the sink would cause such a thing
+      if (!allowNesting) {
+        // a get is always ok to sink
+        if (!get) {
+          assert(expressionStack.size() >= 2);
+          assert(expressionStack[expressionStack.size() - 1] == curr);
+          auto* parent = expressionStack[expressionStack.size() - 2];
+          bool parentIsSet = parent->template is<SetLocal>();
+          // if the parent of this get is a set, we can sink into the set's value,
+          // it would not be nested.
+          if (!parentIsSet) {
+            return;
+          }
+        }
+      }
+      // we can optimize here
+      if (!allowNesting && get && !oneUse) {
+        // if we can't nest 's a copy with multiple uses, then we can't create
+        // a tee, and we can't nop the origin, but we can at least switch to
+        // the copied index, which may make the origin unneeded eventually.
+        curr->index = get->index;
+        anotherCycle = true;
+        return;
+      }
       // sink it, and nop the origin
-      auto* set = (*found->second.item)->cast<SetLocal>();
-      if (firstCycle || getCounter.num[curr->index] == 1) {
-        replaceCurrent(set->value);
+      if (oneUse) {
+        // with just one use, we can sink just the value
+        this->replaceCurrent(set->value);
       } else {
-        replaceCurrent(set);
+        this->replaceCurrent(set);
         assert(!set->isTee());
         set->setTee(true);
       }
@@ -225,7 +254,7 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals>>
     if (set) {
       assert(set->isTee());
       set->setTee(false);
-      replaceCurrent(set);
+      this->replaceCurrent(set);
     }
   }
 
@@ -242,9 +271,11 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals>>
     }
   }
 
+  // a full expression stack is used when !allowNesting, so that we can check if
+  // a sink would cause nesting
   std::vector<Expression*> expressionStack;
 
-  static void visitPre(SimplifyLocals* self, Expression** currp) {
+  static void visitPre(SimplifyLocals<allowTee, allowStructure, allowNesting>* self, Expression** currp) {
     Expression* curr = *currp;
 
     EffectAnalyzer effects(self->getPassOptions());
@@ -252,10 +283,12 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals>>
       self->checkInvalidations(effects);
     }
 
-    self->expressionStack.push_back(curr);
+    if (!allowNesting) {
+      self->expressionStack.push_back(curr);
+    }
   }
 
-  static void visitPost(SimplifyLocals* self, Expression** currp) {
+  static void visitPost(SimplifyLocals<allowTee, allowStructure, allowNesting>* self, Expression** currp) {
     // perform main SetLocal processing here, since we may be the result of
     // replaceCurrent, i.e., the visitor was not called.
     auto* set = (*currp)->dynCast<SetLocal>();
@@ -265,7 +298,7 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals>>
       // store is dead, leave just the value
       auto found = self->sinkables.find(set->index);
       if (found != self->sinkables.end()) {
-        auto* previous = (*found->second.item)->cast<SetLocal>();
+        auto* previous = (*found->second.item)->template cast<SetLocal>();
         assert(!previous->isTee());
         auto* previousValue = previous->value;
         Drop* drop = ExpressionManipulator::convert<SetLocal, Drop>(previous);
@@ -287,7 +320,9 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals>>
       self->sinkables.emplace(std::make_pair(index, SinkableInfo(currp, self->getPassOptions())));
     }
 
-    self->expressionStack.pop_back();
+    if (!allowNesting) {
+      self->expressionStack.pop_back();
+    }
   }
 
   bool canSink(SetLocal* set) {
@@ -308,7 +343,7 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals>>
     auto breaks = std::move(blockBreaks[block->name]);
     blockBreaks.erase(block->name);
     if (breaks.size() == 0) return; // block has no branches TODO we might optimize trivial stuff here too
-    assert(!(*breaks[0].brp)->cast<Break>()->value); // block does not already have a return value (if one break has one, they all do)
+    assert(!(*breaks[0].brp)->template cast<Break>()->value); // block does not already have a return value (if one break has one, they all do)
     // look for a set_local that is present in them all
     bool found = false;
     Index sharedIndex = -1;
@@ -349,8 +384,8 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals>>
       // move break set_local's value to the break
       auto* breakSetLocalPointer = breaks[j].sinkables.at(sharedIndex).item;
       auto* brp = breaks[j].brp;
-      auto* br = (*brp)->cast<Break>();
-      auto* set = (*breakSetLocalPointer)->cast<SetLocal>();
+      auto* br = (*brp)->template cast<Break>();
+      auto* set = (*breakSetLocalPointer)->template cast<SetLocal>();
       if (br->condition) {
         // TODO: optimize
         FindAll<SetLocal> findAll(br->condition);
@@ -361,8 +396,8 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals>>
             // itself, there is any risk
             Nop nop;
             *breakSetLocalPointer = &nop;
-            EffectAnalyzer condition(getPassOptions(), br->condition);
-            EffectAnalyzer value(getPassOptions(), set);
+            EffectAnalyzer condition(this->getPassOptions(), br->condition);
+            EffectAnalyzer value(this->getPassOptions(), set);
             *breakSetLocalPointer = set;
             if (condition.invalidates(value)) {
               // indeed, we can't do this, stop
@@ -384,7 +419,7 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals>>
     }
     // move block set_local's value to the end, in return position, and nop the set
     auto* blockSetLocalPointer = sinkables.at(sharedIndex).item;
-    auto* value = (*blockSetLocalPointer)->cast<SetLocal>()->value;
+    auto* value = (*blockSetLocalPointer)->template cast<SetLocal>()->value;
     block->list[block->list.size() - 1] = value;
     block->type = value->type;
     ExpressionManipulator::nop(*blockSetLocalPointer);
@@ -392,25 +427,25 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals>>
       // move break set_local's value to the break
       auto* breakSetLocalPointer = breaks[j].sinkables.at(sharedIndex).item;
       auto* brp = breaks[j].brp;
-      auto* br = (*brp)->cast<Break>();
+      auto* br = (*brp)->template cast<Break>();
       assert(!br->value);
       // if the break is conditional, then we must set the value here - if the break is not reached, we must still have the new value in the local
-      auto* set = (*breakSetLocalPointer)->cast<SetLocal>();
+      auto* set = (*breakSetLocalPointer)->template cast<SetLocal>();
       if (br->condition) {
         br->value = set;
         set->setTee(true);
-        *breakSetLocalPointer = getModule()->allocator.alloc<Nop>();
+        *breakSetLocalPointer = this->getModule()->allocator.template alloc<Nop>();
         // in addition, as this is a conditional br that now has a value, it now returns a value, so it must be dropped
         br->finalize();
-        *brp = Builder(*getModule()).makeDrop(br);
+        *brp = Builder(*this->getModule()).makeDrop(br);
       } else {
         br->value = set->value;
         ExpressionManipulator::nop(set);
       }
     }
     // finally, create a set_local on the block itself
-    auto* newSetLocal = Builder(*getModule()).makeSetLocal(sharedIndex, block);
-    replaceCurrent(newSetLocal);
+    auto* newSetLocal = Builder(*this->getModule()).makeSetLocal(sharedIndex, block);
+    this->replaceCurrent(newSetLocal);
     sinkables.clear();
     anotherCycle = true;
   }
@@ -446,39 +481,39 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals>>
     }
     // all set, go
     auto *ifTrueItem = ifTrue.at(sharedIndex).item;
-    ifTrueBlock->list[ifTrueBlock->list.size() - 1] = (*ifTrueItem)->cast<SetLocal>()->value;
+    ifTrueBlock->list[ifTrueBlock->list.size() - 1] = (*ifTrueItem)->template cast<SetLocal>()->value;
     ExpressionManipulator::nop(*ifTrueItem);
     ifTrueBlock->finalize();
     assert(ifTrueBlock->type != none);
     auto *ifFalseItem = ifFalse.at(sharedIndex).item;
-    ifFalseBlock->list[ifFalseBlock->list.size() - 1] = (*ifFalseItem)->cast<SetLocal>()->value;
+    ifFalseBlock->list[ifFalseBlock->list.size() - 1] = (*ifFalseItem)->template cast<SetLocal>()->value;
     ExpressionManipulator::nop(*ifFalseItem);
     ifFalseBlock->finalize();
     assert(ifTrueBlock->type != none);
     iff->finalize(); // update type
     assert(iff->type != none);
     // finally, create a set_local on the iff itself
-    auto* newSetLocal = Builder(*getModule()).makeSetLocal(sharedIndex, iff);
+    auto* newSetLocal = Builder(*this->getModule()).makeSetLocal(sharedIndex, iff);
     *currp = newSetLocal;
     anotherCycle = true;
   }
 
   // override scan to add a pre and a post check task to all nodes
-  static void scan(SimplifyLocals* self, Expression** currp) {
+  static void scan(SimplifyLocals<allowTee, allowStructure, allowNesting>* self, Expression** currp) {
     self->pushTask(visitPost, currp);
 
     auto* curr = *currp;
 
     if (curr->is<If>() && curr->cast<If>()->ifFalse) {
       // handle if-elses in a special manner, using the ifStack
-      self->pushTask(SimplifyLocals::doNoteIfElseFalse, currp);
-      self->pushTask(SimplifyLocals::scan, &curr->cast<If>()->ifFalse);
-      self->pushTask(SimplifyLocals::doNoteIfElseTrue, currp);
-      self->pushTask(SimplifyLocals::scan, &curr->cast<If>()->ifTrue);
-      self->pushTask(SimplifyLocals::doNoteIfElseCondition, currp);
-      self->pushTask(SimplifyLocals::scan, &curr->cast<If>()->condition);
+      self->pushTask(SimplifyLocals<allowTee, allowStructure, allowNesting>::doNoteIfElseFalse, currp);
+      self->pushTask(SimplifyLocals<allowTee, allowStructure, allowNesting>::scan, &curr->cast<If>()->ifFalse);
+      self->pushTask(SimplifyLocals<allowTee, allowStructure, allowNesting>::doNoteIfElseTrue, currp);
+      self->pushTask(SimplifyLocals<allowTee, allowStructure, allowNesting>::scan, &curr->cast<If>()->ifTrue);
+      self->pushTask(SimplifyLocals<allowTee, allowStructure, allowNesting>::doNoteIfElseCondition, currp);
+      self->pushTask(SimplifyLocals<allowTee, allowStructure, allowNesting>::scan, &curr->cast<If>()->condition);
     } else {
-      super::scan(self, currp);
+      WalkerPass<LinearExecutionWalker<SimplifyLocals<allowTee, allowStructure, allowNesting>>>::scan(self, currp);
     }
 
     self->pushTask(visitPre, currp);
@@ -500,11 +535,11 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals>>
     do {
       anotherCycle = false;
       // main operation
-      super::doWalkFunction(func);
+      WalkerPass<LinearExecutionWalker<SimplifyLocals<allowTee, allowStructure, allowNesting>>>::doWalkFunction(func);
       // enlarge blocks that were marked, for the next round
       if (blocksToEnlarge.size() > 0) {
         for (auto* block : blocksToEnlarge) {
-          block->list.push_back(getModule()->allocator.alloc<Nop>());
+          block->list.push_back(this->getModule()->allocator.template alloc<Nop>());
         }
         blocksToEnlarge.clear();
         anotherCycle = true;
@@ -512,15 +547,15 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals>>
       // enlarge ifs that were marked, for the next round
       if (ifsToEnlarge.size() > 0) {
         for (auto* iff : ifsToEnlarge) {
-          auto ifTrue = Builder(*getModule()).blockify(iff->ifTrue);
+          auto ifTrue = Builder(*this->getModule()).blockify(iff->ifTrue);
           iff->ifTrue = ifTrue;
-          if (ifTrue->list.size() == 0 || !ifTrue->list.back()->is<Nop>()) {
-            ifTrue->list.push_back(getModule()->allocator.alloc<Nop>());
+          if (ifTrue->list.size() == 0 || !ifTrue->list.back()->template is<Nop>()) {
+            ifTrue->list.push_back(this->getModule()->allocator.template alloc<Nop>());
           }
-          auto ifFalse = Builder(*getModule()).blockify(iff->ifFalse);
+          auto ifFalse = Builder(*this->getModule()).blockify(iff->ifFalse);
           iff->ifFalse = ifFalse;
-          if (ifFalse->list.size() == 0 || !ifFalse->list.back()->is<Nop>()) {
-            ifFalse->list.push_back(getModule()->allocator.alloc<Nop>());
+          if (ifFalse->list.size() == 0 || !ifFalse->list.back()->template is<Nop>()) {
+            ifFalse->list.push_back(this->getModule()->allocator.template alloc<Nop>());
           }
         }
         ifsToEnlarge.clear();
@@ -548,19 +583,23 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals>>
 };
 
 Pass *createSimplifyLocalsPass() {
-  return new SimplifyLocals(true, true);
+  return new SimplifyLocals<true, true>();
 }
 
 Pass *createSimplifyLocalsNoTeePass() {
-  return new SimplifyLocals(false, true);
+  return new SimplifyLocals<false, true>();
 }
 
 Pass *createSimplifyLocalsNoStructurePass() {
-  return new SimplifyLocals(true, false);
+  return new SimplifyLocals<true, false>();
 }
 
 Pass *createSimplifyLocalsNoTeeNoStructurePass() {
-  return new SimplifyLocals(false, false);
+  return new SimplifyLocals<false, false>();
+}
+
+Pass *createSimplifyLocalsNoNestingPass() {
+  return new SimplifyLocals<false, false, false>();
 }
 
 } // namespace wasm
diff --git a/src/passes/pass.cpp b/src/passes/pass.cpp
index 0ec5864e0..ae9484d29 100644
--- a/src/passes/pass.cpp
+++ b/src/passes/pass.cpp
@@ -108,8 +108,9 @@ void PassRegistry::registerPasses() {
   registerPass("reorder-locals", "sorts locals by access frequency", createReorderLocalsPass);
   registerPass("rereloop", "re-optimize control flow using the relooper algorithm", createReReloopPass);
   registerPass("rse", "remove redundant set_locals", createRedundantSetEliminationPass);
-  registerPass("simplify-locals", "miscellaneous locals-related optimizations", createSimplifyLocalsPass);
   registerPass("safe-heap", "instrument loads and stores to check for invalid behavior", createSafeHeapPass);
+  registerPass("simplify-locals", "miscellaneous locals-related optimizations", createSimplifyLocalsPass);
+  registerPass("simplify-locals-nonesting", "miscellaneous locals-related optimizations (no nesting at all; preserves flatness)", createSimplifyLocalsNoNestingPass);
   registerPass("simplify-locals-notee", "miscellaneous locals-related optimizations", createSimplifyLocalsNoTeePass);
   registerPass("simplify-locals-nostructure", "miscellaneous locals-related optimizations", createSimplifyLocalsNoStructurePass);
   registerPass("simplify-locals-notee-nostructure", "miscellaneous locals-related optimizations", createSimplifyLocalsNoTeeNoStructurePass);
diff --git a/src/passes/passes.h b/src/passes/passes.h
index 2eaf1049e..5ab03a439 100644
--- a/src/passes/passes.h
+++ b/src/passes/passes.h
@@ -67,6 +67,7 @@ Pass* createReReloopPass();
 Pass* createRedundantSetEliminationPass();
 Pass* createSafeHeapPass();
 Pass* createSimplifyLocalsPass();
+Pass* createSimplifyLocalsNoNestingPass();
 Pass* createSimplifyLocalsNoTeePass();
 Pass* createSimplifyLocalsNoStructurePass();
 Pass* createSimplifyLocalsNoTeeNoStructurePass();