[Stack Switching] Fix GUFA on continuations

scripts/test/fuzzing.py

@@ -114,6 +114,7 @@
     'vacuum-stack-switching.wast',
     'cont.wast',
     'cont_simple.wast',
+    'gufa-cont.wast',
     'cont_many_unhandled.wast',
     # TODO: fix split_wast() on tricky escaping situations like a string ending
     #       in \\" (the " is not escaped - there is an escaped \ before it)

src/ir/possible-contents.cpp

@@ -1291,14 +1291,40 @@ struct InfoCollector
     // TODO: optimize when possible
     addRoot(curr);
   }
-  void visitResume(Resume* curr) {
-    // TODO: optimize when possible
-    addRoot(curr);
-  }
-  void visitResumeThrow(ResumeThrow* curr) {
+
+  template<typename T> void handleResume(T* curr) {
     // TODO: optimize when possible
     addRoot(curr);
+
+    // Connect handled tags with their branch targets, and materialize non-null
+    // continuation values.
+    auto numTags = curr->handlerTags.size();
+    for (Index tagIndex = 0; tagIndex < numTags; tagIndex++) {
+      auto tag = curr->handlerTags[tagIndex];
+      auto target = curr->handlerBlocks[tagIndex];
+      auto params = getModule()->getTag(tag)->params();
+
+      // Add the values from the tag.
+      for (Index i = 0; i < params.size(); i++) {
+        if (isRelevant(params[i])) {
+          info.links.push_back(
+            {TagLocation{tag, i}, getBreakTargetLocation(target, i)});
+        }
+      }
+
+      // Add the continuation. Its type is determined by the block we break to,
+      // as the last result.
+      auto targetType = findBreakTarget(target)->type;
+      assert(targetType.size() >= 1);
+      auto contType = targetType[targetType.size() - 1];
+      auto location = getRootLocation(contType);
+      info.links.push_back(
+        {location, getBreakTargetLocation(target, params.size())});
+    }
   }
+
+  void visitResume(Resume* curr) { handleResume(curr); }
+  void visitResumeThrow(ResumeThrow* curr) { handleResume(curr); }
   void visitStackSwitch(StackSwitch* curr) {
     // TODO: optimize when possible
     addRoot(curr);

test/lit/passes/gufa-cont.wast

@@ -0,0 +1,132 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: foreach %s %t wasm-opt -all --gufa -S -o - | filecheck %s
+
+(module
+ ;; CHECK:      (type $func (func))
+ (type $func (func))
+ ;; CHECK:      (type $cont (cont $func))
+ (type $cont (cont $func))
+
+ ;; CHECK:      (type $func-i32 (func (result i32)))
+ (type $func-i32 (func (result i32)))
+ ;; CHECK:      (type $cont-i32 (cont $func-i32))
+ (type $cont-i32 (cont $func-i32))
+
+ ;; CHECK:      (type $4 (func (result i32 (ref $cont))))
+
+ ;; CHECK:      (elem declare func $cont $cont-i32)
+
+ ;; CHECK:      (tag $tag (type $func))
+ (tag $tag (type $func))
+
+ ;; CHECK:      (tag $tag-i32 (type $func-i32) (result i32))
+ (tag $tag-i32 (type $func-i32))
+
+ ;; CHECK:      (export "resume" (func $resume))
+
+ ;; CHECK:      (export "resume_throw" (func $resume_throw))
+
+ ;; CHECK:      (export "resume-i32" (func $resume-i32))
+
+ ;; CHECK:      (func $cont (type $func)
+ ;; CHECK-NEXT:  (suspend $tag)
+ ;; CHECK-NEXT: )
+ (func $cont
+  ;; Helper for below.
+  (suspend $tag)
+ )
+
+ ;; CHECK:      (func $resume (type $func)
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (block $block (result (ref $cont))
+ ;; CHECK-NEXT:    (resume $cont (on $tag $block)
+ ;; CHECK-NEXT:     (cont.new $cont
+ ;; CHECK-NEXT:      (ref.func $cont)
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:    (return)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $resume (export "resume")
+  ;; A continuation is created, it suspends, and we handle that. There is
+  ;; nothing to optimize or change here.
+  (drop
+   (block $block (result (ref $cont))
+    (resume $cont (on $tag $block)
+     (cont.new $cont
+      (ref.func $cont)
+     )
+    )
+    (return)
+   )
+  )
+ )
+
+ ;; CHECK:      (func $resume_throw (type $func)
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (block $block (result (ref $cont))
+ ;; CHECK-NEXT:    (resume_throw $cont $tag (on $tag $block)
+ ;; CHECK-NEXT:     (cont.new $cont
+ ;; CHECK-NEXT:      (ref.func $cont)
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:    (return)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $resume_throw (export "resume_throw")
+  ;; As above, but with resume_throw.
+  (drop
+   (block $block (result (ref $cont))
+    (resume_throw $cont $tag (on $tag $block)
+     (cont.new $cont
+      (ref.func $cont)
+     )
+    )
+    (return)
+   )
+  )
+ )
+
+ ;; CHECK:      (func $cont-i32 (type $func-i32) (result i32)
+ ;; CHECK-NEXT:  (suspend $tag)
+ ;; CHECK-NEXT:  (unreachable)
+ ;; CHECK-NEXT:  (unreachable)
+ ;; CHECK-NEXT: )
+ (func $cont-i32 (result i32)
+  ;; Helper for below.
+  (suspend $tag)
+  (unreachable)
+ )
+
+ ;; CHECK:      (func $resume-i32 (type $func)
+ ;; CHECK-NEXT:  (tuple.drop 2
+ ;; CHECK-NEXT:   (block $block (type $4) (result i32 (ref $cont))
+ ;; CHECK-NEXT:    (drop
+ ;; CHECK-NEXT:     (resume $cont-i32 (on $tag-i32 $block)
+ ;; CHECK-NEXT:      (cont.new $cont-i32
+ ;; CHECK-NEXT:       (ref.func $cont-i32)
+ ;; CHECK-NEXT:      )
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:    (return)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $resume-i32 (export "resume-i32")
+  ;; As above, but with more values sent than just the continuation.
+  (tuple.drop 2
+   (block $block (result i32 (ref $cont))
+    (resume $cont-i32 (on $tag-i32 $block)
+     (cont.new $cont-i32
+      (ref.func $cont-i32)
+     )
+    )
+    (return)
+   )
+  )
+ )
+)
+