Asyncify: Instrument indirect calls from functions in add-list or only-list

src/passes/Asyncify.cpp

@@ -265,6 +265,27 @@
 //      If the "only-list" is provided, then *only* the functions in the list
 //      will be instrumented, and nothing else.
 //
+// Note that there are two types of instrumentation that happen for each
+// function: if foo() can be part of a pause/resume operation, then we need to
+// instrument code inside it to support pausing and resuming, but also, we need
+// callers of the function to instrument calls to it. Normally this is already
+// taken care of as the callers need to be instrumented as well anyhow. However,
+// the ignore-indirect option makes things more complicated, since we can't tell
+// where all the calls to foo() are - all we see are indirect calls that do not
+// refer to foo() by name. To make it possible for you to use the add-list or
+// only-list with ignore-indirect, those lists affect *both* kinds of
+// instrumentation. That is, if parent() calls foo() indirectly, and you add
+// parent() to the add-list, then the indirect calls in parent() will be
+// instrumented to support pausing/resuming, even if ignore-indirect is set.
+// Typically you don't need to think about this, and just add all the functions
+// that can be on the stack while pausing - what this means is that when you do
+// so, indirect calls will just work. (The cost is that an instrumented function
+// will check for pausing at all indirect calls, which may be unnecessary in
+// some cases; but this is an instrumented function anyhow, and indirect calls
+// are slower anyhow, so this simple model seems good enough - a more complex
+// model where you can specify "instrument, but not indirect calls from me"
+// would likely have little benefit.)
+//
 // TODO When wasm has GC, extending the live ranges of locals can keep things
 //      alive unnecessarily. We may want to set locals to null at the end
 //      of their original range.
@@ -478,6 +499,7 @@ class ModuleAnalyzer {
     // that call it are instrumented. This is not done for the bottom.
     bool isTopMostRuntime = false;
     bool inRemoveList = false;
+    bool addedFromList = false;
   };
 
   typedef std::map<Function*, Info> Map;
@@ -634,15 +656,22 @@ class ModuleAnalyzer {
       // Only the functions in the only-list can change the state.
       for (auto& func : module.functions) {
         if (!func->imported()) {
-          map[func.get()].canChangeState = onlyList.match(func->name);
+          auto& info = map[func.get()];
+          bool matched = onlyList.match(func->name);
+          info.canChangeState = matched;
+          if (matched) {
+            info.addedFromList = true;
+          }
         }
       }
     }
 
     if (!addListInput.empty()) {
       for (auto& func : module.functions) {
         if (!func->imported() && addList.match(func->name)) {
-          map[func.get()].canChangeState = true;
+          auto& info = map[func.get()];
+          info.canChangeState = true;
+          info.addedFromList = true;
         }
       }
     }
@@ -657,7 +686,7 @@ class ModuleAnalyzer {
     return info.canChangeState && !info.isTopMostRuntime;
   }
 
-  bool canChangeState(Expression* curr) {
+  bool canChangeState(Expression* curr, Function* func) {
     // Look inside to see if we call any of the things we know can change the
     // state.
     // TODO: caching, this is O(N^2)
@@ -683,29 +712,30 @@ class ModuleAnalyzer {
           canChangeState = true;
         }
       }
-      void visitCallIndirect(CallIndirect* curr) {
-        if (canIndirectChangeState) {
-          canChangeState = true;
-        }
-        // TODO optimize the other case, at least by type
-      }
+      void visitCallIndirect(CallIndirect* curr) { hasIndirectCall = true; }
       Module* module;
       ModuleAnalyzer* analyzer;
       Map* map;
-      bool canIndirectChangeState;
+      bool hasIndirectCall = false;
       bool canChangeState = false;
       bool isBottomMostRuntime = false;
     };
     Walker walker;
     walker.module = &module;
     walker.analyzer = this;
     walker.map = &map;
-    walker.canIndirectChangeState = canIndirectChangeState;
     walker.walk(curr);
-    if (walker.isBottomMostRuntime) {
-      walker.canChangeState = false;
+    // An indirect call is normally ignored if we are ignoring indirect calls.
+    // However, see the docs at the top: if the function we are inside was
+    // specifically added by the user (in the only-list or the add-list) then we
+    // instrument indirect calls from it (this allows specifically allowing some
+    // indirect calls but not others).
+    if (walker.hasIndirectCall &&
+        (canIndirectChangeState || map[func].addedFromList)) {
+      walker.canChangeState = true;
     }
-    return walker.canChangeState;
+    // The bottom-most runtime can never change the state.
+    return walker.canChangeState && !walker.isBottomMostRuntime;
   }
 
   FakeGlobalHelper fakeGlobals;
@@ -814,7 +844,7 @@ struct AsyncifyFlow : public Pass {
   Index callIndex = 0;
 
   Expression* process(Expression* curr) {
-    if (!analyzer->canChangeState(curr)) {
+    if (!analyzer->canChangeState(curr, func)) {
       return makeMaybeSkip(curr);
     }
     // The IR is in flat form, which makes this much simpler: there are no
@@ -856,12 +886,13 @@ struct AsyncifyFlow : public Pass {
       Index i = 0;
       auto& list = block->list;
       while (i < list.size()) {
-        if (analyzer->canChangeState(list[i])) {
+        if (analyzer->canChangeState(list[i], func)) {
           list[i] = process(list[i]);
           i++;
         } else {
           Index end = i + 1;
-          while (end < list.size() && !analyzer->canChangeState(list[end])) {
+          while (end < list.size() &&
+                 !analyzer->canChangeState(list[end], func)) {
             end++;
           }
           // We have a range of [i, end) in which the state cannot change,
@@ -886,7 +917,7 @@ struct AsyncifyFlow : public Pass {
     } else if (auto* iff = curr->dynCast<If>()) {
       // The state change cannot be in the condition due to flat form, so it
       // must be in one of the children.
-      assert(!analyzer->canChangeState(iff->condition));
+      assert(!analyzer->canChangeState(iff->condition, func));
       // We must linearize this, which means we pass through both arms if we
       // are rewinding.
       if (!iff->ifFalse) {

test/passes/asyncify_pass-arg=asyncify-addlist@foo_pass-arg=asyncify-ignore-indirect.txt

@@ -0,0 +1,195 @@
+(module
+ (type $none_=>_none (func))
+ (type $i32_=>_none (func (param i32)))
+ (type $none_=>_i32 (func (result i32)))
+ (import "env" "import" (func $import))
+ (memory $0 1 2)
+ (table $0 1 funcref)
+ (global $__asyncify_state (mut i32) (i32.const 0))
+ (global $__asyncify_data (mut i32) (i32.const 0))
+ (export "asyncify_start_unwind" (func $asyncify_start_unwind))
+ (export "asyncify_stop_unwind" (func $asyncify_stop_unwind))
+ (export "asyncify_start_rewind" (func $asyncify_start_rewind))
+ (export "asyncify_stop_rewind" (func $asyncify_stop_rewind))
+ (export "asyncify_get_state" (func $asyncify_get_state))
+ (func $foo
+  (local $0 i32)
+  (local $1 i32)
+  (if
+   (i32.eq
+    (global.get $__asyncify_state)
+    (i32.const 2)
+   )
+   (nop)
+  )
+  (local.set $0
+   (block $__asyncify_unwind (result i32)
+    (block
+     (block
+      (if
+       (i32.eq
+        (global.get $__asyncify_state)
+        (i32.const 2)
+       )
+       (block
+        (i32.store
+         (global.get $__asyncify_data)
+         (i32.add
+          (i32.load
+           (global.get $__asyncify_data)
+          )
+          (i32.const -4)
+         )
+        )
+        (local.set $1
+         (i32.load
+          (i32.load
+           (global.get $__asyncify_data)
+          )
+         )
+        )
+       )
+      )
+      (block
+       (if
+        (i32.eq
+         (global.get $__asyncify_state)
+         (i32.const 0)
+        )
+        (call $nothing)
+       )
+       (if
+        (if (result i32)
+         (i32.eq
+          (global.get $__asyncify_state)
+          (i32.const 0)
+         )
+         (i32.const 1)
+         (i32.eq
+          (local.get $1)
+          (i32.const 0)
+         )
+        )
+        (block
+         (call_indirect (type $none_=>_none)
+          (i32.const 0)
+         )
+         (if
+          (i32.eq
+           (global.get $__asyncify_state)
+           (i32.const 1)
+          )
+          (br $__asyncify_unwind
+           (i32.const 0)
+          )
+         )
+        )
+       )
+      )
+     )
+     (return)
+    )
+   )
+  )
+  (block
+   (i32.store
+    (i32.load
+     (global.get $__asyncify_data)
+    )
+    (local.get $0)
+   )
+   (i32.store
+    (global.get $__asyncify_data)
+    (i32.add
+     (i32.load
+      (global.get $__asyncify_data)
+     )
+     (i32.const 4)
+    )
+   )
+  )
+  (nop)
+ )
+ (func $bar
+  (call $nothing)
+  (call_indirect (type $none_=>_none)
+   (i32.const 0)
+  )
+ )
+ (func $nothing
+  (nop)
+ )
+ (func $asyncify_start_unwind (param $0 i32)
+  (global.set $__asyncify_state
+   (i32.const 1)
+  )
+  (global.set $__asyncify_data
+   (local.get $0)
+  )
+  (if
+   (i32.gt_u
+    (i32.load
+     (global.get $__asyncify_data)
+    )
+    (i32.load offset=4
+     (global.get $__asyncify_data)
+    )
+   )
+   (unreachable)
+  )
+ )
+ (func $asyncify_stop_unwind
+  (global.set $__asyncify_state
+   (i32.const 0)
+  )
+  (if
+   (i32.gt_u
+    (i32.load
+     (global.get $__asyncify_data)
+    )
+    (i32.load offset=4
+     (global.get $__asyncify_data)
+    )
+   )
+   (unreachable)
+  )
+ )
+ (func $asyncify_start_rewind (param $0 i32)
+  (global.set $__asyncify_state
+   (i32.const 2)
+  )
+  (global.set $__asyncify_data
+   (local.get $0)
+  )
+  (if
+   (i32.gt_u
+    (i32.load
+     (global.get $__asyncify_data)
+    )
+    (i32.load offset=4
+     (global.get $__asyncify_data)
+    )
+   )
+   (unreachable)
+  )
+ )
+ (func $asyncify_stop_rewind
+  (global.set $__asyncify_state
+   (i32.const 0)
+  )
+  (if
+   (i32.gt_u
+    (i32.load
+     (global.get $__asyncify_data)
+    )
+    (i32.load offset=4
+     (global.get $__asyncify_data)
+    )
+   )
+   (unreachable)
+  )
+ )
+ (func $asyncify_get_state (result i32)
+  (global.get $__asyncify_state)
+ )
+)

test/passes/asyncify_pass-arg=asyncify-addlist@foo_pass-arg=asyncify-ignore-indirect.wast

@@ -0,0 +1,18 @@
+(module
+  (type $t (func))
+  (memory 1 2)
+  (table 1 funcref)
+  (elem (i32.const 0))
+  (import "env" "import" (func $import))
+  (func $foo ;; doesn't look like it needs instrumentation, but in add list
+    (call $nothing)
+    (call_indirect (type $t) (i32.const 0))
+  )
+  (func $bar ;; doesn't look like it needs instrumentation, and not in add list
+    (call $nothing)
+    (call_indirect (type $t) (i32.const 0))
+  )
+  (func $nothing
+  )
+)
+