ARCCodeMotion: fix two problems in release hoisting:

lib/SILOptimizer/Transforms/ARCCodeMotion.cpp

@@ -618,13 +618,16 @@ class ReleaseBlockState : public BlockState {
   }
 
   /// constructor.
-  ReleaseBlockState(bool IsExit, unsigned size, bool MultiIteration) {
+  ///
+  /// If \p InitOptimistic is true, the block in-bits are initialized to 1
+  /// which enables optimistic data flow evaluation.
+  ReleaseBlockState(bool InitOptimistic, unsigned size) {
     // backward data flow.
     // Initialize to true if we are running optimistic data flow, i.e.
     // MultiIteration is true.
-    BBSetIn.resize(size, MultiIteration);
+    BBSetIn.resize(size, InitOptimistic);
     BBSetOut.resize(size, false);
-    BBMaxSet.resize(size, !IsExit && MultiIteration);
+    BBMaxSet.resize(size, InitOptimistic);
 
     // Genset and Killset are initially empty.
     BBGenSet.resize(size, false);
@@ -735,6 +738,17 @@ bool ReleaseCodeMotionContext::requireIteration() {
 }
 
 void ReleaseCodeMotionContext::initializeCodeMotionDataFlow() {
+  // All blocks which are initialized with 1-bits. These are all blocks which
+  // eventually reach the function exit (return, throw), excluding the
+  // function exit blocks themselves.
+  // Optimistic initialization enables moving releases across loops. On the
+  // other hand, blocks, which never reach the function exit, e.g. infinite
+  // loop blocks, must be excluded. Otherwise we would end up inserting
+  // completely unrelated release instructions in such blocks.
+  llvm::SmallPtrSet<SILBasicBlock *, 32> BlocksInitOptimistically;
+
+  llvm::SmallVector<SILBasicBlock *, 32> Worklist;
+
   // Find all the RC roots in the function.
   for (auto &BB : *F) {
     for (auto &II : BB) {
@@ -750,13 +764,25 @@ void ReleaseCodeMotionContext::initializeCodeMotionDataFlow() {
       RCRootIndex[Root] = RCRootVault.size();
       RCRootVault.insert(Root);
     }
+    if (MultiIteration && BB.getTerminator()->isFunctionExiting())
+      Worklist.push_back(&BB);
+  }
+
+  // Find all blocks from which there is a path to the function exit.
+  // Note: the Worklist is empty if we are not in MultiIteration mode.
+  while (!Worklist.empty()) {
+    SILBasicBlock *BB = Worklist.pop_back_val();
+    for (SILBasicBlock *Pred : BB->getPredecessorBlocks()) {
+      if (BlocksInitOptimistically.insert(Pred).second)
+        Worklist.push_back(Pred);
+    }
   }
 
   // Initialize all the data flow bit vector for all basic blocks.
   for (auto &BB : *F) {
     BlockStates[&BB] = new (BPA.Allocate())
-            ReleaseBlockState(BB.getTerminator()->isFunctionExiting(),
-                              RCRootVault.size(), MultiIteration);
+            ReleaseBlockState(BlocksInitOptimistically.count(&BB) != 0,
+                              RCRootVault.size());
   }
 }
 
@@ -1038,17 +1064,22 @@ class ARCCodeMotion : public SILFunctionTransform {
       return;
 
     DEBUG(llvm::dbgs() << "*** ARCCM on function: " << F->getName() << " ***\n");
+
+    PostOrderAnalysis *POA = PM->getAnalysis<PostOrderAnalysis>();
+
     // Split all critical edges.
     //
     // TODO: maybe we can do this lazily or maybe we should disallow SIL passes
     // to create critical edges.
     bool EdgeChanged = splitAllCriticalEdges(*F, false, nullptr, nullptr);
+    if (EdgeChanged)
+      POA->invalidateFunction(F);
 
-    llvm::SpecificBumpPtrAllocator<BlockState> BPA;
-    auto *PO = PM->getAnalysis<PostOrderAnalysis>()->get(F);
+    auto *PO = POA->get(F);
     auto *AA = PM->getAnalysis<AliasAnalysis>();
     auto *RCFI = PM->getAnalysis<RCIdentityAnalysis>()->get(F);
 
+    llvm::SpecificBumpPtrAllocator<BlockState> BPA;
     bool InstChanged = false;
     if (Kind == Release) {
       // TODO: we should consider Throw block as well, or better we should

test/SILOptimizer/retain_release_code_motion.sil

@@ -443,3 +443,91 @@ bb3:
   %5 = tuple()
   return %5 : $()
 }
+
+// CHECK-LABEL: sil @move_retain_over_loop
+// CHECK:      bb0({{.*}}):
+// CHECK-NEXT:   br bb1
+// CHECK:      bb1:
+// CHECK-NEXT:   cond_br
+// CHECK:      bb2:
+// CHECK:        strong_retain
+// CHECK:        apply
+// CHECK:        strong_release
+// CHECK:        return
+sil @move_retain_over_loop : $@convention(thin) (Builtin.NativeObject) -> () {
+bb0(%0 : $Builtin.NativeObject):
+  strong_retain %0 : $Builtin.NativeObject
+  br bb1
+
+bb1:
+  cond_br undef, bb1, bb2
+
+bb2:
+  %2 = function_ref @blocker : $@convention(thin) () -> ()
+  apply %2() : $@convention(thin) () -> ()
+  strong_release %0 : $Builtin.NativeObject
+  %1 = tuple()
+  return %1 : $()
+}
+
+// CHECK-LABEL: sil @move_release_over_loop
+// CHECK:      bb0{{.*}}:
+// CHECK:        strong_retain
+// CHECK:        apply
+// CHECK:        strong_release
+// CHECK:        br bb1
+// CHECK:      bb1:
+// CHECK-NEXT:   cond_br
+// CHECK:      bb2:
+// CHECK-NEXT:   br bb1
+// CHECK:      bb3:
+// CHECK-NEXT:   tuple
+// CHECK-NEXT:   return
+sil @move_release_over_loop : $@convention(thin) (Builtin.NativeObject) -> () {
+bb0(%0 : $Builtin.NativeObject):
+  strong_retain %0 : $Builtin.NativeObject
+  %2 = function_ref @blocker : $@convention(thin) () -> ()
+  apply %2() : $@convention(thin) () -> ()
+  br bb1
+
+bb1:
+  cond_br undef, bb1, bb2
+
+bb2:
+  strong_release %0 : $Builtin.NativeObject
+  %1 = tuple()
+  return %1 : $()
+}
+
+// CHECK-LABEL: sil @handle_infinite_loop
+// CHECK:      bb0{{.*}}:
+// CHECK-NEXT:   cond_br
+// CHECK:      bb1:
+// CHECK-NOT:    {{(retain|release)}}
+// CHECK:        apply
+// CHECK-NEXT:   br bb2
+// CHECK:      bb2:
+// CHECK-NEXT:   br bb2
+// CHECK:      bb3:
+// CHECK:        strong_retain
+// CHECK:        strong_release
+// CHECK:        return
+sil @handle_infinite_loop : $@convention(thin) (@inout Builtin.NativeObject) -> () {
+bb0(%a : $*Builtin.NativeObject):
+  cond_br undef, bb1, bb3
+
+bb1:
+  %2 = function_ref @blocker : $@convention(thin) () -> ()
+  apply %2() : $@convention(thin) () -> ()
+  br bb2
+
+bb2:
+  br bb2
+
+bb3:
+  %0 = load %a : $*Builtin.NativeObject
+  strong_retain %0 : $Builtin.NativeObject
+  strong_release %0 : $Builtin.NativeObject
+  %1 = tuple()
+  return %1 : $()
+}