yolo

src/ir/return-utils.h

@@ -0,0 +1,46 @@
+/*
+ * 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-traversal.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.
+struct ReturnValueRemover : public PostWalker<ReturnValueRemover> {
+  void visitReturn(Return* curr) {
+    auto* value = curr->value;
+    assert(value);
+    curr->value = nullptr;
+    Builder builder(*module);
+    replaceCurrent(builder.makeSequence(builder.makeDrop(value), curr));
+  }
+
+  void visitFunction(Function* curr) {
+    if (curr->body->type.isConcrete()) {
+      curr->body = Builder(*module).makeDrop(curr->body);
+    }
+  }
+};
+
+} // namespace wasm::ReturnUtils
+
+#endif // wasm_ir_return_h

src/passes/DeadArgumentElimination.cpp

@@ -42,6 +42,7 @@
 #include "ir/find_all.h"
 #include "ir/lubs.h"
 #include "ir/module-utils.h"
+#include "ir/return-utils.h"
 #include "ir/type-updating.h"
 #include "ir/utils.h"
 #include "param-utils.h"
@@ -358,23 +359,7 @@ struct DAE : public Pass {
       }
     }
     // Remove any return values.
-    struct ReturnUpdater : public PostWalker<ReturnUpdater> {
-      Module* module;
-      ReturnUpdater(Function* func, Module* module) : module(module) {
-        walk(func->body);
-      }
-      void visitReturn(Return* curr) {
-        auto* value = curr->value;
-        assert(value);
-        curr->value = nullptr;
-        Builder builder(*module);
-        replaceCurrent(builder.makeSequence(builder.makeDrop(value), curr));
-      }
-    } returnUpdater(func, module);
-    // Remove any value flowing out.
-    if (func->body->type.isConcrete()) {
-      func->body = Builder(*module).makeDrop(func->body);
-    }
+    ReturnUtils::ReturnValueRemover().walkFunctionInModule(func, module);
   }
 
   // Given a function and all the calls to it, see if we can refine the type of

src/passes/Monomorphize.cpp

@@ -92,6 +92,7 @@
 #include "ir/manipulation.h"
 #include "ir/module-utils.h"
 #include "ir/names.h"
+#include "ir/return-utils.h"
 #include "ir/type-updating.h"
 #include "ir/utils.h"
 #include "pass.h"
@@ -480,8 +481,9 @@ struct Monomorphize : public Pass {
         newParams.push_back(operand->type);
       }
     }
-    // TODO: support changes to results.
-    auto newResults = func->getResults();
+    // If we were dropped then we are pulling the drop into the monomorphized
+    // function, which means we return nothing.
+    auto newResults = context.dropped ? Type::none : func->getResults();
     newFunc->type = Signature(Type(newParams), newResults);
 
     // We must update local indexes: the new function has a  potentially
@@ -576,6 +578,10 @@ struct Monomorphize : public Pass {
       newFunc->body = builder.makeBlock(pre);
     }
 
+    if (context.dropped) {
+      ReturnUtils::ReturnValueRemover().walkFunctionInModule(func, &wasm);
+    }
+
     return newFunc;
   }