Fix dominator tree update in array specialization

We need to walk the dominator tree starting at the header and update the
dominator of all nodes outside the loop we hit i.e the dominator tree nodes that
are immediately dominated 'by the loop' instead of only updating dominated exit
blocks.

rdar://34523864
diff --git a/lib/SILOptimizer/LoopTransforms/COWArrayOpt.cpp b/lib/SILOptimizer/LoopTransforms/COWArrayOpt.cpp
index 2d9fcde..b7267e7 100644
--- a/lib/SILOptimizer/LoopTransforms/COWArrayOpt.cpp
+++ b/lib/SILOptimizer/LoopTransforms/COWArrayOpt.cpp
@@ -2219,6 +2219,24 @@
   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);
@@ -2231,22 +2249,19 @@
   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);
@@ -2275,8 +2290,8 @@
                      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));
 
diff --git a/test/SILOptimizer/array_specialize.sil b/test/SILOptimizer/array_specialize.sil
index c39d77d..fdab0ed 100644
--- a/test/SILOptimizer/array_specialize.sil
+++ b/test/SILOptimizer/array_specialize.sil
@@ -109,3 +109,70 @@
 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
+}