1 files changed, 14 insertions, 21 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