[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.

3 files changed, 67 insertions, 67 deletions

diff --git a/src/tools/wasm-split/split-options.cpp b/src/tools/wasm-split/split-options.cpp
index e77957f1f..cada7930b 100644
--- 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
index 105c90c80..8f23928c3 100644
--- 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
index 2b66d1164..1cfe5bad9 100644
--- 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);