[wasm-split] Add an option to skip importing placeholders (#6942)

Wasm-split generally assumes that calls to secondary functions made
before the secondary module has been loaded and instantiated should go
to imported placeholder functions that can be responsible for loading
the secondary module and forwarding the call to the loaded function.
That scheme makes the loading entirely transparent from the
application's point of view, which is not always a good thing. Other
schemes would make it impossible for a secondary function to be called
before the secondary module has been explicitly loaded, in which case
the placeholder functions would never be called. To improve code size
and simplify instantiation under these schemes, add a new
`--no-placeholders` option that skips adding imported placeholder
functions.

Author: tlively

src/ir/module-splitting.cpp

@@ -94,10 +94,9 @@ template<class F> void forEachElement(Module& module, F f) {
     } else if (auto* g = segment->offset->dynCast<GlobalGet>()) {
       base = g->name;
     }
-    ElementUtils::iterElementSegmentFunctionNames(
-      segment, [&](Name& entry, Index i) {
-        f(segment->table, base, offset + i, entry);
-      });
+    for (Index i = 0; i < segment->data.size(); ++i) {
+      f(segment->table, base, offset + i, segment->data[i]);
+    }
   });
 }
 
@@ -209,9 +208,12 @@ TableSlotManager::TableSlotManager(Module& module) : module(module) {
   }
 
   // Initialize funcIndices with the functions already in the table.
-  forEachElement(module, [&](Name table, Name base, Index offset, Name func) {
-    addSlot(func, {table, base, offset});
-  });
+  forEachElement(module,
+                 [&](Name table, Name base, Index offset, Expression* elem) {
+                   if (auto* func = elem->dynCast<RefFunc>()) {
+                     addSlot(func->func, {table, base, offset});
+                   }
+                 });
 }
 
 Table* TableSlotManager::makeTable() {
@@ -693,21 +695,32 @@ void ModuleSplitter::setupTablePatching() {
   // Replace table references to secondary functions with an imported
   // placeholder that encodes the table index in its name:
   // `importNamespace`.`index`.
-  forEachElement(primary, [&](Name, Name, Index index, Name& elem) {
-    if (secondaryFuncs.count(elem)) {
-      placeholderMap[index] = elem;
-      auto* secondaryFunc = secondary.getFunction(elem);
-      replacedElems[index] = secondaryFunc;
-      auto placeholder = std::make_unique<Function>();
-      placeholder->module = config.placeholderNamespace;
-      placeholder->base = std::to_string(index);
-      placeholder->name = Names::getValidFunctionName(
-        primary, std::string("placeholder_") + placeholder->base.toString());
-      placeholder->hasExplicitName = true;
-      placeholder->type = secondaryFunc->type;
-      elem = placeholder->name;
-      primary.addFunction(std::move(placeholder));
+  forEachElement(primary, [&](Name, Name, Index index, Expression*& elem) {
+    auto* ref = elem->dynCast<RefFunc>();
+    if (!ref) {
+      return;
+    }
+    if (!secondaryFuncs.count(ref->func)) {
+      return;
+    }
+    placeholderMap[index] = ref->func;
+    auto* secondaryFunc = secondary.getFunction(ref->func);
+    replacedElems[index] = secondaryFunc;
+    if (!config.usePlaceholders) {
+      // TODO: This can create active element segments with lots of nulls. We
+      // should optimize them like we do data segments with zeros.
+      elem = Builder(primary).makeRefNull(HeapType::nofunc);
+      return;
     }
+    auto placeholder = std::make_unique<Function>();
+    placeholder->module = config.placeholderNamespace;
+    placeholder->base = std::to_string(index);
+    placeholder->name = Names::getValidFunctionName(
+      primary, std::string("placeholder_") + placeholder->base.toString());
+    placeholder->hasExplicitName = true;
+    placeholder->type = secondaryFunc->type;
+    elem = Builder(primary).makeRefFunc(placeholder->name, placeholder->type);
+    primary.addFunction(std::move(placeholder));
   });
 
   if (replacedElems.size() == 0) {

src/ir/module-splitting.h

@@ -52,11 +52,15 @@ struct Config {
   // exists. May or may not include imported functions, which are always kept in
   // the primary module regardless.
   std::set<Name> primaryFuncs;
+  // Whether to import placeholder functions into the primary module that will
+  // be called when a secondary function is called before the secondary module
+  // has been loaded.
+  bool usePlaceholders = true;
   // The namespace from which to import primary functions into the secondary
   // module.
   Name importNamespace = "primary";
   // The namespace from which to import placeholder functions into the primary
-  // module.
+  // module. Ignored if `usePlaceholders` is false.
   Name placeholderNamespace = "placeholder";
   // The prefix to attach to the name of any newly created exports. This can be
   // used to differentiate between "real" exports of the module and exports that

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

@@ -176,6 +176,15 @@ WasmSplitOptions::WasmSplitOptions()
          {Mode::Split},
          Options::Arguments::Zero,
          [&](Options* o, const std::string& argument) { symbolMap = true; })
+    .add(
+      "--no-placeholders",
+      "",
+      "Do not import placeholder functions. Calls to secondary functions will "
+      "fail before the secondary module has been instantiated.",
+      WasmSplitOption,
+      {Mode::Split},
+      Options::Arguments::Zero,
+      [&](Options* o, const std::string& argument) { usePlaceholders = false; })
     .add(
       "--placeholdermap",
       "",

src/tools/wasm-split/split-options.h

@@ -41,6 +41,7 @@ struct WasmSplitOptions : ToolOptions {
   };
   StorageKind storageKind = StorageKind::InGlobals;
 
+  bool usePlaceholders = true;
   bool unescape = false;
   bool verbose = false;
   bool emitBinary = true;

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

@@ -329,6 +329,7 @@ void splitModule(const WasmSplitOptions& options) {
   if (options.exportPrefix.size()) {
     config.newExportPrefix = options.exportPrefix;
   }
+  config.usePlaceholders = options.usePlaceholders;
   config.minimizeNewExportNames = !options.passOptions.debugInfo;
   config.jspi = options.jspi;
   auto splitResults = ModuleSplitting::splitFunctions(wasm, config);

test/lit/help/wasm-split.test

@@ -52,6 +52,11 @@
 ;; CHECK-NEXT:   --symbolmap                          [split] Write a symbol map file for each
 ;; CHECK-NEXT:                                        of the output modules.
 ;; CHECK-NEXT:
+;; CHECK-NEXT:   --no-placeholders                    [split] Do not import placeholder
+;; CHECK-NEXT:                                        functions. Calls to secondary functions
+;; CHECK-NEXT:                                        will fail before the secondary module has
+;; CHECK-NEXT:                                        been instantiated.
+;; CHECK-NEXT:
 ;; CHECK-NEXT:   --placeholdermap                     [split] Write a file mapping placeholder
 ;; CHECK-NEXT:                                        indices to the function names.
 ;; CHECK-NEXT:

test/lit/wasm-split/no-placeholders.wast

@@ -0,0 +1,54 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: wasm-split %s -all --no-placeholders --split-funcs=bar,baz -g -o1 %t.1.wasm -o2 %t.2.wasm
+;; RUN: wasm-dis %t.1.wasm | filecheck %s --check-prefix PRIMARY
+;; RUN: wasm-dis %t.2.wasm | filecheck %s --check-prefix SECONDARY
+
+(module
+ ;; PRIMARY:      (type $0 (func))
+
+ ;; PRIMARY:      (table $0 2 funcref)
+
+ ;; PRIMARY:      (elem $0 (table $0) (i32.const 0) funcref (item (ref.null nofunc)) (item (ref.null nofunc)))
+
+ ;; PRIMARY:      (export "foo" (func $foo))
+
+ ;; PRIMARY:      (export "table" (table $0))
+
+ ;; PRIMARY:      (func $foo
+ ;; PRIMARY-NEXT:  (call_indirect (type $0)
+ ;; PRIMARY-NEXT:   (i32.const 0)
+ ;; PRIMARY-NEXT:  )
+ ;; PRIMARY-NEXT:  (call_indirect (type $0)
+ ;; PRIMARY-NEXT:   (i32.const 1)
+ ;; PRIMARY-NEXT:  )
+ ;; PRIMARY-NEXT: )
+ (func $foo
+  (call $bar)
+  (call $baz)
+ )
+ ;; SECONDARY:      (type $0 (func))
+
+ ;; SECONDARY:      (import "primary" "table" (table $timport$0 2 funcref))
+
+ ;; SECONDARY:      (import "primary" "foo" (func $foo))
+
+ ;; SECONDARY:      (elem $0 (i32.const 0) $bar $baz)
+
+ ;; SECONDARY:      (func $bar
+ ;; SECONDARY-NEXT:  (call $foo)
+ ;; SECONDARY-NEXT:  (call $baz)
+ ;; SECONDARY-NEXT: )
+ (func $bar
+  (call $foo)
+  (call $baz)
+ )
+ ;; SECONDARY:      (func $baz
+ ;; SECONDARY-NEXT:  (call $foo)
+ ;; SECONDARY-NEXT:  (call $bar)
+ ;; SECONDARY-NEXT: )
+ (func $baz
+  (call $foo)
+  (call $bar)
+ )
+)