From 1fe8993ad81900eea4d0f5603dffd239e520c31d Mon Sep 17 00:00:00 2001
From: Akira Hatanaka <ahatanaka@apple.com>
Date: Mon, 8 Nov 2021 21:19:07 -0800
Subject: [PATCH] [ObjC][ARC] Replace uses of ObjC intrinsics that are
 arguments of operand bundle "clang.arc.attachedcall" with ObjC runtime
 functions

The existing code only handles the case where the intrinsic being
rewritten is used as the called function pointer of a call/invoke.
---
 llvm/lib/CodeGen/PreISelIntrinsicLowering.cpp | 16 ++++++++-
 .../PreISelIntrinsicLowering/objc-arc.ll      | 34 +++++++++++++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/PreISelIntrinsicLowering.cpp b/llvm/lib/CodeGen/PreISelIntrinsicLowering.cpp
index 9547fe6f93de..d9c726743db1 100644
--- a/llvm/lib/CodeGen/PreISelIntrinsicLowering.cpp
+++ b/llvm/lib/CodeGen/PreISelIntrinsicLowering.cpp
@@ -13,6 +13,7 @@
 
 #include "llvm/CodeGen/PreISelIntrinsicLowering.h"
 #include "llvm/Analysis/ObjCARCInstKind.h"
+#include "llvm/Analysis/ObjCARCUtil.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
@@ -90,7 +91,20 @@ static bool lowerObjCCall(Function &F, const char *NewFn,
   CallInst::TailCallKind OverridingTCK = getOverridingTailCallKind(F);
 
   for (auto I = F.use_begin(), E = F.use_end(); I != E;) {
-    auto *CI = cast<CallInst>(I->getUser());
+    auto *CB = cast<CallBase>(I->getUser());
+
+    if (CB->getCalledFunction() != &F) {
+      objcarc::ARCInstKind Kind = objcarc::getAttachedARCFunctionKind(CB);
+      (void)Kind;
+      assert((Kind == objcarc::ARCInstKind::RetainRV ||
+              Kind == objcarc::ARCInstKind::ClaimRV) &&
+             "use expected to be the argument of operand bundle "
+             "\"clang.arc.attachedcall\"");
+      I++->set(FCache.getCallee());
+      continue;
+    }
+
+    auto *CI = cast<CallInst>(CB);
     assert(CI->getCalledFunction() && "Cannot lower an indirect call!");
     ++I;
 
diff --git a/llvm/test/Transforms/PreISelIntrinsicLowering/objc-arc.ll b/llvm/test/Transforms/PreISelIntrinsicLowering/objc-arc.ll
index aff01910cd9b..67d8ff697e60 100644
--- a/llvm/test/Transforms/PreISelIntrinsicLowering/objc-arc.ll
+++ b/llvm/test/Transforms/PreISelIntrinsicLowering/objc-arc.ll
@@ -3,6 +3,9 @@
 ; Make sure calls to the objc intrinsics are translated to calls in to the
 ; runtime
 
+declare i8* @foo()
+declare i32 @__gxx_personality_v0(...)
+
 define i8* @test_objc_autorelease(i8* %arg0) {
 ; CHECK-LABEL: test_objc_autorelease
 ; CHECK-NEXT: entry
@@ -153,6 +156,30 @@ entry:
 	ret i8* %0
 }
 
+define void @test_objc_retainAutoreleasedReturnValue_bundle() {
+; CHECK-LABEL: test_objc_retainAutoreleasedReturnValue_bundle(
+; CHECK-NEXT: call i8* @foo() [ "clang.arc.attachedcall"(i8* (i8*)* @objc_retainAutoreleasedReturnValue) ]
+  call i8* @foo() [ "clang.arc.attachedcall"(i8* (i8*)* @llvm.objc.retainAutoreleasedReturnValue) ]
+  ret void
+}
+
+define void @test_objc_retainAutoreleasedReturnValue_bundle_invoke() personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
+; CHECK-LABEL: test_objc_retainAutoreleasedReturnValue_bundle_invoke(
+; CHECK-NEXT: entry
+; CHECK-NEXT: invoke i8* @foo() [ "clang.arc.attachedcall"(i8* (i8*)* @objc_retainAutoreleasedReturnValue) ]
+entry:
+  invoke i8* @foo() [ "clang.arc.attachedcall"(i8* (i8*)* @llvm.objc.retainAutoreleasedReturnValue) ]
+      to label %invoke.cont unwind label %lpad
+
+invoke.cont:
+  ret void
+
+lpad:
+  %1 = landingpad { i8*, i32 }
+          cleanup
+  resume { i8*, i32 } %1
+}
+
 define i8* @test_objc_retainBlock(i8* %arg0) {
 ; CHECK-LABEL: test_objc_retainBlock
 ; CHECK-NEXT: entry
@@ -193,6 +220,13 @@ entry:
   ret i8* %0
 }
 
+define void @test_objc_unsafeClaimAutoreleasedReturnValue_bundle() {
+; CHECK-LABEL: test_objc_unsafeClaimAutoreleasedReturnValue_bundle(
+; CHECK-NEXT: call i8* @foo() [ "clang.arc.attachedcall"(i8* (i8*)* @objc_unsafeClaimAutoreleasedReturnValue) ]
+  call i8* @foo() [ "clang.arc.attachedcall"(i8* (i8*)* @llvm.objc.unsafeClaimAutoreleasedReturnValue) ]
+  ret void
+}
+
 define i8* @test_objc_retainedObject(i8* %arg0) {
 ; CHECK-LABEL: test_objc_retainedObject
 ; CHECK-NEXT: entry