[LVI] Address review comments.

Author: dtcxzyw

llvm/lib/Analysis/LazyValueInfo.cpp

@@ -644,6 +644,13 @@ static void AddNonNullPointersByInstruction(
     AddNonNullPointer(MI->getRawDest(), PtrSet);
     if (MemTransferInst *MTI = dyn_cast<MemTransferInst>(MI))
       AddNonNullPointer(MTI->getRawSource(), PtrSet);
+  } else if (auto *CB = dyn_cast<CallBase>(I)) {
+    for (auto &U : CB->args()) {
+      if (U->getType()->isPointerTy() &&
+          CB->paramHasNonNullAttr(CB->getArgOperandNo(&U),
+                                  /*AllowUndefOrPoison=*/false))
+        AddNonNullPointer(U.get(), PtrSet);
+    }
   }
 }
 
@@ -780,25 +787,12 @@ void LazyValueInfoImpl::intersectAssumeOrGuardBlockValueConstantRange(
   }
 
   if (BBLV.isOverdefined()) {
-    if (PointerType *PTy = dyn_cast<PointerType>(Val->getType())) {
-      // Check whether we're checking at the terminator, and the pointer has
-      // been dereferenced in this block.
-      if (BB->getTerminator() == BBI && isNonNullAtEndOfBlock(Val, BB))
-        BBLV = ValueLatticeElement::getNot(ConstantPointerNull::get(PTy));
-      else {
-        for (Use &U : Val->uses()) {
-          if (auto *CB = dyn_cast<CallBase>(U.getUser())) {
-            if (CB->isArgOperand(&U) &&
-                CB->paramHasNonNullAttr(CB->getArgOperandNo(&U),
-                                        /*AllowUndefOrPoison=*/false) &&
-                isValidAssumeForContext(CB, BBI)) {
-              BBLV = ValueLatticeElement::getNot(ConstantPointerNull::get(PTy));
-              break;
-            }
-          }
-        }
-      }
-    }
+    // Check whether we're checking at the terminator, and the pointer has
+    // been dereferenced in this block.
+    PointerType *PTy = dyn_cast<PointerType>(Val->getType());
+    if (PTy && BB->getTerminator() == BBI &&
+        isNonNullAtEndOfBlock(Val, BB))
+      BBLV = ValueLatticeElement::getNot(ConstantPointerNull::get(PTy));
   }
 }
 

llvm/test/Transforms/CorrelatedValuePropagation/non-null.ll

@@ -349,6 +349,7 @@ define i1 @test_known_nonnull_at_callsite(ptr %src) {
 ; CHECK-LABEL: @test_known_nonnull_at_callsite(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    call void @callee(ptr noundef nonnull [[SRC:%.*]])
+; CHECK-NEXT:    [[NONNULL:%.*]] = icmp eq ptr [[SRC]], null
 ; CHECK-NEXT:    ret i1 false
 ;
 entry:
@@ -361,6 +362,7 @@ define i1 @test_known_nonnull_mixed(ptr %src) {
 ; CHECK-LABEL: @test_known_nonnull_mixed(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    call void @callee2(ptr nonnull [[SRC:%.*]])
+; CHECK-NEXT:    [[NONNULL:%.*]] = icmp eq ptr [[SRC]], null
 ; CHECK-NEXT:    ret i1 false
 ;
 entry:
@@ -373,6 +375,7 @@ define i1 @test_known_nonnull_at_callsite_dereferenceable(ptr %src) {
 ; CHECK-LABEL: @test_known_nonnull_at_callsite_dereferenceable(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    call void @callee(ptr dereferenceable(1) [[SRC:%.*]])
+; CHECK-NEXT:    [[NONNULL:%.*]] = icmp eq ptr [[SRC]], null
 ; CHECK-NEXT:    ret i1 false
 ;
 entry: