Cherrypick SCC with indirect effects

Author: stevenfontanella

src/passes/GlobalEffects.cpp

@@ -21,6 +21,7 @@
 
 #include "ir/effects.h"
 #include "ir/module-utils.h"
+#include "ir/subtypes.h"
 #include "pass.h"
 #include "support/strongly_connected_components.h"
 #include "wasm.h"
@@ -39,6 +40,9 @@ struct FuncInfo {
 
   // Directly-called functions from this function.
   std::unordered_set<Name> calledFunctions;
+
+  // Types that are targets of indirect calls.
+  std::unordered_set<HeapType> indirectCalledTypes;
 };
 
 std::map<Function*, FuncInfo> analyzeFuncs(Module& module,
@@ -83,11 +87,21 @@ std::map<Function*, FuncInfo> analyzeFuncs(Module& module,
             if (auto* call = curr->dynCast<Call>()) {
               // Note the direct call.
               funcInfo.calledFunctions.insert(call->target);
+            } else if (effects.calls && options.closedWorld) {
+              HeapType type;
+              if (auto* callRef = curr->dynCast<CallRef>()) {
+                type = callRef->target->type.getHeapType();
+              } else if (auto* callIndirect = curr->dynCast<CallIndirect>()) {
+                // nullability doesn't matter here
+                // call_indirect is always inexact
+                type = callIndirect->heapType;
+              } else {
+                assert(false && "Unexpected type of call");
+              }
+
+              funcInfo.indirectCalledTypes.insert(type);
             } else if (effects.calls) {
-              // This is an indirect call of some sort, so we must assume the
-              // worst. To do so, clear the effects, which indicates nothing
-              // is known (so anything is possible).
-              // TODO: We could group effects by function type etc.
+              assert(!options.closedWorld);
               funcInfo.effects = UnknownEffects;
             } else {
               // No call here, but update throwing if we see it. (Only do so,
@@ -107,14 +121,45 @@ std::map<Function*, FuncInfo> analyzeFuncs(Module& module,
   return std::move(analysis.map);
 }
 
-std::unordered_map<Function*, std::unordered_set<Function*>>
-buildCallGraph(const Module& module,
-               const std::map<Function*, FuncInfo>& funcInfos) {
-  std::unordered_map<Function*, std::unordered_set<Function*>> callGraph;
-  for (const auto& [func, info] : funcInfos) {
-    for (Name callee : info.calledFunctions) {
-      callGraph[func].insert(module.getFunction(callee));
+using CallGraphNode = std::variant<Function*, HeapType>;
+
+// Build a call graph for indirect and direct calls.
+// key (callee) -> value (caller)
+// Name         -> Name     : callee is called directly by caller
+// Name         -> HeapType : callee is a potential target of a virtual call
+//                            with this HeapType
+// HeapType     -> Name     : callee is indirectly called by caller
+// HeapType     -> HeapType : callee is a subtype of caller If we're
+// running in an open world, we only include Name -> Name edges.
+std::unordered_map<CallGraphNode, std::unordered_set<CallGraphNode>>
+buildCallGraph(Module& module,
+               const std::map<Function*, FuncInfo>& funcInfos,
+               bool closedWorld) {
+  std::unordered_map<CallGraphNode, std::unordered_set<CallGraphNode>>
+    callGraph;
+  std::unordered_set<HeapType> allFunctionTypes;
+  for (const auto& [caller, callerInfo] : funcInfos) {
+    for (Name calleeFunction : callerInfo.calledFunctions) {
+      callGraph[caller].insert(module.getFunction(calleeFunction));
+    }
+    if (!closedWorld) {
+      continue;
+    }
+
+    allFunctionTypes.insert(caller->type.getHeapType());
+    for (HeapType calleeType : callerInfo.indirectCalledTypes) {
+      callGraph[caller].insert(calleeType);
+      allFunctionTypes.insert(calleeType);
     }
+    callGraph[caller->type.getHeapType()].insert(caller);
+  }
+
+  SubTypes subtypes(module);
+  for (HeapType type : allFunctionTypes) {
+    subtypes.iterSubTypes(type, [&callGraph, type](HeapType sub, auto _) {
+      callGraph[type].insert(sub);
+      return true;
+    });
   }
 
   return callGraph;
@@ -123,98 +168,108 @@ buildCallGraph(const Module& module,
 // Propagate effects from callees to callers transitively
 // e.g. if A -> B -> C (A calls B which calls C)
 // Then B inherits effects from C and A inherits effects from both B and C.
-//
-// Generate SCC for the call graph, then traverse it in reverse topological
-// order processing each callee before its callers. When traversing:
-// - Merge all of the effects of functions within the CC
-// - Also merge the (already computed) effects of each callee CC
-// - Add trap effects for potentially recursive call chains
 void propagateEffects(
-  const Module& module,
+  Module& module,
   const PassOptions& passOptions,
   std::map<Function*, FuncInfo>& funcInfos,
-  const std::unordered_map<Function*, std::unordered_set<Function*>>
+  const std::unordered_map<CallGraphNode, std::unordered_set<CallGraphNode>>&
     callGraph) {
+
   struct CallGraphSCCs
-    : SCCs<std::vector<Function*>::const_iterator, CallGraphSCCs> {
+    : SCCs<std::vector<CallGraphNode>::const_iterator, CallGraphSCCs> {
     const std::map<Function*, FuncInfo>& funcInfos;
-    const std::unordered_map<Function*, std::unordered_set<Function*>>&
-      callGraph;
     const Module& module;
+    const std::unordered_map<CallGraphNode, std::unordered_set<CallGraphNode>>&
+      callGraph;
 
     CallGraphSCCs(
-      const std::vector<Function*>& funcs,
+      const std::vector<CallGraphNode>& nodes,
       const std::map<Function*, FuncInfo>& funcInfos,
-      const std::unordered_map<Function*, std::unordered_set<Function*>>&
-        callGraph,
-      const Module& module)
-      : SCCs<std::vector<Function*>::const_iterator, CallGraphSCCs>(
-          funcs.begin(), funcs.end()),
-        funcInfos(funcInfos), callGraph(callGraph), module(module) {}
-
-    void pushChildren(Function* f) {
-      auto callees = callGraph.find(f);
+      Module& module,
+      const std::unordered_map<CallGraphNode,
+                               std::unordered_set<CallGraphNode>>& callGraph)
+      : SCCs<std::vector<CallGraphNode>::const_iterator, CallGraphSCCs>(
+          nodes.begin(), nodes.end()),
+        funcInfos(funcInfos), module(module), callGraph(callGraph) {}
+
+    void pushChildren(CallGraphNode node) {
+      auto callees = callGraph.find(node);
       if (callees == callGraph.end()) {
         return;
       }
-
-      for (auto* callee : callees->second) {
+      for (const auto& callee : callees->second) {
         push(callee);
       }
     }
   };
 
-  std::vector<Function*> allFuncs;
+  std::vector<CallGraphNode> funcs;
+  // We only care about Functions that are roots, not types
   for (auto& [func, info] : funcInfos) {
-    allFuncs.push_back(func);
+    funcs.push_back(func);
   }
-  CallGraphSCCs sccs(allFuncs, funcInfos, callGraph, module);
 
-  std::unordered_map<Function*, int> sccMembers;
+  CallGraphSCCs sccs(funcs, funcInfos, module, callGraph);
+
+  std::unordered_map<CallGraphNode, int> sccMembers;
   std::unordered_map<int, std::optional<EffectAnalyzer>> componentEffects;
 
   int ccIndex = 0;
   for (auto ccIterator : sccs) {
+    std::vector<CallGraphNode> cc(ccIterator.begin(), ccIterator.end());
     ccIndex++;
     std::optional<EffectAnalyzer>& ccEffects = componentEffects[ccIndex];
-    std::vector<Function*> ccFuncs(ccIterator.begin(), ccIterator.end());
-
     ccEffects.emplace(passOptions, module);
 
+    std::vector<Function*> ccFuncs;
+    std::vector<HeapType> ccTypes;
+    for (auto v : cc) {
+      if (auto** func = std::get_if<Function*>(&v)) {
+        ccFuncs.push_back(*func);
+      } else {
+        ccTypes.push_back(std::get<HeapType>(v));
+      }
+    }
+
     for (Function* f : ccFuncs) {
       sccMembers.emplace(f, ccIndex);
     }
+    for (HeapType t : ccTypes) {
+      sccMembers.emplace(t, ccIndex);
+    }
 
     std::unordered_set<int> calleeSccs;
-    for (Function* caller : ccFuncs) {
+    for (const auto& caller : cc) {
       auto callees = callGraph.find(caller);
-      if (callees == callGraph.end()) {
-        continue;
-      }
-      for (auto* callee : callees->second) {
-        calleeSccs.insert(sccMembers.at(callee));
+      if (callees != callGraph.end()) {
+        for (const auto& callee : callees->second) {
+          auto sccIt = sccMembers.find(callee);
+          if (sccIt != sccMembers.end()) {
+            calleeSccs.insert(sccIt->second);
+          }
+        }
       }
     }
 
-    // Merge in effects from callees
     for (int calleeScc : calleeSccs) {
       const auto& calleeComponentEffects = componentEffects.at(calleeScc);
       if (calleeComponentEffects == UnknownEffects) {
         ccEffects = UnknownEffects;
         break;
       }
 
-      else if (ccEffects != UnknownEffects) {
+      else if (ccEffects) {
         ccEffects->mergeIn(*calleeComponentEffects);
       }
     }
 
     // Add trap effects for potential cycles.
-    if (ccFuncs.size() > 1) {
+    if (cc.size() > 1) {
       if (ccEffects != UnknownEffects) {
         ccEffects->trap = true;
       }
-    } else {
+      // A cycle isn't possible for a CC that only contains a type
+    } else if (ccFuncs.size() == 1) {
       auto* func = ccFuncs[0];
       if (funcInfos.at(func).calledFunctions.contains(func->name)) {
         if (ccEffects != UnknownEffects) {
@@ -223,8 +278,7 @@ void propagateEffects(
       }
     }
 
-    // Aggregate effects within this CC
-    if (ccEffects) {
+    if (ccEffects)
       for (Function* f : ccFuncs) {
         const auto& effects = funcInfos.at(f).effects;
         if (effects == UnknownEffects) {
@@ -234,9 +288,7 @@ void propagateEffects(
 
         ccEffects->mergeIn(*effects);
       }
-    }
 
-    // Assign each function's effects to its CC effects.
     for (Function* f : ccFuncs) {
       if (!ccEffects) {
         funcInfos.at(f).effects = UnknownEffects;
@@ -247,27 +299,26 @@ void propagateEffects(
   }
 }
 
-void copyEffectsToFunctions(const std::map<Function*, FuncInfo> funcInfos) {
-  for (auto& [func, info] : funcInfos) {
-    func->effects.reset();
-    if (!info.effects) {
-      continue;
-    }
-
-    func->effects = std::make_shared<EffectAnalyzer>(*info.effects);
-  }
-}
-
 struct GenerateGlobalEffects : public Pass {
   void run(Module* module) override {
     std::map<Function*, FuncInfo> funcInfos =
       analyzeFuncs(*module, getPassOptions());
 
-    auto callGraph = buildCallGraph(*module, funcInfos);
+    auto callGraph =
+      buildCallGraph(*module, funcInfos, getPassOptions().closedWorld);
 
     propagateEffects(*module, getPassOptions(), funcInfos, callGraph);
 
-    copyEffectsToFunctions(funcInfos);
+    // Generate the final data, starting from a blank slate where nothing is
+    // known.
+    for (auto& [func, info] : funcInfos) {
+      func->effects.reset();
+      if (!info.effects) {
+        continue;
+      }
+
+      func->effects = std::make_shared<EffectAnalyzer>(*info.effects);
+    }
   }
 };