[llvm] [ObjCARC] Change autorelease to release when the pool state is not changed between the autorelease and the pool pop (PR #152353)

[via llvm-commits]

https://github.com/AZero13 updated https://github.com/llvm/llvm-project/pull/152353

>From a12776cff505510ae4772f3d626cb6599ef111d2 Mon Sep 17 00:00:00 2001
From: AZero13 <gfunni234 at gmail.com>
Date: Wed, 6 Aug 2025 14:07:16 -0400
Subject: [PATCH] [ObjCARC] Change autorelease to release when the pool state
 is not changed between the autorelease and the pool pop

Co-Authored-By: Jon Roelofs <jroelofs at gmail.com>
---
 llvm/lib/Transforms/ObjCARC/ObjCARCOpts.cpp   | 74 ++++++++++++++++++-
 .../ObjCARC/test_autorelease_pool.ll          |  6 +-
 2 files changed, 72 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Transforms/ObjCARC/ObjCARCOpts.cpp b/llvm/lib/Transforms/ObjCARC/ObjCARCOpts.cpp
index 66a2c7632aadc..12958e4b2af7f 100644
--- a/llvm/lib/Transforms/ObjCARC/ObjCARCOpts.cpp
+++ b/llvm/lib/Transforms/ObjCARC/ObjCARCOpts.cpp
@@ -121,6 +121,62 @@ static const Value *FindSingleUseIdentifiedObject(const Value *Arg) {
 /// \defgroup ARCOpt ARC Optimization.
 /// @{
 
+/// Check if there is an autoreleasePoolPop after the given autorelease
+/// instruction in the same basic block with no intervening calls that
+/// could affect the autorelease pool.
+static bool HasFollowingAutoreleasePoolPop(Instruction *AutoreleaseInst) {
+  assert(IsAutorelease(GetBasicARCInstKind(AutoreleaseInst)));
+  BasicBlock *BB = AutoreleaseInst->getParent();
+
+  // Look forward from the autorelease instruction
+  for (BasicBlock::iterator I = std::next(AutoreleaseInst->getIterator()),
+                            E = BB->end();
+       I != E; ++I) {
+    ARCInstKind Class = GetBasicARCInstKind(&*I);
+
+    switch (Class) {
+    case ARCInstKind::AutoreleasepoolPop:
+      // Found a pool pop - the autorelease will be drained
+      return true;
+
+    case ARCInstKind::AutoreleasepoolPush:
+      // A new pool started - this autorelease won't be drained by a later pop
+      return false;
+
+    case ARCInstKind::CallOrUser:
+    case ARCInstKind::Call:
+      // Unknown call could affect autorelease pool state or return autoreleased
+      // objects
+      return false;
+    case ARCInstKind::Retain:
+    case ARCInstKind::RetainRV:
+    case ARCInstKind::UnsafeClaimRV:
+    case ARCInstKind::RetainBlock:
+    case ARCInstKind::Release:
+    case ARCInstKind::Autorelease:
+    case ARCInstKind::AutoreleaseRV:
+    case ARCInstKind::NoopCast:
+    case ARCInstKind::FusedRetainAutorelease:
+    case ARCInstKind::FusedRetainAutoreleaseRV:
+    case ARCInstKind::LoadWeakRetained:
+    case ARCInstKind::StoreWeak:
+    case ARCInstKind::InitWeak:
+    case ARCInstKind::LoadWeak:
+    case ARCInstKind::MoveWeak:
+    case ARCInstKind::CopyWeak:
+    case ARCInstKind::DestroyWeak:
+    case ARCInstKind::StoreStrong:
+    case ARCInstKind::IntrinsicUser:
+    case ARCInstKind::None:
+      // Everything else is safe to ignore:
+      break;
+    }
+  }
+
+  // Reached end of basic block without finding a pool pop
+  return false;
+}
+
 // TODO: On code like this:
 //
 // objc_retain(%x)
@@ -978,12 +1034,22 @@ void ObjCARCOpt::OptimizeIndividualCallImpl(Function &F, Instruction *Inst,
     break;
   }
 
-  // objc_autorelease(x) -> objc_release(x) if x is otherwise unused.
+  // objc_autorelease(x) -> objc_release(x) if x is otherwise unused
+  // OR if this autorelease is followed by an autoreleasePoolPop.
   if (IsAutorelease(Class) && Inst->use_empty()) {
     CallInst *Call = cast<CallInst>(Inst);
     const Value *Arg = Call->getArgOperand(0);
     Arg = FindSingleUseIdentifiedObject(Arg);
-    if (Arg) {
+    bool ShouldConvert = (Arg != nullptr);
+    const char *Reason = "since x is otherwise unused";
+
+    // Also convert if this autorelease is followed by a pool pop
+    if (!ShouldConvert && HasFollowingAutoreleasePoolPop(Inst)) {
+      ShouldConvert = true;
+      Reason = "since it's followed by autoreleasePoolPop";
+    }
+
+    if (ShouldConvert) {
       Changed = true;
       ++NumAutoreleases;
 
@@ -997,8 +1063,8 @@ void ObjCARCOpt::OptimizeIndividualCallImpl(Function &F, Instruction *Inst,
                            MDNode::get(C, {}));
 
       LLVM_DEBUG(dbgs() << "Replacing autorelease{,RV}(x) with objc_release(x) "
-                           "since x is otherwise unused.\nOld: "
-                        << *Call << "\nNew: " << *NewCall << "\n");
+                        << Reason << ".\nOld: " << *Call
+                        << "\nNew: " << *NewCall << "\n");
 
       EraseInstruction(Call);
       Inst = NewCall;
diff --git a/llvm/test/Transforms/ObjCARC/test_autorelease_pool.ll b/llvm/test/Transforms/ObjCARC/test_autorelease_pool.ll
index 896717f92146f..150837f922c39 100644
--- a/llvm/test/Transforms/ObjCARC/test_autorelease_pool.ll
+++ b/llvm/test/Transforms/ObjCARC/test_autorelease_pool.ll
@@ -44,7 +44,7 @@ define void @test_multiple_autoreleases() {
 ; CHECK-NEXT:    call void @use_object(ptr [[OBJ1]])
 ; CHECK-NEXT:    [[TMP1:%.*]] = call ptr @llvm.objc.autorelease(ptr [[OBJ1]]) #[[ATTR0]]
 ; CHECK-NEXT:    call void @use_object(ptr [[OBJ2]])
-; CHECK-NEXT:    [[TMP2:%.*]] = call ptr @llvm.objc.autorelease(ptr [[OBJ2]]) #[[ATTR0]]
+; CHECK-NEXT:    call void @llvm.objc.release(ptr [[OBJ2]]) #[[ATTR0]], !clang.imprecise_release [[META0]]
 ; CHECK-NEXT:    call void @llvm.objc.autoreleasePoolPop(ptr [[POOL]]) #[[ATTR0]]
 ; CHECK-NEXT:    ret void
 ;
@@ -211,9 +211,7 @@ define void @test_complex_shadowing() {
 ; CHECK-NEXT:    [[OBJ3:%.*]] = call ptr @create_object()
 ; CHECK-NEXT:    call void @llvm.objc.release(ptr [[OBJ1]]) #[[ATTR0]], !clang.imprecise_release [[META0]]
 ; CHECK-NEXT:    call void @llvm.objc.release(ptr [[OBJ2]]) #[[ATTR0]], !clang.imprecise_release [[META0]]
-; CHECK-NEXT:    [[INNER2_POOL:%.*]] = call ptr @llvm.objc.autoreleasePoolPush() #[[ATTR0]]
-; CHECK-NEXT:    [[TMP1:%.*]] = call ptr @llvm.objc.autorelease(ptr [[OBJ3]]) #[[ATTR0]]
-; CHECK-NEXT:    call void @llvm.objc.autoreleasePoolPop(ptr [[INNER2_POOL]]) #[[ATTR0]]
+; CHECK-NEXT:    call void @llvm.objc.release(ptr [[OBJ3]]) #[[ATTR0]], !clang.imprecise_release [[META0]]
 ; CHECK-NEXT:    ret void
 ;
   %obj1 = call ptr @create_object()