[ObjCARC] Traverse chain downwards to replace uses of argument passed to ObjC library call with call return

[ahatanak]

Summary

ObjCARCContract::runOnFunction tries to replace uses of an argument passed to an objective-c library call with the call
return value. For example, in the following IR, it replaces uses of argument %9 and uses of the values discovered traversing the chain upwards (%7 and %8) with the call return %10, if they are dominated by the call to @objc_autoreleaseReturnValue. This transformation enables code-gen to tail-call the call to @objc_autoreleaseReturnValue, which is necessary to enable auto release return value optimization.

%7 = tail call i8* @objc_loadWeakRetained(i8** %6)
%8 = bitcast i8* %7 to %0*
%9 = bitcast %0* %8 to i8*
%10 = tail call i8* @objc_autoreleaseReturnValue(i8* %9)
ret %0* %8

Since r276727, llvm started removing redundant bitcasts and as a result started feeding the following IR to ARC contraction:

%7 = tail call i8* @objc_loadWeakRetained(i8** %6)
%8 = bitcast i8* %7 to %0*
%9 = tail call i8* @objc_autoreleaseReturnValue(i8* %7)
ret %0* %8

ARC contraction no longer does the optimization described above since it only traverses the chain upwards and fails to recognize that the function return can be replaced by the call return from @objc_autoreleaseReturnValue. This patch changes ObjCARCContract::runOnFunction to traverse the chain downwards too and replace uses of bitcasts with the call return.

rdar://problem/28011339

lib/Transforms/ObjCARC/ObjCARCContract.cpp

for (inst_iterator I = inst_begin(&F), E = inst_end(&F); I != E;) {
     if (tryToPeepholeInstruction(F, Inst, I, DependingInstructions, Visited,
                                  TailOkForStoreStrongs))
       continue;
 
     // Otherwise, try to undo objc-arc-expand.
 
     // Don't use GetArgRCIdentityRoot because we don't want to look through bitcasts
     // and such; to do the replacement, the argument must have type i8*.
-    Value *Arg = cast<CallInst>(Inst)->getArgOperand(0);
 
-    // TODO: Change this to a do-while.
-    for (;;) {
+    // Function for replacing uses of Arg dominated by Inst.
+    auto ReplaceArgUses = [Inst, this](Value *Arg) {
       // If we're compiling bugpointed code, don't get in trouble.
       if (!isa<Instruction>(Arg) && !isa<Argument>(Arg))
-        break;
+        return;
+
       // Look through the uses of the pointer.
       for (Value::use_iterator UI = Arg->use_begin(), UE = Arg->use_end();
            UI != UE; ) {
         // Increment UI now, because we may unlink its element.
         Use &U = *UI++;
         unsigned OperandNo = U.getOperandNo();
 
         // If the call's return value dominates a use of the call's argument
auto ReplaceArgUses = [Inst, this](Value *Arg) {
           } else {
             if (Replacement->getType() != UseTy)
               Replacement = new BitCastInst(Replacement, UseTy, "",
                                             cast<Instruction>(U.getUser()));
             U.set(Replacement);
           }
         }
       }
+    };
+
+
+    Value *Arg = cast<CallInst>(Inst)->getArgOperand(0), *OrigArg = Arg;
+
+    // TODO: Change this to a do-while.
+    for (;;) {
+      ReplaceArgUses(Arg);
 
       // If Arg is a no-op casted pointer, strip one level of casts and iterate.
       if (const BitCastInst *BI = dyn_cast<BitCastInst>(Arg))
         Arg = BI->getOperand(0);
       else if (isa<GEPOperator>(Arg) &&
                cast<GEPOperator>(Arg)->hasAllZeroIndices())
         Arg = cast<GEPOperator>(Arg)->getPointerOperand();
       else if (isa<GlobalAlias>(Arg) &&
                !cast<GlobalAlias>(Arg)->isInterposable())
         Arg = cast<GlobalAlias>(Arg)->getAliasee();
       else
         break;
     }
+
+    // Replace bitcast users of Arg that are dominated by Inst.
+    SmallVector<BitCastInst *, 2> BitCastUsers;
+
+    // Add all bitcast users of the function argument first.
+    for (User *U : OrigArg->users())
+      if (auto *BC = dyn_cast<BitCastInst>(U))
+        BitCastUsers.push_back(BC);
+
+    // Replace the bitcasts with the call return. Iterate until list is empty.
+    while (!BitCastUsers.empty()) {
+      auto *BC = BitCastUsers.pop_back_val();
+      for (User *U : BC->users())
+        if (auto *B = dyn_cast<BitCastInst>(U))
+          BitCastUsers.push_back(B);
+
+      ReplaceArgUses(BC);
+    }
   }
 
   // If this function has no escaping allocas or suspicious vararg usage,
   // objc_storeStrong calls can be marked with the "tail" keyword.
   if (TailOkForStoreStrongs)
     for (CallInst *CI : StoreStrongCalls)
       CI->setTailCall();
   StoreStrongCalls.clear();

test/Transforms/ObjCARC/contract-replace-arg-use.ll

+; RUN: opt -objc-arc-contract -S < %s | FileCheck %s
+
+declare i8* @objc_autoreleaseReturnValue(i8*)
+declare i8* @foo1()
+
+; Check that ARC contraction replaces the function return with the value
+; returned by @objc_autoreleaseReturnValue.
+
+; CHECK: %[[V0:[0-9]+]] = tail call i8* @objc_autoreleaseReturnValue(
+; CHECK: %[[V1:[0-9]+]] = bitcast i8* %[[V0]] to i32*
+; CHECK: ret i32* %[[V1]]
+
+define i32* @autoreleaseRVTailCall() {
+  %1 = call i8* @foo1()
+  %2 = bitcast i8* %1 to i32*
+  %3 = tail call i8* @objc_autoreleaseReturnValue(i8* %1)
+  ret i32* %2
+}