1 files changed, 155 insertions, 0 deletions
diff --git a/src/passes/EncloseWorld.cpp b/src/passes/EncloseWorld.cpp
new file mode 100644
index 000000000..5c6b70546
--- /dev/null
+++ b/src/passes/EncloseWorld.cpp
@@ -0,0 +1,155 @@
+/*
+ * 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.
+ */
+
+//
+// "Closes" the world, in the sense of making it more compatible with the
+// --closed-world flag, in a potentially destructive manner. This is mainly
+// useful for fuzzing (in that a random module is usually very incomptable with
+// closed world, with most types being public and hence unoptimizable, but
+// running this pass makes as many as we can fully private).
+//
+// The fixup we do is to find references sent out/received in, and to
+// externalize / internalize them. For example, this export:
+//
+//  (func $refs (export "refs") (param $x (ref $X)) (result (ref $Y))
+//
+// would have the following function exported in its place:
+//
+//  (func $refs-closed (export "refs") (param $x externref) (result externref)
+//    (extern.convert_any
+//      (call $refs
+//        (ref.cast (ref $X)
+//          (any.convert_extern
+//            (local.get $x))))))
+//
+
+#include "ir/names.h"
+#include "pass.h"
+#include "wasm-builder.h"
+#include "wasm.h"
+
+namespace wasm {
+
+namespace {
+
+struct EncloseWorld : public Pass {
+  void run(Module* module) override {
+    // Handle exports.
+    // TODO: Non-function exports.
+    std::vector<std::unique_ptr<Export>> newExports;
+    for (auto& ex : module->exports) {
+      if (ex->kind == ExternalKind::Function) {
+        auto* func = module->getFunction(ex->value);
+        // If this opens up types, replace it with an enclosed stub.
+        if (opensTypes(func)) {
+          auto stubName = makeStubStubForExport(func, module);
+          ex->value = stubName;
+        }
+      }
+    }
+    for (auto& ex : newExports) {
+      module->addExport(std::move(ex));
+    }
+
+    // TODO: Handle imports.
+  }
+
+private:
+  // Whether a type is an "open" ref, that is, a type that closed-world would
+  // consider to keep things public and prevent some amount of closed-world
+  // optimizations.
+  bool isOpenRef(Type t) {
+    // Only externref keeps things closed, and we must ignore things that
+    // cannot be converted to/from it (like funcrefs), so we can just check for
+    // the top type being any.
+    return t.isRef() && t.getHeapType().getTop() == HeapType::any;
+  }
+
+  // Whether a function causes types to be open.
+  bool opensTypes(Function* func) {
+    for (const auto& param : func->getParams()) {
+      if (isOpenRef(param)) {
+        return true;
+      }
+    }
+    // TODO: Handle tuple results.
+    return isOpenRef(func->getResults());
+  }
+
+  // Make an enclosed stub function for an exported function, and return its
+  // name.
+  Name makeStubStubForExport(Function* func, Module* module) {
+    // Pick a valid name for the stub we are about to create.
+    auto stubName = Names::getValidFunctionName(
+      *module, std::string("stub$") + func->name.toString());
+
+    // Create the stub.
+    Builder builder(*module);
+
+    // The stub's body is just a call to the original function, but with some
+    // conversions to/from externref.
+    std::vector<Expression*> params;
+
+    auto externref = Type(HeapType::ext, Nullable);
+
+    // Handle params.
+    std::vector<Type> stubParams;
+    for (const auto& param : func->getParams()) {
+      if (!isOpenRef(param)) {
+        // A normal parameter. Just pass it to the original function.
+        auto* get = builder.makeLocalGet(stubParams.size(), param);
+        params.push_back(get);
+        stubParams.push_back(param);
+      } else {
+        // A type we must fix up: receive as an externref and then internalize
+        // and cast before sending to the original function.
+        auto* get = builder.makeLocalGet(stubParams.size(), externref);
+        auto* interned = builder.makeRefAs(AnyConvertExtern, get);
+        // This cast may be trivial, but we leave it to the optimizer to remove.
+        auto* cast = builder.makeRefCast(interned, param);
+        params.push_back(cast);
+        stubParams.push_back(externref);
+      }
+    }
+
+    auto* call = builder.makeCall(func->name, params, func->getResults());
+
+    // Generate the stub's type.
+    auto oldResults = func->getResults();
+    Type resultsType = isOpenRef(oldResults) ? externref : oldResults;
+    auto type = Signature(Type(stubParams), resultsType);
+
+    // Handle the results and make the body.
+    Expression* body;
+    if (!isOpenRef(oldResults)) {
+      // Just use the call.
+      body = call;
+    } else {
+      // Fix up the call's result.
+      body = builder.makeRefAs(ExternConvertAny, call);
+    }
+
+    module->addFunction(builder.makeFunction(stubName, type, {}, body));
+
+    return stubName;
+  }
+};
+
+} // anonymous namespace
+
+Pass* createEncloseWorldPass() { return new EncloseWorld(); }
+
+} // namespace wasm