Allow --keepfuncs and --splitfuncs to be use alongside a profile data (#6322)

There are times after collecting a profile, we wish to manually include
specific functions into the primary module.
It could be due to non-deterministic profiling or functions for error
scenarios (e.g. _trap).

This PR helps to unlock this workflow by honoring both the
`--keep-funcs` flag as well as the `--profile` flag

2 files changed, 38 insertions, 25 deletions

diff --git a/src/tools/wasm-split/split-options.cpp b/src/tools/wasm-split/split-options.cpp
index b166b575c..9a9351998 100644
--- a/src/tools/wasm-split/split-options.cpp
+++ b/src/tools/wasm-split/split-options.cpp
@@ -129,7 +129,7 @@ WasmSplitOptions::WasmSplitOptions()
     .add("--keep-funcs",
          "",
          "Comma-separated list of functions to keep in the primary module. The "
-         "rest will be split out. Cannot be used with --profile or "
+         "rest will be split out. Can be used alongside --profile and "
          "--split-funcs. You can also pass a file with one function per line "
          "by passing @filename.",
          WasmSplitOption,
@@ -141,8 +141,9 @@ WasmSplitOptions::WasmSplitOptions()
     .add("--split-funcs",
          "",
          "Comma-separated list of functions to split out to the secondary "
-         "module. The rest will be kept. Cannot be used with --profile or "
-         "--keep-funcs. You can also pass a file with one function per line "
+         "module. The rest will be kept. Can be used alongside --profile and "
+         "--keep-funcs. This takes precedence over other split options. "
+         "You can also pass a file with one function per line "
          "by passing @filename.",
          WasmSplitOption,
          {Mode::Split},
@@ -421,18 +422,6 @@ bool WasmSplitOptions::validate() {
     }
   }
 
-  if (mode == Mode::Split) {
-    if (profileFile.size() && keepFuncs.size()) {
-      fail("Cannot use both --profile and --keep-funcs.");
-    }
-    if (profileFile.size() && splitFuncs.size()) {
-      fail("Cannot use both --profile and --split-funcs.");
-    }
-    if (keepFuncs.size() && splitFuncs.size()) {
-      fail("Cannot use both --keep-funcs and --split-funcs.");
-    }
-  }
-
   return valid;
 }
 
diff --git a/src/tools/wasm-split/wasm-split.cpp b/src/tools/wasm-split/wasm-split.cpp
index d7dc19d67..abb1646c1 100644
--- a/src/tools/wasm-split/wasm-split.cpp
+++ b/src/tools/wasm-split/wasm-split.cpp
@@ -212,30 +212,35 @@ void splitModule(const WasmSplitOptions& options) {
   parseInput(wasm, options);
 
   std::set<Name> keepFuncs;
+  std::set<Name> splitFuncs;
 
   if (options.profileFile.size()) {
     // Use the profile to set `keepFuncs`.
     uint64_t hash = hashFile(options.inputFiles[0]);
-    std::set<Name> splitFuncs;
     getFunctionsToKeepAndSplit(
       wasm, hash, options.profileFile, keepFuncs, splitFuncs);
-  } else if (options.keepFuncs.size()) {
+  }
+
+  if (options.keepFuncs.size()) {
     // Use the explicitly provided `keepFuncs`.
     for (auto& func : options.keepFuncs) {
       if (!options.quiet && wasm.getFunctionOrNull(func) == nullptr) {
         std::cerr << "warning: function " << func << " does not exist\n";
+        continue;
       }
+
       keepFuncs.insert(func);
+      splitFuncs.erase(func);
     }
-  } else if (options.splitFuncs.size()) {
+  }
+
+  if (options.splitFuncs.size()) {
     // Use the explicitly provided `splitFuncs`.
-    for (auto& func : wasm.functions) {
-      keepFuncs.insert(func->name);
-    }
     for (auto& func : options.splitFuncs) {
       auto* function = wasm.getFunctionOrNull(func);
       if (!options.quiet && function == nullptr) {
         std::cerr << "warning: function " << func << " does not exist\n";
+        continue;
       }
       if (function && function->imported()) {
         if (!options.quiet) {
@@ -243,20 +248,39 @@ void splitModule(const WasmSplitOptions& options) {
                     << "\n";
         }
       } else {
+        if (!options.quiet && keepFuncs.count(func) > 0) {
+          std::cerr
+            << "warning: function " << func
+            << " was to be kept in primary module. "
+            << "However it will now be split out into secondary module.\n";
+        }
+
+        splitFuncs.insert(func);
         keepFuncs.erase(func);
       }
     }
-  }
 
-  if (options.jspi) {
-    // The load secondary module function must be kept in the main module.
-    keepFuncs.insert(ModuleSplitting::LOAD_SECONDARY_MODULE);
+    if (keepFuncs.empty()) {
+      // could be the case where every function has been split out
+      // or when `splitFuncs` is used standalone, which is the case we'll cover
+      // here
+      for (auto& func : wasm.functions) {
+        if (splitFuncs.count(func->name) == 0) {
+          keepFuncs.insert(func->name);
+        }
+      }
+    }
   }
 
   if (!options.quiet && keepFuncs.size() == 0) {
     std::cerr << "warning: not keeping any functions in the primary module\n";
   }
 
+  if (options.jspi) {
+    // The load secondary module function must be kept in the main module.
+    keepFuncs.insert(ModuleSplitting::LOAD_SECONDARY_MODULE);
+  }
+
   // If warnings are enabled, check that any functions are being split out.
   if (!options.quiet) {
     std::set<Name> splitFuncs;