Allow --keepfuncs and --splitfuncs to be use alongside a profile data

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;
 }
 

src/tools/wasm-split/wasm-split.cpp

@@ -212,51 +212,75 @@ 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) {
           std::cerr << "warning: cannot split out imported function " << func
                     << "\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;

test/lit/help/wasm-split.test

@@ -30,17 +30,18 @@
 ;; CHECK-NEXT:
 ;; CHECK-NEXT:   --keep-funcs                         [split] Comma-separated list of functions
 ;; CHECK-NEXT:                                        to keep in the primary module. The rest
-;; CHECK-NEXT:                                        will be split out. Cannot be used with
-;; CHECK-NEXT:                                        --profile or --split-funcs. You can also
+;; CHECK-NEXT:                                        will be split out. Can be used alongside
+;; CHECK-NEXT:                                        --profile and --split-funcs. You can also
 ;; CHECK-NEXT:                                        pass a file with one function per line by
 ;; CHECK-NEXT:                                        passing @filename.
 ;; CHECK-NEXT:
 ;; CHECK-NEXT:   --split-funcs                        [split] Comma-separated list of functions
 ;; CHECK-NEXT:                                        to split out to the secondary module. The
-;; CHECK-NEXT:                                        rest will be kept. Cannot be used with
-;; CHECK-NEXT:                                        --profile or --keep-funcs. You can also
-;; CHECK-NEXT:                                        pass a file with one function per line by
-;; CHECK-NEXT:                                        passing @filename.
+;; CHECK-NEXT:                                        rest will be kept. Can be used alongside
+;; CHECK-NEXT:                                        --profile and --keep-funcs. This takes
+;; CHECK-NEXT:                                        precedence over other split options. You
+;; CHECK-NEXT:                                        can also pass a file with one function
+;; CHECK-NEXT:                                        per line by passing @filename.
 ;; CHECK-NEXT:
 ;; CHECK-NEXT:   --primary-output,-o1                 [split] Output file for the primary
 ;; CHECK-NEXT:                                        module.

test/lit/wasm-split/basic.wast

@@ -47,6 +47,12 @@
 ;; RUN: wasm-dis %t.split-bar.1.wasm | filecheck %s --check-prefix KEEP-FOO-PRIMARY
 ;; RUN: wasm-dis %t.split-bar.2.wasm | filecheck %s --check-prefix KEEP-FOO-SECONDARY
 
+;; Check workflow where --split-funcs supersede --keep-funcs
+;; RUN: wasm-split %s --export-prefix='%' -g -o1 %t.split-bar.1.wasm -o2 %t.split-bar.2.wasm --keep-funcs=@%S/both.txt --split-funcs=bar -v 2>&1 \
+;; RUN:     | filecheck %s --check-prefix SPLIT-BAR-SUPERSEDE
+;; RUN: wasm-dis %t.split-bar.1.wasm | filecheck %s --check-prefix KEEP-FOO-PRIMARY
+;; RUN: wasm-dis %t.split-bar.2.wasm | filecheck %s --check-prefix KEEP-FOO-SECONDARY
+
 (module
  (table $table 1 1 funcref)
  (elem (i32.const 0) $foo)
@@ -170,3 +176,7 @@
 ;; KEEP-BOTH-SECONDARY:      (module
 ;; KEEP-BOTH-SECONDARY-NEXT:  (import "primary" "%table" (table $table 1 1 funcref))
 ;; KEEP-BOTH-SECONDARY-NEXT: )
+
+;; SPLIT-BAR-SUPERSEDE: warning: function bar was to be kept in primary module. However it will now be split out into secondary module.
+;; SPLIT-BAR-SUPERSEDE: Keeping functions: foo{{$}}
+;; SPLIT-BAR-SUPERSEDE: Splitting out functions: bar{{$}}

test/lit/wasm-split/invalid-options.wast

@@ -52,18 +52,6 @@
 ;; RUN: not wasm-split %s --merge-profiles -g 2>&1 \
 ;; RUN:   | filecheck %s --check-prefix MERGE-DEBUGINFO
 
-;; --profile cannot be used with --keep-funcs
-;; RUN: not wasm-split %s --profile=foo --keep-funcs=foo 2>&1 \
-;; RUN:   | filecheck %s --check-prefix PROFILE-KEEP
-
-;; --profile cannot be used with --split-funcs
-;; RUN: not wasm-split %s --profile=foo --split-funcs=foo 2>&1 \
-;; RUN:   | filecheck %s --check-prefix PROFILE-SPLIT
-
-;; --keep-funcs cannot be used with --split-funcs
-;; RUN: not wasm-split %s  --keep-funcs=foo --split-funcs=foo 2>&1 \
-;; RUN:   | filecheck %s --check-prefix KEEP-SPLIT
-
 ;; INSTRUMENT-PROFILE: error: Option --profile cannot be used in instrument mode.
 
 ;; INSTRUMENT-OUT1: error: Option --primary-output cannot be used in instrument mode.
@@ -90,10 +78,4 @@
 
 ;; MERGE-DEBUGINFO: error: Option --debuginfo cannot be used in merge-profiles mode.
 
-;; PROFILE-KEEP: error: Cannot use both --profile and --keep-funcs.
-
-;; PROFILE-SPLIT: error: Cannot use both --profile and --split-funcs.
-
-;; KEEP-SPLIT: error: Cannot use both --keep-funcs and --split-funcs.
-
 (module)

test/lit/wasm-split/profile-guided.wast

@@ -25,6 +25,12 @@
 ;; RUN: wasm-split -all %s --profile=%t.none.prof -v -o1 %t.none.1.wasm -o2 %t.none.2.wasm \
 ;; RUN:   | filecheck %s --check-prefix NONE
 
+;; RUN: wasm-split -all %s --profile=%t.bar.prof --keep-funcs=uncalled -v -o1 %t.bar.1.wasm -o2 %t.bar.2.wasm \
+;; RUN:   | filecheck %s --check-prefix PROFILE_KEEP
+
+;; RUN: wasm-split -all %s --profile=%t.both.prof --split-funcs=shared_callee -v -o1 %t.both.1.wasm -o2 %t.both.2.wasm 2>&1 \
+;; RUN:   | filecheck %s --check-prefix PROFILE_SPLIT
+
 ;; =================================
 ;; Do it all again using --in-memory
 ;; =================================
@@ -52,6 +58,12 @@
 ;; RUN: wasm-split -all %s --profile=%t.none.prof -v -o1 %t.none.1.wasm -o2 %t.none.2.wasm \
 ;; RUN:   | filecheck %s --check-prefix NONE
 
+;; RUN: wasm-split -all %s --profile=%t.bar.prof --keep-funcs=uncalled -v -o1 %t.bar.1.wasm -o2 %t.bar.2.wasm \
+;; RUN:   | filecheck %s --check-prefix PROFILE_KEEP
+
+;; RUN: wasm-split -all %s --profile=%t.both.prof --split-funcs=shared_callee -v -o1 %t.both.1.wasm -o2 %t.both.2.wasm 2>&1 \
+;; RUN:   | filecheck %s --check-prefix PROFILE_SPLIT
+
 ;; =======
 ;; Results
 ;; =======
@@ -68,6 +80,13 @@
 ;; NONE: Keeping functions:
 ;; NONE: Splitting out functions: bar, bar_callee, deep_foo_callee, foo, foo_callee, shared_callee, uncalled
 
+;; PROFILE_KEEP: Keeping functions: bar, bar_callee, shared_callee, uncalled
+;; PROFILE_KEEP: Splitting out functions: deep_foo_callee, foo, foo_callee
+
+;; PROFILE_SPLIT: warning: function shared_callee was to be kept in primary module. However it will now be split out into secondary module.
+;; PROFILE_SPLIT: Keeping functions: bar, bar_callee, deep_foo_callee, foo, foo_callee
+;; PROFILE_SPLIT: Splitting out functions: shared_callee, uncalled
+
 (module
   (memory $mem 1 1 shared)
   (export "memory" (memory $mem))