[CaptureTracking][NFCI] Expose capture tracking logic

[jdoerfert]

Summary

The logic exposed by this patch via llvm::DetermineUseCaptureKind was
part of llvm::PointerMayBeCaptured. In the Attributor we want to keep
track of the work list items but still reuse the logic if a use might
capture a value. A follow up for the Attributor removes ~100 lines of
code and complexity while making future handling of simplified values
possible.

[nikic]

LGTM

llvm/include/llvm/Analysis/CaptureTracking.h

119

"the user" -> "the transitive uses"? We're inspecting the user either way.

[sidorovd]

llvm/include/llvm/Analysis/CaptureTracking.h

113

Hi, while PASSTHROUGH seems to be a good name for the enum value, it so happened, that it is also a macro defined in wingdi.h (which is included in windows.h). Can we consider renaming it or we should guard/undef the macro in the downstream?

[jdoerfert]

llvm/include/llvm/Analysis/CaptureTracking.h

113

Rename it here and in all use sites. I guess we could add UCK_ ad prefix, or make it something like USER_MAY_CAPTURE.

llvm/include/llvm/Analysis/CaptureTracking.h

 //===----- llvm/Analysis/CaptureTracking.h - Pointer capture ----*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains routines that help determine which pointers are captured.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_CAPTURETRACKING_H
 #define LLVM_ANALYSIS_CAPTURETRACKING_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLFunctionalExtras.h"
 
 namespace llvm {
 
   class Value;
   class Use;
   class DataLayout;
   class Instruction;
   class DominatorTree;
struct CaptureTracker {
     virtual bool captured(const Use *U) = 0;
 
     /// isDereferenceableOrNull - Overload to allow clients with additional
     /// knowledge about pointer dereferenceability to provide it and thereby
     /// avoid conservative responses when a pointer is compared to null.
     virtual bool isDereferenceableOrNull(Value *O, const DataLayout &DL);
   };
 
+  /// Types of use capture kinds, see \p DetermineUseCaptureKind.
+  enum class UseCaptureKind {
+    NO_CAPTURE,
+    MAY_CAPTURE,
+    PASSTHROUGH,
+  };
+
+  /// Determine what kind of capture behaviour \p U may exhibit.
+  ///
+  /// A use can be no-capture, a use can potentially capture, or a use can be
+  /// passthrough such that the uses of the user or \p U should be inspected.
+  /// The \p IsDereferenceableOrNull callback is used to rule out capturing for
+  /// certain comparisons.
+  UseCaptureKind
+  DetermineUseCaptureKind(const Use &U,
+                          llvm::function_ref<bool(Value *, const DataLayout &)>
+                              IsDereferenceableOrNull);
+
   /// PointerMayBeCaptured - Visit the value and the values derived from it and
   /// find values which appear to be capturing the pointer value. This feeds
   /// results into and is controlled by the CaptureTracker object.
   /// MaxUsesToExplore specifies how many uses the analysis should explore for
   /// one value before giving up due too "too many uses". If MaxUsesToExplore
   /// is zero, a default value is assumed.
   void PointerMayBeCaptured(const Value *V, CaptureTracker *Tracker,
                             unsigned MaxUsesToExplore = 0);

llvm/lib/Analysis/CaptureTracking.cpp

Instruction *llvm::FindEarliestCapture(const Value *V, Function &F,
   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!");
-  if (MaxUsesToExplore == 0)
-    MaxUsesToExplore = DefaultMaxUsesToExplore;
-
-  SmallVector<const Use *, 20> Worklist;
-  Worklist.reserve(getDefaultMaxUsesToExploreForCaptureTracking());
-  SmallSet<const Use *, 20> Visited;
-
-  auto AddUses = [&](const Value *V) {
-    unsigned Count = 0;
-    for (const Use &U : V->uses()) {
-      // If there are lots of uses, conservatively say that the value
-      // is captured to avoid taking too much compile time.
-      if (Count++ >= MaxUsesToExplore) {
-        Tracker->tooManyUses();
-        return false;
-      }
-      if (!Visited.insert(&U).second)
-        continue;
-      if (!Tracker->shouldExplore(&U))
-        continue;
-      Worklist.push_back(&U);
-    }
-    return true;
-  };
-  if (!AddUses(V))
-    return;
-
-  while (!Worklist.empty()) {
-    const Use *U = Worklist.pop_back_val();
-    Instruction *I = cast<Instruction>(U->getUser());
+UseCaptureKind llvm::DetermineUseCaptureKind(
+    const Use &U,
+    function_ref<bool(Value *, const DataLayout &)> IsDereferenceableOrNull) {
+  Instruction *I = cast<Instruction>(U.getUser());
 
-    switch (I->getOpcode()) {
-    case Instruction::Call:
-    case Instruction::Invoke: {
-      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())
-        break;
+      return UseCaptureKind::NO_CAPTURE;
 
-      // The pointer is not captured if returned pointer is not captured.
-      // NOTE: CaptureTracking users should not assume that only functions
-      // marked with nocapture do not capture. This means that places like
-      // getUnderlyingObject in ValueTracking or DecomposeGEPExpression
-      // in BasicAA also need to know about this property.
-      if (isIntrinsicReturningPointerAliasingArgumentWithoutCapturing(Call,
-                                                                      true)) {
-        if (!AddUses(Call))
-          return;
-        break;
-      }
+    if (isIntrinsicReturningPointerAliasingArgumentWithoutCapturing(Call, true))
+      return UseCaptureKind::PASSTHROUGH;
 
-      // Volatile operations effectively capture the memory location that they
-      // load and store to.
-      if (auto *MI = dyn_cast<MemIntrinsic>(Call))
-        if (MI->isVolatile())
-          if (Tracker->captured(U))
-            return;
+        return UseCaptureKind::MAY_CAPTURE;
 
-      // Calling a function pointer does not in itself cause the pointer to
-      // be captured.  This is a subtle point considering that (for example)
-      // the callee might return its own address.  It is analogous to saying
-      // that loading a value from a pointer does not cause the pointer to be
-      // captured, even though the loaded value might be the pointer itself
-      // (think of self-referential objects).
-      if (Call->isCallee(U))
-        break;
+    if (Call->isCallee(&U))
+      return UseCaptureKind::NO_CAPTURE;
 
-      // Not captured if only passed via 'nocapture' arguments.
-      if (Call->isDataOperand(U) &&
-          !Call->doesNotCapture(Call->getDataOperandNo(U))) {
-        // The parameter is not marked 'nocapture' - captured.
-        if (Tracker->captured(U))
-          return;
-      }
-      break;
-    }
-    case Instruction::Load:
-      // Volatile loads make the address observable.
-      if (cast<LoadInst>(I)->isVolatile())
-        if (Tracker->captured(U))
-          return;
-      break;
-    case Instruction::VAArg:
-      // "va-arg" from a pointer does not cause it to be captured.
-      break;
-    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())
-        if (Tracker->captured(U))
-          return;
-      break;
-    case Instruction::AtomicRMW: {
-      // atomicrmw conceptually includes both a load and store from
-      // the same location.
-      // As with a store, the location being accessed is not captured,
-      // but the value being stored is.
-      // Volatile stores make the address observable.
-      auto *ARMWI = cast<AtomicRMWInst>(I);
-      if (U->getOperandNo() == 1 || ARMWI->isVolatile())
-        if (Tracker->captured(U))
-          return;
-      break;
-    }
-    case Instruction::AtomicCmpXchg: {
-      // cmpxchg conceptually includes both a load and store from
-      // the same location.
-      // As with a store, the location being accessed is not captured,
-      // but the value being stored is.
-      // Volatile stores make the address observable.
-      auto *ACXI = cast<AtomicCmpXchgInst>(I);
-      if (U->getOperandNo() == 1 || U->getOperandNo() == 2 ||
-          ACXI->isVolatile())
-        if (Tracker->captured(U))
-          return;
-      break;
-    }
-    case Instruction::BitCast:
-    case Instruction::GetElementPtr:
-    case Instruction::PHI:
-    case Instruction::Select:
-    case Instruction::AddrSpaceCast:
-      // The original value is not captured via this if the new value isn't.
-      if (!AddUses(I))
-        return;
-      break;
-    case Instruction::ICmp: {
-      unsigned Idx = U->getOperandNo();
-      unsigned OtherIdx = 1 - Idx;
-      if (auto *CPN = dyn_cast<ConstantPointerNull>(I->getOperand(OtherIdx))) {
-        // 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 (isNoAliasCall(U->get()->stripPointerCasts()))
-            break;
-        if (!I->getFunction()->nullPointerIsDefined()) {
-          auto *O = I->getOperand(Idx)->stripPointerCastsSameRepresentation();
-          // Comparing a dereferenceable_or_null pointer against null cannot
-          // lead to pointer escapes, because if it is not null it must be a
-          // valid (in-bounds) pointer.
-          if (Tracker->isDereferenceableOrNull(O, I->getModule()->getDataLayout()))
-            break;
-        }
-      }
-      // Comparison against value stored in global variable. Given the pointer
-      // does not escape, its value cannot be guessed and stored separately in a
-      // global variable.
-      auto *LI = dyn_cast<LoadInst>(I->getOperand(OtherIdx));
-      if (LI && isa<GlobalVariable>(LI->getPointerOperand()))
-        break;
-      // Otherwise, be conservative. There are crazy ways to capture pointers
-      // using comparisons.
-      if (Tracker->captured(U))
-        return;
-      break;
-    }
-    default:
-      // Something else - be conservative and say it is captured.
+    if (Call->isDataOperand(&U) &&
+        !Call->doesNotCapture(Call->getDataOperandNo(&U))) {
+      return UseCaptureKind::MAY_CAPTURE;
+    return UseCaptureKind::NO_CAPTURE;
+      return UseCaptureKind::MAY_CAPTURE;
+    return UseCaptureKind::NO_CAPTURE;
+    return UseCaptureKind::NO_CAPTURE;
+    if (U.getOperandNo() == 0 || cast<StoreInst>(I)->isVolatile())
+      return UseCaptureKind::MAY_CAPTURE;
+    return UseCaptureKind::NO_CAPTURE;
+    if (U.getOperandNo() == 1 || ARMWI->isVolatile())
+      return UseCaptureKind::MAY_CAPTURE;
+    return UseCaptureKind::NO_CAPTURE;
+    if (U.getOperandNo() == 1 || U.getOperandNo() == 2 || ACXI->isVolatile())
+      return UseCaptureKind::MAY_CAPTURE;
+    return UseCaptureKind::NO_CAPTURE;
+    return UseCaptureKind::PASSTHROUGH;
+    unsigned Idx = U.getOperandNo();
+        if (isNoAliasCall(U.get()->stripPointerCasts()))
+          return UseCaptureKind::NO_CAPTURE;
+        const DataLayout &DL = I->getModule()->getDataLayout();
+        if (IsDereferenceableOrNull && IsDereferenceableOrNull(O, DL))
+          return UseCaptureKind::NO_CAPTURE;
+      return UseCaptureKind::NO_CAPTURE;
+    return UseCaptureKind::MAY_CAPTURE;
+    return UseCaptureKind::MAY_CAPTURE;
+  }
+}
+
+void llvm::PointerMayBeCaptured(const Value *V, CaptureTracker *Tracker,
+                                unsigned MaxUsesToExplore) {
+  assert(V->getType()->isPointerTy() && "Capture is for pointers only!");
+  if (MaxUsesToExplore == 0)
+    MaxUsesToExplore = DefaultMaxUsesToExplore;
+
+  SmallVector<const Use *, 20> Worklist;
+  Worklist.reserve(getDefaultMaxUsesToExploreForCaptureTracking());
+  SmallSet<const Use *, 20> Visited;
+
+  auto AddUses = [&](const Value *V) {
+    unsigned Count = 0;
+    for (const Use &U : V->uses()) {
+      // If there are lots of uses, conservatively say that the value
+      // is captured to avoid taking too much compile time.
+      if (Count++ >= MaxUsesToExplore) {
+        Tracker->tooManyUses();
+        return false;
+      }
+      if (!Visited.insert(&U).second)
+        continue;
+      if (!Tracker->shouldExplore(&U))
+        continue;
+      Worklist.push_back(&U);
+    }
+    return true;
+  };
+  if (!AddUses(V))
+    return;
+
+  auto IsDereferenceableOrNull = [Tracker](Value *V, const DataLayout &DL) {
+    return Tracker->isDereferenceableOrNull(V, DL);
+  };
+  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;
-      break;
+      continue;
+    case UseCaptureKind::PASSTHROUGH:
+      if (!AddUses(U->getUser()))
+        return;
+      continue;
     }
   }
 
   // All uses examined.
 }
 
 bool llvm::isNonEscapingLocalObject(
     const Value *V, SmallDenseMap<const Value *, bool, 8> *IsCapturedCache) {