[Outlining] Fix outlining control flow

Changes the controlFlowQueue used in stringify-walker to push values of Expression*,  This ensures that we walk the Wasm module in the same order, regardless of whether the control flow expression is outlined.

Reviewers: tlively

Reviewed By: tlively

Pull Request: https://github.com/WebAssembly/binaryen/pull/6139

4 files changed, 208 insertions, 11 deletions

diff --git a/src/passes/Outlining.cpp b/src/passes/Outlining.cpp
index f799d2c7a..51d000745 100644
--- a/src/passes/Outlining.cpp
+++ b/src/passes/Outlining.cpp
@@ -100,8 +100,20 @@ struct ReconstructStringifyWalker
       existingBuilder.push(curr->iff->condition);
       ASSERT_OK(existingBuilder.visitIfStart(curr->iff));
       DBG(desc = "If Start for ");
+    } else if (reason.getElseStart()) {
+      ASSERT_OK(existingBuilder.visitElse());
+      DBG(desc = "Else Start at ");
     } else if (reason.getEnd()) {
       ASSERT_OK(existingBuilder.visitEnd());
+      // Outlining performs an unnested walk of the Wasm module, visiting
+      // each scope one at a time. IRBuilder, in contrast, expects to
+      // visit several nested scopes at a time. Thus, calling end() finalizes
+      // the control flow and places it on IRBuilder's internal stack, ready for
+      // the enclosing scope to consume its expressions off the stack. Since
+      // outlining walks unnested, the enclosing scope never arrives to retrieve
+      // its expressions off the stack, so we must call build() after visitEnd()
+      // to clear the internal stack IRBuilder manages.
+      ASSERT_OK(existingBuilder.build());
       DBG(desc = "End for ");
     } else {
       DBG(desc = "addUniqueSymbol for unimplemented control flow ");
diff --git a/src/passes/stringify-walker-impl.h b/src/passes/stringify-walker-impl.h
index 8ed166d75..cca2a3405 100644
--- a/src/passes/stringify-walker-impl.h
+++ b/src/passes/stringify-walker-impl.h
@@ -43,7 +43,7 @@ template<typename SubType>
 inline void StringifyWalker<SubType>::scan(SubType* self, Expression** currp) {
   Expression* curr = *currp;
   if (Properties::isControlFlowStructure(curr)) {
-    self->controlFlowQueue.push(currp);
+    self->controlFlowQueue.push(curr);
     self->pushTask(doVisitExpression, currp);
     // The if-condition is a value child consumed by the if control flow, which
     // makes the if-condition a true sibling rather than part of its contents in
@@ -60,9 +60,8 @@ inline void StringifyWalker<SubType>::scan(SubType* self, Expression** currp) {
 // of control flow.
 template<typename SubType> void StringifyWalker<SubType>::dequeueControlFlow() {
   auto& queue = controlFlowQueue;
-  Expression** currp = queue.front();
+  Expression* curr = queue.front();
   queue.pop();
-  Expression* curr = *currp;
 
   // TODO: Issue #5796, Make a ControlChildIterator
   switch (curr->_id) {
@@ -80,7 +79,7 @@ template<typename SubType> void StringifyWalker<SubType>::dequeueControlFlow() {
       addUniqueSymbol(SeparatorReason::makeIfStart(iff));
       Super::walk(iff->ifTrue);
       if (iff->ifFalse) {
-        addUniqueSymbol(SeparatorReason::makeElseStart(iff));
+        addUniqueSymbol(SeparatorReason::makeElseStart());
         Super::walk(iff->ifFalse);
       }
       addUniqueSymbol(SeparatorReason::makeEnd());
diff --git a/src/passes/stringify-walker.h b/src/passes/stringify-walker.h
index eb4a4482b..6f50f58d4 100644
--- a/src/passes/stringify-walker.h
+++ b/src/passes/stringify-walker.h
@@ -82,9 +82,7 @@ struct StringifyWalker
       If* iff;
     };
 
-    struct ElseStart {
-      If* iff;
-    };
+    struct ElseStart {};
 
     struct LoopStart {
       Loop* loop;
@@ -119,8 +117,8 @@ struct StringifyWalker
     static SeparatorReason makeIfStart(If* iff) {
       return SeparatorReason(IfStart{iff});
     }
-    static SeparatorReason makeElseStart(If* iff) {
-      return SeparatorReason(ElseStart{iff});
+    static SeparatorReason makeElseStart() {
+      return SeparatorReason(ElseStart{});
     }
     static SeparatorReason makeLoopStart(Loop* loop) {
       return SeparatorReason(LoopStart{loop});
@@ -170,7 +168,11 @@ struct StringifyWalker
     return o << "~~~Undefined in operator<< overload~~~";
   }
 
-  std::queue<Expression**> controlFlowQueue;
+  // To ensure control flow children are walked consistently during outlining,
+  // we push a copy of the control flow expression. This avoids an issue where
+  // control flow no longer points to the same expression after being
+  // outlined into a new function.
+  std::queue<Expression*> controlFlowQueue;
 
   /*
    * To initiate the walk, subclasses should call walkModule with a pointer to
diff --git a/test/lit/passes/outlining.wast b/test/lit/passes/outlining.wast
index ad712b31a..98cb5a316 100644
--- a/test/lit/passes/outlining.wast
+++ b/test/lit/passes/outlining.wast
@@ -238,7 +238,59 @@
   )
 )
 
-;; Tests that outlining works correctly with If control flow
+;; Tests that outlining works correctly with if-condition
+(module
+  ;; CHECK:      (type $0 (func))
+
+  ;; CHECK:      (type $1 (func (result i32)))
+
+  ;; CHECK:      (global $global$1 (mut i32) (i32.const 100))
+  (global $global$1 (mut i32) (i32.const 100))
+  ;; CHECK:      (func $outline$ (type $1) (result i32)
+  ;; CHECK-NEXT:  (i32.eqz
+  ;; CHECK-NEXT:   (global.get $global$1)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+
+  ;; CHECK:      (func $a (type $0)
+  ;; CHECK-NEXT:  (if
+  ;; CHECK-NEXT:   (call $outline$)
+  ;; CHECK-NEXT:   (global.set $global$1
+  ;; CHECK-NEXT:    (i32.const 15)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $a
+    (if
+      (i32.eqz
+        (global.get $global$1)
+      )
+      (global.set $global$1
+        (i32.const 15)
+      )
+    )
+  )
+  ;; CHECK:      (func $b (type $0)
+  ;; CHECK-NEXT:  (if
+  ;; CHECK-NEXT:   (call $outline$)
+  ;; CHECK-NEXT:   (global.set $global$1
+  ;; CHECK-NEXT:    (i32.const 20)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $b
+    (if
+      (i32.eqz
+        (global.get $global$1)
+      )
+      (global.set $global$1
+        (i32.const 20)
+      )
+    )
+  )
+)
+
+;; Outline if-true.
 (module
   ;; CHECK:      (type $0 (func (param i32)))
 
@@ -288,6 +340,138 @@
   )
 )
 
+;; Outline if-false.
+(module
+  ;; CHECK:      (type $0 (func))
+
+  ;; CHECK:      (global $global$1 (mut i32) (i32.const 100))
+  (global $global$1 (mut i32) (i32.const 100))
+  ;; CHECK:      (func $outline$ (type $0)
+  ;; CHECK-NEXT:  (global.set $global$1
+  ;; CHECK-NEXT:   (i32.const 100)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+
+  ;; CHECK:      (func $a (type $0)
+  ;; CHECK-NEXT:  (if
+  ;; CHECK-NEXT:   (i32.eqz
+  ;; CHECK-NEXT:    (global.get $global$1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (global.set $global$1
+  ;; CHECK-NEXT:    (i32.const 15)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (call $outline$)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $a
+    (if
+      (i32.eqz
+        (global.get $global$1)
+      )
+      (global.set $global$1
+        (i32.const 15)
+      )
+      (block
+        (global.set $global$1
+          (i32.const 100)
+        )
+      )
+    )
+  )
+  ;; CHECK:      (func $b (type $0)
+  ;; CHECK-NEXT:  (if
+  ;; CHECK-NEXT:   (i32.ctz
+  ;; CHECK-NEXT:    (global.get $global$1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (global.set $global$1
+  ;; CHECK-NEXT:    (i32.const 30)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (call $outline$)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $b
+    (if
+      (i32.ctz
+        (global.get $global$1)
+      )
+      (global.set $global$1
+        (i32.const 30)
+      )
+      (block
+        (global.set $global$1
+          (i32.const 100)
+        )
+      )
+    )
+  )
+)
+
+;; Outline if control flow, with matching if-condition, if-true, if-false
+;; TODO: Ideally outlining would keep the if-true and if-false inline in
+;; $outline$, instead of moving them to another outlined function ($outline$_3
+;; & $outline$_4) because of the unique symbol between the if-condition and
+;; if-true and the unique symbol between if-true and if-false.
+(module
+  ;; CHECK:      (type $0 (func))
+
+  ;; CHECK:      (global $global$1 (mut i32) (i32.const 100))
+  (global $global$1 (mut i32) (i32.const 100))
+  ;; CHECK:      (func $outline$ (type $0)
+  ;; CHECK-NEXT:  (if
+  ;; CHECK-NEXT:   (i32.eqz
+  ;; CHECK-NEXT:    (global.get $global$1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (call $outline$_3)
+  ;; CHECK-NEXT:   (call $outline$_4)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+
+  ;; CHECK:      (func $outline$_3 (type $0)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i32.const 10)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+
+  ;; CHECK:      (func $outline$_4 (type $0)
+  ;; CHECK-NEXT:  (global.set $global$1
+  ;; CHECK-NEXT:   (i32.const 20)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+
+  ;; CHECK:      (func $a (type $0)
+  ;; CHECK-NEXT:  (call $outline$)
+  ;; CHECK-NEXT: )
+  (func $a
+    (if
+      (i32.eqz
+        (global.get $global$1)
+      )
+      (drop
+        (i32.const 10)
+      )
+      (global.set $global$1
+        (i32.const 20)
+      )
+    )
+  )
+  ;; CHECK:      (func $b (type $0)
+  ;; CHECK-NEXT:  (call $outline$)
+  ;; CHECK-NEXT: )
+  (func $b
+    (if
+      (i32.eqz
+        (global.get $global$1)
+      )
+      (drop
+        (i32.const 10)
+      )
+      (global.set $global$1
+        (i32.const 20)
+      )
+    )
+  )
+)
+
 ;; Tests that local.get instructions are correctly filtered from being outlined.
 (module
   ;; CHECK:      (type $0 (func (param i32)))