[Stack Switching] RemoveUnusedModuleElements: Fix closed world analysis of continuation functions

scripts/test/fuzzing.py

@@ -119,6 +119,7 @@
     'cont_export.wast',
     'cont_export_throw.wast',
     'type-merging-cont.wast',
+    'remove-unused-module-elements-cont.wast',
     # TODO: fix split_wast() on tricky escaping situations like a string ending
     #       in \\" (the " is not escaped - there is an escaped \ before it)
     'string-lifting-section.wast',

src/passes/RemoveUnusedModuleElements.cpp

@@ -190,6 +190,16 @@ struct Noter : public PostWalker<Noter, UnifiedExpressionVisitor<Noter>> {
     auto type = curr->ref->type.getHeapType();
     noteStructField(StructField{type, curr->index});
   }
+
+  void visitContNew(ContNew* curr) {
+    // The function reference that is passed in here will be called, just as if
+    // we were a call_ref, except at a potentially later time.
+    if (!curr->func->type.isRef()) {
+      return;
+    }
+
+    noteCallRef(curr->func->type.getHeapType());
+  }
 };
 
 // Analyze a module to find what things are referenced and what things are used.

test/lit/passes/remove-unused-module-elements-cont.wast

@@ -0,0 +1,136 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: foreach %s %t wasm-opt --remove-unused-module-elements --closed-world -all -S -o - | filecheck %s
+;; RUN: foreach %s %t wasm-opt --remove-unused-module-elements                -all -S -o - | filecheck %s --check-prefix OPEN_WORLD
+
+;; Test both open world (default) and closed world. In a closed world we can do
+;; more with function refs, as we assume nothing calls them on the outside, so
+;; if no calls exist to the right type, the function is not reached.
+
+(module
+ ;; CHECK:      (type $func (func))
+ ;; OPEN_WORLD:      (type $func (func))
+ (type $func (func))
+ ;; CHECK:      (type $cont (cont $func))
+ ;; OPEN_WORLD:      (type $cont (cont $func))
+ (type $cont (cont $func))
+
+ ;; CHECK:      (elem declare func $func)
+
+ ;; CHECK:      (tag $tag (type $func))
+ ;; OPEN_WORLD:      (elem declare func $func)
+
+ ;; OPEN_WORLD:      (tag $tag (type $func))
+ (tag $tag (type $func))
+
+ ;; CHECK:      (export "run" (func $run))
+
+ ;; CHECK:      (func $func (type $func)
+ ;; CHECK-NEXT:  (nop)
+ ;; CHECK-NEXT: )
+ ;; OPEN_WORLD:      (export "run" (func $run))
+
+ ;; OPEN_WORLD:      (func $func (type $func)
+ ;; OPEN_WORLD-NEXT:  (nop)
+ ;; OPEN_WORLD-NEXT: )
+ (func $func
+  ;; This function is only ever referred to by a ref.func that is passed into
+  ;; a cont.new. We should not modify this to unreachable, which we would do if
+  ;; we didn't realize it will execute (in closed world).
+  (nop)
+ )
+
+ ;; CHECK:      (func $run (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 $func)
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:    (return)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ ;; OPEN_WORLD:      (func $run (type $func)
+ ;; OPEN_WORLD-NEXT:  (drop
+ ;; OPEN_WORLD-NEXT:   (block $block (result (ref $cont))
+ ;; OPEN_WORLD-NEXT:    (resume $cont (on $tag $block)
+ ;; OPEN_WORLD-NEXT:     (cont.new $cont
+ ;; OPEN_WORLD-NEXT:      (ref.func $func)
+ ;; OPEN_WORLD-NEXT:     )
+ ;; OPEN_WORLD-NEXT:    )
+ ;; OPEN_WORLD-NEXT:    (return)
+ ;; OPEN_WORLD-NEXT:   )
+ ;; OPEN_WORLD-NEXT:  )
+ ;; OPEN_WORLD-NEXT: )
+ (func $run (export "run")
+  ;; No changes are expected here.
+  (drop
+   (block $block (result (ref $cont))
+    (resume $cont (on $tag $block)
+     (cont.new $cont
+      (ref.func $func)
+     )
+    )
+    (return)
+   )
+  )
+ )
+)
+
+(module
+ ;; CHECK:      (type $func (func))
+ ;; OPEN_WORLD:      (type $func (func))
+ (type $func (func))
+ ;; CHECK:      (type $cont (cont $func))
+ ;; OPEN_WORLD:      (type $cont (cont $func))
+ (type $cont (cont $func))
+
+ ;; CHECK:      (tag $tag (type $func))
+ ;; OPEN_WORLD:      (tag $tag (type $func))
+ (tag $tag (type $func))
+
+ ;; CHECK:      (export "run" (func $run))
+
+ ;; CHECK:      (func $run (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.null nofunc)
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:    (return)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ ;; OPEN_WORLD:      (export "run" (func $run))
+
+ ;; OPEN_WORLD:      (func $run (type $func)
+ ;; OPEN_WORLD-NEXT:  (drop
+ ;; OPEN_WORLD-NEXT:   (block $block (result (ref $cont))
+ ;; OPEN_WORLD-NEXT:    (resume $cont (on $tag $block)
+ ;; OPEN_WORLD-NEXT:     (cont.new $cont
+ ;; OPEN_WORLD-NEXT:      (ref.null nofunc)
+ ;; OPEN_WORLD-NEXT:     )
+ ;; OPEN_WORLD-NEXT:    )
+ ;; OPEN_WORLD-NEXT:    (return)
+ ;; OPEN_WORLD-NEXT:   )
+ ;; OPEN_WORLD-NEXT:  )
+ ;; OPEN_WORLD-NEXT: )
+ (func $run (export "run")
+  ;; We should not error on a null function reference in the cont.new.
+  (drop
+   (block $block (result (ref $cont))
+    (resume $cont (on $tag $block)
+     (cont.new $cont
+      (ref.null $func)
+     )
+    )
+    (return)
+   )
+  )
+ )
+)
+