[ObjC][ARC] Prevent moving objc_retain calls past objc_release calls

that release the retained object

This patch fixes what looks like a longstanding bug in ARC optimizer
where it reverses the order of objc_retain calls and objc_release calls
that retain and release the same object.

The code in ARC optimizer that is responsible for code motion takes the
following steps:

1. Traverse the CFG bottom-up and determine how far up objc_release
   calls can be moved. Determine the insertion points for the
   objc_release calls, but don't actually move them.
2. Traverse the CFG top-down and determine how far down objc_retain
   calls can be moved. Determine the insertion points for the
   objc_retain calls, but don't actually move them.
3. Try to move the objc_retain and objc_release calls if they can't be
   removed.

The problem is that the insertion points for the objc_retain calls are
determined in step 2 without taking into consideration the insertion
points for objc_release calls determined in step 1, so the order of an
objc_retain call and an objc_release call can be reversed, which is
incorrect, even though each step is correct in isolation.

To fix this bug, this patch teaches the top-down traversal step to take
into consideration the insertion points for objc_release calls
determined in the bottom-up traversal step. Code motion for an
objc_retain call is disabled if there is a possibility that it can be
moved past an objc_release call that releases the retained object.

rdar://79292791

Differential Revision: https://reviews.llvm.org/D104953

Author: ahatanaka

llvm/lib/Transforms/ObjCARC/ObjCARCOpts.cpp

@@ -532,12 +532,15 @@ class ObjCARCOpt {
   bool VisitBottomUp(BasicBlock *BB,
                      DenseMap<const BasicBlock *, BBState> &BBStates,
                      BlotMapVector<Value *, RRInfo> &Retains);
-  bool VisitInstructionTopDown(Instruction *Inst,
-                               DenseMap<Value *, RRInfo> &Releases,
-                               BBState &MyStates);
-  bool VisitTopDown(BasicBlock *BB,
-                    DenseMap<const BasicBlock *, BBState> &BBStates,
-                    DenseMap<Value *, RRInfo> &Releases);
+  bool VisitInstructionTopDown(
+      Instruction *Inst, DenseMap<Value *, RRInfo> &Releases, BBState &MyStates,
+      const DenseMap<const Instruction *, SmallPtrSet<const Value *, 2>>
+          &ReleaseInsertPtToRCIdentityRoots);
+  bool VisitTopDown(
+      BasicBlock *BB, DenseMap<const BasicBlock *, BBState> &BBStates,
+      DenseMap<Value *, RRInfo> &Releases,
+      const DenseMap<const Instruction *, SmallPtrSet<const Value *, 2>>
+          &ReleaseInsertPtToRCIdentityRoots);
   bool Visit(Function &F, DenseMap<const BasicBlock *, BBState> &BBStates,
              BlotMapVector<Value *, RRInfo> &Retains,
              DenseMap<Value *, RRInfo> &Releases);
@@ -1495,14 +1498,62 @@ bool ObjCARCOpt::VisitBottomUp(BasicBlock *BB,
   return NestingDetected;
 }
 
-bool
-ObjCARCOpt::VisitInstructionTopDown(Instruction *Inst,
-                                    DenseMap<Value *, RRInfo> &Releases,
-                                    BBState &MyStates) {
+// Fill ReleaseInsertPtToRCIdentityRoots, which is a map from insertion points
+// to the set of RC identity roots that would be released by the release calls
+// moved to the insertion points.
+static void collectReleaseInsertPts(
+    const BlotMapVector<Value *, RRInfo> &Retains,
+    DenseMap<const Instruction *, SmallPtrSet<const Value *, 2>>
+        &ReleaseInsertPtToRCIdentityRoots) {
+  for (auto &P : Retains) {
+    // Retains is a map from an objc_retain call to a RRInfo of the RC identity
+    // root of the call. Get the RC identity root of the objc_retain call.
+    Instruction *Retain = cast<Instruction>(P.first);
+    Value *Root = GetRCIdentityRoot(Retain->getOperand(0));
+    // Collect all the insertion points of the objc_release calls that release
+    // the RC identity root of the objc_retain call.
+    for (const Instruction *InsertPt : P.second.ReverseInsertPts)
+      ReleaseInsertPtToRCIdentityRoots[InsertPt].insert(Root);
+  }
+}
+
+// Get the RC identity roots from an insertion point of an objc_release call.
+// Return nullptr if the passed instruction isn't an insertion point.
+static const SmallPtrSet<const Value *, 2> *
+getRCIdentityRootsFromReleaseInsertPt(
+    const Instruction *InsertPt,
+    const DenseMap<const Instruction *, SmallPtrSet<const Value *, 2>>
+        &ReleaseInsertPtToRCIdentityRoots) {
+  auto I = ReleaseInsertPtToRCIdentityRoots.find(InsertPt);
+  if (I == ReleaseInsertPtToRCIdentityRoots.end())
+    return nullptr;
+  return &I->second;
+}
+
+bool ObjCARCOpt::VisitInstructionTopDown(
+    Instruction *Inst, DenseMap<Value *, RRInfo> &Releases, BBState &MyStates,
+    const DenseMap<const Instruction *, SmallPtrSet<const Value *, 2>>
+        &ReleaseInsertPtToRCIdentityRoots) {
   bool NestingDetected = false;
   ARCInstKind Class = GetARCInstKind(Inst);
   const Value *Arg = nullptr;
 
+  // Make sure a call to objc_retain isn't moved past insertion points of calls
+  // to objc_release.
+  if (const SmallPtrSet<const Value *, 2> *Roots =
+          getRCIdentityRootsFromReleaseInsertPt(
+              Inst, ReleaseInsertPtToRCIdentityRoots))
+    for (auto *Root : *Roots) {
+      TopDownPtrState &S = MyStates.getPtrTopDownState(Root);
+      // Disable code motion if the current position is S_Retain to prevent
+      // moving the objc_retain call past objc_release calls. If it's
+      // S_CanRelease or larger, it's not necessary to disable code motion as
+      // the insertion points that prevent the objc_retain call from moving down
+      // should have been set already.
+      if (S.GetSeq() == S_Retain)
+        S.SetCFGHazardAfflicted(true);
+    }
+
   LLVM_DEBUG(dbgs() << "        Class: " << Class << "\n");
 
   switch (Class) {
@@ -1565,10 +1616,11 @@ ObjCARCOpt::VisitInstructionTopDown(Instruction *Inst,
   return NestingDetected;
 }
 
-bool
-ObjCARCOpt::VisitTopDown(BasicBlock *BB,
-                         DenseMap<const BasicBlock *, BBState> &BBStates,
-                         DenseMap<Value *, RRInfo> &Releases) {
+bool ObjCARCOpt::VisitTopDown(
+    BasicBlock *BB, DenseMap<const BasicBlock *, BBState> &BBStates,
+    DenseMap<Value *, RRInfo> &Releases,
+    const DenseMap<const Instruction *, SmallPtrSet<const Value *, 2>>
+        &ReleaseInsertPtToRCIdentityRoots) {
   LLVM_DEBUG(dbgs() << "\n== ObjCARCOpt::VisitTopDown ==\n");
   bool NestingDetected = false;
   BBState &MyStates = BBStates[BB];
@@ -1608,7 +1660,8 @@ ObjCARCOpt::VisitTopDown(BasicBlock *BB,
   for (Instruction &Inst : *BB) {
     LLVM_DEBUG(dbgs() << "    Visiting " << Inst << "\n");
 
-    NestingDetected |= VisitInstructionTopDown(&Inst, Releases, MyStates);
+    NestingDetected |= VisitInstructionTopDown(
+        &Inst, Releases, MyStates, ReleaseInsertPtToRCIdentityRoots);
 
     // Bail out if the number of pointers being tracked becomes too large so
     // that this pass can complete in a reasonable amount of time.
@@ -1728,10 +1781,15 @@ bool ObjCARCOpt::Visit(Function &F,
       return false;
   }
 
+  DenseMap<const Instruction *, SmallPtrSet<const Value *, 2>>
+      ReleaseInsertPtToRCIdentityRoots;
+  collectReleaseInsertPts(Retains, ReleaseInsertPtToRCIdentityRoots);
+
   // Use reverse-postorder for top-down.
   bool TopDownNestingDetected = false;
   for (BasicBlock *BB : llvm::reverse(PostOrder)) {
-    TopDownNestingDetected |= VisitTopDown(BB, BBStates, Releases);
+    TopDownNestingDetected |=
+        VisitTopDown(BB, BBStates, Releases, ReleaseInsertPtToRCIdentityRoots);
     if (DisableRetainReleasePairing)
       return false;
   }

llvm/test/Transforms/ObjCARC/code-motion.ll

@@ -39,26 +39,32 @@ define void @test2() {
   ret void
 }
 
-; ARC optimizer shouldn't reverse the order of retains and releases in
-; if.then in @test3 and @test4.
+; Check that code motion is disabled in @test3 and @test4.
+; Previously, ARC optimizer would move the release past the retain.
+
+; if.then:
+;   call void @readOnlyFunc(i8* %obj, i8* null)
+;   call void @llvm.objc.release(i8* %obj) #1, !clang.imprecise_release !2
+;   %1 = add i32 1, 2
+;   %2 = tail call i8* @llvm.objc.retain(i8* %obj)
+;
+; Ideally, the retain/release pairs in BB if.then should be removed.
 
 define void @test3(i8* %obj, i1 %cond) {
 ; CHECK-LABEL: @test3(
+; CHECK-NEXT:    [[TMP2:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ:%.*]])
 ; CHECK-NEXT:    br i1 [[COND:%.*]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
 ; CHECK:       if.then:
-; CHECK-NEXT:    call void @readOnlyFunc(i8* [[OBJ:%.*]], i8* null)
-; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ]]) {{.*}}, !clang.imprecise_release !2
+; CHECK-NEXT:    call void @readOnlyFunc(i8* [[OBJ]], i8* null)
 ; CHECK-NEXT:    [[TMP1:%.*]] = add i32 1, 2
-; CHECK-NEXT:    [[TMP2:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ]])
 ; CHECK-NEXT:    call void @alterRefCount()
 ; CHECK-NEXT:    br label [[JOIN:%.*]]
 ; CHECK:       if.else:
-; CHECK-NEXT:    [[TMP3:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ]])
 ; CHECK-NEXT:    call void @alterRefCount()
 ; CHECK-NEXT:    call void @use(i8* [[OBJ]])
-; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ]]) {{.*}}, !clang.imprecise_release !2
 ; CHECK-NEXT:    br label [[JOIN]]
 ; CHECK:       join:
+; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ]]) {{.*}}, !clang.imprecise_release !2
 ; CHECK-NEXT:    ret void
 ;
   %v0 = call i8* @llvm.objc.retain(i8* %obj)
@@ -82,26 +88,22 @@ join:
 
 define void @test4(i8* %obj0, i8* %obj1, i1 %cond) {
 ; CHECK-LABEL: @test4(
+; CHECK-NEXT:    [[TMP3:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ0:%.*]])
+; CHECK-NEXT:    [[TMP2:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ1:%.*]])
 ; CHECK-NEXT:    br i1 [[COND:%.*]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
 ; CHECK:       if.then:
-; CHECK-NEXT:    call void @readOnlyFunc(i8* [[OBJ0:%.*]], i8* [[OBJ1:%.*]])
-; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ1]]) {{.*}}, !clang.imprecise_release !2
-; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ0]]) {{.*}}, !clang.imprecise_release !2
+; CHECK-NEXT:    call void @readOnlyFunc(i8* [[OBJ0]], i8* [[OBJ1]])
 ; CHECK-NEXT:    [[TMP1:%.*]] = add i32 1, 2
-; CHECK-NEXT:    [[TMP2:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ1]])
-; CHECK-NEXT:    [[TMP3:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ0]])
 ; CHECK-NEXT:    call void @alterRefCount()
 ; CHECK-NEXT:    br label [[JOIN:%.*]]
 ; CHECK:       if.else:
-; CHECK-NEXT:    [[TMP4:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ1]])
-; CHECK-NEXT:    [[TMP5:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ0]])
 ; CHECK-NEXT:    call void @alterRefCount()
 ; CHECK-NEXT:    call void @use(i8* [[OBJ0]])
-; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ0]]) {{.*}}, !clang.imprecise_release !2
 ; CHECK-NEXT:    call void @use(i8* [[OBJ1]])
-; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ1]]) {{.*}}, !clang.imprecise_release !2
 ; CHECK-NEXT:    br label [[JOIN]]
 ; CHECK:       join:
+; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ0]]) {{.*}}, !clang.imprecise_release !2
+; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ1]]) {{.*}}, !clang.imprecise_release !2
 ; CHECK-NEXT:    ret void
 ;
   %v0 = call i8* @llvm.objc.retain(i8* %obj0)
@@ -131,27 +133,25 @@ join:
 
 define void @test5(i8* %obj, i1 %cond0, i1 %cond1) {
 ; CHECK-LABEL: @test5(
+; CHECK-NEXT:    [[V0:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ:%.*]])
 ; CHECK-NEXT:    br i1 [[COND0:%.*]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
 ; CHECK:       if.then:
-; CHECK-NEXT:    call void @readOnlyFunc(i8* [[OBJ:%.*]], i8* null)
-; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ]])
+; CHECK-NEXT:    call void @readOnlyFunc(i8* [[OBJ]], i8* null)
 ; CHECK-NEXT:    br i1 [[COND1:%.*]], label [[IF_THEN2:%.*]], label [[IF_ELSE2:%.*]]
 ; CHECK:       if.then2:
 ; CHECK-NEXT:    br label [[BB1:%.*]]
 ; CHECK:       if.else2:
 ; CHECK-NEXT:    br label [[BB1]]
 ; CHECK:       bb1:
 ; CHECK-NEXT:    [[TMP1:%.*]] = add i32 1, 2
-; CHECK-NEXT:    [[TMP2:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ]])
 ; CHECK-NEXT:    call void @alterRefCount()
 ; CHECK-NEXT:    br label [[JOIN:%.*]]
 ; CHECK:       if.else:
-; CHECK-NEXT:    [[TMP3:%.*]] = tail call i8* @llvm.objc.retain(i8* [[OBJ]])
 ; CHECK-NEXT:    call void @alterRefCount()
 ; CHECK-NEXT:    call void @use(i8* [[OBJ]])
-; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ]])
 ; CHECK-NEXT:    br label [[JOIN]]
 ; CHECK:       join:
+; CHECK-NEXT:    call void @llvm.objc.release(i8* [[OBJ]])
 ; CHECK-NEXT:    ret void
 ;
   %v0 = call i8* @llvm.objc.retain(i8* %obj)