Merge pull request #67920 from tshortli/unavailable-decl-opt-complete-resilient-enum-switch

SILOptimizer: Remove switch cases matching unavailable enum elements

Author: tshortli

include/swift/AST/Decl.h

@@ -650,13 +650,16 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
                         /// address diversified ptrauth qualified field.
                         IsNonTrivialPtrAuth : 1);
 
-  SWIFT_INLINE_BITFIELD(EnumDecl, NominalTypeDecl, 2+1,
+  SWIFT_INLINE_BITFIELD(EnumDecl, NominalTypeDecl, 2+1+1,
     /// True if the enum has cases and at least one case has associated values.
     HasAssociatedValues : 2,
     /// True if the enum has at least one case that has some availability
     /// attribute.  A single bit because it's lazily computed along with the
     /// HasAssociatedValues bit.
-    HasAnyUnavailableValues : 1
+    HasAnyUnavailableValues : 1,
+    /// True if \c isAvailableDuringLowering() is false for any cases. Lazily
+    /// computed along with HasAssociatedValues.
+    HasAnyUnavailableDuringLoweringValues : 1
   );
 
   SWIFT_INLINE_BITFIELD(ModuleDecl, TypeDecl, 1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1,
@@ -1237,7 +1240,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 +4210,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 +4228,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.
@@ -4279,6 +4297,11 @@ class EnumDecl final : public NominalTypeDecl {
   /// Note that this is false for enums with absolutely no cases.
   bool hasPotentiallyUnavailableCaseValue() const;
 
+  /// True if \c isAvailableDuringLowering() is false for any cases.
+  ///
+  /// Note that this is false for enums with absolutely no cases.
+  bool hasCasesUnavailableDuringLowering() const;
+
   /// True if the enum has cases.
   bool hasCases() const {
     return !getAllElements().empty();

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/AST/Availability.cpp

@@ -312,6 +312,9 @@ bool Decl::isAvailableDuringLowering() const {
       UnavailableDeclOptimization::Complete)
     return true;
 
+  if (hasClangNode())
+    return true;
+
   return !isUnconditionallyUnavailable(this);
 }
 

lib/AST/Decl.cpp

@@ -5944,14 +5944,13 @@ EnumCaseDecl *EnumCaseDecl::create(SourceLoc CaseLoc,
 }
 
 bool EnumDecl::hasPotentiallyUnavailableCaseValue() const {
-  switch (static_cast<AssociatedValueCheck>(Bits.EnumDecl.HasAssociatedValues)) {
-    case AssociatedValueCheck::Unchecked:
-      // Compute below
-      this->hasOnlyCasesWithoutAssociatedValues();
-      LLVM_FALLTHROUGH;
-    default:
-      return static_cast<bool>(Bits.EnumDecl.HasAnyUnavailableValues);
-  }
+  (void)this->hasOnlyCasesWithoutAssociatedValues(); // Prime the cache
+  return static_cast<bool>(Bits.EnumDecl.HasAnyUnavailableValues);
+}
+
+bool EnumDecl::hasCasesUnavailableDuringLowering() const {
+  (void)this->hasOnlyCasesWithoutAssociatedValues(); // Prime the cache
+  return static_cast<bool>(Bits.EnumDecl.HasAnyUnavailableDuringLoweringValues);
 }
 
 bool EnumDecl::hasOnlyCasesWithoutAssociatedValues() const {
@@ -5970,6 +5969,7 @@ bool EnumDecl::hasOnlyCasesWithoutAssociatedValues() const {
   }
 
   bool hasAnyUnavailableValues = false;
+  bool hasAnyUnavailableDuringLoweringValues = false;
   bool hasAssociatedValues = false;
 
   for (auto elt : getAllElements()) {
@@ -5980,13 +5980,18 @@ bool EnumDecl::hasOnlyCasesWithoutAssociatedValues() const {
       }
     }
 
+    if (!elt->isAvailableDuringLowering())
+      hasAnyUnavailableDuringLoweringValues = true;
+
     if (elt->hasAssociatedValues())
       hasAssociatedValues = true;
   }
 
   EnumDecl *enumDecl = const_cast<EnumDecl *>(this);
 
   enumDecl->Bits.EnumDecl.HasAnyUnavailableValues = hasAnyUnavailableValues;
+  enumDecl->Bits.EnumDecl.HasAnyUnavailableDuringLoweringValues =
+      hasAnyUnavailableDuringLoweringValues;
   enumDecl->Bits.EnumDecl.HasAssociatedValues = static_cast<unsigned>(
       hasAssociatedValues ? AssociatedValueCheck::HasAssociatedValues
                           : AssociatedValueCheck::NoAssociatedValues);

lib/DriverTool/sil_opt_main.cpp

@@ -499,6 +499,21 @@ struct SILOptOptions {
   llvm::cl::opt<bool> DebugDiagnosticNames = llvm::cl::opt<bool>(
       "debug-diagnostic-names",
       llvm::cl::desc("Include diagnostic names when printing"));
+
+  llvm::cl::opt<swift::UnavailableDeclOptimization>
+      UnavailableDeclOptimization =
+          llvm::cl::opt<swift::UnavailableDeclOptimization>(
+              "unavailable-decl-optimization",
+              llvm::cl::desc("Optimization mode for unavailable declarations"),
+              llvm::cl::values(
+                  clEnumValN(swift::UnavailableDeclOptimization::None, "none",
+                             "Don't optimize unavailable decls"),
+                  clEnumValN(swift::UnavailableDeclOptimization::Stub, "stub",
+                             "Lower unavailable functions to stubs"),
+                  clEnumValN(
+                      swift::UnavailableDeclOptimization::Complete, "complete",
+                      "Eliminate unavailable decls from lowered SIL/IR")),
+              llvm::cl::init(swift::UnavailableDeclOptimization::None));
 };
 
 /// Regular expression corresponding to the value given in one of the
@@ -646,6 +661,9 @@ int sil_opt_main(ArrayRef<const char *> argv, void *MainAddr) {
 
   Invocation.getLangOptions().EnableCXXInterop = options.EnableCxxInterop;
 
+  Invocation.getLangOptions().UnavailableDeclOptimizationMode =
+      options.UnavailableDeclOptimization;
+
   Invocation.getDiagnosticOptions().VerifyMode =
     options.VerifyMode ? DiagnosticOptions::Verify : DiagnosticOptions::NoVerify;
 

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,96 @@ 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;
+
+  EnumDecl *ED = SWI->getOperand(0)->getType().getEnumOrBoundGenericEnum();
+  assert(ED && "operand is not an enum");
+
+  // No need to check the instruction if all elements are available.
+  if (!ED->hasCasesUnavailableDuringLowering())
+    return false;
+
+  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 +1026,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 +1172,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 +1195,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.

lib/Serialization/ModuleFileSharedCore.cpp

@@ -531,9 +531,8 @@ std::string serialization::StatusToString(Status S) {
   case Status::TargetIncompatible: return "TargetIncompatible";
   case Status::TargetTooNew: return "TargetTooNew";
   case Status::SDKMismatch: return "SDKMismatch";
-  default:
-    llvm_unreachable("The switch should cover all cases");
   }
+  llvm_unreachable("The switch should cover all cases");
 }
 
 bool serialization::isSerializedAST(StringRef data) {