[BasicAA] Improve escape source analysis for return-only captures

llvm/include/llvm/Analysis/AliasAnalysis.h

@@ -888,9 +888,9 @@ bool isIdentifiedFunctionLocal(const Value *V);
 /// can be inbounds w.r.t the actual underlying object.
 bool isBaseOfObject(const Value *V);
 
-/// Returns true if the pointer is one which would have been considered an
-/// escape by isNonEscapingLocalObject.
-bool isEscapeSource(const Value *V);
+/// Returns true if the pointer cannot alias a non-escaping local object,
+/// except via the values returned in the MayAlias argument.
+bool isEscapeSource(const Value *V, SmallVectorImpl<Value *> &MayAlias);
 
 /// Return true if Object memory is not visible after an unwind, in the sense
 /// that program semantics cannot depend on Object containing any particular

llvm/include/llvm/IR/InstrTypes.h

@@ -1692,10 +1692,11 @@ class CallBase : public Instruction {
     return capturesNothing(getCaptureInfo(OpNo));
   }
 
-  /// Returns whether the call has an argument that has an attribute like
+  /// Collects call arguments that have an attribute like
   /// captures(ret: address, provenance), where the return capture components
   /// are not a subset of the other capture components.
-  bool hasArgumentWithAdditionalReturnCaptureComponents() const;
+  void getArgumentsWithAdditionalReturnCaptureComponents(
+      SmallVectorImpl<Value *> &Args) const;
 
   /// Determine whether this argument is passed by value.
   bool isByValArgument(unsigned ArgNo) const {

llvm/lib/Analysis/AliasAnalysis.cpp

@@ -840,15 +840,16 @@ bool llvm::isBaseOfObject(const Value *V) {
   return (isa<AllocaInst>(V) || isa<GlobalVariable>(V));
 }
 
-bool llvm::isEscapeSource(const Value *V) {
+bool llvm::isEscapeSource(const Value *V, SmallVectorImpl<Value *> &MayAlias) {
   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();
+    // attribute may alias with the corresponding argument. If it doesn't, then
+    // it is an escape source.
+    CB->getArgumentsWithAdditionalReturnCaptureComponents(MayAlias);
+    return true;
   }
 
   // The load case works because isNonEscapingLocalObject considers all

llvm/lib/Analysis/BasicAliasAnalysis.cpp

@@ -1606,11 +1606,21 @@ AliasResult BasicAAResult::aliasCheck(const Value *V1, LocationSize V1Size,
     // temporary store the nocapture argument's value in a temporary memory
     // location if that memory location doesn't escape. Or it may pass a
     // nocapture value to other functions as long as they don't capture it.
-    if (isEscapeSource(O1) && AAQI.CA->isNotCapturedBefore(
-                                  O2, dyn_cast<Instruction>(O1), /*OrAt*/ true))
-      return AliasResult::NoAlias;
-    if (isEscapeSource(O2) && AAQI.CA->isNotCapturedBefore(
-                                  O1, dyn_cast<Instruction>(O2), /*OrAt*/ true))
+    auto EscapeSourceDoesNotAliasNonEscapingLocalObject = [&](const Value *O1,
+                                                              const Value *O2) {
+      SmallVector<Value *> MayAlias;
+      if (!isEscapeSource(O1, MayAlias) ||
+          !AAQI.CA->isNotCapturedBefore(O2, dyn_cast<Instruction>(O1),
+                                        /*OrAt=*/true))
+        return false;
+      for (Value *V : MayAlias)
+        if (AAQI.AAR.alias(MemoryLocation::getBeforeOrAfter(O2),
+                           MemoryLocation::getBeforeOrAfter(V), AAQI))
+          return false;
+      return true;
+    };
+    if (EscapeSourceDoesNotAliasNonEscapingLocalObject(O1, O2) ||
+        EscapeSourceDoesNotAliasNonEscapingLocalObject(O2, O1))
       return AliasResult::NoAlias;
   }
 

llvm/lib/IR/Instructions.cpp

@@ -713,7 +713,8 @@ CaptureInfo CallBase::getCaptureInfo(unsigned OpNo) const {
   return OBU.isDeoptOperandBundle() ? CaptureInfo::none() : CaptureInfo::all();
 }
 
-bool CallBase::hasArgumentWithAdditionalReturnCaptureComponents() const {
+void CallBase::getArgumentsWithAdditionalReturnCaptureComponents(
+    SmallVectorImpl<Value *> &Args) const {
   for (unsigned I = 0, E = arg_size(); I < E; ++I) {
     if (!getArgOperand(I)->getType()->isPointerTy())
       continue;
@@ -722,9 +723,8 @@ bool CallBase::hasArgumentWithAdditionalReturnCaptureComponents() const {
     if (auto *Fn = dyn_cast<Function>(getCalledOperand()))
       CI &= Fn->getAttributes().getParamAttrs(I).getCaptureInfo();
     if (capturesAnything(CI.getRetComponents() & ~CI.getOtherComponents()))
-      return true;
+      Args.push_back(getArgOperand(I));
   }
-  return false;
 }
 
 //===----------------------------------------------------------------------===//

llvm/lib/Transforms/Utils/InlineFunction.cpp

@@ -1273,7 +1273,8 @@ static void AddAliasScopeMetadata(CallBase &CB, ValueToValueMapTy &VMap,
           UsesAliasingPtr = true;
         }
 
-        if (isEscapeSource(V)) {
+        SmallVector<Value *> MayAlias;
+        if (isEscapeSource(V, MayAlias) && MayAlias.empty()) {
           // An escape source can only alias with a noalias argument if it has
           // been captured beforehand.
           RequiresNoCaptureBefore = true;

llvm/test/Analysis/BasicAA/captures.ll

@@ -40,3 +40,16 @@ define void @address_capture_and_full_capture() {
   load i32, ptr %a
   ret void
 }
+
+declare ptr @capture_ret(ptr, ptr)
+
+; CHECK-LABEL: capture_ret_only
+; CHECK: MayAlias:	i8* %a, i8* %ret
+; CHECK: NoAlias:	i8* %b, i8* %ret
+define void @capture_ret_only(ptr noalias %a, ptr noalias %b) {
+  %ret = call ptr @capture_ret(ptr captures(ret: address, provenance) %a, ptr captures(none) %b)
+  load i8, ptr %ret
+  load i8, ptr %a
+  load i8, ptr %b
+  ret void
+}