[wasm-split] Add an --initial-table option (#3454)

Adds an option to wasm-split to allow the user to specify the initial table size
for both instrumenting and splitting use cases. In the short term this will be
used in Emscripten SPLIT_MODULE + dynamic linking workflow to ensure that the
expected table size baked into the JS works for both the instrumented and the
primary split modules. In the longer term this may be replaced with a more
elegant mechanism for making the JS works in both cases.

2 files changed, 51 insertions, 0 deletions

diff --git a/src/tools/wasm-split.cpp b/src/tools/wasm-split.cpp
index 32a726791..5b44e5f0e 100644
--- a/src/tools/wasm-split.cpp
+++ b/src/tools/wasm-split.cpp
@@ -66,6 +66,11 @@ struct WasmSplitOptions : ToolOptions {
   std::string placeholderNamespace;
   std::string exportPrefix;
 
+  // A hack to ensure the split and instrumented modules have the same table
+  // size when using Emscripten's SPLIT_MODULE mode with dynamic linking. TODO:
+  // Figure out a more elegant solution for that use case and remove this.
+  int initialTableSize = -1;
+
   WasmSplitOptions();
   bool validate();
   void parse(int argc, const char* argv[]);
@@ -179,6 +184,16 @@ WasmSplitOptions::WasmSplitOptions()
          [&](Options* o, const std::string& arguments) {
            passOptions.debugInfo = true;
          })
+    .add("--initial-table",
+         "",
+         "A hack to ensure the split and instrumented modules have the same "
+         "table size when using Emscripten's SPLIT_MODULE mode with dynamic "
+         "linking. TODO: Figure out a more elegant solution for that use "
+         "case and remove this.",
+         Options::Arguments::One,
+         [&](Options* o, const std::string& argument) {
+           initialTableSize = std::stoi(argument);
+         })
     .add_positional(
       "INFILE",
       Options::Arguments::One,
@@ -453,6 +468,24 @@ uint64_t hashFile(const std::string& filename) {
   return uint64_t(digest);
 }
 
+void adjustTableSize(Module& wasm, int initialSize) {
+  if (initialSize < 0) {
+    return;
+  }
+  if (!wasm.table.exists) {
+    Fatal() << "--initial-table used but there is no table";
+  }
+  if ((uint64_t)initialSize < wasm.table.initial) {
+    Fatal() << "Specified initial table size too small, should be at least "
+            << wasm.table.initial;
+  }
+  if ((uint64_t)initialSize > wasm.table.max) {
+    Fatal() << "Specified initial table size larger than max table size "
+            << wasm.table.max;
+  }
+  wasm.table.initial = initialSize;
+}
+
 void instrumentModule(Module& wasm, const WasmSplitOptions& options) {
   // Check that the profile export name is not already taken
   if (wasm.getExportOrNull(options.profileExport) != nullptr) {
@@ -463,6 +496,8 @@ void instrumentModule(Module& wasm, const WasmSplitOptions& options) {
   PassRunner runner(&wasm, options.passOptions);
   Instrumenter(options.profileExport, moduleHash).run(&runner, &wasm);
 
+  adjustTableSize(wasm, options.initialTableSize);
+
   // Write the output modules
   ModuleWriter writer;
   writer.setBinary(options.emitBinary);
@@ -600,6 +635,9 @@ void splitModule(Module& wasm, const WasmSplitOptions& options) {
   std::unique_ptr<Module> secondary =
     ModuleSplitting::splitFunctions(wasm, config);
 
+  adjustTableSize(wasm, options.initialTableSize);
+  adjustTableSize(*secondary, options.initialTableSize);
+
   // Write the output modules
   ModuleWriter writer;
   writer.setBinary(options.emitBinary);
diff --git a/test/lit/wasm-split/initial-table.wast b/test/lit/wasm-split/initial-table.wast
new file mode 100644
index 000000000..534f7f3a9
--- /dev/null
+++ b/test/lit/wasm-split/initial-table.wast
@@ -0,0 +1,13 @@
+;; Test that the --initial-table flag works as expected
+
+;; RUN: wasm-split %s --instrument --initial-table=1234 -S | filecheck %s
+
+;; RUN: wasm-split %s -g -o1 %t.1.wasm -o2 %t.2.wasm --initial-table=1234
+;; RUN: wasm-dis %t.1.wasm | filecheck %s
+;; RUN: wasm-dis %t.2.wasm | filecheck %s
+
+;; CHECK: (table $table 1234 funcref)
+
+(module
+ (table $table 3 funcref)
+)