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

clang/test/CodeGen/allow-ubsan-check.c

@@ -86,7 +86,7 @@ int div(int x, int y) {
 }
 
 // CHECK-LABEL: define dso_local i32 @null(
-// CHECK-SAME: ptr noundef readonly [[X:%.*]]) local_unnamed_addr #[[ATTR0]] {
+// CHECK-SAME: ptr noundef readonly captures(address_is_null) [[X:%.*]]) local_unnamed_addr #[[ATTR0]] {
 // CHECK-NEXT:  [[ENTRY:.*:]]
 // CHECK-NEXT:    [[TMP0:%.*]] = icmp eq ptr [[X]], null, !nosanitize [[META2]]
 //
@@ -102,7 +102,7 @@ int div(int x, int y) {
 // CHECK-NEXT:    ret i32 [[TMP2]]
 //
 // TR-LABEL: define dso_local i32 @null(
-// TR-SAME: ptr noundef readonly [[X:%.*]]) local_unnamed_addr #[[ATTR0]] {
+// TR-SAME: ptr noundef readonly captures(address_is_null) [[X:%.*]]) local_unnamed_addr #[[ATTR0]] {
 // TR-NEXT:  [[ENTRY:.*:]]
 // TR-NEXT:    [[TMP0:%.*]] = icmp eq ptr [[X]], null, !nosanitize [[META2]]
 // TR-NEXT:    [[TMP1:%.*]] = tail call i1 @llvm.allow.ubsan.check(i8 29), !nosanitize [[META2]]
@@ -116,7 +116,7 @@ int div(int x, int y) {
 // TR-NEXT:    ret i32 [[TMP2]]
 //
 // REC-LABEL: define dso_local i32 @null(
-// REC-SAME: ptr noundef readonly [[X:%.*]]) local_unnamed_addr #[[ATTR0]] {
+// REC-SAME: ptr noundef readonly captures(address_is_null) [[X:%.*]]) local_unnamed_addr #[[ATTR0]] {
 // REC-NEXT:  [[ENTRY:.*:]]
 // REC-NEXT:    [[TMP0:%.*]] = icmp eq ptr [[X]], null, !nosanitize [[META2]]
 // REC-NEXT:    [[TMP1:%.*]] = tail call i1 @llvm.allow.ubsan.check(i8 29), !nosanitize [[META2]]

clang/test/CodeGenCXX/RelativeVTablesABI/dynamic-cast.cpp

@@ -3,7 +3,7 @@
 
 // RUN: %clang_cc1 %s -triple=aarch64-unknown-fuchsia -O3 -o - -emit-llvm | FileCheck %s
 
-// CHECK:      define{{.*}} ptr @_Z6upcastP1B(ptr noundef readnone returned %b) local_unnamed_addr
+// CHECK:      define{{.*}} ptr @_Z6upcastP1B(ptr noundef readnone returned captures(ret: address, provenance) %b) local_unnamed_addr
 // CHECK-NEXT: entry:
 // CHECK-NEXT:   ret ptr %b
 // CHECK-NEXT: }
@@ -22,12 +22,12 @@
 
 // CHECK: declare ptr @__dynamic_cast(ptr, ptr, ptr, i64) local_unnamed_addr
 
-// CHECK:      define{{.*}} ptr @_Z8selfcastP1B(ptr noundef readnone returned %b) local_unnamed_addr
+// CHECK:      define{{.*}} ptr @_Z8selfcastP1B(ptr noundef readnone returned captures(ret: address, provenance) %b) local_unnamed_addr
 // CHECK-NEXT: entry
 // CHECK-NEXT:   ret ptr %b
 // CHECK-NEXT: }
 
-// CHECK: define{{.*}} ptr @_Z9void_castP1B(ptr noundef readonly %b) local_unnamed_addr
+// CHECK: define{{.*}} ptr @_Z9void_castP1B(ptr noundef readonly captures(address_is_null, ret: address, provenance) %b) local_unnamed_addr
 // CHECK-NEXT: entry:
 // CHECK-NEXT:   [[isnull:%[0-9]+]] = icmp eq ptr %b, null
 // CHECK-NEXT:   br i1 [[isnull]], label %[[dynamic_cast_end:[a-z0-9._]+]], label %[[dynamic_cast_notnull:[a-z0-9._]+]]

clang/test/CodeGenCXX/RelativeVTablesABI/type-info.cpp

@@ -24,7 +24,7 @@
 // CHECK-NEXT:   ret ptr @_ZTS1A
 // CHECK-NEXT: }
 
-// CHECK:      define{{.*}} i1 @_Z5equalP1A(ptr noundef readonly %a) local_unnamed_addr
+// CHECK:      define{{.*}} i1 @_Z5equalP1A(ptr noundef readonly captures(address_is_null) %a) local_unnamed_addr
 // CHECK-NEXT: entry:
 // CHECK-NEXT:   [[isnull:%[0-9]+]] = icmp eq ptr %a, null
 // CHECK-NEXT:   br i1 [[isnull]], label %[[bad_typeid:[a-z0-9._]+]], label %[[end:[a-z0-9.+]+]]

clang/test/CodeGenOpenCL/amdgcn-buffer-rsrc-type.cl

@@ -22,7 +22,7 @@ __amdgpu_buffer_rsrc_t getBuffer(void *p) {
 }
 
 // CHECK-LABEL: define {{[^@]+}}@consumeBufferPtr
-// CHECK-SAME: (ptr addrspace(5) noundef readonly [[P:%.*]]) local_unnamed_addr #[[ATTR0]] {
+// CHECK-SAME: (ptr addrspace(5) noundef readonly captures(address) [[P:%.*]]) local_unnamed_addr #[[ATTR0]] {
 // CHECK-NEXT:  entry:
 // CHECK-NEXT:    [[TOBOOL_NOT:%.*]] = icmp eq ptr addrspace(5) [[P]], addrspacecast (ptr null to ptr addrspace(5))
 // CHECK-NEXT:    br i1 [[TOBOOL_NOT]], label [[IF_END:%.*]], label [[IF_THEN:%.*]]
@@ -39,7 +39,7 @@ void consumeBufferPtr(__amdgpu_buffer_rsrc_t *p) {
 }
 
 // CHECK-LABEL: define {{[^@]+}}@test
-// CHECK-SAME: (ptr addrspace(5) noundef readonly [[A:%.*]]) local_unnamed_addr #[[ATTR0]] {
+// CHECK-SAME: (ptr addrspace(5) noundef readonly captures(address) [[A:%.*]]) local_unnamed_addr #[[ATTR0]] {
 // CHECK-NEXT:  entry:
 // CHECK-NEXT:    [[TMP0:%.*]] = load i32, ptr addrspace(5) [[A]], align 16, !tbaa [[TBAA8:![0-9]+]]
 // CHECK-NEXT:    [[TOBOOL_NOT:%.*]] = icmp eq i32 [[TMP0]], 0

clang/test/CodeGenOpenCL/as_type.cl

@@ -67,7 +67,7 @@ int3 f8(char16 x) {
   return __builtin_astype(x, int3);
 }
 
-//CHECK: define{{.*}} spir_func noundef ptr addrspace(1) @addr_cast(ptr noundef readnone %[[x:.*]])
+//CHECK: define{{.*}} spir_func noundef ptr addrspace(1) @addr_cast(ptr noundef readnone captures(ret: address, provenance) %[[x:.*]])
 //CHECK: %[[cast:.*]] ={{.*}} addrspacecast ptr %[[x]] to ptr addrspace(1)
 //CHECK: ret ptr addrspace(1) %[[cast]]
 global int* addr_cast(int *x) {

llvm/include/llvm/Analysis/CaptureTracking.h

@@ -14,11 +14,13 @@
 #define LLVM_ANALYSIS_CAPTURETRACKING_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/Support/ModRef.h"
 
 namespace llvm {
 
   class Value;
   class Use;
+  class CaptureInfo;
   class DataLayout;
   class Instruction;
   class DominatorTree;
@@ -71,10 +73,47 @@ namespace llvm {
                                    bool ReturnCaptures, const DominatorTree &DT,
                                    unsigned MaxUsesToExplore = 0);
 
+  /// Capture information for a specific Use.
+  struct UseCaptureInfo {
+    /// Components captured by this use.
+    CaptureComponents UseCC;
+    /// Components captured by the return value of the user of this Use.
+    CaptureComponents ResultCC;
+
+    UseCaptureInfo(CaptureComponents UseCC,
+                   CaptureComponents ResultCC = CaptureComponents::None)
+        : UseCC(UseCC), ResultCC(ResultCC) {}
+
+    static UseCaptureInfo passthrough() {
+      return UseCaptureInfo(CaptureComponents::None, CaptureComponents::All);
+    }
+
+    bool isPassthrough() const {
+      return capturesNothing(UseCC) && capturesAnything(ResultCC);
+    }
+
+    operator CaptureComponents() const { return UseCC | ResultCC; }
+  };
+
   /// This callback is used in conjunction with PointerMayBeCaptured. In
   /// addition to the interface here, you'll need to provide your own getters
   /// to see whether anything was captured.
   struct CaptureTracker {
+    /// Action returned from captures().
+    enum Action {
+      /// Stop the traversal.
+      Stop,
+      /// Continue traversal, and also follow the return value of the user if
+      /// it has additional capture components (that is, if it has capture
+      /// components in Ret that are not part of Other).
+      Continue,
+      /// Continue traversal, but do not follow the return value of the user,
+      /// even if it has additional capture components. Should only be used if
+      /// captures() has already taken the potential return captures into
+      /// account.
+      ContinueIgnoringReturn,
+    };
+
     virtual ~CaptureTracker();
 
     /// tooManyUses - The depth of traversal has breached a limit. There may be
@@ -88,32 +127,31 @@ namespace llvm {
     /// U->getUser() is always an Instruction.
     virtual bool shouldExplore(const Use *U);
 
-    /// captured - Information about the pointer was captured by the user of
-    /// use U. Return true to stop the traversal or false to continue looking
-    /// for more capturing instructions.
-    virtual bool captured(const Use *U) = 0;
+    /// Use U directly captures CI.UseCC and additionally CI.ResultCC
+    /// through the return value of the user of U.
+    ///
+    /// Return one of Stop, Continue or ContinueIgnoringReturn to control
+    /// further traversal.
+    virtual Action captured(const Use *U, UseCaptureInfo CI) = 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 returned UseCaptureInfo contains the components captured directly
+  /// by the use (UseCC) and the components captured through the return value
+  /// of the user (ResultCC).
+  ///
+  /// \p Base is the starting value of the capture analysis, which is
+  /// relevant for address_is_null captures.
   /// The \p IsDereferenceableOrNull callback is used to rule out capturing for
   /// certain comparisons.
-  UseCaptureKind
-  DetermineUseCaptureKind(const Use &U,
+  UseCaptureInfo
+  DetermineUseCaptureKind(const Use &U, const Value *Base,
                           llvm::function_ref<bool(Value *, const DataLayout &)>
                               IsDereferenceableOrNull);
 

llvm/include/llvm/IR/InstrTypes.h

@@ -1692,6 +1692,11 @@ class CallBase : public Instruction {
     return capturesNothing(getCaptureInfo(OpNo));
   }
 
+  /// Returns whether the call has an argument that has an attribute like
+  /// captures(ret: address, provenance), where the return capture components
+  /// are not a subset of the other capture components.
+  bool hasArgumentWithAdditionalReturnCaptureComponents() const;
+
   /// Determine whether this argument is passed by value.
   bool isByValArgument(unsigned ArgNo) const {
     return paramHasAttr(ArgNo, Attribute::ByVal);

llvm/include/llvm/Support/ModRef.h

@@ -326,6 +326,10 @@ inline bool capturesFullProvenance(CaptureComponents CC) {
   return (CC & CaptureComponents::Provenance) == CaptureComponents::Provenance;
 }
 
+inline bool capturesAll(CaptureComponents CC) {
+  return CC == CaptureComponents::All;
+}
+
 raw_ostream &operator<<(raw_ostream &OS, CaptureComponents CC);
 
 /// Represents which components of the pointer may be captured in which
@@ -350,6 +354,15 @@ class CaptureInfo {
   /// Create CaptureInfo that may capture all components of the pointer.
   static CaptureInfo all() { return CaptureInfo(CaptureComponents::All); }
 
+  /// Create CaptureInfo that may only capture via the return value.
+  static CaptureInfo
+  retOnly(CaptureComponents RetComponents = CaptureComponents::All) {
+    return CaptureInfo(CaptureComponents::None, RetComponents);
+  }
+
+  /// Whether the pointer is only captured via the return value.
+  bool isRetOnly() const { return capturesNothing(OtherComponents); }
+
   /// Get components potentially captured by the return value.
   CaptureComponents getRetComponents() const { return RetComponents; }
 

llvm/lib/Analysis/AliasAnalysis.cpp

@@ -635,7 +635,13 @@ ModRefInfo AAResults::callCapturesBefore(const Instruction *I,
     // Only look at the no-capture or byval pointer arguments.  If this
     // pointer were passed to arguments that were neither of these, then it
     // couldn't be no-capture.
-    if (!(*CI)->getType()->isPointerTy() || !Call->doesNotCapture(ArgNo))
+    if (!(*CI)->getType()->isPointerTy())
+      continue;
+
+    // Make sure we still check captures(ret: address, provenance) arguments,
+    // as these wouldn't be treated as a capture at the call-site.
+    CaptureInfo Captures = Call->getCaptureInfo(ArgNo);
+    if (!capturesNothing(Captures.getOtherComponents()))
       continue;
 
     AliasResult AR =
@@ -834,9 +840,15 @@ bool llvm::isBaseOfObject(const Value *V) {
 }
 
 bool llvm::isEscapeSource(const Value *V) {
-  if (auto *CB = dyn_cast<CallBase>(V))
-    return !isIntrinsicReturningPointerAliasingArgumentWithoutCapturing(CB,
-                                                                        true);
+  if (auto *CB = dyn_cast<CallBase>(V)) {
+    if (isIntrinsicReturningPointerAliasingArgumentWithoutCapturing(CB, true))
+      return false;
+
+    // The return value of a function with a captures(ret: address, provenance)
+    // attribute is not necessarily an escape source. The return value may
+    // alias with a non-escaping object.
+    return !CB->hasArgumentWithAdditionalReturnCaptureComponents();
+  }
 
   // The load case works because isNonEscapingLocalObject considers all
   // stores to be escapes (it passes true for the StoreCaptures argument