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.