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

Relative to the previous attempt, this adjusts isEscapeSource()
to not treat calls with captures(ret: address, provenance) or similar
arguments as escape sources. This addresses the miscompile reported at:
#125880 (comment)

The implementation uses a helper function on CallBase to make this
check a bit more efficient (e.g. by skipping the byval checks) as
checking attributes on all arguments if fairly expensive.

------

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.

Author: nikic

Diff for: 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]]

Diff for: 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._]+]]

Diff for: 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.+]+]]

Diff for: 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

Diff for: 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) {

Diff for: 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;
@@ -77,10 +79,47 @@ namespace llvm {
                                    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
@@ -94,32 +133,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);
 

Diff for: 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);

Diff for: 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; }
 

Diff for: llvm/lib/Analysis/AliasAnalysis.cpp

@@ -835,9 +835,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