From 5f040ec8a5a4b4c2f7ef500286b379caf98500bd Mon Sep 17 00:00:00 2001
From: Arnold Schwaighofer <aschwaighofer@apple.com>
Date: Mon, 25 Sep 2017 20:04:58 -0700
Subject: [PATCH] [4.0] 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
---
 .../LoopTransforms/COWArrayOpt.cpp            | 39 +++++++----
 test/SILOptimizer/array_specialize.sil        | 67 +++++++++++++++++++
 2 files changed, 94 insertions(+), 12 deletions(-)

diff --git a/lib/SILOptimizer/LoopTransforms/COWArrayOpt.cpp b/lib/SILOptimizer/LoopTransforms/COWArrayOpt.cpp
index 359dec86da82a..0a8d4f2a08252 100644
--- a/lib/SILOptimizer/LoopTransforms/COWArrayOpt.cpp
+++ b/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));
 
diff --git a/test/SILOptimizer/array_specialize.sil b/test/SILOptimizer/array_specialize.sil
index c39d77db75014..fdab0eda9f05a 100644
--- a/test/SILOptimizer/array_specialize.sil
+++ b/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
+}