Be more aggressive with ObjCARCContract.cpp

If Inst is a retain, we can be more aggressive by checking it

Author: AZero13

llvm/lib/Transforms/ObjCARC/ObjCARCContract.cpp

@@ -243,13 +243,21 @@ static StoreInst *findSafeStoreForStoreStrongContraction(LoadInst *Load,
 
     // Ok, now we know we have not seen a store yet.
 
-    // If Inst is a retain, we don't care about it as it doesn't prevent moving
-    // the load to the store.
-    //
-    // TODO: This is one area where the optimization could be made more
-    // aggressive.
-    if (IsRetain(Class))
-      continue;
+    // If Inst is a retain, we can be more aggressive by checking if:
+    // 1. The retain is on an unrelated object (different RC identity root)
+    // 2. The retain's argument doesn't alias with our load location
+    // In these cases, we can safely move past the retain.
+    if (IsRetain(Class)) {
+      Value *RetainArg = GetArgRCIdentityRoot(Inst);
+      Value *LoadRoot = GetRCIdentityRoot(Load);
+      if (RetainArg != LoadRoot && !AA->alias(MemoryLocation::get(Load), 
+          MemoryLocation::getBeforeOrAfter(RetainArg))) {
+        continue;
+      }
+      // If the retain is on the same object or we can't prove no aliasing,
+      // be conservative and bail out
+      return nullptr;
+    }
 
     // See if Inst can write to our load location, if it can not, just ignore
     // the instruction.

llvm/test/Transforms/ObjCARC/contract.ll

@@ -32,7 +32,7 @@ entry:
   ret void
 }
 
-; Merge objc_retain and objc_autorelease into objc_retainAutorelease.
+; Merge objc_retain and objc_autorelease into objc_retainAutoreleasedReturnValue.
 
 ; CHECK-LABEL: define void @test2(
 ; CHECK: tail call ptr @llvm.objc.retainAutorelease(ptr %x) [[NUW:#[0-9]+]]
@@ -238,3 +238,34 @@ declare void @llvm.objc.clang.arc.use(...) nounwind
 declare void @llvm.objc.clang.arc.noop.use(...) nounwind
 
 ; CHECK: attributes [[NUW]] = { nounwind }
+
+; CHECK-LABEL: define void @test_aggressive_retain_opt(
+; CHECK: tail call ptr @llvm.objc.retainAutoreleasedReturnValue(ptr %p)
+; CHECK: }
+define void @test_aggressive_retain_opt() {
+  %p = call ptr @returner()
+  %unrelated = call ptr @returner()
+  tail call ptr @llvm.objc.retain(ptr %unrelated) nounwind
+  tail call ptr @llvm.objc.retain(ptr %p) nounwind
+  ret void
+}
+
+; CHECK-LABEL: define void @test_aggressive_retain_opt_alias(
+; CHECK: tail call ptr @llvm.objc.retain(ptr %p)
+; CHECK: }
+define void @test_aggressive_retain_opt_alias(ptr %p, ptr %q) {
+  store ptr %p, ptr %q
+  tail call ptr @llvm.objc.retain(ptr %p) nounwind
+  ret void
+}
+
+; CHECK-LABEL: define void @test_aggressive_retain_opt_no_alias(
+; CHECK: tail call ptr @llvm.objc.retainAutoreleasedReturnValue(ptr %p)
+; CHECK: }
+define void @test_aggressive_retain_opt_no_alias() {
+  %p = call ptr @returner()
+  %q = alloca ptr
+  store ptr null, ptr %q
+  tail call ptr @llvm.objc.retain(ptr %p) nounwind
+  ret void
+}