1 files changed, 86 insertions, 7 deletions

diff --git a/llvm/lib/Analysis/CaptureTracking.cpp b/llvm/lib/Analysis/CaptureTracking.cpp
index 5fe4f9befc86..8955658cb9e7 100644
--- a/llvm/lib/Analysis/CaptureTracking.cpp
+++ b/llvm/lib/Analysis/CaptureTracking.cpp
@@ -98,10 +98,10 @@ namespace {
   /// as the given instruction and the use.
   struct CapturesBefore : public CaptureTracker {
 
-    CapturesBefore(bool ReturnCaptures, const Instruction *I, const DominatorTree *DT,
-                   bool IncludeI)
-      : BeforeHere(I), DT(DT),
-        ReturnCaptures(ReturnCaptures), IncludeI(IncludeI), Captured(false) {}
+    CapturesBefore(bool ReturnCaptures, const Instruction *I,
+                   const DominatorTree *DT, bool IncludeI, const LoopInfo *LI)
+        : BeforeHere(I), DT(DT), ReturnCaptures(ReturnCaptures),
+          IncludeI(IncludeI), Captured(false), LI(LI) {}
 
     void tooManyUses() override { Captured = true; }
 
@@ -115,7 +115,7 @@ namespace {
         return true;
 
       // Check whether there is a path from I to BeforeHere.
-      return !isPotentiallyReachable(I, BeforeHere, nullptr, DT);
+      return !isPotentiallyReachable(I, BeforeHere, nullptr, DT, LI);
     }
 
     bool captured(const Use *U) override {
@@ -140,6 +140,68 @@ namespace {
     bool IncludeI;
 
     bool Captured;
+
+    const LoopInfo *LI;
+  };
+
+  /// Find the 'earliest' instruction before which the pointer is known not to
+  /// be captured. Here an instruction A is considered earlier than instruction
+  /// B, if A dominates B. If 2 escapes do not dominate each other, the
+  /// terminator of the common dominator is chosen. If not all uses cannot be
+  /// analyzed, the earliest escape is set to the first instruction in the
+  /// function entry block.
+  // NOTE: Users have to make sure instructions compared against the earliest
+  // escape are not in a cycle.
+  struct EarliestCaptures : public CaptureTracker {
+
+    EarliestCaptures(bool ReturnCaptures, Function &F, const DominatorTree &DT)
+        : DT(DT), ReturnCaptures(ReturnCaptures), Captured(false), F(F) {}
+
+    void tooManyUses() override {
+      Captured = true;
+      EarliestCapture = &*F.getEntryBlock().begin();
+    }
+
+    bool captured(const Use *U) override {
+      Instruction *I = cast<Instruction>(U->getUser());
+      if (isa<ReturnInst>(I) && !ReturnCaptures)
+        return false;
+
+      if (!EarliestCapture) {
+        EarliestCapture = I;
+      } else if (EarliestCapture->getParent() == I->getParent()) {
+        if (I->comesBefore(EarliestCapture))
+          EarliestCapture = I;
+      } else {
+        BasicBlock *CurrentBB = I->getParent();
+        BasicBlock *EarliestBB = EarliestCapture->getParent();
+        if (DT.dominates(EarliestBB, CurrentBB)) {
+          // EarliestCapture already comes before the current use.
+        } else if (DT.dominates(CurrentBB, EarliestBB)) {
+          EarliestCapture = I;
+        } else {
+          // Otherwise find the nearest common dominator and use its terminator.
+          auto *NearestCommonDom =
+              DT.findNearestCommonDominator(CurrentBB, EarliestBB);
+          EarliestCapture = NearestCommonDom->getTerminator();
+        }
+      }
+      Captured = true;
+
+      // Return false to continue analysis; we need to see all potential
+      // captures.
+      return false;
+    }
+
+    Instruction *EarliestCapture = nullptr;
+
+    const DominatorTree &DT;
+
+    bool ReturnCaptures;
+
+    bool Captured;
+
+    Function &F;
   };
 }
 
@@ -183,7 +245,8 @@ bool llvm::PointerMayBeCaptured(const Value *V,
 bool llvm::PointerMayBeCapturedBefore(const Value *V, bool ReturnCaptures,
                                       bool StoreCaptures, const Instruction *I,
                                       const DominatorTree *DT, bool IncludeI,
-                                      unsigned MaxUsesToExplore) {
+                                      unsigned MaxUsesToExplore,
+                                      const LoopInfo *LI) {
   assert(!isa<GlobalValue>(V) &&
          "It doesn't make sense to ask whether a global is captured.");
 
@@ -194,7 +257,7 @@ bool llvm::PointerMayBeCapturedBefore(const Value *V, bool ReturnCaptures,
   // TODO: See comment in PointerMayBeCaptured regarding what could be done
   // with StoreCaptures.
 
-  CapturesBefore CB(ReturnCaptures, I, DT, IncludeI);
+  CapturesBefore CB(ReturnCaptures, I, DT, IncludeI, LI);
   PointerMayBeCaptured(V, &CB, MaxUsesToExplore);
   if (CB.Captured)
     ++NumCapturedBefore;
@@ -203,6 +266,22 @@ bool llvm::PointerMayBeCapturedBefore(const Value *V, bool ReturnCaptures,
   return CB.Captured;
 }
 
+Instruction *llvm::FindEarliestCapture(const Value *V, Function &F,
+                                       bool ReturnCaptures, bool StoreCaptures,
+                                       const DominatorTree &DT,
+                                       unsigned MaxUsesToExplore) {
+  assert(!isa<GlobalValue>(V) &&
+         "It doesn't make sense to ask whether a global is captured.");
+
+  EarliestCaptures CB(ReturnCaptures, F, DT);
+  PointerMayBeCaptured(V, &CB, MaxUsesToExplore);
+  if (CB.Captured)
+    ++NumCapturedBefore;
+  else
+    ++NumNotCapturedBefore;
+  return CB.EarliestCapture;
+}
+
 void llvm::PointerMayBeCaptured(const Value *V, CaptureTracker *Tracker,
                                 unsigned MaxUsesToExplore) {
   assert(V->getType()->isPointerTy() && "Capture is for pointers only!");