SILOptimizer: Remove switch cases matching unavailable enum elements.

Unavailable enum elements cannot be instantiated at runtime without invoking
UB. Therefore the optimizer can consider a basic block unreachable if its only
predecessor is a block that terminates in a switch instruction matching an
unavailable enum element. Furthermore, removing the switch instruction cases
that refer to unavailable enum elements is _mandatory_ when
`-unavailable-decl-optimization=complete` is specified because otherwise
lowered IR for these instructions could refer to enum tag accessors that will
not be lowered, resulting in a failure during linking.

Resolves rdar://113872720.

Author: tshortli

include/swift/AST/Decl.h

@@ -1237,7 +1237,10 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
   /// Returns true if this declaration should be considered available during
   /// SIL/IR lowering. A declaration would not be available during lowering if,
   /// for example, it is annotated as unavailable with `@available` and
-  /// optimization settings require that it be omitted.
+  /// optimization settings require that it be omitted. Canonical SIL must not
+  /// contain any references to declarations that are unavailable during
+  /// lowering because the resulting IR could reference non-existent symbols
+  /// and fail to link.
   bool isAvailableDuringLowering() const;
 
   /// Returns true if ABI compatibility stubs must be emitted for the given
@@ -4204,6 +4207,11 @@ class EnumDecl final : public NominalTypeDecl {
   /// A range for iterating the elements of an enum.
   using ElementRange = DowncastFilterRange<EnumElementDecl, DeclRange>;
 
+  /// A range for iterating the elements of an enum that are available during
+  /// lowering.
+  using ElementRangeForLowering = OptionalTransformRange<
+      ElementRange, AvailableDuringLoweringDeclFilter<EnumElementDecl>>;
+
   /// A range for iterating the cases of an enum.
   using CaseRange = DowncastFilterRange<EnumCaseDecl, DeclRange>;
 
@@ -4217,6 +4225,13 @@ class EnumDecl final : public NominalTypeDecl {
     return std::distance(eltRange.begin(), eltRange.end());
   }
 
+  /// Returns a range that iterates over all the elements of an enum for which
+  /// \c isAvailableDuringLowering() is true.
+  ElementRangeForLowering getAllElementsForLowering() const {
+    return ElementRangeForLowering(
+        getAllElements(), AvailableDuringLoweringDeclFilter<EnumElementDecl>());
+  }
+
   /// If this enum has a unique element, return it. A unique element can
   /// either hold a value or not, and the existence of one unique element does
   /// not imply the existence or non-existence of the opposite unique element.

include/swift/SIL/SILInstruction.h

@@ -6583,7 +6583,7 @@ class SelectEnumInstBase : public BaseTy {
     }
 
     llvm::SmallPtrSet<EnumElementDecl *, 4> unswitchedElts;
-    for (auto elt : decl->getAllElements())
+    for (auto elt : decl->getAllElementsForLowering())
       unswitchedElts.insert(elt);
 
     for (unsigned i = 0, e = this->getNumCases(); i != e; ++i) {
@@ -9812,7 +9812,7 @@ class SwitchEnumInstBase : public BaseTy {
     assert(decl && "switch_enum operand is not an enum");
 
     SmallPtrSet<EnumElementDecl *, 4> unswitchedElts;
-    for (auto elt : decl->getAllElements())
+    for (auto elt : decl->getAllElementsForLowering())
       unswitchedElts.insert(elt);
 
     for (unsigned i = 0, e = getNumCases(); i != e; ++i) {

lib/SIL/IR/SILType.cpp

@@ -940,6 +940,12 @@ TypeBase::replaceSubstitutedSILFunctionTypesWithUnsubstituted(SILModule &M) cons
 bool SILType::isEffectivelyExhaustiveEnumType(SILFunction *f) {
   EnumDecl *decl = getEnumOrBoundGenericEnum();
   assert(decl && "Called for a non enum type");
+
+  // Since unavailable enum elements cannot be referenced in canonical SIL,
+  // enums with these elements cannot be treated as exhaustive types.
+  if (decl->hasCasesUnavailableDuringLowering())
+    return false;
+
   return decl->isEffectivelyExhaustive(f->getModule().getSwiftModule(),
                                        f->getResilienceExpansion());
 }

lib/SIL/Verifier/SILVerifier.cpp

@@ -4906,6 +4906,14 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
       requireSameType(
           result->getType(), SEO->getType(),
           "select_enum case operand must match type of instruction");
+
+      // In canonical SIL, select instructions must not cover any enum elements
+      // that are unavailable.
+      if (F.getModule().getStage() >= SILStage::Canonical) {
+        require(elt->isAvailableDuringLowering(),
+                "select_enum dispatches on enum element that is unavailable "
+                "during lowering.");
+      }
     }
 
     // If the select is non-exhaustive, we require a default.
@@ -4994,6 +5002,15 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
                 "arguments");
         return;
       }
+
+      // In canonical SIL, switch instructions must not cover any enum elements
+      // that are unavailable.
+      if (F.getModule().getStage() >= SILStage::Canonical) {
+        require(elt->isAvailableDuringLowering(),
+                "switch_enum dispatches on enum element that is unavailable "
+                "during lowering.");
+      }
+
       // Check for a valid switch result type.
       if (dest->getArguments().size() == 1) {
         SILType eltArgTy = uTy.getEnumElementType(elt, F.getModule(),
@@ -5098,6 +5115,14 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
               "more than once");
       unswitchedElts.erase(elt);
 
+      // In canonical SIL, switch instructions must not cover any enum elements
+      // that are unavailable.
+      if (F.getModule().getStage() >= SILStage::Canonical) {
+        require(elt->isAvailableDuringLowering(),
+                "switch_enum_addr dispatches on enum element that is "
+                "unavailable during lowering.");
+      }
+
       // The destination BB must not have BB arguments.
       require(dest->getArguments().empty(),
               "switch_enum_addr destination must take no BB args");

lib/SILOptimizer/Mandatory/DiagnoseUnreachable.cpp

@@ -51,6 +51,7 @@ enum class UnreachableKind {
   FoldedBranch,
   FoldedSwitchEnum,
   NoreturnCall,
+  UnavailableEnumElement,
 };
 
 /// Information about a folded conditional branch instruction: it's location
@@ -867,6 +868,92 @@ static bool simplifyBlocksWithCallsToNoReturn(SILBasicBlock &BB,
   return true;
 }
 
+/// Replaces  `switch_enum` / `switch_enum_addr` instructions that have cases
+/// matching unavailable enum elements with new instructions that have those
+/// cases removed since unavailable enum elements cannot be instantiated at
+/// run time.
+///
+/// If this pass removes a case, it will add a default case that traps if there
+/// was not already an existing default case. This artificial default case
+/// serves two purposes. The first is to ensure the new switch instruction still
+/// "covers" the entire enum as required by SIL verification (this analysis
+/// could be changed in the future to ignore unavailable elements). Secondly,
+/// the trapping default case may help developers catch issues at runtime where
+/// UB has been invoked and an unavailable element has been instantiated
+/// unexpectedly. Ideally, this debugging aid would be available consistently,
+/// but since we cannot leave any references to the unavailable elements in SIL
+/// it is not possible to transform switch instructions that already have
+/// default cases in such a way that unavailable elements would be detected.
+static bool eliminateSwitchDispatchOnUnavailableElements(
+    SILBasicBlock &BB, UnreachableUserCodeReportingState *State) {
+  SwitchEnumTermInst SWI(BB.getTerminator());
+  if (!SWI)
+    return false;
+
+  // FIXME: Return early if EnumDecl doesn't have any elements that are
+  // unavailable.
+
+  ASTContext &ctx = BB.getModule().getASTContext();
+  SILLocation Loc = SWI->getLoc();
+  bool DidRemoveUnavailableCase = false;
+
+  SmallVector<std::pair<EnumElementDecl *, SILBasicBlock *>, 4> NewCaseBBs;
+  for (unsigned i : range(SWI.getNumCases())) {
+    auto CaseBB = SWI.getCase(i);
+    auto availableAtr = CaseBB.first->getAttrs().getUnavailable(ctx);
+
+    if (availableAtr && availableAtr->isUnconditionallyUnavailable()) {
+      // Mark the basic block as potentially unreachable.
+      SILBasicBlock *UnreachableBlock = CaseBB.second;
+      if (!State->PossiblyUnreachableBlocks.contains(UnreachableBlock)) {
+        // If this is the first time we see this unreachable block, store it.
+        State->PossiblyUnreachableBlocks.insert(UnreachableBlock);
+        State->MetaMap.insert(
+            {UnreachableBlock,
+             UnreachableInfo{UnreachableKind::UnavailableEnumElement, Loc,
+                             true}});
+      }
+      DidRemoveUnavailableCase = true;
+    } else {
+      NewCaseBBs.push_back(CaseBB);
+    }
+  }
+
+  if (!DidRemoveUnavailableCase)
+    return false;
+
+  // Since at least one case was removed, we need to add a default case that
+  // traps if there isn't already an existing default case. The resulting SIL
+  // will have a structure that matches what SILGen would have produced for a
+  // switch statment that was written in source with unavailable cases
+  // unhandled.
+  SILBasicBlock *DefaultBB = SWI.getDefaultBBOrNull().getPtrOrNull();
+  bool DidCreateDefault = false;
+  if (!DefaultBB) {
+    DefaultBB = BB.getParent()->createBasicBlock();
+    SILLocation genLoc = Loc.asAutoGenerated();
+    SILBuilder B(DefaultBB);
+    B.createUnconditionalFail(genLoc, "unexpected enum value");
+    B.createUnreachable(genLoc);
+    DidCreateDefault = true;
+  }
+
+  if (isa<SwitchEnumInst>(*SWI)) {
+    auto *SEI = SILBuilderWithScope(SWI).createSwitchEnum(
+        Loc, SWI.getOperand(), DefaultBB, NewCaseBBs);
+
+    if (DidCreateDefault)
+      SEI->createDefaultResult();
+  } else {
+    assert(isa<SwitchEnumAddrInst>(*SWI) && "unknown switch_enum instruction");
+    SILBuilderWithScope(SWI).createSwitchEnumAddr(Loc, SWI.getOperand(),
+                                                  DefaultBB, NewCaseBBs);
+  }
+  SWI->eraseFromParent();
+
+  return true;
+}
+
 /// Issue an "unreachable code" diagnostic if the blocks contains or
 /// leads to another block that contains user code.
 ///
@@ -935,6 +1022,11 @@ static bool diagnoseUnreachableBlock(
         }
         break;
       }
+      case (UnreachableKind::UnavailableEnumElement): {
+        // Don't diagnose blocks removed because they were only reachable from
+        // a case matching an unavailable enum element.
+        break;
+      }
       }
 
       // Record that we've reported this unreachable block to avoid duplicates
@@ -1076,6 +1168,11 @@ static void diagnoseUnreachable(SILFunction &Fn) {
     // function.
     if (simplifyBlocksWithCallsToNoReturn(BB, &State))
       continue;
+
+    // Remove the cases of switch_enum / switch_enum_addr instructions that
+    // match unavailable enum elements.
+    if (eliminateSwitchDispatchOnUnavailableElements(BB, &State))
+      continue;
   }
 
   // Remove unreachable blocks.
@@ -1094,6 +1191,10 @@ static void diagnoseUnreachable(SILFunction &Fn) {
     // function.
     if (simplifyBlocksWithCallsToNoReturn(BB, &State))
       continue;
+    // Remove the cases of switch_enum / switch_enum_addr instructions that
+    // match unavailable enum elements.
+    if (eliminateSwitchDispatchOnUnavailableElements(BB, &State))
+      continue;
   }
 
   // Remove unreachable blocks.

test/Inputs/resilient_enum.swift

@@ -208,3 +208,28 @@ public enum ResilientIndirectEnum {
   // -2
   indirect case B(ResilientIndirectEnum, ResilientIndirectEnum)
 }
+
+public enum ResilientEnumWithUnavailableCase {
+  case available
+
+  @available(*, unavailable)
+  case unavailable
+}
+
+public enum ResilientEnumWithUnavailableCaseAndPayload {
+  @available(*, unavailable)
+  case int(_ i: UnavailableResilientInt)
+
+  case double(_ d: ResilientDouble)
+}
+
+extension ResilientEnumWithUnavailableCaseAndPayload {
+  @_alwaysEmitIntoClient
+  public var intValue: Int {
+    switch self {
+    case .int(let i): return i.i
+    case .double(let d): return Int(d.d)
+    default: return -1
+    }
+  }
+}

test/Inputs/resilient_struct.swift

@@ -110,3 +110,12 @@ public struct Container {
 public struct PairContainer {
   public var pair : (Container, Container)
 }
+
+@available(*, unavailable)
+public struct UnavailableResilientInt {
+  public let i: Int
+
+  public init(i: Int) {
+    self.i = i
+  }
+}

test/Interpreter/enum_resilience.swift

@@ -11,6 +11,8 @@
 
 // RUN: %target-run %t/main %t/%target-library-name(resilient_struct) %t/%target-library-name(resilient_enum)
 
+// Test against libraries built with -whole-module-optimization.
+
 // RUN: %target-build-swift-dylib(%t/%target-library-name(resilient_struct_wmo)) -enable-library-evolution %S/../Inputs/resilient_struct.swift -emit-module -emit-module-path %t/resilient_struct.swiftmodule -module-name resilient_struct -whole-module-optimization
 // RUN: %target-codesign %t/%target-library-name(resilient_struct_wmo)
 
@@ -22,6 +24,20 @@
 
 // RUN: %target-run %t/main2 %t/%target-library-name(resilient_struct_wmo) %t/%target-library-name(resilient_enum_wmo)
 
+// Test with -unavailable-decl-optimization=complete.
+
+// RUN: %target-build-swift-dylib(%t/%target-library-name(resilient_struct_udoc)) -enable-library-evolution %S/../Inputs/resilient_struct.swift -emit-module -emit-module-path %t/resilient_struct.swiftmodule -module-name resilient_struct -unavailable-decl-optimization=complete
+// RUN: %target-codesign %t/%target-library-name(resilient_struct_udoc)
+
+// RUN: %target-build-swift-dylib(%t/%target-library-name(resilient_enum_udoc)) -enable-library-evolution %S/../Inputs/resilient_enum.swift -emit-module -emit-module-path %t/resilient_enum.swiftmodule -module-name resilient_enum -I%t -L%t -lresilient_struct_udoc -unavailable-decl-optimization=complete
+// RUN: %target-codesign %t/%target-library-name(resilient_enum_udoc)
+
+// RUN: %target-build-swift %s -L %t -I %t -lresilient_struct_udoc -lresilient_enum_udoc -o %t/main3 %target-rpath(%t)
+// RUN: %target-codesign %t/main3
+
+// RUN: %target-run %t/main3 %t/%target-library-name(resilient_struct_udoc) %t/%target-library-name(resilient_enum_udoc)
+
+
 // REQUIRES: executable_test
 
 import StdlibUnittest
@@ -491,4 +507,31 @@ ResilientEnumTestSuite.test("ResilientEnumSingleCase") {
   // This used to crash.
   test(Status(fst: .only(nested: Nested(str: "foobar", r: ResilientInt(i: 1))), snd: false))
 }
+
+ResilientEnumTestSuite.test("ResilientEnumWithUnavailableCase") {
+  let a: [ResilientEnumWithUnavailableCase] = [.available]
+
+  let b: [Int] = a.map {
+    switch $0 {
+    case .available:
+      return 0
+    case .unavailable:
+      return 1
+    default:
+      return -1
+    }
+  }
+
+  expectEqual(b, [0])
+}
+
+ResilientEnumTestSuite.test("ResilientEnumWithUnavailableCaseAndPayload") {
+  let a: [ResilientEnumWithUnavailableCaseAndPayload] =
+      [.double(ResilientDouble(d: 42.0))]
+
+  let b: [Int] = a.map { return $0.intValue }
+
+  expectEqual(b, [42])
+}
+
 runAllTests()

test/SILOptimizer/definite_init_unavailable_enum_element.swift

@@ -0,0 +1,43 @@
+// RUN: %target-swift-frontend -emit-sil -unavailable-decl-optimization=none -verify %s
+// RUN: %target-swift-frontend -emit-sil -unavailable-decl-optimization=complete -verify %s
+
+public enum HasUnavailableElement {
+  case available
+
+  @available(*, unavailable)
+  case unavailable
+}
+
+public func markUsed<T>(_ t: T) {}
+
+public func testSwitch(_ e: HasUnavailableElement) {
+  switch e {
+  case .available:
+    ()
+  case .unavailable:
+    let x: Int // expected-note {{constant defined here}}
+    markUsed(x) // expected-error {{constant 'x' used before being initialized}}
+  }
+}
+
+public func testIfCase(_ e: HasUnavailableElement) {
+  if case .unavailable = e {
+    let x: Int // expected-note {{constant defined here}}
+    markUsed(x) // expected-error {{constant 'x' used before being initialized}}
+  }
+}
+
+public func testInitSwitch() {
+  struct S {
+    var x: Int // expected-note {{'self.x' not initialized}}
+
+    init(_ e: HasUnavailableElement) {
+      switch e {
+      case .available:
+        x = 1
+      case .unavailable:
+        ()
+      }
+    } // expected-error {{return from initializer without initializing all stored properties}}
+  }
+}