Update IRBuilder to visit control flow correctly (#6124)

Besides If, no control flow structure consumes values from the stack. Fix a
bug in IRBuilder that was causing it to pop control flow children. Also fix a
follow on bug in outlining where it did not make the If condition available on
the stack when starting to visit an If. This required making push() part of
the public API of IRBuilder.

As a drive-by, also add helpful debug logging to IRBuilder.

Co-authored-by: Ashley Nelson <nashley@google.com>

4 files changed, 136 insertions, 5 deletions

diff --git a/src/passes/Outlining.cpp b/src/passes/Outlining.cpp
index c7644c339..f799d2c7a 100644
--- a/src/passes/Outlining.cpp
+++ b/src/passes/Outlining.cpp
@@ -93,6 +93,11 @@ struct ReconstructStringifyWalker
       ASSERT_OK(existingBuilder.visitBlockStart(curr->block));
       DBG(desc = "Block Start for ");
     } else if (auto curr = reason.getIfStart()) {
+      // IR builder needs the condition of the If pushed onto the builder before
+      // visitIfStart(), which will expect to be able to pop the condition.
+      // This is always okay to do because the correct condition was installed
+      // onto the If when the outer scope was visited.
+      existingBuilder.push(curr->iff->condition);
       ASSERT_OK(existingBuilder.visitIfStart(curr->iff));
       DBG(desc = "If Start for ");
     } else if (reason.getEnd()) {
diff --git a/src/wasm-ir-builder.h b/src/wasm-ir-builder.h
index e02d623f8..0089f547f 100644
--- a/src/wasm-ir-builder.h
+++ b/src/wasm-ir-builder.h
@@ -52,6 +52,10 @@ public:
   // initialized to initialize the child fields and refinalize it.
   [[nodiscard]] Result<> visit(Expression*);
 
+  // Like visit, but pushes the expression onto the stack as-is without popping
+  // any children or refinalization.
+  void push(Expression*);
+
   // Handle the boundaries of control flow structures. Users may choose to use
   // the corresponding `makeXYZ` function below instead of `visitXYZStart`, but
   // either way must call `visitEnd` and friends at the appropriate times.
@@ -187,7 +191,7 @@ public:
 
   // Private functions that must be public for technical reasons.
   [[nodiscard]] Result<> visitExpression(Expression*);
-  [[nodiscard]] Result<> visitBlock(Block*);
+  [[nodiscard]] Result<> visitIf(If*);
   [[nodiscard]] Result<> visitReturn(Return*);
   [[nodiscard]] Result<> visitStructNew(StructNew*);
   [[nodiscard]] Result<> visitArrayNew(ArrayNew*);
@@ -308,7 +312,7 @@ private:
       WASM_UNREACHABLE("unexpected scope kind");
     }
     Name getOriginalLabel() {
-      if (getFunction()) {
+      if (std::get_if<NoScope>(&scope) || getFunction()) {
         return Name{};
       }
       if (auto* block = getBlock()) {
@@ -381,7 +385,6 @@ private:
   [[nodiscard]] Result<Name> getLabelName(Index label);
   [[nodiscard]] Result<Index> addScratchLocal(Type);
   [[nodiscard]] Result<Expression*> pop();
-  void push(Expression*);
 
   struct HoistedVal {
     // The index in the stack of the original value-producing expression.
@@ -401,6 +404,8 @@ private:
 
   [[nodiscard]] Result<Expression*> getBranchValue(Name labelName,
                                                    std::optional<Index> label);
+
+  void dump();
 };
 
 } // namespace wasm
diff --git a/src/wasm/wasm-ir-builder.cpp b/src/wasm/wasm-ir-builder.cpp
index 433145d13..a4a4ec5d0 100644
--- a/src/wasm/wasm-ir-builder.cpp
+++ b/src/wasm/wasm-ir-builder.cpp
@@ -17,9 +17,18 @@
 #include <cassert>
 
 #include "ir/names.h"
+#include "ir/properties.h"
 #include "ir/utils.h"
 #include "wasm-ir-builder.h"
 
+#define IR_BUILDER_DEBUG 0
+
+#if IR_BUILDER_DEBUG
+#define DBG(statement) statement
+#else
+#define DBG(statement)
+#endif
+
 using namespace std::string_literals;
 
 namespace wasm {
@@ -144,6 +153,9 @@ void IRBuilder::push(Expression* expr) {
     scope.unreachable = true;
   }
   scope.exprStack.push_back(expr);
+
+  DBG(std::cerr << "After pushing " << ShallowExpression(expr) << ":\n");
+  DBG(dump());
 }
 
 Result<Expression*> IRBuilder::pop() {
@@ -185,7 +197,55 @@ Result<Expression*> IRBuilder::build() {
   return expr;
 }
 
+void IRBuilder::dump() {
+#if IR_BUILDER_DEBUG
+  std::cerr << "Scope stack";
+  if (func) {
+    std::cerr << " in function $" << func->name;
+  }
+  std::cerr << ":\n";
+
+  for (auto& scope : scopeStack) {
+    std::cerr << "  scope ";
+    if (std::get_if<ScopeCtx::NoScope>(&scope.scope)) {
+      std::cerr << "none";
+    } else if (auto* f = std::get_if<ScopeCtx::FuncScope>(&scope.scope)) {
+      std::cerr << "func " << f->func->name;
+    } else if (std::get_if<ScopeCtx::BlockScope>(&scope.scope)) {
+      std::cerr << "block";
+    } else if (std::get_if<ScopeCtx::IfScope>(&scope.scope)) {
+      std::cerr << "if";
+    } else if (std::get_if<ScopeCtx::ElseScope>(&scope.scope)) {
+      std::cerr << "else";
+    } else if (std::get_if<ScopeCtx::LoopScope>(&scope.scope)) {
+      std::cerr << "loop";
+    } else {
+      WASM_UNREACHABLE("unexpected scope");
+    }
+
+    if (auto name = scope.getOriginalLabel()) {
+      std::cerr << " (original label: " << name << ")";
+    }
+
+    if (scope.label) {
+      std::cerr << " (label: " << scope.label << ")";
+    }
+
+    if (scope.unreachable) {
+      std::cerr << " (unreachable)";
+    }
+
+    std::cerr << ":\n";
+
+    for (auto* expr : scope.exprStack) {
+      std::cerr << "    " << ShallowExpression(expr) << "\n";
+    }
+  }
+#endif // IR_BUILDER_DEBUG
+}
+
 Result<> IRBuilder::visit(Expression* curr) {
+  // Call either `visitExpression` or an expression-specific override.
   auto val = UnifiedExpressionVisitor<IRBuilder, Result<>>::visit(curr);
   CHECK_ERR(val);
   if (auto* block = curr->dynCast<Block>()) {
@@ -202,6 +262,12 @@ Result<> IRBuilder::visit(Expression* curr) {
 // Handle the common case of instructions with a constant number of children
 // uniformly.
 Result<> IRBuilder::visitExpression(Expression* curr) {
+  if (Properties::isControlFlowStructure(curr)) {
+    // Control flow structures (besides `if`, handled separately) do not consume
+    // stack values.
+    return Ok{};
+  }
+
 #define DELEGATE_ID curr->_id
 #define DELEGATE_START(id) [[maybe_unused]] auto* expr = curr->cast<id>();
 #define DELEGATE_FIELD_CHILD(id, field)                                        \
@@ -238,8 +304,12 @@ Result<> IRBuilder::visitExpression(Expression* curr) {
   return Ok{};
 }
 
-Result<> IRBuilder::visitBlock(Block* curr) {
-  // No children; pushing and finalizing will be handled by `visit`.
+Result<> IRBuilder::visitIf(If* curr) {
+  // Only the condition is popped from the stack. The ifTrue and ifFalse are
+  // self-contained so we do not modify them.
+  auto cond = pop();
+  CHECK_ERR(cond);
+  curr->condition = *cond;
   return Ok{};
 }
 
diff --git a/test/lit/passes/outlining.wast b/test/lit/passes/outlining.wast
index ffd5eb081..a080162ca 100644
--- a/test/lit/passes/outlining.wast
+++ b/test/lit/passes/outlining.wast
@@ -243,3 +243,54 @@
     )
   )
 )
+
+
+;; Tests that outlining works correctly with If control flow
+(module
+  ;; CHECK:      (type $0 (func (param i32)))
+
+  ;; CHECK:      (type $1 (func))
+
+  ;; CHECK:      (func $outline$
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i32.const 10)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+
+  ;; CHECK:      (func $a (param $0 i32)
+  ;; CHECK-NEXT:  (if
+  ;; CHECK-NEXT:   (i32.eqz
+  ;; CHECK-NEXT:    (local.get $0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (call $outline$)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $a (param i32)
+    (if
+      (i32.eqz
+        (local.get 0)
+      )
+      (drop
+        (i32.const 10)
+      )
+    )
+  )
+  ;; CHECK:      (func $b (param $0 i32)
+  ;; CHECK-NEXT:  (if
+  ;; CHECK-NEXT:   (i32.eqz
+  ;; CHECK-NEXT:    (local.get $0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (call $outline$)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $b (param i32)
+    (if
+      (i32.eqz
+        (local.get 0)
+      )
+      (drop
+        (i32.const 10)
+      )
+    )
+  )
+)