4 files changed, 116 insertions, 14 deletions

diff --git a/src/ast_utils.h b/src/ast_utils.h
index e785a8daa..f883d5f66 100644
--- a/src/ast_utils.h
+++ b/src/ast_utils.h
@@ -28,18 +28,26 @@ namespace wasm {
 struct BreakSeeker : public PostWalker<BreakSeeker, Visitor<BreakSeeker>> {
   Name target; // look for this one XXX looking by name may fall prey to duplicate names
   Index found;
+  WasmType valueType;
 
-  BreakSeeker(Name target) : target(target), found(false) {}
+  BreakSeeker(Name target) : target(target), found(0) {}
+
+  void noteFound(Expression* value) {
+    found++;
+    if (found == 1) valueType = unreachable;
+    if (!value) valueType = none;
+    else if (value->type != unreachable) valueType = value->type;
+  }
 
   void visitBreak(Break *curr) {
-    if (curr->name == target) found++;
+    if (curr->name == target) noteFound(curr->value);
   }
 
   void visitSwitch(Switch *curr) {
     for (auto name : curr->targets) {
-      if (name == target) found++;
+      if (name == target) noteFound(curr->value);
     }
-    if (curr->default_ == target) found++;
+    if (curr->default_ == target) noteFound(curr->value);
   }
 
   static bool has(Expression* tree, Name target) {
diff --git a/src/passes/Vacuum.cpp b/src/passes/Vacuum.cpp
index 592f5a063..0fdd12dec 100644
--- a/src/passes/Vacuum.cpp
+++ b/src/passes/Vacuum.cpp
@@ -157,7 +157,12 @@ struct Vacuum : public WalkerPass<ExpressionStackWalker<Vacuum, Visitor<Vacuum>>
         Break* br = list[z - skip]->dynCast<Break>();
         Switch* sw = list[z - skip]->dynCast<Switch>();
         if ((br && !br->condition) || sw) {
+          auto* last = list.back();
           list.resize(z - skip + 1);
+          // if we removed the last one, and it was a return value, it must be returned
+          if (list.back() != last && isConcreteWasmType(last->type)) {
+            list.push_back(last);
+          }
           needResize = false;
           break;
         }
@@ -222,18 +227,32 @@ struct Vacuum : public WalkerPass<ExpressionStackWalker<Vacuum, Visitor<Vacuum>>
       auto* last = block->list.back();
       if (isConcreteWasmType(last->type)) {
         assert(block->type == last->type);
-        block->list.back() = last = optimize(last, false);
+        last = optimize(last, false);
         if (!last) {
-          block->list.pop_back();
-          // we don't need the drop anymore, let's see what we have left in the block
-          if (block->list.size() > 1) {
-            replaceCurrent(block);
-          } else if (block->list.size() == 1) {
-            replaceCurrent(block->list[0]);
-          } else {
-            ExpressionManipulator::nop(curr);
+          // we may be able to remove this, if there are no brs
+          bool canPop = true;
+          if (block->name.is()) {
+            BreakSeeker breakSeeker(block->name);
+            Expression* temp = block;
+            breakSeeker.walk(temp);
+            if (breakSeeker.found && breakSeeker.valueType != none) {
+              canPop = false;
+            }
+          }
+          if (canPop) {
+            block->list.back() = last;
+            block->list.pop_back();
+            block->type = none;
+            // we don't need the drop anymore, let's see what we have left in the block
+            if (block->list.size() > 1) {
+              replaceCurrent(block);
+            } else if (block->list.size() == 1) {
+              replaceCurrent(block->list[0]);
+            } else {
+              ExpressionManipulator::nop(curr);
+            }
+            return;
           }
-          return;
         }
       }
     }
diff --git a/test/passes/vacuum.txt b/test/passes/vacuum.txt
index 2d274a6d7..04a04efc0 100644
--- a/test/passes/vacuum.txt
+++ b/test/passes/vacuum.txt
@@ -181,4 +181,30 @@
   (func $drop-get-global (type $0)
     (call $drop-get-global)
   )
+  (func $relooperJumpThreading1 (type $0)
+    (local $$vararg_ptr5 i32)
+    (local $$11 i32)
+    (loop $while-in$1
+      (drop
+        (block $jumpthreading$outer$8
+          (block $jumpthreading$inner$8
+            (br $jumpthreading$outer$8
+              (i32.const 0)
+            )
+          )
+          (i32.store
+            (get_local $$vararg_ptr5)
+            (get_local $$11)
+          )
+          (i32.const 0)
+        )
+      )
+    )
+  )
+  (func $relooperJumpThreading2 (type $0)
+    (nop)
+  )
+  (func $relooperJumpThreading3 (type $0)
+    (nop)
+  )
 )
diff --git a/test/passes/vacuum.wast b/test/passes/vacuum.wast
index 5d443380f..84f3eeb6f 100644
--- a/test/passes/vacuum.wast
+++ b/test/passes/vacuum.wast
@@ -363,4 +363,53 @@
       )
     )
   )
+  (func $relooperJumpThreading1
+    (local $$vararg_ptr5 i32)
+    (local $$11 i32)
+    (loop $while-in$1
+      (drop
+        (block $jumpthreading$outer$8
+          (block $jumpthreading$inner$8
+            (br $jumpthreading$outer$8 ;; the rest is dead in the outer block, but be careful to leave the return value!
+              (i32.const 0)
+            )
+          )
+          (i32.store
+            (get_local $$vararg_ptr5)
+            (get_local $$11)
+          )
+          (i32.const 0)
+        )
+      )
+    )
+  )
+  (func $relooperJumpThreading2
+    (loop $while-in$1
+      (drop
+        (block $jumpthreading$outer$8
+          (block $jumpthreading$inner$8
+            (br $jumpthreading$outer$8
+              (i32.const 0)
+            )
+          )
+          (i32.const 0)
+        )
+      )
+    )
+  )
+  (func $relooperJumpThreading3
+    (loop $while-in$1
+      (drop
+        (block $jumpthreading$outer$8
+          (br $jumpthreading$outer$8 ;; code after this is dead, can kill it, but preserve the return value at the end!
+            (i32.const 0)
+          )
+          (drop (i32.const 3))
+          (drop (i32.const 2))
+          (drop (i32.const 1))
+          (i32.const 0)
+        )
+      )
+    )
+  )
 )