[CaptureTracking][FunctionAttrs] Add support for CaptureInfo (#125880)

This extends CaptureTracking to support inferring non-trivial
CaptureInfos. The focus of this patch is to only support FunctionAttrs,
other users of CaptureTracking will be updated in followups.

The key API changes here are:

* DetermineUseCaptureKind() now returns a UseCaptureInfo where the UseCC
component specifies what is captured at that Use and the ResultCC
component specifies what may be captured via the return value of the
User. Usually only one or the other will be used (corresponding to
previous MAY_CAPTURE or PASSTHROUGH results), but both may be set for
call captures.
* The CaptureTracking::captures() extension point is passed this
UseCaptureInfo as well and then can decide what to do with it by
returning an Action, which is one of: Stop: stop traversal.
ContinueIgnoringReturn: continue traversal but don't follow the
instruction return value. Continue: continue traversal and follow the
instruction return value if it has additional CaptureComponents.

For now, this patch retains the (unsound) special logic for comparison
of null with a dereferenceable pointer. I'd like to switch key code to
take advantage of address/address_is_null before dropping it.

This PR mainly intends to introduce necessary API changes and basic
inference support, there are various possible improvements marked with
TODOs.

1 files changed, 74 insertions, 53 deletions

diff --git a/llvm/lib/Analysis/CaptureTracking.cpp b/llvm/lib/Analysis/CaptureTracking.cpp
index 49baf2e..5120b91 100644
--- a/llvm/lib/Analysis/CaptureTracking.cpp
+++ b/llvm/lib/Analysis/CaptureTracking.cpp
@@ -81,14 +81,15 @@ struct SimpleCaptureTracker : public CaptureTracker {
     Captured = true;
   }
 
-  bool captured(const Use *U) override {
+  Action captured(const Use *U, UseCaptureInfo CI) override {
+    // TODO(captures): Use UseCaptureInfo.
     if (isa<ReturnInst>(U->getUser()) && !ReturnCaptures)
-      return false;
+      return ContinueIgnoringReturn;
 
     LLVM_DEBUG(dbgs() << "Captured by: " << *U->getUser() << "\n");
 
     Captured = true;
-    return true;
+    return Stop;
   }
 
   bool ReturnCaptures;
@@ -122,19 +123,21 @@ struct CapturesBefore : public CaptureTracker {
     return !isPotentiallyReachable(I, BeforeHere, nullptr, DT, LI);
   }
 
-  bool captured(const Use *U) override {
+  Action captured(const Use *U, UseCaptureInfo CI) override {
+    // TODO(captures): Use UseCaptureInfo.
     Instruction *I = cast<Instruction>(U->getUser());
     if (isa<ReturnInst>(I) && !ReturnCaptures)
-      return false;
+      return ContinueIgnoringReturn;
 
     // Check isSafeToPrune() here rather than in shouldExplore() to avoid
     // an expensive reachability query for every instruction we look at.
     // Instead we only do one for actual capturing candidates.
     if (isSafeToPrune(I))
-      return false;
+      // If the use is not reachable, the instruction result isn't either.
+      return ContinueIgnoringReturn;
 
     Captured = true;
-    return true;
+    return Stop;
   }
 
   const Instruction *BeforeHere;
@@ -166,10 +169,11 @@ struct EarliestCaptures : public CaptureTracker {
     EarliestCapture = &*F.getEntryBlock().begin();
   }
 
-  bool captured(const Use *U) override {
+  Action captured(const Use *U, UseCaptureInfo CI) override {
+    // TODO(captures): Use UseCaptureInfo.
     Instruction *I = cast<Instruction>(U->getUser());
     if (isa<ReturnInst>(I) && !ReturnCaptures)
-      return false;
+      return ContinueIgnoringReturn;
 
     if (!EarliestCapture)
       EarliestCapture = I;
@@ -177,9 +181,10 @@ struct EarliestCaptures : public CaptureTracker {
       EarliestCapture = DT.findNearestCommonDominator(EarliestCapture, I);
     Captured = true;
 
-    // Return false to continue analysis; we need to see all potential
-    // captures.
-    return false;
+    // Continue analysis, as we need to see all potential captures. However,
+    // we do not need to follow the instruction result, as this use will
+    // dominate any captures made through the instruction result..
+    return ContinueIgnoringReturn;
   }
 
   Instruction *EarliestCapture = nullptr;
@@ -274,25 +279,26 @@ Instruction *llvm::FindEarliestCapture(const Value *V, Function &F,
   return CB.EarliestCapture;
 }
 
-UseCaptureKind llvm::DetermineUseCaptureKind(
-    const Use &U,
+UseCaptureInfo llvm::DetermineUseCaptureKind(
+    const Use &U, const Value *Base,
     function_ref<bool(Value *, const DataLayout &)> IsDereferenceableOrNull) {
   Instruction *I = dyn_cast<Instruction>(U.getUser());
 
   // TODO: Investigate non-instruction uses.
   if (!I)
-    return UseCaptureKind::MAY_CAPTURE;
+    return CaptureComponents::All;
 
   switch (I->getOpcode()) {
   case Instruction::Call:
   case Instruction::Invoke: {
+    // TODO(captures): Make this more precise.
     auto *Call = cast<CallBase>(I);
     // Not captured if the callee is readonly, doesn't return a copy through
     // its return value and doesn't unwind (a readonly function can leak bits
     // by throwing an exception or not depending on the input value).
     if (Call->onlyReadsMemory() && Call->doesNotThrow() &&
         Call->getType()->isVoidTy())
-      return UseCaptureKind::NO_CAPTURE;
+      return CaptureComponents::None;
 
     // The pointer is not captured if returned pointer is not captured.
     // NOTE: CaptureTracking users should not assume that only functions
@@ -300,13 +306,13 @@ UseCaptureKind llvm::DetermineUseCaptureKind(
     // getUnderlyingObject in ValueTracking or DecomposeGEPExpression
     // in BasicAA also need to know about this property.
     if (isIntrinsicReturningPointerAliasingArgumentWithoutCapturing(Call, true))
-      return UseCaptureKind::PASSTHROUGH;
+      return UseCaptureInfo::passthrough();
 
     // Volatile operations effectively capture the memory location that they
     // load and store to.
     if (auto *MI = dyn_cast<MemIntrinsic>(Call))
       if (MI->isVolatile())
-        return UseCaptureKind::MAY_CAPTURE;
+        return CaptureComponents::All;
 
     // Calling a function pointer does not in itself cause the pointer to
     // be captured.  This is a subtle point considering that (for example)
@@ -315,30 +321,27 @@ UseCaptureKind llvm::DetermineUseCaptureKind(
     // captured, even though the loaded value might be the pointer itself
     // (think of self-referential objects).
     if (Call->isCallee(&U))
-      return UseCaptureKind::NO_CAPTURE;
+      return CaptureComponents::None;
 
     // Not captured if only passed via 'nocapture' arguments.
     assert(Call->isDataOperand(&U) && "Non-callee must be data operand");
-    if (!Call->doesNotCapture(Call->getDataOperandNo(&U))) {
-      // The parameter is not marked 'nocapture' - captured.
-      return UseCaptureKind::MAY_CAPTURE;
-    }
-    return UseCaptureKind::NO_CAPTURE;
+    CaptureInfo CI = Call->getCaptureInfo(Call->getDataOperandNo(&U));
+    return UseCaptureInfo(CI.getOtherComponents(), CI.getRetComponents());
   }
   case Instruction::Load:
     // Volatile loads make the address observable.
     if (cast<LoadInst>(I)->isVolatile())
-      return UseCaptureKind::MAY_CAPTURE;
-    return UseCaptureKind::NO_CAPTURE;
+      return CaptureComponents::All;
+    return CaptureComponents::None;
   case Instruction::VAArg:
     // "va-arg" from a pointer does not cause it to be captured.
-    return UseCaptureKind::NO_CAPTURE;
+    return CaptureComponents::None;
   case Instruction::Store:
     // Stored the pointer - conservatively assume it may be captured.
     // Volatile stores make the address observable.
     if (U.getOperandNo() == 0 || cast<StoreInst>(I)->isVolatile())
-      return UseCaptureKind::MAY_CAPTURE;
-    return UseCaptureKind::NO_CAPTURE;
+      return CaptureComponents::All;
+    return CaptureComponents::None;
   case Instruction::AtomicRMW: {
     // atomicrmw conceptually includes both a load and store from
     // the same location.
@@ -347,8 +350,8 @@ UseCaptureKind llvm::DetermineUseCaptureKind(
     // Volatile stores make the address observable.
     auto *ARMWI = cast<AtomicRMWInst>(I);
     if (U.getOperandNo() == 1 || ARMWI->isVolatile())
-      return UseCaptureKind::MAY_CAPTURE;
-    return UseCaptureKind::NO_CAPTURE;
+      return CaptureComponents::All;
+    return CaptureComponents::None;
   }
   case Instruction::AtomicCmpXchg: {
     // cmpxchg conceptually includes both a load and store from
@@ -358,31 +361,35 @@ UseCaptureKind llvm::DetermineUseCaptureKind(
     // Volatile stores make the address observable.
     auto *ACXI = cast<AtomicCmpXchgInst>(I);
     if (U.getOperandNo() == 1 || U.getOperandNo() == 2 || ACXI->isVolatile())
-      return UseCaptureKind::MAY_CAPTURE;
-    return UseCaptureKind::NO_CAPTURE;
+      return CaptureComponents::All;
+    return CaptureComponents::None;
   }
   case Instruction::GetElementPtr:
     // AA does not support pointers of vectors, so GEP vector splats need to
     // be considered as captures.
     if (I->getType()->isVectorTy())
-      return UseCaptureKind::MAY_CAPTURE;
-    return UseCaptureKind::PASSTHROUGH;
+      return CaptureComponents::All;
+    return UseCaptureInfo::passthrough();
   case Instruction::BitCast:
   case Instruction::PHI:
   case Instruction::Select:
   case Instruction::AddrSpaceCast:
     // The original value is not captured via this if the new value isn't.
-    return UseCaptureKind::PASSTHROUGH;
+    return UseCaptureInfo::passthrough();
   case Instruction::ICmp: {
     unsigned Idx = U.getOperandNo();
     unsigned OtherIdx = 1 - Idx;
-    if (auto *CPN = dyn_cast<ConstantPointerNull>(I->getOperand(OtherIdx))) {
+    if (isa<ConstantPointerNull>(I->getOperand(OtherIdx)) &&
+        cast<ICmpInst>(I)->isEquality()) {
+      // TODO(captures): Remove these special cases once we make use of
+      // captures(address_is_null).
+
       // Don't count comparisons of a no-alias return value against null as
       // captures. This allows us to ignore comparisons of malloc results
       // with null, for example.
-      if (CPN->getType()->getAddressSpace() == 0)
+      if (U->getType()->getPointerAddressSpace() == 0)
         if (isNoAliasCall(U.get()->stripPointerCasts()))
-          return UseCaptureKind::NO_CAPTURE;
+          return CaptureComponents::None;
       if (!I->getFunction()->nullPointerIsDefined()) {
         auto *O = I->getOperand(Idx)->stripPointerCastsSameRepresentation();
         // Comparing a dereferenceable_or_null pointer against null cannot
@@ -390,17 +397,23 @@ UseCaptureKind llvm::DetermineUseCaptureKind(
         // valid (in-bounds) pointer.
         const DataLayout &DL = I->getDataLayout();
         if (IsDereferenceableOrNull && IsDereferenceableOrNull(O, DL))
-          return UseCaptureKind::NO_CAPTURE;
+          return CaptureComponents::None;
       }
+
+      // Check whether this is a comparison of the base pointer against
+      // null.
+      if (U.get() == Base)
+        return CaptureComponents::AddressIsNull;
     }
 
     // Otherwise, be conservative. There are crazy ways to capture pointers
-    // using comparisons.
-    return UseCaptureKind::MAY_CAPTURE;
+    // using comparisons. However, only the address is captured, not the
+    // provenance.
+    return CaptureComponents::Address;
   }
   default:
     // Something else - be conservative and say it is captured.
-    return UseCaptureKind::MAY_CAPTURE;
+    return CaptureComponents::All;
   }
 }
 
@@ -438,18 +451,26 @@ void llvm::PointerMayBeCaptured(const Value *V, CaptureTracker *Tracker,
   };
   while (!Worklist.empty()) {
     const Use *U = Worklist.pop_back_val();
-    switch (DetermineUseCaptureKind(*U, IsDereferenceableOrNull)) {
-    case UseCaptureKind::NO_CAPTURE:
-      continue;
-    case UseCaptureKind::MAY_CAPTURE:
-      if (Tracker->captured(U))
-        return;
-      continue;
-    case UseCaptureKind::PASSTHROUGH:
-      if (!AddUses(U->getUser()))
+    UseCaptureInfo CI = DetermineUseCaptureKind(*U, V, IsDereferenceableOrNull);
+    if (capturesAnything(CI.UseCC)) {
+      switch (Tracker->captured(U, CI)) {
+      case CaptureTracker::Stop:
         return;
-      continue;
+      case CaptureTracker::ContinueIgnoringReturn:
+        continue;
+      case CaptureTracker::Continue:
+        // Fall through to passthrough handling, but only if ResultCC contains
+        // additional components that UseCC does not. We assume that a
+        // capture at this point will be strictly more constraining than a
+        // later capture from following the return value.
+        if (capturesNothing(CI.ResultCC & ~CI.UseCC))
+          continue;
+        break;
+      }
     }
+    // TODO(captures): We could keep track of ResultCC for the users.
+    if (capturesAnything(CI.ResultCC) && !AddUses(U->getUser()))
+      return;
   }
 
   // All uses examined.