Fix address_is_null inference

nikic · nikic · commit 6dd18a6aac73 · 2025-02-12T16:43:56.000+01:00
diff --git a/llvm/include/llvm/Analysis/CaptureTracking.h b/llvm/include/llvm/Analysis/CaptureTracking.h
@@ -130,8 +130,11 @@ namespace llvm {
   /// which components may be captured by following uses of the user of \p U.
   /// The \p IsDereferenceableOrNull callback is used to rule out capturing for
   /// certain comparisons.
+  ///
+  /// \p Base is the starting value of the capture analysis, which is
+  /// relevant for address_is_null captures.
   CaptureInfo
-  DetermineUseCaptureKind(const Use &U,
+  DetermineUseCaptureKind(const Use &U, const Value *Base,
                           llvm::function_ref<bool(Value *, const DataLayout &)>
                               IsDereferenceableOrNull);
 
diff --git a/llvm/lib/Analysis/CaptureTracking.cpp b/llvm/lib/Analysis/CaptureTracking.cpp
@@ -280,7 +280,7 @@ Instruction *llvm::FindEarliestCapture(const Value *V, Function &F,
 }
 
 CaptureInfo llvm::DetermineUseCaptureKind(
-    const Use &U,
+    const Use &U, const Value *Base,
     function_ref<bool(Value *, const DataLayout &)> IsDereferenceableOrNull) {
   Instruction *I = dyn_cast<Instruction>(U.getUser());
 
@@ -378,14 +378,15 @@ CaptureInfo llvm::DetermineUseCaptureKind(
   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 CaptureInfo::none();
       if (!I->getFunction()->nullPointerIsDefined()) {
@@ -397,7 +398,16 @@ CaptureInfo llvm::DetermineUseCaptureKind(
         if (IsDereferenceableOrNull && IsDereferenceableOrNull(O, DL))
           return CaptureInfo::none();
       }
-      return CaptureInfo::otherOnly(CaptureComponents::AddressIsNull);
+
+      // Check whether this is a comparison of the base pointer against
+      // null. We can also strip inbounds GEPs, as inbounds preserves
+      // the null-ness of the pointer.
+      Value *Stripped = U.get();
+      if (!NullPointerIsDefined(I->getFunction(),
+                                U->getType()->getPointerAddressSpace()))
+        Stripped = Stripped->stripInBoundsOffsets();
+      if (Stripped == Base)
+        return CaptureInfo::otherOnly(CaptureComponents::AddressIsNull);
     }
 
     // Otherwise, be conservative. There are crazy ways to capture pointers
@@ -445,7 +455,7 @@ void llvm::PointerMayBeCaptured(const Value *V, CaptureTracker *Tracker,
   };
   while (!Worklist.empty()) {
     const Use *U = Worklist.pop_back_val();
-    CaptureInfo CI = DetermineUseCaptureKind(*U, IsDereferenceableOrNull);
+    CaptureInfo CI = DetermineUseCaptureKind(*U, V, IsDereferenceableOrNull);
     if (capturesNothing(CI))
       continue;
     CaptureComponents OtherCC = CI.getOtherComponents();
diff --git a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
@@ -3971,7 +3971,8 @@ struct AANoAliasCallSiteArgument final : AANoAliasImpl {
       //       is CGSCC runs. For those we would need to "allow" AANoCapture for
       //       a value in the module slice.
       // TODO(captures): Make this more precise.
-      CaptureInfo CI = DetermineUseCaptureKind(U, IsDereferenceableOrNull);
+      CaptureInfo CI =
+          DetermineUseCaptureKind(U, /*Base=*/nullptr, IsDereferenceableOrNull);
       if (capturesNothing(CI))
         return true;
       if (CI.isRetOnly()) {
@@ -6018,7 +6019,8 @@ ChangeStatus AANoCaptureImpl::updateImpl(Attributor &A) {
 
   auto UseCheck = [&](const Use &U, bool &Follow) -> bool {
     // TODO(captures): Make this more precise.
-    CaptureInfo CI = DetermineUseCaptureKind(U, IsDereferenceableOrNull);
+    CaptureInfo CI =
+        DetermineUseCaptureKind(U, /*Base=*/nullptr, IsDereferenceableOrNull);
     if (capturesNothing(CI))
       return true;
     if (CI.isRetOnly()) {
@@ -12149,7 +12151,7 @@ struct AAGlobalValueInfoFloating : public AAGlobalValueInfo {
     auto UsePred = [&](const Use &U, bool &Follow) -> bool {
       Uses.insert(&U);
       // TODO(captures): Make this more precise.
-      CaptureInfo CI = DetermineUseCaptureKind(U, nullptr);
+      CaptureInfo CI = DetermineUseCaptureKind(U, /*Base=*/nullptr, nullptr);
       if (capturesAnything(CI) && CI.isRetOnly()) {
         Follow = true;
         return true;
diff --git a/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp b/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp
@@ -1550,7 +1550,8 @@ bool MemCpyOptPass::performStackMoveOptzn(Instruction *Load, Instruction *Store,
         }
         if (!Visited.insert(&U).second)
           continue;
-        CaptureInfo CI = DetermineUseCaptureKind(U, IsDereferenceableOrNull);
+        CaptureInfo CI =
+            DetermineUseCaptureKind(U, AI, IsDereferenceableOrNull);
         // TODO(captures): Make this more precise.
         if (capturesAnything(CI)) {
           if (CI.isRetOnly()) {
diff --git a/llvm/test/Transforms/FunctionAttrs/nocapture.ll b/llvm/test/Transforms/FunctionAttrs/nocapture.ll
@@ -811,6 +811,46 @@ define i1 @nocaptureInboundsGEPICmpRev(ptr %x) {
   ret i1 %2
 }
 
+define i1 @notInboundsGEPICmp(ptr %x) {
+; FNATTRS: Function Attrs: mustprogress nofree norecurse nosync nounwind willreturn memory(none)
+; FNATTRS-LABEL: define i1 @notInboundsGEPICmp
+; FNATTRS-SAME: (ptr readnone captures(address) [[X:%.*]]) #[[ATTR0]] {
+; FNATTRS-NEXT:    [[TMP1:%.*]] = getelementptr i32, ptr [[X]], i32 5
+; FNATTRS-NEXT:    [[TMP2:%.*]] = icmp eq ptr [[TMP1]], null
+; FNATTRS-NEXT:    ret i1 [[TMP2]]
+;
+; ATTRIBUTOR: Function Attrs: mustprogress nofree norecurse nosync nounwind willreturn memory(none)
+; ATTRIBUTOR-LABEL: define i1 @notInboundsGEPICmp
+; ATTRIBUTOR-SAME: (ptr nofree readnone [[X:%.*]]) #[[ATTR0]] {
+; ATTRIBUTOR-NEXT:    [[TMP1:%.*]] = getelementptr i32, ptr [[X]], i32 5
+; ATTRIBUTOR-NEXT:    [[TMP2:%.*]] = icmp eq ptr [[TMP1]], null
+; ATTRIBUTOR-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = getelementptr i32, ptr %x, i32 5
+  %2 = icmp eq ptr %1, null
+  ret i1 %2
+}
+
+define i1 @inboundsGEPICmpNullPointerDefined(ptr %x) null_pointer_is_valid {
+; FNATTRS: Function Attrs: mustprogress nofree norecurse nosync nounwind null_pointer_is_valid willreturn memory(none)
+; FNATTRS-LABEL: define i1 @inboundsGEPICmpNullPointerDefined
+; FNATTRS-SAME: (ptr readnone captures(address) [[X:%.*]]) #[[ATTR16:[0-9]+]] {
+; FNATTRS-NEXT:    [[TMP1:%.*]] = getelementptr i32, ptr [[X]], i32 5
+; FNATTRS-NEXT:    [[TMP2:%.*]] = icmp eq ptr [[TMP1]], null
+; FNATTRS-NEXT:    ret i1 [[TMP2]]
+;
+; ATTRIBUTOR: Function Attrs: mustprogress nofree norecurse nosync nounwind null_pointer_is_valid willreturn memory(none)
+; ATTRIBUTOR-LABEL: define i1 @inboundsGEPICmpNullPointerDefined
+; ATTRIBUTOR-SAME: (ptr nofree readnone [[X:%.*]]) #[[ATTR12:[0-9]+]] {
+; ATTRIBUTOR-NEXT:    [[TMP1:%.*]] = getelementptr i32, ptr [[X]], i32 5
+; ATTRIBUTOR-NEXT:    [[TMP2:%.*]] = icmp eq ptr [[TMP1]], null
+; ATTRIBUTOR-NEXT:    ret i1 [[TMP2]]
+;
+  %1 = getelementptr i32, ptr %x, i32 5
+  %2 = icmp eq ptr %1, null
+  ret i1 %2
+}
+
 define i1 @nocaptureDereferenceableOrNullICmp(ptr dereferenceable_or_null(4) %x) {
 ; FNATTRS: Function Attrs: mustprogress nofree norecurse nosync nounwind willreturn memory(none)
 ; FNATTRS-LABEL: define noundef i1 @nocaptureDereferenceableOrNullICmp
@@ -831,13 +871,13 @@ define i1 @nocaptureDereferenceableOrNullICmp(ptr dereferenceable_or_null(4) %x)
 define i1 @captureDereferenceableOrNullICmp(ptr dereferenceable_or_null(4) %x) null_pointer_is_valid {
 ; FNATTRS: Function Attrs: mustprogress nofree norecurse nosync nounwind null_pointer_is_valid willreturn memory(none)
 ; FNATTRS-LABEL: define noundef i1 @captureDereferenceableOrNullICmp
-; FNATTRS-SAME: (ptr readnone captures(address_is_null) dereferenceable_or_null(4) [[X:%.*]]) #[[ATTR16:[0-9]+]] {
+; FNATTRS-SAME: (ptr readnone captures(address_is_null) dereferenceable_or_null(4) [[X:%.*]]) #[[ATTR16]] {
 ; FNATTRS-NEXT:    [[TMP1:%.*]] = icmp eq ptr [[X]], null
 ; FNATTRS-NEXT:    ret i1 [[TMP1]]
 ;
 ; ATTRIBUTOR: Function Attrs: mustprogress nofree norecurse nosync nounwind null_pointer_is_valid willreturn memory(none)
 ; ATTRIBUTOR-LABEL: define i1 @captureDereferenceableOrNullICmp
-; ATTRIBUTOR-SAME: (ptr nofree readnone dereferenceable_or_null(4) [[X:%.*]]) #[[ATTR12:[0-9]+]] {
+; ATTRIBUTOR-SAME: (ptr nofree readnone dereferenceable_or_null(4) [[X:%.*]]) #[[ATTR12]] {
 ; ATTRIBUTOR-NEXT:    [[TMP1:%.*]] = icmp eq ptr [[X]], null
 ; ATTRIBUTOR-NEXT:    ret i1 [[TMP1]]
 ;
