4 files changed, 72 insertions, 31 deletions

diff --git a/src/passes/RemoveUnusedBrs.cpp b/src/passes/RemoveUnusedBrs.cpp
index 6e9c9c5df..5d236211b 100644
--- a/src/passes/RemoveUnusedBrs.cpp
+++ b/src/passes/RemoveUnusedBrs.cpp
@@ -30,6 +30,34 @@
 
 namespace wasm {
 
+// Grab a slice out of a block, replacing it with nops, and returning
+// either another block with the contents (if more than 1) or a single
+// expression.
+// This does not finalize the input block; it leaves that for the caller.
+static Expression*
+stealSlice(Builder& builder, Block* input, Index from, Index to) {
+  Expression* ret;
+  if (to == from + 1) {
+    // just one
+    ret = input->list[from];
+  } else {
+    auto* block = builder.makeBlock();
+    for (Index i = from; i < to; i++) {
+      block->list.push_back(input->list[i]);
+    }
+    block->finalize();
+    ret = block;
+  }
+  if (to == input->list.size()) {
+    input->list.resize(from);
+  } else {
+    for (Index i = from; i < to; i++) {
+      input->list[i] = builder.makeNop();
+    }
+  }
+  return ret;
+}
+
 // to turn an if into a br-if, we must be able to reorder the
 // condition and possible value, and the possible value must
 // not have side effects (as they would run unconditionally)
@@ -431,7 +459,7 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
           // we need the ifTrue to break, so it cannot reach the code we want to
           // move
           if (iff->ifTrue->type == Type::unreachable) {
-            iff->ifFalse = builder.stealSlice(block, i + 1, list.size());
+            iff->ifFalse = stealSlice(builder, block, i + 1, list.size());
             iff->finalize();
             block->finalize();
             return true;
@@ -476,13 +504,13 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
 
           if (iff->ifTrue->type == Type::unreachable) {
             iff->ifFalse = blockifyMerge(
-              iff->ifFalse, builder.stealSlice(block, i + 1, list.size()));
+              iff->ifFalse, stealSlice(builder, block, i + 1, list.size()));
             iff->finalize();
             block->finalize();
             return true;
           } else if (iff->ifFalse->type == Type::unreachable) {
             iff->ifTrue = blockifyMerge(
-              iff->ifTrue, builder.stealSlice(block, i + 1, list.size()));
+              iff->ifTrue, stealSlice(builder, block, i + 1, list.size()));
             iff->finalize();
             block->finalize();
             return true;
@@ -517,7 +545,8 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
               list[i] =
                 builder.makeIf(brIf->condition,
                                builder.makeBreak(brIf->name),
-                               builder.stealSlice(block, i + 1, list.size()));
+                               stealSlice(builder, block, i + 1, list.size()));
+              block->finalize();
               return true;
             }
           }
diff --git a/src/wasm-builder.h b/src/wasm-builder.h
index 3ca3d9daf..134d42a78 100644
--- a/src/wasm-builder.h
+++ b/src/wasm-builder.h
@@ -1022,33 +1022,6 @@ public:
     return block;
   }
 
-  // Grab a slice out of a block, replacing it with nops, and returning
-  // either another block with the contents (if more than 1) or a single
-  // expression
-  Expression* stealSlice(Block* input, Index from, Index to) {
-    Expression* ret;
-    if (to == from + 1) {
-      // just one
-      ret = input->list[from];
-    } else {
-      auto* block = wasm.allocator.alloc<Block>();
-      for (Index i = from; i < to; i++) {
-        block->list.push_back(input->list[i]);
-      }
-      block->finalize();
-      ret = block;
-    }
-    if (to == input->list.size()) {
-      input->list.resize(from);
-    } else {
-      for (Index i = from; i < to; i++) {
-        input->list[i] = wasm.allocator.alloc<Nop>();
-      }
-    }
-    input->finalize();
-    return ret;
-  }
-
   // Drop an expression if it has a concrete type
   Expression* dropIfConcretelyTyped(Expression* curr) {
     if (!curr->type.isConcrete()) {
diff --git a/test/passes/remove-unused-brs_all-features.txt b/test/passes/remove-unused-brs_all-features.txt
index 4bf1cfb33..3365ebb75 100644
--- a/test/passes/remove-unused-brs_all-features.txt
+++ b/test/passes/remove-unused-brs_all-features.txt
@@ -1,6 +1,7 @@
 (module
  (type $struct (struct (field (ref null $vector))))
  (type $vector (array (mut i32)))
+ (type $none_=>_f64 (func (result f64)))
  (type $none_=>_ref?|$struct| (func (result (ref null $struct))))
  (func $foo (result (ref null $struct))
   (if (result (ref null $struct))
@@ -15,4 +16,21 @@
    (ref.null $struct)
   )
  )
+ (func $test-prefinalize (result f64)
+  (loop $loop (result f64)
+   (if (result f64)
+    (i32.const 1)
+    (f64.const 0)
+    (block $block (result f64)
+     (nop)
+     (br_if $loop
+      (i32.eqz
+       (i32.const 0)
+      )
+     )
+     (unreachable)
+    )
+   )
+  )
+ )
 )
diff --git a/test/passes/remove-unused-brs_all-features.wast b/test/passes/remove-unused-brs_all-features.wast
index 4d1a337c6..3703bb9f2 100644
--- a/test/passes/remove-unused-brs_all-features.wast
+++ b/test/passes/remove-unused-brs_all-features.wast
@@ -16,4 +16,25 @@
    (ref.null $struct)
   )
  )
+ (func $test-prefinalize (result f64)
+  (loop $loop (result f64)
+   (block $block (result f64)
+    (drop
+     (br_if $block
+      (f64.const 0)
+      (i32.const 1)
+     )
+    )
+    (if
+     (i32.const 0)
+     (unreachable)
+    )
+    ;; this will be moved from $block into the if right before it. we must be
+    ;; careful to properly finalize() things, as if we finalize the block too
+    ;; early - before the if - then the block ends in a none type, which is
+    ;; invalid.
+    (br $loop)
+   )
+  )
+ )
 )