remove final elements in vacuum carefully: we must preserve a return value if there is one

2 files changed, 41 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;
         }
       }
     }