5 files changed, 94 insertions, 1 deletions
diff --git a/llvm/lib/Transforms/ObjCARC/ObjCARC.h b/llvm/lib/Transforms/ObjCARC/ObjCARC.h
index cd9b3d96a14..745dac88619 100644
--- a/llvm/lib/Transforms/ObjCARC/ObjCARC.h
+++ b/llvm/lib/Transforms/ObjCARC/ObjCARC.h
@@ -82,6 +82,26 @@ static inline const Instruction *getreturnRVOperand(const Instruction &Inst,
   return dyn_cast<InvokeInst>(Opnd);
 }
 
+/// Return the list of PHI nodes that are equivalent to PN.
+template<class PHINodeTy, class VectorTy>
+void getEquivalentPHIs(PHINodeTy &PN, VectorTy &PHIList) {
+  auto *BB = PN.getParent();
+  for (auto &P : BB->phis()) {
+    if (&P == &PN) // Do not add PN to the list.
+      continue;
+    unsigned I = 0, E = PN.getNumIncomingValues();
+    for (; I < E; ++I) {
+      auto *BB = PN.getIncomingBlock(I);
+      auto *PNOpnd = PN.getIncomingValue(I)->stripPointerCasts();
+      auto *POpnd = P.getIncomingValueForBlock(BB)->stripPointerCasts();
+      if (PNOpnd != POpnd)
+        break;
+    }
+    if (I == E)
+      PHIList.push_back(&P);
+  }
+}
+
 } // end namespace objcarc
 } // end namespace llvm
 
diff --git a/llvm/lib/Transforms/ObjCARC/ObjCARCContract.cpp b/llvm/lib/Transforms/ObjCARC/ObjCARCContract.cpp
index c4e61218f3f..5deb39449e9 100644
--- a/llvm/lib/Transforms/ObjCARC/ObjCARCContract.cpp
+++ b/llvm/lib/Transforms/ObjCARC/ObjCARCContract.cpp
@@ -618,8 +618,17 @@ bool ObjCARCContract::runOnFunction(Function &F) {
       else if (isa<GlobalAlias>(Arg) &&
                !cast<GlobalAlias>(Arg)->isInterposable())
         Arg = cast<GlobalAlias>(Arg)->getAliasee();
-      else
+      else {
+        // If Arg is a PHI node, get PHIs that are equivalent to it and replace
+        // their uses.
+        if (PHINode *PN = dyn_cast<PHINode>(Arg)) {
+          SmallVector<Value *, 1> PHIList;
+          getEquivalentPHIs(*PN, PHIList);
+          for (Value *PHI : PHIList)
+            ReplaceArgUses(PHI);
+        }
         break;
+      }
     }
 
     // Replace bitcast users of Arg that are dominated by Inst.
diff --git a/llvm/lib/Transforms/ObjCARC/ObjCARCOpts.cpp b/llvm/lib/Transforms/ObjCARC/ObjCARCOpts.cpp
index 99ed6863c22..ecec85444b1 100644
--- a/llvm/lib/Transforms/ObjCARC/ObjCARCOpts.cpp
+++ b/llvm/lib/Transforms/ObjCARC/ObjCARCOpts.cpp
@@ -652,6 +652,11 @@ void ObjCARCOpt::OptimizeAutoreleaseRVCall(Function &F,
 
   SmallVector<const Value *, 2> Users;
   Users.push_back(Ptr);
+
+  // Add PHIs that are equivalent to Ptr to Users.
+  if (const PHINode *PN = dyn_cast<PHINode>(Ptr))
+    getEquivalentPHIs(*PN, Users);
+
   do {
     Ptr = Users.pop_back_val();
     for (const User *U : Ptr->users()) {
diff --git a/llvm/test/Transforms/ObjCARC/contract-replace-arg-use.ll b/llvm/test/Transforms/ObjCARC/contract-replace-arg-use.ll
index 26b0cd9a42a..4cff9f7fc09 100644
--- a/llvm/test/Transforms/ObjCARC/contract-replace-arg-use.ll
+++ b/llvm/test/Transforms/ObjCARC/contract-replace-arg-use.ll
@@ -6,6 +6,7 @@ declare i8* @foo1()
 ; Check that ARC contraction replaces the function return with the value
 ; returned by @objc_autoreleaseReturnValue.
 
+; CHECK-LABEL: define i32* @autoreleaseRVTailCall(
 ; CHECK: %[[V0:[0-9]+]] = tail call i8* @objc_autoreleaseReturnValue(
 ; CHECK: %[[V1:[0-9]+]] = bitcast i8* %[[V0]] to i32*
 ; CHECK: ret i32* %[[V1]]
@@ -16,3 +17,30 @@ define i32* @autoreleaseRVTailCall() {
   %3 = tail call i8* @objc_autoreleaseReturnValue(i8* %1)
   ret i32* %2
 }
+
+declare i32* @foo2(i32);
+
+; CHECK-LABEL: define i32* @autoreleaseRVTailCallPhi(
+; CHECK: %[[PHIVAL:.*]] = phi i8* [ %{{.*}}, %bb1 ], [ %{{.*}}, %bb2 ]
+; CHECK: %[[RETVAL:.*]] = phi i32* [ %{{.*}}, %bb1 ], [ %{{.*}}, %bb2 ]
+; CHECK: %[[V4:.*]] = tail call i8* @objc_autoreleaseReturnValue(i8* %[[PHIVAL]])
+; CHECK: %[[V0:.*]] = bitcast i8* %[[V4]] to i32*
+; CHECK: ret i32* %[[V0]]
+
+define i32* @autoreleaseRVTailCallPhi(i1 %cond) {
+entry:
+  br i1 %cond, label %bb1, label %bb2
+bb1:
+  %v0 = call i32* @foo2(i32 1)
+  %v1 = bitcast i32* %v0 to i8*
+  br label %bb3
+bb2:
+  %v2 = call i32* @foo2(i32 2)
+  %v3 = bitcast i32* %v2 to i8*
+  br label %bb3
+bb3:
+  %phival = phi i8* [ %v1, %bb1 ], [ %v3, %bb2 ]
+  %retval = phi i32* [ %v0, %bb1 ], [ %v2, %bb2 ]
+  %v4 = tail call i8* @objc_autoreleaseReturnValue(i8* %phival)
+  ret i32* %retval
+}
diff --git a/llvm/test/Transforms/ObjCARC/rv.ll b/llvm/test/Transforms/ObjCARC/rv.ll
index d84a875684f..425f86cafb3 100644
--- a/llvm/test/Transforms/ObjCARC/rv.ll
+++ b/llvm/test/Transforms/ObjCARC/rv.ll
@@ -333,6 +333,37 @@ bb2:
   ret i8* %v2
 }
 
+declare i32* @func27(i32);
+
+; Check that ObjCARCOpt::OptimizeAutoreleaseRVCall doesn't turn a call to
+; @objc_autoreleaseReturnValue into a call to @objc_autorelease when a return
+; instruction uses a value equivalent to @objc_autoreleaseReturnValue's operand.
+; In the code below, %phival and %retval are considered equivalent.
+
+; CHECK-LABEL: define i32* @test27(
+; CHECK: %[[PHIVAL:.*]] = phi i8* [ %{{.*}}, %bb1 ], [ %{{.*}}, %bb2 ]
+; CHECK: %[[RETVAL:.*]] = phi i32* [ %{{.*}}, %bb1 ], [ %{{.*}}, %bb2 ]
+; CHECK: tail call i8* @objc_autoreleaseReturnValue(i8* %[[PHIVAL]])
+; CHECK: ret i32* %[[RETVAL]]
+
+define i32* @test27(i1 %cond) {
+entry:
+  br i1 %cond, label %bb1, label %bb2
+bb1:
+  %v0 = call i32* @func27(i32 1)
+  %v1 = bitcast i32* %v0 to i8*
+  br label %bb3
+bb2:
+  %v2 = call i32* @func27(i32 2)
+  %v3 = bitcast i32* %v2 to i8*
+  br label %bb3
+bb3:
+  %phival = phi i8* [ %v1, %bb1 ], [ %v3, %bb2 ]
+  %retval = phi i32* [ %v0, %bb1 ], [ %v2, %bb2 ]
+  %v4 = tail call i8* @objc_autoreleaseReturnValue(i8* %phival)
+  ret i32* %retval
+}
+
 !0 = !{}
 
 ; CHECK: attributes [[NUW]] = { nounwind }