3 files changed, 139 insertions, 0 deletions
diff --git a/src/ir/CMakeLists.txt b/src/ir/CMakeLists.txt
index 996daa564..45b08702d 100644
--- a/src/ir/CMakeLists.txt
+++ b/src/ir/CMakeLists.txt
@@ -17,6 +17,7 @@ set(ir_SOURCES
   LocalGraph.cpp
   LocalStructuralDominance.cpp
   ReFinalize.cpp
+  return-utils.cpp
   stack-utils.cpp
   table-utils.cpp
   type-updating.cpp
diff --git a/src/ir/return-utils.cpp b/src/ir/return-utils.cpp
new file mode 100644
index 000000000..20b3a194b
--- /dev/null
+++ b/src/ir/return-utils.cpp
@@ -0,0 +1,99 @@
+/*
+ * Copyright 2024 WebAssembly Community Group participants
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "ir/return-utils.h"
+#include "ir/module-utils.h"
+#include "wasm-builder.h"
+#include "wasm-traversal.h"
+#include "wasm.h"
+
+namespace wasm::ReturnUtils {
+
+namespace {
+
+struct ReturnValueRemover : public PostWalker<ReturnValueRemover> {
+  void visitReturn(Return* curr) {
+    auto* value = curr->value;
+    assert(value);
+    curr->value = nullptr;
+    Builder builder(*getModule());
+    replaceCurrent(builder.makeSequence(builder.makeDrop(value), curr));
+  }
+
+  void visitCall(Call* curr) { handleReturnCall(curr); }
+  void visitCallIndirect(CallIndirect* curr) { handleReturnCall(curr); }
+  void visitCallRef(CallRef* curr) { handleReturnCall(curr); }
+
+  template<typename T> void handleReturnCall(T* curr) {
+    if (curr->isReturn) {
+      Fatal() << "Cannot remove return_calls in ReturnValueRemover";
+    }
+  }
+
+  void visitFunction(Function* curr) {
+    if (curr->body->type.isConcrete()) {
+      curr->body = Builder(*getModule()).makeDrop(curr->body);
+    }
+  }
+};
+
+} // anonymous namespace
+
+void removeReturns(Function* func, Module& wasm) {
+  ReturnValueRemover().walkFunctionInModule(func, &wasm);
+}
+
+std::unordered_map<Function*, bool> findReturnCallers(Module& wasm) {
+  ModuleUtils::ParallelFunctionAnalysis<bool> analysis(
+    wasm, [&](Function* func, bool& hasReturnCall) {
+      if (func->imported()) {
+        return;
+      }
+
+      struct Finder : PostWalker<Finder> {
+        bool hasReturnCall = false;
+
+        void visitCall(Call* curr) {
+          if (curr->isReturn) {
+            hasReturnCall = true;
+          }
+        }
+        void visitCallIndirect(CallIndirect* curr) {
+          if (curr->isReturn) {
+            hasReturnCall = true;
+          }
+        }
+        void visitCallRef(CallRef* curr) {
+          if (curr->isReturn) {
+            hasReturnCall = true;
+          }
+        }
+      } finder;
+
+      finder.walk(func->body);
+      hasReturnCall = finder.hasReturnCall;
+    });
+
+  // Convert to an unordered map for fast lookups. TODO: Avoid a copy here.
+  std::unordered_map<Function*, bool> ret;
+  ret.reserve(analysis.map.size());
+  for (auto& [k, v] : analysis.map) {
+    ret[k] = v;
+  }
+  return ret;
+}
+
+} // namespace wasm::ReturnUtils
diff --git a/src/ir/return-utils.h b/src/ir/return-utils.h
new file mode 100644
index 000000000..a5214ba01
--- /dev/null
+++ b/src/ir/return-utils.h
@@ -0,0 +1,39 @@
+/*
+ * Copyright 2024 WebAssembly Community Group participants
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef wasm_ir_return_h
+#define wasm_ir_return_h
+
+#include "wasm.h"
+
+namespace wasm::ReturnUtils {
+
+// Removes values from both explicit returns and implicit ones (values that flow
+// from the body). This is useful after changing a function's type to no longer
+// return anything.
+//
+// This does *not* handle return calls, and will error on them. Removing a
+// return call may change the semantics of the program, so we do not do it
+// automatically here.
+void removeReturns(Function* func, Module& wasm);
+
+// Return a map of every function to whether it does a return call.
+using ReturnCallersMap = std::unordered_map<Function*, bool>;
+ReturnCallersMap findReturnCallers(Module& wasm);
+
+} // namespace wasm::ReturnUtils
+
+#endif // wasm_ir_return_h