From 60755cf327b6e41bceade0db7bb2c574840bc5d3 Mon Sep 17 00:00:00 2001
From: Marcin Kolny <mkolny@amazon.com>
Date: Fri, 16 Aug 2024 20:00:58 +0100
Subject: [PATCH] Add a pass to remove exports from the WASM bundle

This is helpful when e.g. there's a WASM shared library with all possible
functions exported but the application only needs a subset of those exports.
Removing exports in that case will enable further optimizers to remove more
code from the file.
---
 src/passes/CMakeLists.txt           |  1 +
 src/passes/RemoveExports.cpp        | 55 +++++++++++++++++++++++++++++
 src/passes/pass.cpp                 |  1 +
 src/passes/passes.h                 |  1 +
 test/lit/help/wasm-opt.test         |  2 ++
 test/lit/passes/remove-exports.wast | 33 +++++++++++++++++
 6 files changed, 93 insertions(+)
 create mode 100644 src/passes/RemoveExports.cpp
 create mode 100644 test/lit/passes/remove-exports.wast

diff --git a/src/passes/CMakeLists.txt b/src/passes/CMakeLists.txt
index 703b27e056d..bbff18ada04 100644
--- a/src/passes/CMakeLists.txt
+++ b/src/passes/CMakeLists.txt
@@ -95,6 +95,7 @@ set(passes_SOURCES
   StripTargetFeatures.cpp
   TraceCalls.cpp
   RedundantSetElimination.cpp
+  RemoveExports.cpp
   RemoveImports.cpp
   RemoveMemory.cpp
   RemoveNonJSOps.cpp
diff --git a/src/passes/RemoveExports.cpp b/src/passes/RemoveExports.cpp
new file mode 100644
index 00000000000..f396a2cc26e
--- /dev/null
+++ b/src/passes/RemoveExports.cpp
@@ -0,0 +1,55 @@
+/*
+ * 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.
+ */
+
+//
+// Removes function exports. This is particularly helpful for shared
+// libraries where not all the exported functions are being used for
+// a specific use case.
+//
+
+#include <unordered_map>
+
+#include "pass.h"
+#include "support/string.h"
+
+namespace wasm {
+
+struct RemoveExports : public Pass {
+  void run(Module* module) override {
+    auto arg =
+      getArgument("remove-exports",
+                  "RemoveExports usage: wasm-opt "
+                  "--remove-exports=EXPORT_NAME_1[,EXPORT_NAME_2[,...]]");
+
+    auto argsItems = String::Split(arg, ",");
+    auto exportsToRemove =
+      std::unordered_set<std::string>(argsItems.begin(), argsItems.end());
+
+    auto it = module->exports.begin();
+    while (it != module->exports.end()) {
+      if ((*it)->kind == ExternalKind::Function &&
+          exportsToRemove.count((*it)->name.toString())) {
+        it = module->exports.erase(it);
+      } else {
+        ++it;
+      }
+    }
+  }
+};
+
+Pass* createRemoveExportsPass() { return new RemoveExports(); }
+
+} // namespace wasm
diff --git a/src/passes/pass.cpp b/src/passes/pass.cpp
index 3f84ee60464..80896b1f94a 100644
--- a/src/passes/pass.cpp
+++ b/src/passes/pass.cpp
@@ -390,6 +390,7 @@ void PassRegistry::registerPasses() {
   registerPass("remove-imports",
                "removes imports and replaces them with nops",
                createRemoveImportsPass);
+  registerPass("remove-exports", "remove function exports", createRemoveExportsPass);
   registerPass(
     "remove-memory", "removes memory segments", createRemoveMemoryPass);
   registerPass("remove-unused-brs",
diff --git a/src/passes/passes.h b/src/passes/passes.h
index bd3aabf9f7e..73b08dc825b 100644
--- a/src/passes/passes.h
+++ b/src/passes/passes.h
@@ -125,6 +125,7 @@ Pass* createPrintFeaturesPass();
 Pass* createPrintFunctionMapPass();
 Pass* createPropagateGlobalsGloballyPass();
 Pass* createRemoveNonJSOpsPass();
+Pass* createRemoveExportsPass();
 Pass* createRemoveImportsPass();
 Pass* createRemoveMemoryPass();
 Pass* createRemoveUnusedBrsPass();
diff --git a/test/lit/help/wasm-opt.test b/test/lit/help/wasm-opt.test
index a48b8e303e4..fc9aee6287a 100644
--- a/test/lit/help/wasm-opt.test
+++ b/test/lit/help/wasm-opt.test
@@ -103,6 +103,8 @@
 ;; CHECK-NEXT:   --cfp                                         propagate constant struct field
 ;; CHECK-NEXT:                                                 values
 ;; CHECK-NEXT:
+;; CHECK-NEXT:   --remove-exports                              remove function exports
+;; CHECK-NEXT:
 ;; CHECK-NEXT:   --cfp-reftest                                 propagate constant struct field
 ;; CHECK-NEXT:                                                 values, using ref.test
 ;; CHECK-NEXT:
diff --git a/test/lit/passes/remove-exports.wast b/test/lit/passes/remove-exports.wast
new file mode 100644
index 00000000000..c18664ff015
--- /dev/null
+++ b/test/lit/passes/remove-exports.wast
@@ -0,0 +1,33 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: wasm-opt --remove-exports=foo,bar %s -S -o - | filecheck %s
+
+(module
+
+  ;; CHECK:      (type $0 (func))
+
+  ;; CHECK:      (export "baz" (func $baz))
+
+  ;; CHECK:      (func $foo
+  ;; CHECK-NEXT:  (unreachable)
+  ;; CHECK-NEXT: )
+  (func $foo
+   (unreachable)
+  )
+  ;; CHECK:      (func $bar
+  ;; CHECK-NEXT:  (unreachable)
+  ;; CHECK-NEXT: )
+  (func $bar
+   (unreachable)
+  )
+  ;; CHECK:      (func $baz
+  ;; CHECK-NEXT:  (unreachable)
+  ;; CHECK-NEXT: )
+  (func $baz
+   (unreachable)
+  )
+
+  (export "foo" (func $foo))
+  (export "bar" (func $bar))
+  (export "baz" (func $baz))
+)