Flattening rewrite (#1201)

Rename flatten-control-flow to flatten, which now flattens everything, not just control flow, so e.g.

(i32.add
  (call $x)
  (call $y)
)
==>
(block
  (set_local $temp_x (call $x))
  (set_local $temp_y (call $y))
  (i32.add
    (get_local $x)
    (get_local $y)
  )
)

This uses more locals than before, but is much simpler and avoids a bunch of corner cases and fuzz bugs the old one hit. We can optimize later if necessary.

Author: kripken

auto_update_tests.py

@@ -7,7 +7,7 @@
     ASM2WASM, MOZJS, S2WASM, WASM_SHELL, WASM_OPT, WASM_AS, WASM_DIS,
     WASM_CTOR_EVAL, WASM_MERGE, WASM_REDUCE, WASM2ASM,
     BINARYEN_INSTALL_DIR, has_shell_timeout)
-from scripts.test.wasm2asm import tests, spec_tests, extra_tests
+from scripts.test.wasm2asm import tests, spec_tests, extra_tests, assert_tests
 
 
 print '[ processing and updating testcases... ]\n'
@@ -64,19 +64,6 @@
     expected_file = os.path.join('test', dot_s_dir, wasm)
     with open(expected_file, 'w') as o: o.write(actual)
 
-'''
-for wasm in ['address.wast']:#os.listdir(os.path.join('test', 'spec')):
-  if wasm.endswith('.wast'):
-    print '..', wasm
-    asm = wasm.replace('.wast', '.2asm.js')
-    proc = subprocess.Popen([os.path.join('bin', 'wasm2asm'), os.path.join('test', 'spec', wasm)], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-    actual, err = proc.communicate()
-    assert proc.returncode == 0, err
-    assert err == '', 'bad err:' + err
-    expected_file = os.path.join('test', asm)
-    open(expected_file, 'w').write(actual)
-'''
-
 for t in sorted(os.listdir(os.path.join('test', 'print'))):
   if t.endswith('.wast'):
     print '..', t
@@ -281,6 +268,22 @@
   out = run_command(cmd)
   with open(expected_file, 'w') as o: o.write(out)
 
+for wasm in assert_tests:
+  print '..', wasm
+
+  asserts = os.path.basename(wasm).replace('.wast.asserts', '.asserts.js')
+  traps = os.path.basename(wasm).replace('.wast.asserts', '.traps.js')
+  asserts_expected_file = os.path.join('test', asserts)
+  traps_expected_file = os.path.join('test', traps)
+
+  cmd = WASM2ASM + [os.path.join('test', wasm), '--allow-asserts']
+  out = run_command(cmd)
+  with open(asserts_expected_file, 'w') as o: o.write(out)
+
+  cmd += ['--pedantic']
+  out = run_command(cmd)
+  with open(traps_expected_file, 'w') as o: o.write(out)
+
 if has_shell_timeout():
   print '\n[ checking wasm-reduce ]\n'
 

src/passes/CMakeLists.txt

@@ -7,7 +7,7 @@ SET(passes_SOURCES
   DeadCodeElimination.cpp
   DuplicateFunctionElimination.cpp
   ExtractFunction.cpp
-  FlattenControlFlow.cpp
+  Flatten.cpp
   Inlining.cpp
   LegalizeJSInterface.cpp
   LocalCSE.cpp

src/passes/Flatten.cpp

@@ -0,0 +1,337 @@
+/*
+ * Copyright 2017 WebAssembly Community Group participants
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+//
+// Flattens code, removing nesting.e.g. an if return value would be
+// converted to a local
+//
+//  (i32.add
+//    (if (..condition..)
+//      (..if true..)
+//      (..if false..)
+//    )
+//    (i32.const 1)
+//  )
+// =>
+//  (if (..condition..)
+//    (set_local $temp
+//      (..if true..)
+//    )
+//    (set_local $temp
+//      (..if false..)
+//    )
+//  )
+//  (i32.add
+//    (get_local $temp)
+//    (i32.const 1)
+//  )
+//
+// Formally, this pass flattens in the precise sense of
+// making the AST have these properties:
+//
+//  1. The operands of an instruction must be a get_local or a const.
+//     anything else is written to a local earlier.
+//  2. Disallow block, loop, and if return values, i.e., do not use
+//     control flow to pass around values.
+//
+
+#include <wasm.h>
+#include <pass.h>
+#include <wasm-builder.h>
+#include <ast_utils.h>
+#include <ast/effects.h>
+
+namespace wasm {
+
+// We use the following algorithm: we maintain a list of "preludes", code
+// that runs right before an expression. When we visit an expression we
+// must handle it and its preludes. If the expression has side effects,
+// we reduce it to a get_local and add a prelude for that. We then handle
+// the preludes, by moving them to the parent or handling them directly.
+// we can move them to the parent if the parent is not a control flow
+// structure. Otherwise, if the parent is a control flow structure, it
+// will incorporate the preludes of its children accordingly.
+// As a result, when we reach a node, we know its children have no
+// side effects (they have been moved to a prelude), or we are a
+// control flow structure (which allows children with side effects,
+// e.g. a return as a block element).
+// Once exception is that we allow an (unreachable) node, which is used
+// when we move something unreachable to another place, and need a
+// placeholder. We will never reach that (unreachable) anyhow
+struct Flatten : public WalkerPass<ExpressionStackWalker<Flatten, UnifiedExpressionVisitor<Flatten>>> {
+  bool isFunctionParallel() override { return true; }
+
+  Pass* create() override { return new Flatten; }
+
+  // For each expression, a bunch of expressions that should execute right before it
+  std::unordered_map<Expression*, std::vector<Expression*>> preludes;
+
+  // Break values are sent through a temp local
+  std::unordered_map<Name, Index> breakTemps;
+
+  void visitExpression(Expression* curr) {
+    std::vector<Expression*> ourPreludes;
+    Builder builder(*getModule());
+
+    if (isControlFlowStructure(curr)) {
+      // handle control flow explicitly. our children do not have control flow,
+      // but they do have preludes which we need to set up in the right place
+      assert(preludes.find(curr) == preludes.end()); // no one should have given us preludes, they are on the children
+      if (auto* block = curr->dynCast<Block>()) {
+        // make a new list, where each item's preludes are added before it
+        ExpressionList newList(getModule()->allocator);
+        for (auto* item : block->list) {
+          auto iter = preludes.find(item);
+          if (iter != preludes.end()) {
+            auto& itemPreludes = iter->second;
+            for (auto* prelude : itemPreludes) {
+              newList.push_back(prelude);
+            }
+            itemPreludes.clear();
+          }
+          newList.push_back(item);
+        }
+        block->list.swap(newList);
+        // remove a block return value
+        auto type = block->type;
+        if (isConcreteWasmType(type)) {
+          // if there is a temp index for breaking to the block, use that
+          Index temp;
+          auto iter = breakTemps.find(block->name);
+          if (iter != breakTemps.end()) {
+            temp = iter->second;
+          } else {
+            temp = builder.addVar(getFunction(), type);
+          }
+          auto*& last = block->list.back();
+          if (isConcreteWasmType(last->type)) {
+            last = builder.makeSetLocal(temp, last);
+          }
+          block->finalize(none);
+          // and we leave just a get of the value
+          auto* rep = builder.makeGetLocal(temp, type);
+          replaceCurrent(rep);
+          // the whole block is now a prelude
+          ourPreludes.push_back(block);
+        }
+        // the block now has no return value, and may have become unreachable
+        block->finalize(none);
+      } else if (auto* iff = curr->dynCast<If>()) {
+        // condition preludes go before the entire if
+        auto* rep = getPreludesWithExpression(iff->condition, iff);
+        // arm preludes go in the arms. we must also remove an if value
+        auto* originalIfTrue = iff->ifTrue;
+        auto* originalIfFalse = iff->ifFalse;
+        auto type = iff->type;
+        Expression* prelude = nullptr;
+        if (isConcreteWasmType(type)) {
+          Index temp = builder.addVar(getFunction(), type);
+          if (isConcreteWasmType(iff->ifTrue->type)) {
+            iff->ifTrue = builder.makeSetLocal(temp, iff->ifTrue);
+          }
+          if (iff->ifFalse && isConcreteWasmType(iff->ifFalse->type)) {
+            iff->ifFalse = builder.makeSetLocal(temp, iff->ifFalse);
+          }
+          // the whole if (+any preludes from the condition) is now a prelude
+          prelude = rep;
+          // and we leave just a get of the value
+          rep = builder.makeGetLocal(temp, type);
+        }
+        iff->ifTrue = getPreludesWithExpression(originalIfTrue, iff->ifTrue);
+        if (iff->ifFalse) iff->ifFalse = getPreludesWithExpression(originalIfFalse, iff->ifFalse);
+        iff->finalize();
+        if (prelude) {
+          ReFinalizeNode().visit(prelude);
+          ourPreludes.push_back(prelude);
+        }
+        replaceCurrent(rep);
+      } else if (auto* loop = curr->dynCast<Loop>()) {
+        // remove a loop value
+        Expression* rep = loop;
+        auto* originalBody = loop->body;
+        auto type = loop->type;
+        if (isConcreteWasmType(type)) {
+          Index temp = builder.addVar(getFunction(), type);
+          loop->body = builder.makeSetLocal(temp, loop->body);
+          // and we leave just a get of the value
+          rep = builder.makeGetLocal(temp, type);
+          // the whole if is now a prelude
+          ourPreludes.push_back(loop);
+          loop->type = none;
+        }
+        loop->body = getPreludesWithExpression(originalBody, loop->body);
+        loop->finalize();
+        replaceCurrent(rep);
+      } else {
+        WASM_UNREACHABLE();
+      }
+    } else {
+      // for anything else, there may be existing preludes
+      auto iter = preludes.find(curr);
+      if (iter != preludes.end()) {
+        ourPreludes.swap(iter->second);
+      }
+      // special handling
+      if (auto* br = curr->dynCast<Break>()) {
+        if (br->value) {
+          auto type = br->value->type;
+          if (isConcreteWasmType(type)) {
+            // we are sending a value. use a local instead
+            Index temp = getTempForBreakTarget(br->name, type);
+            ourPreludes.push_back(builder.makeSetLocal(temp, br->value));
+            if (br->condition) {
+              // the value must also flow out
+              ourPreludes.push_back(br);
+              if (isConcreteWasmType(br->type)) {
+                replaceCurrent(builder.makeGetLocal(temp, type));
+              } else {
+                assert(br->type == unreachable);
+                replaceCurrent(builder.makeUnreachable());
+              }
+            }
+            br->value = nullptr;
+            br->finalize();
+          } else {
+            assert(type == unreachable);
+            // we don't need the br at all
+            replaceCurrent(br->value);
+          }
+        }
+      } else if (auto* sw = curr->dynCast<Switch>()) {
+        if (sw->value) {
+          auto type = sw->value->type;
+          if (isConcreteWasmType(type)) {
+            // we are sending a value. use a local instead
+            Index temp = builder.addVar(getFunction(), type);
+            ourPreludes.push_back(builder.makeSetLocal(temp, sw->value));
+            // we don't know which break target will be hit - assign to them all
+            std::set<Name> names;
+            for (auto target : sw->targets) {
+              names.insert(target);
+            }
+            names.insert(sw->default_);
+            for (auto name : names) {
+              ourPreludes.push_back(builder.makeSetLocal(
+                getTempForBreakTarget(name, type),
+                builder.makeGetLocal(temp, type)
+              ));
+            }
+            sw->value = nullptr;
+            sw->finalize();
+          } else {
+            assert(type == unreachable);
+            // we don't need the br at all
+            replaceCurrent(sw->value);
+          }
+        }
+      }
+    }
+    // continue for general handling of everything, control flow or otherwise
+    curr = getCurrent(); // we may have replaced it
+    // we have changed children
+    ReFinalizeNode().visit(curr);
+    // handle side effects and control flow, things we need to be
+    // in the prelude. note that we must handle anything here, not just
+    // side effects, as a sibling after us may have side effect for us,
+    // and thus we need to move in anticipation of that (e.g., we are
+    // a get, and a later sibling is a tee - if just the tee moves,
+    // that is bade) TODO optimize
+    if (isControlFlowStructure(curr) || EffectAnalyzer(getPassOptions(), curr).hasAnything()) {
+      // we need to move the side effects to the prelude
+      if (curr->type == unreachable) {
+        ourPreludes.push_back(curr);
+        replaceCurrent(builder.makeUnreachable());
+      } else if (curr->type == none) {
+        if (!curr->is<Nop>()) {
+          ourPreludes.push_back(curr);
+          replaceCurrent(builder.makeNop());
+        }
+      } else {
+        // use a local
+        auto type = curr->type;
+        Index temp = builder.addVar(getFunction(), type);
+        ourPreludes.push_back(builder.makeSetLocal(temp, curr));
+        replaceCurrent(builder.makeGetLocal(temp, type));
+      }
+    }
+    // next, finish up: migrate our preludes if we can
+    if (!ourPreludes.empty()) {
+      auto* parent = getParent();
+      if (parent && !isControlFlowStructure(parent)) {
+        auto& parentPreludes = preludes[parent];
+        for (auto* prelude : ourPreludes) {
+          parentPreludes.push_back(prelude);
+        }
+      } else {
+        // keep our preludes, parent will handle them
+        preludes[getCurrent()].swap(ourPreludes);
+      }
+    }
+  }
+
+  void visitFunction(Function* curr) {
+    auto* originalBody = curr->body;
+    // if the body is a block with a result, turn that into a return
+    if (isConcreteWasmType(curr->body->type)) {
+      curr->body = Builder(*getModule()).makeReturn(curr->body);
+    }
+    // the body may have preludes
+    curr->body = getPreludesWithExpression(originalBody, curr->body);
+  }
+
+private:
+  bool isControlFlowStructure(Expression* curr) {
+    return curr->is<Block>() || curr->is<If>() || curr->is<Loop>();
+  }
+
+  // gets an expression, either by itself, or in a block with its
+  // preludes (which we use up) before it
+  Expression* getPreludesWithExpression(Expression* curr) {
+    return getPreludesWithExpression(curr, curr);
+  }
+
+  // gets an expression, either by itself, or in a block with some
+  // preludes (which we use up) for another expression before it
+  Expression* getPreludesWithExpression(Expression* preluder, Expression* after) {
+    auto iter = preludes.find(preluder);
+    if (iter == preludes.end()) return after;
+    // we have preludes
+    auto& thePreludes = iter->second;
+    auto* ret = Builder(*getModule()).makeBlock(thePreludes);
+    thePreludes.clear();
+    ret->list.push_back(after);
+    ret->finalize();
+    return ret;
+  }
+
+  // get the temp local to be used for breaks to that target. allocates
+  // one if there isn't one yet
+  Index getTempForBreakTarget(Name name, WasmType type) {
+    auto iter = breakTemps.find(name);
+    if (iter != breakTemps.end()) {
+      return iter->second;
+    } else {
+      return breakTemps[name] = Builder(*getModule()).addVar(getFunction(), type);
+    }
+  }
+};
+
+Pass *createFlattenPass() {
+  return new Flatten();
+}
+
+} // namespace wasm
+