[CaptureTracking] Don't let comparisons against null escape inbounds pointers

Pointers that are in-bounds (either through dereferenceable_or_null or
thorough a getelementptr inbounds) cannot be captured with a comparison
against null. There is no way to construct a pointer that is still in
bounds but also NULL.

This helps safe languages that insert null checks before load/store
instructions. Without this patch, almost all pointers would be
considered captured even for simple loads. With this patch, an icmp with
null will not be seen as escaping as long as certain conditions are met.

There was a lot of discussion about this patch. See the Phabricator
thread for detals.

Differential Revision: https://reviews.llvm.org/D60047

llvm-svn: 362900

2 files changed, 51 insertions, 5 deletions

diff --git a/llvm/lib/Analysis/CaptureTracking.cpp b/llvm/lib/Analysis/CaptureTracking.cpp
index d2a0920413c..adaa83a6c44 100644
--- a/llvm/lib/Analysis/CaptureTracking.cpp
+++ b/llvm/lib/Analysis/CaptureTracking.cpp
@@ -330,14 +330,32 @@ void llvm::PointerMayBeCaptured(const Value *V, CaptureTracker *Tracker,
       AddUses(I);
       break;
     case Instruction::ICmp: {
-      // 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 (ConstantPointerNull *CPN =
-          dyn_cast<ConstantPointerNull>(I->getOperand(1)))
+      if (auto *CPN = dyn_cast<ConstantPointerNull>(I->getOperand(1))) {
+        // 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(V->stripPointerCasts()))
             break;
+        if (!I->getFunction()->nullPointerIsDefined()) {
+          auto *O = I->getOperand(0)->stripPointerCastsSameRepresentation();
+          // An inbounds GEP can either be a valid pointer (pointing into
+          // or to the end of an allocation), or be null in the default
+          // address space. So for an inbounds GEPs there is no way to let
+          // the pointer escape using clever GEP hacking because doing so
+          // would make the pointer point outside of the allocated object
+          // and thus make the GEP result a poison value.
+          if (auto *GEP = dyn_cast<GetElementPtrInst>(O))
+            if (GEP->isInBounds())
+              break;
+          // Comparing a dereferenceable_or_null argument against null
+          // cannot lead to pointer escapes, because if it is not null it
+          // must be a valid (in-bounds) pointer.
+          bool CanBeNull;
+          if (O->getPointerDereferenceableBytes(I->getModule()->getDataLayout(), CanBeNull))
+            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.
diff --git a/llvm/test/Transforms/FunctionAttrs/nocapture.ll b/llvm/test/Transforms/FunctionAttrs/nocapture.ll
index 0c3fc0af725..b109249674b 100644
--- a/llvm/test/Transforms/FunctionAttrs/nocapture.ll
+++ b/llvm/test/Transforms/FunctionAttrs/nocapture.ll
@@ -253,5 +253,33 @@ define void @captureStrip(i8* %p) {
   ret void
 }
 
+; CHECK: define i1 @captureICmp(i32* readnone %x)
+define i1 @captureICmp(i32* %x) {
+  %1 = icmp eq i32* %x, null
+  ret i1 %1
+}
+
+; CHECK: define i1 @nocaptureInboundsGEPICmp(i32* nocapture readnone %x)
+define i1 @nocaptureInboundsGEPICmp(i32* %x) {
+  %1 = getelementptr inbounds i32, i32* %x, i32 5
+  %2 = bitcast i32* %1 to i8*
+  %3 = icmp eq i8* %2, null
+  ret i1 %3
+}
+
+; CHECK: define i1 @nocaptureDereferenceableOrNullICmp(i32* nocapture readnone dereferenceable_or_null(4) %x)
+define i1 @nocaptureDereferenceableOrNullICmp(i32* dereferenceable_or_null(4) %x) {
+  %1 = bitcast i32* %x to i8*
+  %2 = icmp eq i8* %1, null
+  ret i1 %2
+}
+
+; CHECK: define i1 @captureDereferenceableOrNullICmp(i32* readnone dereferenceable_or_null(4) %x)
+define i1 @captureDereferenceableOrNullICmp(i32* dereferenceable_or_null(4) %x) "null-pointer-is-valid"="true" {
+  %1 = bitcast i32* %x to i8*
+  %2 = icmp eq i8* %1, null
+  ret i1 %2
+}
+
 declare i8* @llvm.launder.invariant.group.p0i8(i8*)
 declare i8* @llvm.strip.invariant.group.p0i8(i8*)