Account for addressed funcs

Author: stevenfontanella

src/passes/GlobalEffects.cpp

@@ -22,6 +22,7 @@
 #include <ranges>
 
 #include "ir/effects.h"
+#include "ir/element-utils.h"
 #include "ir/module-utils.h"
 #include "pass.h"
 #include "support/strongly_connected_components.h"
@@ -46,8 +47,55 @@ struct FuncInfo {
   std::unordered_set<HeapType> indirectCalledTypes;
 };
 
+std::unordered_set<Function*> getAddressedFuncs(Module& module) {
+  struct AddressedFuncsWalker
+    : PostWalker<AddressedFuncsWalker,
+                 UnifiedExpressionVisitor<AddressedFuncsWalker>> {
+    const Module& module;
+    std::unordered_set<Function*>& addressedFuncs;
+
+    AddressedFuncsWalker(const Module& module,
+                         std::unordered_set<Function*>& addressedFuncs)
+      : module(module), addressedFuncs(addressedFuncs) {}
+
+    void visitExpression(Expression* curr) {
+      if (auto* refFunc = curr->dynCast<RefFunc>()) {
+        addressedFuncs.insert(module.getFunction(refFunc->func));
+      }
+    }
+  };
+
+  std::unordered_set<Function*> addressedFuncs;
+  AddressedFuncsWalker walker(module, addressedFuncs);
+  walker.walkModule(&module);
+
+  ModuleUtils::iterImportedFunctions(
+    module, [&addressedFuncs, &module](Function* import) {
+      addressedFuncs.insert(module.getFunction(import->name));
+    });
+
+  for (const auto& export_ : module.exports) {
+    if (export_->kind != ExternalKind::Function) {
+      continue;
+    }
+
+    // TODO: internal or external? I think internal
+    // This might be why we failed to lookup the function earlier
+    // Maybe we can use Function* after all
+    addressedFuncs.insert(module.getFunction(*export_->getInternalName()));
+  }
+
+  ElementUtils::iterAllElementFunctionNames(
+    &module, [&addressedFuncs, &module](Name func) {
+      addressedFuncs.insert(module.getFunction(func));
+    });
+
+  return addressedFuncs;
+}
+
 std::map<Function*, FuncInfo> analyzeFuncs(Module& module,
                                            const PassOptions& passOptions) {
+  std::unordered_set<Name> addressedFuncs;
   ModuleUtils::ParallelFunctionAnalysis<FuncInfo> analysis(
     module, [&](Function* func, FuncInfo& funcInfo) {
       if (func->imported()) {
@@ -78,10 +126,14 @@ std::map<Function*, FuncInfo> analyzeFuncs(Module& module,
           const PassOptions& options;
           FuncInfo& funcInfo;
 
+          std::unordered_set<Name>& addressedFuncs;
+
           CallScanner(Module& wasm,
                       const PassOptions& options,
-                      FuncInfo& funcInfo)
-            : wasm(wasm), options(options), funcInfo(funcInfo) {}
+                      FuncInfo& funcInfo,
+                      std::unordered_set<Name>& addressedFuncs)
+            : wasm(wasm), options(options), funcInfo(funcInfo),
+              addressedFuncs(addressedFuncs) {}
 
           void visitExpression(Expression* curr) {
             ShallowEffectAnalyzer effects(options, wasm, curr);
@@ -114,7 +166,7 @@ std::map<Function*, FuncInfo> analyzeFuncs(Module& module,
             }
           }
         };
-        CallScanner scanner(module, passOptions, funcInfo);
+        CallScanner scanner(module, passOptions, funcInfo, addressedFuncs);
         scanner.walkFunction(func);
       }
     });
@@ -146,6 +198,7 @@ using CallGraph =
 
 CallGraph buildCallGraph(const Module& module,
                          const std::map<Function*, FuncInfo>& funcInfos,
+                         const std::unordered_set<Function*> addressedFuncs,
                          bool closedWorld) {
   CallGraph callGraph;
 
@@ -170,7 +223,9 @@ CallGraph buildCallGraph(const Module& module,
     }
 
     // Type -> Function
-    callGraph[caller->type.getHeapType()].insert(caller);
+    if (addressedFuncs.contains(caller)) {
+      callGraph[caller->type.getHeapType()].insert(caller);
+    }
   }
 
   // Type -> Type
@@ -325,11 +380,12 @@ void copyEffectsToFunctions(const std::map<Function*, FuncInfo>& funcInfos) {
 
 struct GenerateGlobalEffects : public Pass {
   void run(Module* module) override {
-    std::map<Function*, FuncInfo> funcInfos =
-      analyzeFuncs(*module, getPassOptions());
+    auto funcInfos = analyzeFuncs(*module, getPassOptions());
+
+    auto addressedFuncs = getAddressedFuncs(*module);
 
-    auto callGraph =
-      buildCallGraph(*module, funcInfos, getPassOptions().closedWorld);
+    auto callGraph = buildCallGraph(
+      *module, funcInfos, addressedFuncs, getPassOptions().closedWorld);
 
     propagateEffects(*module, getPassOptions(), funcInfos, callGraph);
 

test/lit/passes/global-effects-closed-world.wast

@@ -344,15 +344,12 @@
   )
 
   ;; CHECK:      (func $f (type $1) (param $ref (ref $only-has-effects-in-not-addressable-function))
-  ;; CHECK-NEXT:  (call $calls-type-with-effects-but-not-addressable
-  ;; CHECK-NEXT:   (local.get $ref)
-  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (nop)
   ;; CHECK-NEXT: )
   (func $f (param $ref (ref $only-has-effects-in-not-addressable-function))
     ;; The type $has-effects-but-not-exported doesn't have an address because
     ;; it's not exported and it's never the target of a ref.func.
-    ;; We should be able to determine that $ref can only point to $nop.
-    ;; TODO: Only aggregate effects from functions that are addressed.
+    ;; So the call_ref has no potential targets and thus no effects.
     (call $calls-type-with-effects-but-not-addressable (local.get $ref))
   )
 )