1 files changed, 21 insertions, 1 deletions
diff --git a/src/passes/CodeFolding.cpp b/src/passes/CodeFolding.cpp
index 5a18937bb..dc2301f9c 100644
--- a/src/passes/CodeFolding.cpp
+++ b/src/passes/CodeFolding.cpp
@@ -172,7 +172,7 @@ struct CodeFolding : public WalkerPass<ControlFlowWalker<CodeFolding>> {
     }
   }
 
-  void visitReturn(Return* curr) {
+  void handleReturn(Expression* curr) {
     if (!controlFlowStack.empty()) {
       // we can easily optimize if we are at the end of the parent block
       Block* parent = controlFlowStack.back()->dynCast<Block>();
@@ -186,6 +186,26 @@ struct CodeFolding : public WalkerPass<ControlFlowWalker<CodeFolding>> {
     returnTails.push_back(Tail(curr, getCurrentPointer()));
   }
 
+  void visitReturn(Return* curr) { handleReturn(curr); }
+
+  void visitCall(Call* curr) {
+    if (curr->isReturn) {
+      handleReturn(curr);
+    }
+  }
+
+  void visitCallIndirect(CallIndirect* curr) {
+    if (curr->isReturn) {
+      handleReturn(curr);
+    }
+  }
+
+  void visitCallRef(CallRef* curr) {
+    if (curr->isReturn) {
+      handleReturn(curr);
+    }
+  }
+
   void visitBlock(Block* curr) {
     if (curr->list.empty()) {
       return;