[wasm-split] Simplify handling of --keep-funcs and --split-funcs (#6948)

Maintain the invariant that every defined functions belongs to either the set of kept functions or the set of split functions. Functions are kept by default except when --keep-funcs is specified without --split-funcs on the command line. This is mostly NFC except that it changes the default behavior when no arguments are specified on the command line to keep all functions. This will simplify a follow-on PR that switches from passing the kept functions to the module splitting utility to passing the split functions.
WebAssembly · Sep 17, 2024 · f9b64c8 · f9b64c8
1 parent 0da6d3e
commit f9b64c8
Show file tree

Hide file tree

Showing 6 changed files with 74 additions and 75 deletions.
diff --git a/src/tools/wasm-split/split-options.cpp b/src/tools/wasm-split/split-options.cpp
@@ -147,6 +147,7 @@ WasmSplitOptions::WasmSplitOptions()
          Options::Arguments::One,
          [&](Options* o, const std::string& argument) {
            keepFuncs = parseNameList(argument);
+           hasKeepFuncs = true;
          })
     .add("--split-funcs",
          "",
@@ -160,6 +161,7 @@ WasmSplitOptions::WasmSplitOptions()
          Options::Arguments::One,
          [&](Options* o, const std::string& argument) {
            splitFuncs = parseNameList(argument);
+           hasSplitFuncs = true;
          })
     .add(
       "--manifest",

diff --git a/src/tools/wasm-split/split-options.h b/src/tools/wasm-split/split-options.h
@@ -58,6 +58,8 @@ struct WasmSplitOptions : ToolOptions {
 
   std::set<Name> keepFuncs;
   std::set<Name> splitFuncs;
+  bool hasKeepFuncs = false;
+  bool hasSplitFuncs = false;
 
   std::vector<std::string> inputFiles;
   std::string output;

diff --git a/src/tools/wasm-split/wasm-split.cpp b/src/tools/wasm-split/wasm-split.cpp
@@ -216,6 +216,7 @@ void splitModule(const WasmSplitOptions& options) {
   Module wasm;
   parseInput(wasm, options);
 
+  // All defined functions will be in one set or the other.
   std::set<Name> keepFuncs;
   std::set<Name> splitFuncs;
 
@@ -224,57 +225,53 @@ void splitModule(const WasmSplitOptions& options) {
     uint64_t hash = hashFile(options.inputFiles[0]);
     getFunctionsToKeepAndSplit(
       wasm, hash, options.profileFile, keepFuncs, splitFuncs);
+  } else {
+    // Normally the default is to keep each function, but if --keep-funcs is the
+    // only thing specified, then all other functions will be split.
+    bool defaultSplit = options.hasKeepFuncs && !options.hasSplitFuncs;
+    if (defaultSplit) {
+      ModuleUtils::iterDefinedFunctions(
+        wasm, [&](Function* func) { splitFuncs.insert(func->name); });
+    } else {
+      ModuleUtils::iterDefinedFunctions(
+        wasm, [&](Function* func) { keepFuncs.insert(func->name); });
+    }
   }
 
-  if (options.keepFuncs.size()) {
-    // Use the explicitly provided `keepFuncs`.
-    for (auto& func : options.keepFuncs) {
-      if (!options.quiet && wasm.getFunctionOrNull(func) == nullptr) {
+  // Use the explicitly provided `keepFuncs`.
+  for (auto& func : options.keepFuncs) {
+    if (!wasm.getFunctionOrNull(func)) {
+      if (!options.quiet) {
         std::cerr << "warning: function " << func << " does not exist\n";
-        continue;
       }
-
-      keepFuncs.insert(func);
-      splitFuncs.erase(func);
+      continue;
     }
+    keepFuncs.insert(func);
+    splitFuncs.erase(func);
   }
 
-  if (options.splitFuncs.size()) {
-    // Use the explicitly provided `splitFuncs`.
-    for (auto& func : options.splitFuncs) {
-      auto* function = wasm.getFunctionOrNull(func);
-      if (!options.quiet && function == nullptr) {
+  // Use the explicitly provided `splitFuncs`.
+  for (auto& func : options.splitFuncs) {
+    auto* function = wasm.getFunctionOrNull(func);
+    if (!function) {
+      if (!options.quiet) {
         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);
       }
+      continue;
     }
-
-    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 (function->imported()) {
+      if (!options.quiet) {
+        std::cerr << "warning: cannot split out imported function " << func
+                  << "\n";
       }
+      continue;
     }
+    if (!options.quiet && options.keepFuncs.count(func)) {
+      std::cerr << "warning: function " << func
+                << " was to be both kept and split. It will be split.\n";
+    }
+    splitFuncs.insert(func);
+    keepFuncs.erase(func);
   }
 
   if (!options.quiet && keepFuncs.size() == 0) {
@@ -284,43 +281,42 @@ void splitModule(const WasmSplitOptions& options) {
   if (options.jspi) {
     // The load secondary module function must be kept in the main module.
     keepFuncs.insert(ModuleSplitting::LOAD_SECONDARY_MODULE);
+    splitFuncs.erase(ModuleSplitting::LOAD_SECONDARY_MODULE);
   }
 
   // If warnings are enabled, check that any functions are being split out.
-  if (!options.quiet) {
-    std::set<Name> splitFuncs;
-    ModuleUtils::iterDefinedFunctions(wasm, [&](Function* func) {
-      if (keepFuncs.count(func->name) == 0) {
-        splitFuncs.insert(func->name);
-      }
-    });
-
-    if (splitFuncs.size() == 0) {
-      std::cerr
-        << "warning: not splitting any functions out to the secondary module\n";
-    }
+  if (!options.quiet && splitFuncs.size() == 0) {
+    std::cerr
+      << "warning: not splitting any functions out to the secondary module\n";
+  }
 
-    // Dump the kept and split functions if we are verbose
-    if (options.verbose) {
-      auto printCommaSeparated = [&](auto funcs) {
-        for (auto it = funcs.begin(); it != funcs.end(); ++it) {
-          if (it != funcs.begin()) {
-            std::cout << ", ";
-          }
-          std::cout << *it;
+  // Dump the kept and split functions if we are verbose.
+  if (options.verbose) {
+    auto printCommaSeparated = [&](auto funcs) {
+      for (auto it = funcs.begin(); it != funcs.end(); ++it) {
+        if (it != funcs.begin()) {
+          std::cout << ", ";
         }
-      };
+        std::cout << *it;
+      }
+    };
 
-      std::cout << "Keeping functions: ";
-      printCommaSeparated(keepFuncs);
-      std::cout << "\n";
+    std::cout << "Keeping functions: ";
+    printCommaSeparated(keepFuncs);
+    std::cout << "\n";
 
-      std::cout << "Splitting out functions: ";
-      printCommaSeparated(splitFuncs);
-      std::cout << "\n";
-    }
+    std::cout << "Splitting out functions: ";
+    printCommaSeparated(splitFuncs);
+    std::cout << "\n";
   }
 
+#ifndef NDEBUG
+  // Check that all defined functions are in one set or the other.
+  ModuleUtils::iterDefinedFunctions(wasm, [&](Function* func) {
+    assert(keepFuncs.count(func->name) || splitFuncs.count(func->name));
+  });
+#endif // NDEBUG
+
   // Actually perform the splitting
   ModuleSplitting::Config config;
   config.primaryFuncs = std::move(keepFuncs);

diff --git a/test/lit/wasm-split/basic.wast b/test/lit/wasm-split/basic.wast
@@ -1,11 +1,6 @@
 ;; RUN: not wasm-split %s --export-prefix='%' -g -o1 %t.none.1.wasm -o2 %t.none.2.wasm [email protected] -v 2>&1 \
 ;; RUN:     | filecheck %s --check-prefix FAILED-TO-OPEN
 
-;; RUN: wasm-split %s --export-prefix='%' -g -o1 %t.none.1.wasm -o2 %t.none.2.wasm -v 2>&1 \
-;; RUN:     | filecheck %s --check-prefix KEEP-NONE
-;; RUN: wasm-dis %t.none.1.wasm | filecheck %s --check-prefix KEEP-NONE-PRIMARY
-;; RUN: wasm-dis %t.none.2.wasm | filecheck %s --check-prefix KEEP-NONE-SECONDARY
-
 ;; RUN: wasm-split %s --export-prefix='%' -g -o1 %t.none.1.wasm -o2 %t.none.2.wasm --keep-funcs=@%S/none.txt -v 2>&1 \
 ;; RUN:     | filecheck %s --check-prefix KEEP-NONE
 ;; RUN: wasm-dis %t.none.1.wasm | filecheck %s --check-prefix KEEP-NONE-PRIMARY
@@ -31,6 +26,11 @@
 ;; RUN: wasm-dis %t.bar.1.wasm | filecheck %s --check-prefix KEEP-BAR-PRIMARY
 ;; RUN: wasm-dis %t.bar.2.wasm | filecheck %s --check-prefix KEEP-BAR-SECONDARY
 
+;; RUN: wasm-split %s --export-prefix='%' -g -o1 %t.none.1.wasm -o2 %t.none.2.wasm -v 2>&1 \
+;; RUN:     | filecheck %s --check-prefix KEEP-BOTH
+;; RUN: wasm-dis %t.none.1.wasm | filecheck %s --check-prefix KEEP-BOTH-PRIMARY
+;; RUN: wasm-dis %t.none.2.wasm | filecheck %s --check-prefix KEEP-BOTH-SECONDARY
+
 ;; RUN: wasm-split %s --export-prefix='%' -g -o1 %t.both.1.wasm -o2 %t.both.2.wasm --keep-funcs=foo,bar -v 2>&1 \
 ;; RUN:     | filecheck %s --check-prefix KEEP-BOTH
 ;; RUN: wasm-dis %t.both.1.wasm | filecheck %s --check-prefix KEEP-BOTH-PRIMARY
@@ -176,6 +176,6 @@
 ;; KEEP-BOTH-SECONDARY:      (module
 ;; 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: warning: function bar was to be both kept and split. It will be split.
 ;; SPLIT-BAR-SUPERSEDE: Keeping functions: foo{{$}}
 ;; SPLIT-BAR-SUPERSEDE: Splitting out functions: bar{{$}}
diff --git a/test/lit/wasm-split/jspi-secondary-export.wast b/test/lit/wasm-split/jspi-secondary-export.wast
@@ -1,5 +1,5 @@
 ;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
-;; RUN: wasm-split %s --export-prefix='%' -g -o1 %t.1.wasm -o2 %t.2.wasm --jspi
+;; RUN: wasm-split %s --export-prefix='%' -g -o1 %t.1.wasm -o2 %t.2.wasm --jspi --split-funcs=foo,bar
 ;; RUN: wasm-dis %t.1.wasm | filecheck %s --check-prefix PRIMARY
 ;; RUN: wasm-dis %t.2.wasm | filecheck %s --check-prefix SECONDARY
 

diff --git a/test/lit/wasm-split/profile-guided.wast b/test/lit/wasm-split/profile-guided.wast
@@ -83,7 +83,6 @@
 ;; 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