vacuum the top of functions too

1 files changed, 12 insertions, 0 deletions

diff --git a/src/passes/Vacuum.cpp b/src/passes/Vacuum.cpp
index 6d6f1490a..1524927ec 100644
--- a/src/passes/Vacuum.cpp
+++ b/src/passes/Vacuum.cpp
@@ -212,6 +212,18 @@ struct Vacuum : public WalkerPass<PostWalker<Vacuum, Visitor<Vacuum>>> {
 
     self->pushTask(visitPre, currp);
   }
+
+  void visitFunction(Function* curr) {
+    auto* optimized = optimize(curr->body, curr->result != none);
+    if (optimized) {
+      curr->body = optimized;
+    } else {
+      ExpressionManipulator::nop(curr->body);
+    }
+    if (curr->result == none && !EffectAnalyzer(curr->body).hasSideEffects()) {
+      ExpressionManipulator::nop(curr->body);
+    }
+  }
 };
 
 static RegisterPass<Vacuum> registerPass("vacuum", "removes obviously unneeded code");