[wasm-split] Add an --asyncify option (#4513)

Add an option for running the asyncify transformation on the primary module
emitted by wasm-split. The idea is that the placeholder functions should be able
to unwind the stack while the secondary module is asynchronously loaded, then
once the placeholder functions have been patched out by the secondary module the
stack should be rewound and end up in the correct secondary function.

5 files changed, 199 insertions, 0 deletions

diff --git a/src/tools/wasm-split/split-options.cpp b/src/tools/wasm-split/split-options.cpp
index bd6d6f09a..b5929aad7 100644
--- a/src/tools/wasm-split/split-options.cpp
+++ b/src/tools/wasm-split/split-options.cpp
@@ -191,6 +191,15 @@ WasmSplitOptions::WasmSplitOptions()
            placeholderNamespace = argument;
          })
     .add(
+      "--asyncify",
+      "",
+      "Transform the module to support unwinding the stack from placeholder "
+      "functions and rewinding it once the secondary module has been loaded.",
+      WasmSplitOption,
+      {Mode::Split},
+      Options::Arguments::Zero,
+      [&](Options* o, const std::string& argument) { asyncify = true; })
+    .add(
       "--export-prefix",
       "",
       "An identifying prefix to prepend to new export names created "
diff --git a/src/tools/wasm-split/split-options.h b/src/tools/wasm-split/split-options.h
index d52f215cd..16e7b75e6 100644
--- a/src/tools/wasm-split/split-options.h
+++ b/src/tools/wasm-split/split-options.h
@@ -43,6 +43,7 @@ struct WasmSplitOptions : ToolOptions {
   bool emitBinary = true;
   bool symbolMap = false;
   bool placeholderMap = false;
+  bool asyncify = false;
 
   // TODO: Remove this. See the comment in wasm-binary.h.
   bool emitModuleNames = false;
diff --git a/src/tools/wasm-split/wasm-split.cpp b/src/tools/wasm-split/wasm-split.cpp
index 2732502ad..871d6bc95 100644
--- a/src/tools/wasm-split/wasm-split.cpp
+++ b/src/tools/wasm-split/wasm-split.cpp
@@ -280,6 +280,16 @@ void splitModule(const WasmSplitOptions& options) {
   adjustTableSize(wasm, options.initialTableSize);
   adjustTableSize(*secondary, options.initialTableSize);
 
+  // Run asyncify on the primary module
+  if (options.asyncify) {
+    PassOptions passOptions;
+    passOptions.optimizeLevel = 1;
+    passOptions.arguments.insert({"asyncify-ignore-imports", ""});
+    PassRunner runner(&wasm, passOptions);
+    runner.add("asyncify");
+    runner.run();
+  }
+
   if (options.symbolMap) {
     writeSymbolMap(wasm, options.primaryOutput + ".symbols");
     writeSymbolMap(*secondary, options.secondaryOutput + ".symbols");
diff --git a/test/lit/help/wasm-split.test b/test/lit/help/wasm-split.test
index cd275c79f..6319f475f 100644
--- a/test/lit/help/wasm-split.test
+++ b/test/lit/help/wasm-split.test
@@ -59,6 +59,11 @@
 ;; CHECK-NEXT:                                        import placeholder functions into the
 ;; CHECK-NEXT:                                        primary module.
 ;; CHECK-NEXT:
+;; CHECK-NEXT:   --asyncify                           [split] Transform the module to support
+;; CHECK-NEXT:                                        unwinding the stack from placeholder
+;; CHECK-NEXT:                                        functions and rewinding it once the
+;; CHECK-NEXT:                                        secondary module has been loaded.
+;; CHECK-NEXT:
 ;; CHECK-NEXT:   --export-prefix                      [split] An identifying prefix to prepend
 ;; CHECK-NEXT:                                        to new export names created by module
 ;; CHECK-NEXT:                                        splitting.
diff --git a/test/lit/wasm-split/asyncify.wast b/test/lit/wasm-split/asyncify.wast
new file mode 100644
index 000000000..4894f58a7
--- /dev/null
+++ b/test/lit/wasm-split/asyncify.wast
@@ -0,0 +1,174 @@
+;; RUN: wasm-split %s --export-prefix='%' -g -o1 %t.1.wasm -o2 %t.2.wasm --keep-funcs=foo --asyncify
+;; RUN: wasm-dis %t.1.wasm | filecheck %s --check-prefix PRIMARY
+;; RUN: wasm-dis %t.2.wasm | filecheck %s --check-prefix SECONDARY
+
+;; Check that the --asyncify option instruments the primary module but not the
+;; secondary module.
+
+(module
+ (func $foo (param i32) (result i32)
+  (call $bar (i32.const 0))
+ )
+ (func $bar (param i32) (result i32)
+  (call $foo (i32.const 1))
+ )
+)
+
+;; PRIMARY:      (module
+;; PRIMARY-NEXT:  (type $i32_=>_i32 (func (param i32) (result i32)))
+;; PRIMARY-NEXT:  (type $i32_=>_none (func (param i32)))
+;; PRIMARY-NEXT:  (type $none_=>_none (func))
+;; PRIMARY-NEXT:  (type $none_=>_i32 (func (result i32)))
+;; PRIMARY-NEXT:  (import "placeholder" "0" (func $placeholder_0 (param i32) (result i32)))
+;; PRIMARY-NEXT:  (global $global$0 (mut i32) (i32.const 0))
+;; PRIMARY-NEXT:  (global $global$1 (mut i32) (i32.const 0))
+;; PRIMARY-NEXT:  (memory $0 1 1)
+;; PRIMARY-NEXT:  (table $0 1 funcref)
+;; PRIMARY-NEXT:  (elem (i32.const 0) $placeholder_0)
+;; PRIMARY-NEXT:  (export "%foo" (func $foo))
+;; PRIMARY-NEXT:  (export "%table" (table $0))
+;; PRIMARY-NEXT:  (export "asyncify_start_unwind" (func $asyncify_start_unwind))
+;; PRIMARY-NEXT:  (export "asyncify_stop_unwind" (func $asyncify_stop_unwind))
+;; PRIMARY-NEXT:  (export "asyncify_start_rewind" (func $asyncify_start_rewind))
+;; PRIMARY-NEXT:  (export "asyncify_stop_rewind" (func $asyncify_stop_rewind))
+;; PRIMARY-NEXT:  (export "asyncify_get_state" (func $asyncify_get_state))
+;; PRIMARY-NEXT:  (func $foo (param $0 i32) (result i32)
+;; PRIMARY-NEXT:   (local $1 i32)
+;; PRIMARY-NEXT:   (if
+;; PRIMARY-NEXT:    (i32.eq
+;; PRIMARY-NEXT:     (global.get $global$0)
+;; PRIMARY-NEXT:     (i32.const 2)
+;; PRIMARY-NEXT:    )
+;; PRIMARY-NEXT:    (block
+;; PRIMARY-NEXT:     (i32.store
+;; PRIMARY-NEXT:      (global.get $global$1)
+;; PRIMARY-NEXT:      (i32.sub
+;; PRIMARY-NEXT:       (i32.load
+;; PRIMARY-NEXT:        (global.get $global$1)
+;; PRIMARY-NEXT:       )
+;; PRIMARY-NEXT:       (i32.const 4)
+;; PRIMARY-NEXT:      )
+;; PRIMARY-NEXT:     )
+;; PRIMARY-NEXT:     (local.set $0
+;; PRIMARY-NEXT:      (i32.load
+;; PRIMARY-NEXT:       (i32.load
+;; PRIMARY-NEXT:        (global.get $global$1)
+;; PRIMARY-NEXT:       )
+;; PRIMARY-NEXT:      )
+;; PRIMARY-NEXT:     )
+;; PRIMARY-NEXT:    )
+;; PRIMARY-NEXT:   )
+;; PRIMARY-NEXT:   (local.set $1
+;; PRIMARY-NEXT:    (block $label$2 (result i32)
+;; PRIMARY-NEXT:     (if
+;; PRIMARY-NEXT:      (i32.eqz
+;; PRIMARY-NEXT:       (select
+;; PRIMARY-NEXT:        (if (result i32)
+;; PRIMARY-NEXT:         (i32.eq
+;; PRIMARY-NEXT:          (global.get $global$0)
+;; PRIMARY-NEXT:          (i32.const 2)
+;; PRIMARY-NEXT:         )
+;; PRIMARY-NEXT:         (block (result i32)
+;; PRIMARY-NEXT:          (i32.store
+;; PRIMARY-NEXT:           (global.get $global$1)
+;; PRIMARY-NEXT:           (i32.sub
+;; PRIMARY-NEXT:            (i32.load
+;; PRIMARY-NEXT:             (global.get $global$1)
+;; PRIMARY-NEXT:            )
+;; PRIMARY-NEXT:            (i32.const 4)
+;; PRIMARY-NEXT:           )
+;; PRIMARY-NEXT:          )
+;; PRIMARY-NEXT:          (i32.load
+;; PRIMARY-NEXT:           (i32.load
+;; PRIMARY-NEXT:            (global.get $global$1)
+;; PRIMARY-NEXT:           )
+;; PRIMARY-NEXT:          )
+;; PRIMARY-NEXT:         )
+;; PRIMARY-NEXT:         (local.get $1)
+;; PRIMARY-NEXT:        )
+;; PRIMARY-NEXT:        (i32.const 0)
+;; PRIMARY-NEXT:        (global.get $global$0)
+;; PRIMARY-NEXT:       )
+;; PRIMARY-NEXT:      )
+;; PRIMARY-NEXT:      (block
+;; PRIMARY-NEXT:       (local.set $1
+;; PRIMARY-NEXT:        (call_indirect (type $i32_=>_i32)
+;; PRIMARY-NEXT:         (i32.const 0)
+;; PRIMARY-NEXT:         (i32.const 0)
+;; PRIMARY-NEXT:        )
+;; PRIMARY-NEXT:       )
+;; PRIMARY-NEXT:       (drop
+;; PRIMARY-NEXT:        (br_if $label$2
+;; PRIMARY-NEXT:         (i32.const 0)
+;; PRIMARY-NEXT:         (i32.eq
+;; PRIMARY-NEXT:          (global.get $global$0)
+;; PRIMARY-NEXT:          (i32.const 1)
+;; PRIMARY-NEXT:         )
+;; PRIMARY-NEXT:        )
+;; PRIMARY-NEXT:       )
+;; PRIMARY-NEXT:       (local.set $0
+;; PRIMARY-NEXT:        (local.get $1)
+;; PRIMARY-NEXT:       )
+;; PRIMARY-NEXT:      )
+;; PRIMARY-NEXT:     )
+;; PRIMARY-NEXT:     (if
+;; PRIMARY-NEXT:      (i32.eqz
+;; PRIMARY-NEXT:       (global.get $global$0)
+;; PRIMARY-NEXT:      )
+;; PRIMARY-NEXT:      (return
+;; PRIMARY-NEXT:       (local.get $0)
+;; PRIMARY-NEXT:      )
+;; PRIMARY-NEXT:     )
+;; PRIMARY-NEXT:     (unreachable)
+;; PRIMARY-NEXT:    )
+;; PRIMARY-NEXT:   )
+;; PRIMARY-NEXT:   (i32.store
+;; PRIMARY-NEXT:    (i32.load
+;; PRIMARY-NEXT:     (global.get $global$1)
+;; PRIMARY-NEXT:    )
+;; PRIMARY-NEXT:    (local.get $1)
+;; PRIMARY-NEXT:   )
+;; PRIMARY-NEXT:   (i32.store
+;; PRIMARY-NEXT:    (global.get $global$1)
+;; PRIMARY-NEXT:    (i32.add
+;; PRIMARY-NEXT:     (i32.load
+;; PRIMARY-NEXT:      (global.get $global$1)
+;; PRIMARY-NEXT:     )
+;; PRIMARY-NEXT:     (i32.const 4)
+;; PRIMARY-NEXT:    )
+;; PRIMARY-NEXT:   )
+;; PRIMARY-NEXT:   (i32.store
+;; PRIMARY-NEXT:    (i32.load
+;; PRIMARY-NEXT:     (global.get $global$1)
+;; PRIMARY-NEXT:    )
+;; PRIMARY-NEXT:    (local.get $0)
+;; PRIMARY-NEXT:   )
+;; PRIMARY-NEXT:   (i32.store
+;; PRIMARY-NEXT:    (global.get $global$1)
+;; PRIMARY-NEXT:    (i32.add
+;; PRIMARY-NEXT:     (i32.load
+;; PRIMARY-NEXT:      (global.get $global$1)
+;; PRIMARY-NEXT:     )
+;; PRIMARY-NEXT:     (i32.const 4)
+;; PRIMARY-NEXT:    )
+;; PRIMARY-NEXT:   )
+;; PRIMARY-NEXT:   (i32.const 0)
+;; PRIMARY-NEXT:  )
+;; PRIMARY:       (func $asyncify_start_unwind (param $0 i32)
+;; PRIMARY:       (func $asyncify_stop_unwind
+;; PRIMARY:       (func $asyncify_start_rewind (param $0 i32)
+;; PRIMARY:       (func $asyncify_stop_rewind
+;; PRIMARY:       (func $asyncify_get_state (result i32)
+;; PRIMARY:      )
+
+;; SECONDARY:      (module
+;; SECONDARY-NEXT:  (type $i32_=>_i32 (func (param i32) (result i32)))
+;; SECONDARY-NEXT:  (import "primary" "%table" (table $timport$0 1 funcref))
+;; SECONDARY-NEXT:  (import "primary" "%foo" (func $foo (param i32) (result i32)))
+;; SECONDARY-NEXT:  (elem (i32.const 0) $bar)
+;; SECONDARY-NEXT:  (func $bar (param $0 i32) (result i32)
+;; SECONDARY-NEXT:   (call $foo
+;; SECONDARY-NEXT:    (i32.const 1)
+;; SECONDARY-NEXT:   )
+;; SECONDARY-NEXT:  )
+;; SECONDARY-NEXT: )