[4.0] Fix dominator tree update in array specialization

lib/SILOptimizer/LoopTransforms/COWArrayOpt.cpp

@@ -2213,6 +2213,24 @@ static void replaceArrayPropsCall(SILBuilder &B, ArraySemanticsCall C) {
   C.removeCall();
 }
 
+/// Collects all loop dominated blocks outside the loop that are immediately
+/// dominated by the loop.
+static void
+collectImmediateLoopDominatedBlocks(const SILLoop *Lp, DominanceInfoNode *Node,
+                                    SmallVectorImpl<SILBasicBlock *> &Blocks) {
+  SILBasicBlock *BB = Node->getBlock();
+
+  // Base case: First loop dominated block outside of loop.
+  if (!Lp->contains(BB)) {
+    Blocks.push_back(BB);
+    return;
+  }
+
+  // Loop contains the basic block. Look at immediately dominated nodes.
+  for (auto *Child : *Node)
+    collectImmediateLoopDominatedBlocks(Lp, Child, Blocks);
+}
+
 void ArrayPropertiesSpecializer::specializeLoopNest() {
   auto *Lp = getLoop();
   assert(Lp);
@@ -2225,22 +2243,19 @@ void ArrayPropertiesSpecializer::specializeLoopNest() {
   auto *CheckBlock = splitBasicBlockAndBranch(B,
       HoistableLoopPreheader->getTerminator(), DomTree, nullptr);
 
-  // Get the exit blocks of the original loop.
   auto *Header = CheckBlock->getSingleSuccessorBlock();
   assert(Header);
 
-  // Our loop info is not really completely valid anymore since the cloner does
-  // not update it. However, exit blocks of the original loop are still valid.
+  // Collect all loop dominated blocks (e.g exit blocks could be among them). We
+  // need to update their dominator.
+  SmallVector<SILBasicBlock *, 16> LoopDominatedBlocks;
+  collectImmediateLoopDominatedBlocks(Lp, DomTree->getNode(Header),
+                                      LoopDominatedBlocks);
+
+  // Collect all exit blocks.
   SmallVector<SILBasicBlock *, 16> ExitBlocks;
   Lp->getExitBlocks(ExitBlocks);
 
-  // Collect the exit blocks dominated by the loop - they will be dominated by
-  // the check block.
-  SmallVector<SILBasicBlock *, 16> ExitBlocksDominatedByPreheader;
-  for (auto *ExitBlock: ExitBlocks)
-    if (DomTree->dominates(CheckBlock, ExitBlock))
-      ExitBlocksDominatedByPreheader.push_back(ExitBlock);
-
   // Split the preheader before the first instruction.
   SILBasicBlock *NewPreheader =
     splitBasicBlockAndBranch(B, &*CheckBlock->begin(), DomTree, nullptr);
@@ -2269,8 +2284,8 @@ void ArrayPropertiesSpecializer::specializeLoopNest() {
                      IsFastNativeArray, ClonedPreheader, NewPreheader);
   CheckBlock->getTerminator()->eraseFromParent();
 
-  // Fixup the exit blocks. They are now dominated by the check block.
-  for (auto *BB : ExitBlocksDominatedByPreheader)
+  // Fixup the loop dominated blocks. They are now dominated by the check block.
+  for (auto *BB : LoopDominatedBlocks)
     DomTree->changeImmediateDominator(DomTree->getNode(BB),
                                       DomTree->getNode(CheckBlock));
 

test/SILOptimizer/array_specialize.sil

@@ -109,3 +109,70 @@ bb4:
 bb5(%9 : $Error):
   throw %9 : $Error
 }
+
+sil @dominator_update_outside_non_exit_block : $@convention(thin) (@inout MyArray<MyClass>, @inout Builtin.Int1) -> Builtin.Int1 {
+bb0(%0 : $*MyArray<MyClass>, %1 : $*Builtin.Int1):
+  %3 = load %0 : $*MyArray<MyClass>
+  br bb1
+
+bb1:
+  %2 = function_ref @arrayPropertyIsNative : $@convention(method) (@owned MyArray<MyClass>) -> Bool
+  %4 = load %1 : $*Builtin.Int1
+  retain_value %3 : $MyArray<MyClass>
+  %5 = apply %2(%3) : $@convention(method) (@owned MyArray<MyClass>) -> Bool
+  cond_br %4, bb2, bb4
+
+bb2:
+ cond_br undef, bb3, bb5
+
+bb3:
+ %6 = integer_literal $Builtin.Int1, -1
+ cond_br %6, bb1, bb7
+
+bb4: // Exit block; b1 dom b4
+ cond_br undef, bb5, bb6
+
+bb5: // Exit Block; b1 dom b4
+ br bb6
+
+bb6: // Non-exit Dominated by bb1
+ br bb7
+
+bb7:
+  return %4 : $Builtin.Int1
+}
+
+sil @dominator_update_outside_non_exit_block_2 : $@convention(thin) (@inout MyArray<MyClass>, @inout Builtin.Int1) -> Builtin.Int1 {
+bb0(%0 : $*MyArray<MyClass>, %1 : $*Builtin.Int1):
+  %3 = load %0 : $*MyArray<MyClass>
+  br bb1
+
+bb1:
+  %2 = function_ref @arrayPropertyIsNative : $@convention(method) (@owned MyArray<MyClass>) -> Bool
+  %4 = load %1 : $*Builtin.Int1
+  retain_value %3 : $MyArray<MyClass>
+  %5 = apply %2(%3) : $@convention(method) (@owned MyArray<MyClass>) -> Bool
+  cond_br %4, bb2, bb4
+
+bb2:
+ cond_br undef, bb3, bb5
+
+bb3:
+ %6 = integer_literal $Builtin.Int1, -1
+ cond_br %6, bb1, bb7
+
+bb4: // Exit block; b1 dom b4
+ cond_br undef, bb5, bb6
+
+bb5: // Exit Block; b1 dom b4
+ br bb6
+
+bb6: // Non-exit Dominated by bb1
+ br bb8
+
+bb7: // Exit dominated by bb3
+ br bb8
+
+bb8: // Non-exit dominated by bb1
+  return %4 : $Builtin.Int1
+}