2 files changed, 48 insertions, 23 deletions

diff --git a/src/ir/possible-contents.cpp b/src/ir/possible-contents.cpp
index 4755c9904..fd186a74f 100644
--- a/src/ir/possible-contents.cpp
+++ b/src/ir/possible-contents.cpp
@@ -418,6 +418,20 @@ struct InfoCollector
     addRoot(
       curr,
       PossibleContents::literal(Literal(curr->func, curr->type.getHeapType())));
+
+    // The presence of a RefFunc indicates the function may be called
+    // indirectly, so add the relevant connections for this particular function.
+    // We do so here in the RefFunc so that we only do it for functions that
+    // actually have a RefFunc.
+    auto* func = getModule()->getFunction(curr->func);
+    for (Index i = 0; i < func->getParams().size(); i++) {
+      info.links.push_back(
+        {SignatureParamLocation{func->type, i}, LocalLocation{func, i, 0}});
+    }
+    for (Index i = 0; i < func->getResults().size(); i++) {
+      info.links.push_back(
+        {ResultLocation{func, i}, SignatureResultLocation{func->type, i}});
+    }
   }
   void visitRefEq(RefEq* curr) {
     addRoot(curr);
@@ -616,9 +630,6 @@ struct InfoCollector
     handleIndirectCall(curr, curr->heapType);
   }
   void visitCallRef(CallRef* curr) {
-    // TODO: Optimize like RefCast etc.: the values reaching us depend on the
-    //       possible values of |target| (which might be nothing, or might be a
-    //       constant function).
     handleIndirectCall(curr, curr->target->type);
   }
 
@@ -1295,24 +1306,6 @@ Flower::Flower(Module& wasm) : wasm(wasm) {
   }
 
 #ifdef POSSIBLE_CONTENTS_DEBUG
-  std::cout << "func phase\n";
-#endif
-
-  // Connect function parameters to their signature, so that any indirect call
-  // of that signature will reach them.
-  // TODO: find which functions are even taken by reference
-  for (auto& func : wasm.functions) {
-    for (Index i = 0; i < func->getParams().size(); i++) {
-      links.insert(getIndexes({SignatureParamLocation{func->type, i},
-                               LocalLocation{func.get(), i, 0}}));
-    }
-    for (Index i = 0; i < func->getResults().size(); i++) {
-      links.insert(getIndexes({ResultLocation{func.get(), i},
-                               SignatureResultLocation{func->type, i}}));
-    }
-  }
-
-#ifdef POSSIBLE_CONTENTS_DEBUG
   std::cout << "struct phase\n";
 #endif
 
diff --git a/test/lit/passes/gufa-refs.wast b/test/lit/passes/gufa-refs.wast
index d8d12e9f0..24ebda381 100644
--- a/test/lit/passes/gufa-refs.wast
+++ b/test/lit/passes/gufa-refs.wast
@@ -3995,10 +3995,10 @@
 
 ;; call_ref types
 (module
-  ;; CHECK:      (type $none_=>_i32 (func_subtype (result i32) func))
-
   ;; CHECK:      (type $i1 (func_subtype (param i32) func))
   (type $i1 (func (param i32)))
+  ;; CHECK:      (type $none_=>_i32 (func_subtype (result i32) func))
+
   ;; CHECK:      (type $i2 (func_subtype (param i32) func))
   (type $i2 (func (param i32)))
 
@@ -4007,6 +4007,9 @@
   ;; CHECK:      (import "a" "b" (func $import (result i32)))
   (import "a" "b" (func $import (result i32)))
 
+  ;; CHECK:      (global $func (ref func) (ref.func $reffed-in-global-code))
+  (global $func (ref func) (ref.func $reffed-in-global-code))
+
   ;; CHECK:      (elem declare func $reffed1 $reffed2)
 
   ;; CHECK:      (func $reffed1 (type $i1) (param $x i32)
@@ -4022,6 +4025,35 @@
     )
   )
 
+  ;; CHECK:      (func $not-reffed (type $i1) (param $x i32)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (unreachable)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $not-reffed (type $i1) (param $x i32)
+    ;; This function has the same type as the previous one, but it is never
+    ;; taken by reference, which means the call_refs below do not affect it. As
+    ;; there are no other calls, this local.get can be turned into an
+    ;; unreachable.
+    (drop
+      (local.get $x)
+    )
+  )
+
+  ;; CHECK:      (func $reffed-in-global-code (type $i1) (param $x i32)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i32.const 42)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $reffed-in-global-code (type $i1) (param $x i32)
+    ;; The only ref to this function is in global code, so this tests that we
+    ;; scan that properly. This can be optimized like $reffed, that is, we can
+    ;; infer 42 here.
+    (drop
+      (local.get $x)
+    )
+  )
+
   ;; CHECK:      (func $reffed2 (type $i2) (param $x i32)
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (local.get $x)