[Attributor] Create a call site position for AACalledges

[kuter]

Summary

This patch adds a call site position for AACallEdges, this
allows us to ask questions about which functions a specific
CallBase might call.

[kuter]

Small fix with change tracking.

[jdoerfert]

LG

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

struct AANoUndefCallSiteReturned final
     : AACallSiteReturnedFromReturned<AANoUndef, AANoUndefImpl> {
   AANoUndefCallSiteReturned(const IRPosition &IRP, Attributor &A)
       : AACallSiteReturnedFromReturned<AANoUndef, AANoUndefImpl>(IRP, A) {}
 
   /// See AbstractAttribute::trackStatistics()
   void trackStatistics() const override { STATS_DECLTRACK_CSRET_ATTR(noundef) }
 };
 
-struct AACallEdgesFunction : public AACallEdges {
-  AACallEdgesFunction(const IRPosition &IRP, Attributor &A)
-      : AACallEdges(IRP, A) {}
+struct AACallEdgesImpl : public AACallEdges {
+  AACallEdgesImpl(const IRPosition &IRP, Attributor &A) : AACallEdges(IRP, A) {}
 
-  /// See AbstractAttribute::updateImpl(...).
-  ChangeStatus updateImpl(Attributor &A) override {
-    ChangeStatus Change = ChangeStatus::UNCHANGED;
-    bool OldHasUnknownCallee = HasUnknownCallee;
-    bool OldHasUnknownCalleeNonAsm = HasUnknownCalleeNonAsm;
+  virtual const SetVector<Function *> &getOptimisticEdges() const override {
+    return CalledFunctions;
+  }
+
+  virtual bool hasUnknownCallee() const override { return HasUnknownCallee; }
+
+  virtual bool hasNonAsmUnknownCallee() const override {
+    return HasUnknownCalleeNonAsm;
+  }
 
-    auto AddCalledFunction = [&](Function *Fn) {
-      if (CalledFunctions.insert(Fn)) {
-        Change = ChangeStatus::CHANGED;
-        LLVM_DEBUG(dbgs() << "[AACallEdges] New call edge: " << Fn->getName()
-                          << "\n");
-      }
-    };
+  const std::string getAsStr() const override {
+    return "CallEdges[" + std::to_string(HasUnknownCallee) + "," +
+           std::to_string(CalledFunctions.size()) + "]";
+  }
+
+  void trackStatistics() const override {}
+
+protected:
+  ChangeStatus addCalledFunction(Function *Fn) {
+    ChangeStatus Change = ChangeStatus::UNCHANGED;
+    return Change;
+  }
+
+  ChangeStatus setHasUnknownCallee(bool NonAsm) {
+    ChangeStatus Change = ChangeStatus::UNCHANGED;
+    if (!HasUnknownCallee)
+      Change = ChangeStatus::CHANGED;
+    if (NonAsm && !HasUnknownCalleeNonAsm)
+      Change = ChangeStatus::CHANGED;
+    HasUnknownCalleeNonAsm |= NonAsm;
+    HasUnknownCallee = true;
+
+    return Change;
+  }
+
+private:
+  /// Optimistic set of functions that might be called by this position.
+  SetVector<Function *> CalledFunctions;
+
+  /// Is there any call with a unknown callee.
+  bool HasUnknownCallee = false;
+
+  /// Is there any call with a unknown callee, excluding any inline asm.
+  bool HasUnknownCalleeNonAsm = false;
 
+struct AACallEdgesCallSite : public AACallEdgesImpl {
+  AACallEdgesCallSite(const IRPosition &IRP, Attributor &A)
+      : AACallEdgesImpl(IRP, A) {}
+  /// See AbstractAttribute::updateImpl(...).
+  ChangeStatus updateImpl(Attributor &A) override {
+    ChangeStatus Change = ChangeStatus::UNCHANGED;
+
     auto VisitValue = [&](Value &V, const Instruction *CtxI, bool &HasUnknown,
                           bool Stripped) -> bool {
       if (Function *Fn = dyn_cast<Function>(&V)) {
-        AddCalledFunction(Fn);
+        addCalledFunction(Fn);
       } else {
         LLVM_DEBUG(dbgs() << "[AACallEdges] Unrecognized value: " << V << "\n");
-        HasUnknown = true;
-        HasUnknownCalleeNonAsm = true;
+        setHasUnknownCallee(true);
       }
 
       // Explore all values.
       return true;
     };
 
     // Process any value that we might call.
-    auto ProcessCalledOperand = [&](Value *V, Instruction *Ctx) {
-      if (!genericValueTraversal<bool>(
-              A, IRPosition::value(*V), *this, HasUnknownCallee, VisitValue,
-              nullptr, false)) {
+    auto ProcessCalledOperand = [&](Value *V) {
+      bool DummyValue = false;
+      if (!genericValueTraversal<bool>(A, IRPosition::value(*V), *this,
+                                       DummyValue, VisitValue, nullptr,
+                                       false)) {
         // If we haven't gone through all values, assume that there are unknown
         // callees.
-        HasUnknownCallee = true;
-        HasUnknownCalleeNonAsm = true;
+        Change = Change | setHasUnknownCallee(true);
       }
     };
 
-    auto ProcessCallInst = [&](Instruction &Inst) {
-      CallBase &CB = static_cast<CallBase &>(Inst);
-      if (CB.isInlineAsm()) {
-        HasUnknownCallee = true;
-        return true;
-      }
+    CallBase *CB = static_cast<CallBase *>(getCtxI());
+
+    if (CB->isInlineAsm()) {
+      Change = Change | setHasUnknownCallee(false);
+      return Change;
 
-      // Process callee metadata if available.
-      if (auto *MD = Inst.getMetadata(LLVMContext::MD_callees)) {
-        for (auto &Op : MD->operands()) {
-          Function *Callee = mdconst::extract_or_null<Function>(Op);
-          if (Callee)
-            AddCalledFunction(Callee);
-        }
-        // Callees metadata grantees that the called function is one of its
-        // operands, So we are done.
-        return true;
-      }
+    if (auto *MD = getCtxI()->getMetadata(LLVMContext::MD_callees)) {
+          Change = Change | addCalledFunction(Callee);
+      return Change;
 
-      // The most simple case.
-      ProcessCalledOperand(CB.getCalledOperand(), &Inst);
+    ProcessCalledOperand(CB->getCalledOperand());
 
-      // Process callback functions.
-      SmallVector<const Use *, 4u> CallbackUses;
-      AbstractCallSite::getCallbackUses(CB, CallbackUses);
-      for (const Use *U : CallbackUses)
-        ProcessCalledOperand(U->get(), &Inst);
+    AbstractCallSite::getCallbackUses(*CB, CallbackUses);
+      ProcessCalledOperand(U->get());
+
+    return Change;
+  }
+};
+
+struct AACallEdgesFunction : public AACallEdgesImpl {
+  AACallEdgesFunction(const IRPosition &IRP, Attributor &A)
+      : AACallEdgesImpl(IRP, A) {}
+
+  /// See AbstractAttribute::updateImpl(...).
+  ChangeStatus updateImpl(Attributor &A) override {
+    ChangeStatus Change = ChangeStatus::UNCHANGED;
+
+    auto ProcessCallInst = [&](Instruction &Inst) {
+      CallBase &CB = static_cast<CallBase &>(Inst);
+
+      auto &CBEdges = A.getAAFor<AACallEdges>(
+          *this, IRPosition::callsite_function(CB), DepClassTy::REQUIRED);
+      if (CBEdges.hasNonAsmUnknownCallee())
+        Change = Change | setHasUnknownCallee(false);
+      if (CBEdges.hasUnknownCallee())
+        Change = Change | setHasUnknownCallee(true);
+
+      for (Function *F : CBEdges.getOptimisticEdges())
+        Change = Change | addCalledFunction(F);
 
       return true;
     };
 
     // Visit all callable instructions.
     bool UsedAssumedInformation = false;
     if (!A.checkForAllCallLikeInstructions(ProcessCallInst, *this,
                                            UsedAssumedInformation)) {
       // If we haven't looked at all call like instructions, assume that there
       // are unknown callees.
-      HasUnknownCallee = true;
-      HasUnknownCalleeNonAsm = true;
+      setHasUnknownCallee(true);
     }
 
-    // Track changes.
-    if (OldHasUnknownCallee != HasUnknownCallee ||
-        OldHasUnknownCalleeNonAsm != HasUnknownCalleeNonAsm)
-      Change = ChangeStatus::CHANGED;
-
     return Change;
   }
-
-  virtual const SetVector<Function *> &getOptimisticEdges() const override {
-    return CalledFunctions;
-  };
-
-  virtual bool hasUnknownCallee() const override { return HasUnknownCallee; }
-
-  virtual bool hasNonAsmUnknownCallee() const override {
-    return HasUnknownCalleeNonAsm;
-  }
-
-  const std::string getAsStr() const override {
-    return "CallEdges[" + std::to_string(HasUnknownCallee) + "," +
-           std::to_string(CalledFunctions.size()) + "]";
-  }
-
-  void trackStatistics() const override {}
-
-  /// Optimistic set of functions that might be called by this function.
-  SetVector<Function *> CalledFunctions;
-
-  /// Is there any call with a unknown callee.
-  bool HasUnknownCallee = false;
-
-  /// Is there any call with a unknown callee, excluding any inline asm.
-  bool HasUnknownCalleeNonAsm = false;
 };
 
 struct AAFunctionReachabilityFunction : public AAFunctionReachability {
   AAFunctionReachabilityFunction(const IRPosition &IRP, Attributor &A)
       : AAFunctionReachability(IRP, A) {}
 
   bool canReach(Attributor &A, Function *Fn) const override {
     // Assume that we can reach any function if we can reach a call with

 
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoUnwind)
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoSync)
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoRecurse)
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAWillReturn)
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoReturn)
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAReturnedValues)
 CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAMemoryLocation)
+CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AACallEdges)
 
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANonNull)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoAlias)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAPrivatizablePtr)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AADereferenceable)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAAlign)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoCapture)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAValueConstantRange)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAPotentialValues)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoUndef)
 
 CREATE_ALL_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAValueSimplify)
 CREATE_ALL_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAIsDead)
 CREATE_ALL_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoFree)
 
 CREATE_FUNCTION_ONLY_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAHeapToStack)
 CREATE_FUNCTION_ONLY_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAReachability)
 CREATE_FUNCTION_ONLY_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAUndefinedBehavior)
-CREATE_FUNCTION_ONLY_ABSTRACT_ATTRIBUTE_FOR_POSITION(AACallEdges)
 CREATE_FUNCTION_ONLY_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAFunctionReachability)
 
 CREATE_NON_RET_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAMemoryBehavior)
 
 #undef CREATE_FUNCTION_ONLY_ABSTRACT_ATTRIBUTE_FOR_POSITION
 #undef CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION
 #undef CREATE_NON_RET_ABSTRACT_ATTRIBUTE_FOR_POSITION
 #undef CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION
 #undef CREATE_ALL_ABSTRACT_ATTRIBUTE_FOR_POSITION
 #undef SWITCH_PK_CREATE
 #undef SWITCH_PK_INV