Allow import mutable globals used in Asyncify pass (#4427)

This PR is part of the solution to emscripten-core/emscripten#15594.

emscripten Asyncify won't work properly in side modules, because the
globals, __asyncify_state and __asyncify_data, are not synchronized
between main-module and side-modules.

A new pass arg, asyncify-side-module, is added to make
__asyncify_state and __asyncify_data imported in the instrumented
wasm.

2 files changed, 131 insertions, 10 deletions

diff --git a/src/passes/Asyncify.cpp b/src/passes/Asyncify.cpp
index 7b03741d9..fc867f32e 100644
--- a/src/passes/Asyncify.cpp
+++ b/src/passes/Asyncify.cpp
@@ -296,6 +296,7 @@
 //      of their original range.
 //
 
+#include "asmjs/shared-constants.h"
 #include "cfg/liveness-traversal.h"
 #include "ir/effects.h"
 #include "ir/find_all.h"
@@ -1454,6 +1455,8 @@ struct Asyncify : public Pass {
       runner->options.getArgumentOrDefault("asyncify-asserts", "") != "";
     auto verbose =
       runner->options.getArgumentOrDefault("asyncify-verbose", "") != "";
+    auto sideModule =
+      runner->options.getArgumentOrDefault("asyncify-side-module", "") != "";
 
     removeList = handleBracketingOperators(removeList);
     addList = handleBracketingOperators(addList);
@@ -1488,7 +1491,7 @@ struct Asyncify : public Pass {
                             verbose);
 
     // Add necessary globals before we emit code to use them.
-    addGlobals(module);
+    addGlobals(module, sideModule);
 
     // Instrument the flow of code, adding code instrumentation and
     // skips for when rewinding. We do this on flat IR so that it is
@@ -1543,16 +1546,28 @@ struct Asyncify : public Pass {
   }
 
 private:
-  void addGlobals(Module* module) {
+  void addGlobals(Module* module, bool imported) {
     Builder builder(*module);
-    module->addGlobal(builder.makeGlobal(ASYNCIFY_STATE,
-                                         Type::i32,
-                                         builder.makeConst(int32_t(0)),
-                                         Builder::Mutable));
-    module->addGlobal(builder.makeGlobal(ASYNCIFY_DATA,
-                                         Type::i32,
-                                         builder.makeConst(int32_t(0)),
-                                         Builder::Mutable));
+
+    auto asyncifyState = builder.makeGlobal(ASYNCIFY_STATE,
+                                            Type::i32,
+                                            builder.makeConst(int32_t(0)),
+                                            Builder::Mutable);
+    if (imported) {
+      asyncifyState->module = ENV;
+      asyncifyState->base = ASYNCIFY_STATE;
+    }
+    module->addGlobal(std::move(asyncifyState));
+
+    auto asyncifyData = builder.makeGlobal(ASYNCIFY_DATA,
+                                           Type::i32,
+                                           builder.makeConst(int32_t(0)),
+                                           Builder::Mutable);
+    if (imported) {
+      asyncifyData->module = ENV;
+      asyncifyData->base = ASYNCIFY_DATA;
+    }
+    module->addGlobal(std::move(asyncifyData));
   }
 
   void addFunctions(Module* module) {
diff --git a/test/lit/passes/asyncify_pass-arg=asyncify-side-module.wast b/test/lit/passes/asyncify_pass-arg=asyncify-side-module.wast
new file mode 100644
index 000000000..66332e3c4
--- /dev/null
+++ b/test/lit/passes/asyncify_pass-arg=asyncify-side-module.wast
@@ -0,0 +1,106 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; NOTE: This test was ported using port_test.py and could be cleaned up.
+
+;; RUN: foreach %s %t wasm-opt --enable-mutable-globals --asyncify --pass-arg=asyncify-side-module -S -o - | filecheck %s
+
+(module
+)
+;; CHECK:      (type $i32_=>_none (func (param i32)))
+
+;; CHECK:      (type $none_=>_none (func))
+
+;; CHECK:      (type $none_=>_i32 (func (result i32)))
+
+;; CHECK:      (import "env" "__asyncify_state" (global $__asyncify_state (mut i32)))
+
+;; CHECK:      (import "env" "__asyncify_data" (global $__asyncify_data (mut i32)))
+
+;; CHECK:      (memory $0 1 1)
+
+;; CHECK:      (export "asyncify_start_unwind" (func $asyncify_start_unwind))
+
+;; CHECK:      (export "asyncify_stop_unwind" (func $asyncify_stop_unwind))
+
+;; CHECK:      (export "asyncify_start_rewind" (func $asyncify_start_rewind))
+
+;; CHECK:      (export "asyncify_stop_rewind" (func $asyncify_stop_rewind))
+
+;; CHECK:      (export "asyncify_get_state" (func $asyncify_get_state))
+
+;; CHECK:      (func $asyncify_start_unwind (param $0 i32)
+;; CHECK-NEXT:  (global.set $__asyncify_state
+;; CHECK-NEXT:   (i32.const 1)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (global.set $__asyncify_data
+;; CHECK-NEXT:   (local.get $0)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (if
+;; CHECK-NEXT:   (i32.gt_u
+;; CHECK-NEXT:    (i32.load
+;; CHECK-NEXT:     (global.get $__asyncify_data)
+;; CHECK-NEXT:    )
+;; CHECK-NEXT:    (i32.load offset=4
+;; CHECK-NEXT:     (global.get $__asyncify_data)
+;; CHECK-NEXT:    )
+;; CHECK-NEXT:   )
+;; CHECK-NEXT:   (unreachable)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT: )
+
+;; CHECK:      (func $asyncify_stop_unwind
+;; CHECK-NEXT:  (global.set $__asyncify_state
+;; CHECK-NEXT:   (i32.const 0)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (if
+;; CHECK-NEXT:   (i32.gt_u
+;; CHECK-NEXT:    (i32.load
+;; CHECK-NEXT:     (global.get $__asyncify_data)
+;; CHECK-NEXT:    )
+;; CHECK-NEXT:    (i32.load offset=4
+;; CHECK-NEXT:     (global.get $__asyncify_data)
+;; CHECK-NEXT:    )
+;; CHECK-NEXT:   )
+;; CHECK-NEXT:   (unreachable)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT: )
+
+;; CHECK:      (func $asyncify_start_rewind (param $0 i32)
+;; CHECK-NEXT:  (global.set $__asyncify_state
+;; CHECK-NEXT:   (i32.const 2)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (global.set $__asyncify_data
+;; CHECK-NEXT:   (local.get $0)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (if
+;; CHECK-NEXT:   (i32.gt_u
+;; CHECK-NEXT:    (i32.load
+;; CHECK-NEXT:     (global.get $__asyncify_data)
+;; CHECK-NEXT:    )
+;; CHECK-NEXT:    (i32.load offset=4
+;; CHECK-NEXT:     (global.get $__asyncify_data)
+;; CHECK-NEXT:    )
+;; CHECK-NEXT:   )
+;; CHECK-NEXT:   (unreachable)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT: )
+
+;; CHECK:      (func $asyncify_stop_rewind
+;; CHECK-NEXT:  (global.set $__asyncify_state
+;; CHECK-NEXT:   (i32.const 0)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (if
+;; CHECK-NEXT:   (i32.gt_u
+;; CHECK-NEXT:    (i32.load
+;; CHECK-NEXT:     (global.get $__asyncify_data)
+;; CHECK-NEXT:    )
+;; CHECK-NEXT:    (i32.load offset=4
+;; CHECK-NEXT:     (global.get $__asyncify_data)
+;; CHECK-NEXT:    )
+;; CHECK-NEXT:   )
+;; CHECK-NEXT:   (unreachable)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT: )
+
+;; CHECK:      (func $asyncify_get_state (result i32)
+;; CHECK-NEXT:  (global.get $__asyncify_state)
+;; CHECK-NEXT: )