[Attributor][FIX] Handle non-pointers when following uses

When we follow uses, e.g., in AAMemoryBehavior or AANoCapture, we need
to make sure the value is a pointer before we ask for abstract
attributes only valid for pointers. This happens because we follow
pointers through calls that do not capture but may return the value.

3 files changed, 40 insertions, 7 deletions

diff --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index 46e5c16..c65925b 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -5074,10 +5074,14 @@ bool AAMemoryBehaviorFloating::followUsersOfUseIn(Attributor &A, const Use *U,
   // general capturing of the underlying argument. The reason is that the
   // call might the argument "through return", which we allow and for which we
   // need to check call users.
-  unsigned ArgNo = ICS.getArgumentNo(U);
-  const auto &ArgNoCaptureAA =
-      A.getAAFor<AANoCapture>(*this, IRPosition::callsite_argument(ICS, ArgNo));
-  return !ArgNoCaptureAA.isAssumedNoCapture();
+  if (U->get()->getType()->isPointerTy()) {
+    unsigned ArgNo = ICS.getArgumentNo(U);
+    const auto &ArgNoCaptureAA = A.getAAFor<AANoCapture>(
+        *this, IRPosition::callsite_argument(ICS, ArgNo));
+    return !ArgNoCaptureAA.isAssumedNoCapture();
+  }
+
+  return true;
 }
 
 void AAMemoryBehaviorFloating::analyzeUseIn(Attributor &A, const Use *U,
@@ -5123,9 +5127,12 @@ void AAMemoryBehaviorFloating::analyzeUseIn(Attributor &A, const Use *U,
 
     // Adjust the possible access behavior based on the information on the
     // argument.
-    unsigned ArgNo = ICS.getArgumentNo(U);
-    const IRPosition &ArgPos = IRPosition::callsite_argument(ICS, ArgNo);
-    const auto &MemBehaviorAA = A.getAAFor<AAMemoryBehavior>(*this, ArgPos);
+    IRPosition Pos;
+    if (U->get()->getType()->isPointerTy())
+      Pos = IRPosition::callsite_argument(ICS, ICS.getArgumentNo(U));
+    else
+      Pos = IRPosition::callsite_function(ICS);
+    const auto &MemBehaviorAA = A.getAAFor<AAMemoryBehavior>(*this, Pos);
     // "assumed" has at most the same bits as the MemBehaviorAA assumed
     // and at least "known".
     intersectAssumedBits(MemBehaviorAA.getAssumed());
diff --git a/llvm/test/Transforms/Attributor/nocapture-1.ll b/llvm/test/Transforms/Attributor/nocapture-1.ll
index 135e916..78a2ba6 100644
--- a/llvm/test/Transforms/Attributor/nocapture-1.ll
+++ b/llvm/test/Transforms/Attributor/nocapture-1.ll
@@ -342,5 +342,19 @@ entry:
   ret i8* %p
 }
 
+declare i8* @maybe_returned_ptr(i8* readonly %ptr) readonly nounwind
+declare i8 @maybe_returned_val(i8* %ptr) readonly nounwind
+declare void @val_use(i8 %ptr) readonly nounwind
+
+; FIXME: Both pointers should be nocapture
+define void @ptr_uses(i8* %ptr, i8* %wptr) {
+; CHECK: define void @ptr_uses(i8* %ptr, i8* nocapture nonnull writeonly dereferenceable(1) %wptr)
+  %call_ptr = call i8* @maybe_returned_ptr(i8* %ptr)
+  %call_val = call i8 @maybe_returned_val(i8* %call_ptr)
+  call void @val_use(i8 %call_val)
+  store i8 0, i8* %wptr
+  ret void
+}
+
 declare i8* @llvm.launder.invariant.group.p0i8(i8*)
 declare i8* @llvm.strip.invariant.group.p0i8(i8*)
diff --git a/llvm/test/Transforms/Attributor/readattrs.ll b/llvm/test/Transforms/Attributor/readattrs.ll
index aa272a3..e23dc92 100644
--- a/llvm/test/Transforms/Attributor/readattrs.ll
+++ b/llvm/test/Transforms/Attributor/readattrs.ll
@@ -195,3 +195,15 @@ define void @testbyval(i8* %read_only) {
   ret void
 }
 ;}
+
+declare i8* @maybe_returned_ptr(i8* readonly %ptr) readonly nounwind
+declare i8 @maybe_returned_val(i8* %ptr) readonly nounwind
+declare void @val_use(i8 %ptr) readonly nounwind
+
+define void @ptr_uses(i8* %ptr) {
+; ATTRIBUTOR: define void @ptr_uses(i8* nocapture readonly %ptr)
+  %call_ptr = call i8* @maybe_returned_ptr(i8* %ptr)
+  %call_val = call i8 @maybe_returned_val(i8* %call_ptr)
+  call void @val_use(i8 %call_val)
+  ret void
+}