Fix crash in EpilogueArcAnalysis

EpilogueARCState for unreachable blocks can be non-existent. Fix the
EpilogueARCContext::getState method and its users.

Author: meg-gupta

include/swift/SILOptimizer/Analysis/EpilogueARCAnalysis.h

@@ -83,8 +83,14 @@ class EpilogueARCContext {
   /// The exit blocks of the function.
   llvm::SmallPtrSet<SILBasicBlock *, 2> ExitBlocks;
 
-  EpilogueARCBlockState &getState(SILBasicBlock *BB) {
-    return IndexToStateMap[*PO->getPONumber(BB)];
+  /// Returns the EpligoyeARCBlockState for \p BB. If \p BB is unreachable,
+  /// returns None
+  Optional<EpilogueARCBlockState *> getState(SILBasicBlock *BB) {
+    // poNumber will be None for unreachable blocks
+    auto poNumber = PO->getPONumber(BB);
+    if (poNumber.hasValue())
+      return &IndexToStateMap[*poNumber];
+    return None;
   }
 
   /// Return true if this is a function exiting block this epilogue ARC
@@ -113,7 +119,10 @@ class EpilogueARCContext {
   }
 
   SILValue getArg(SILBasicBlock *BB) {
-    SILValue A = getState(BB).LocalArg;
+    auto state = getState(BB);
+    if (!state)
+      return SILValue();
+    SILValue A = state.getValue()->LocalArg;
     if (A)
       return A;
     return Arg;

lib/SILOptimizer/Analysis/EpilogueARCAnalysis.cpp

@@ -48,7 +48,9 @@ void EpilogueARCContext::initializeDataflow() {
       for (auto *X : A->getParent()->getPredecessorBlocks()) {
         // Try to find the predecessor edge-value.
         SILValue IA = A->getIncomingPhiValue(X);
-        getState(X).LocalArg = IA;
+        auto state = getState(X);
+        if (state.hasValue())
+          state.getValue()->LocalArg = IA;
 
         // Maybe the edge value is another SILArgument.
         ToProcess.push_back(IA);
@@ -64,15 +66,19 @@ bool EpilogueARCContext::convergeDataflow() {
     Changed = false;
     // Iterate until the data flow converges.
     for (SILBasicBlock *B : PO->getPostOrder()) {
-      auto &BS = getState(B);
+      // Since the basic block is in PO, it is reachable and will always have a
+      // state
+      auto *BS = getState(B).getValue();
       // Merge in all the successors.
       bool BBSetOut = false;
       if (!B->succ_empty()) {
         auto Iter = B->succ_begin();
-        BBSetOut = getState(*Iter).BBSetIn;
+        // Since the basic block is reachable, its successors are reachable, and
+        // will always have a state.
+        BBSetOut = getState(*Iter).getValue()->BBSetIn;
         Iter = std::next(Iter);
         for (auto E = B->succ_end(); Iter != E; ++Iter) {
-          BBSetOut &= getState(*Iter).BBSetIn;
+          BBSetOut &= getState(*Iter).getValue()->BBSetIn;
         }
       } else if (isExitBlock(B)) {
         // We set the BBSetOut for exit blocks.
@@ -100,8 +106,8 @@ bool EpilogueARCContext::convergeDataflow() {
       }
 
       // Update BBSetIn.
-      Changed |= (BS.BBSetIn != BBSetOut);
-      BS.BBSetIn = BBSetOut;
+      Changed |= (BS->BBSetIn != BBSetOut);
+      BS->BBSetIn = BBSetOut;
     }
   } while (Changed);
   return true;
@@ -122,10 +128,12 @@ bool EpilogueARCContext::computeEpilogueARC() {
       // not have an epilogue ARC instruction, which means the data flow has 
       // failed.
       auto Iter = B->succ_begin();
-      auto Base = getState(*Iter).BBSetIn;
+      // Since basic block B is in PO, its successors will be reachable and will
+      // always have a state
+      auto Base = getState(*Iter).getValue()->BBSetIn;
       Iter = std::next(Iter);
       for (auto E = B->succ_end(); Iter != E; ++Iter) {
-        if (getState(*Iter).BBSetIn != Base)
+        if (getState(*Iter).getValue()->BBSetIn != Base)
           return false;
       }
       BBSetOut = Base;

test/SILOptimizer/epilogue_arc_analysis_bugs.sil

@@ -0,0 +1,36 @@
+// RUN: %target-sil-opt -arc-sequence-opts %s
+
+import Builtin
+import Swift
+
+class Klass {
+}
+
+// Test to make sure there is no crash in EpilogueArcAnalysis, while querying the state of unreachable block bb1
+sil @foo : $@convention(thin) (@owned Klass) -> () {
+bb0(%arg : $Klass):
+  strong_retain %arg : $Klass
+  %0 = enum $Optional<Klass>, #Optional.none!enumelt
+  br bb2(%0 : $Optional<Klass>)
+
+bb1: // unreachable
+  %1 = enum $Optional<Klass>, #Optional.some!enumelt, undef : $Klass
+  br bb2(%1 : $Optional<Klass>)
+
+bb2(%2 : $Optional<Klass>):
+  switch_enum %2 : $Optional<Klass>, case #Optional.some!enumelt: bb3, case #Optional.none!enumelt: bb4
+
+bb3:
+  br bb5
+
+bb4:
+  br bb5
+
+bb5:
+ release_value %2 : $Optional<Klass>
+ strong_release %arg : $Klass
+ strong_release %arg : $Klass
+ %res = tuple ()
+ return %res : $()
+}
+