Merge pull request #67636 from tshortli/enum-unavailable-element-derived-equatable-hashable

Sema: Fix unavailable enum element cases in derived hashable/equatable impls

Author: tshortli

lib/AST/ASTPrinter.cpp

@@ -5659,9 +5659,10 @@ void PrintAST::visitCaseStmt(CaseStmt *CS) {
                [&] { Printer << ", "; });
   }
   Printer << ":";
-  Printer.printNewline();
 
-  printASTNodes((cast<BraceStmt>(CS->getBody())->getElements()));
+  if (!printASTNodes((cast<BraceStmt>(CS->getBody())->getElements())))
+    Printer.printNewline();
+  indent();
 }
 
 void PrintAST::visitFailStmt(FailStmt *stmt) {

lib/Sema/DerivedConformanceComparable.cpp

@@ -342,5 +342,7 @@ void DerivedConformance::tryDiagnoseFailedComparableDerivation(
                          rawType, nominal->getDeclaredInterfaceType(),
                          comparableProto->getDeclaredInterfaceType());
     }
+    // FIXME: Diagnose potentially unavailable enum elements that are preventing
+    // Comparable synthesis.
   }
 }

lib/Sema/DerivedConformanceEquatableHashable.cpp

@@ -176,6 +176,13 @@ deriveBodyEquatable_enum_hasAssociatedValues_eq(AbstractFunctionDecl *eqDecl,
   for (auto elt : enumDecl->getAllElements()) {
     ++elementCount;
 
+    if (auto *unavailableElementCase =
+            DerivedConformance::unavailableEnumElementCaseStmt(
+                enumType, elt, eqDecl, /*subPatternCount=*/2)) {
+      cases.push_back(unavailableElementCase);
+      continue;
+    }
+
     // .<elt>(let l0, let l1, ...)
     SmallVector<VarDecl*, 3> lhsPayloadVars;
     auto lhsSubpattern = DerivedConformance::enumElementPayloadSubpattern(elt, 'l', eqDecl,
@@ -692,6 +699,13 @@ deriveBodyHashable_enum_hasAssociatedValues_hashInto(
   // For each enum element, generate a case statement that binds the associated
   // values so that they can be fed to the hasher.
   for (auto elt : enumDecl->getAllElements()) {
+    if (auto *unavailableElementCase =
+            DerivedConformance::unavailableEnumElementCaseStmt(enumType, elt,
+                                                               hashIntoDecl)) {
+      cases.push_back(unavailableElementCase);
+      continue;
+    }
+
     // case .<elt>(let a0, let a1, ...):
     SmallVector<VarDecl*, 3> payloadVars;
     SmallVector<ASTNode, 3> statements;

lib/Sema/DerivedConformances.cpp

@@ -493,6 +493,19 @@ DerivedConformance::createBuiltinCall(ASTContext &ctx,
   return call;
 }
 
+CallExpr *DerivedConformance::createDiagnoseUnavailableCodeReachedCallExpr(
+    ASTContext &ctx) {
+  FuncDecl *diagnoseDecl = ctx.getDiagnoseUnavailableCodeReached();
+  auto diagnoseDeclRefExpr =
+      new (ctx) DeclRefExpr(diagnoseDecl, DeclNameLoc(), true);
+  diagnoseDeclRefExpr->setType(diagnoseDecl->getInterfaceType());
+  auto argList = ArgumentList::createImplicit(ctx, {});
+  auto callExpr = CallExpr::createImplicit(ctx, diagnoseDeclRefExpr, argList);
+  callExpr->setType(ctx.getNeverType());
+  callExpr->setThrows(false);
+  return callExpr;
+}
+
 AccessorDecl *DerivedConformance::
 addGetterToReadOnlyDerivedProperty(VarDecl *property,
                                    Type propertyContextType) {
@@ -761,6 +774,13 @@ DeclRefExpr *DerivedConformance::convertEnumToIndex(SmallVectorImpl<ASTNode> &st
   unsigned index = 0;
   SmallVector<ASTNode, 4> cases;
   for (auto elt : enumDecl->getAllElements()) {
+    if (auto *unavailableElementCase =
+            DerivedConformance::unavailableEnumElementCaseStmt(enumType, elt,
+                                                               funcDecl)) {
+      cases.push_back(unavailableElementCase);
+      continue;
+    }
+
     // generate: case .<Case>:
     auto pat = new (C)
         EnumElementPattern(TypeExpr::createImplicit(enumType, C), SourceLoc(),
@@ -907,6 +927,52 @@ Pattern *DerivedConformance::enumElementPayloadSubpattern(
   return ParenPattern::createImplicit(C, letPattern);
 }
 
+CaseStmt *DerivedConformance::unavailableEnumElementCaseStmt(
+    Type enumType, EnumElementDecl *elt, DeclContext *parentDC,
+    unsigned subPatternCount) {
+  assert(subPatternCount > 0);
+
+  ASTContext &C = parentDC->getASTContext();
+  auto availableAttr = elt->getAttrs().getUnavailable(C);
+  if (!availableAttr)
+    return nullptr;
+
+  if (!availableAttr->isUnconditionallyUnavailable())
+    return nullptr;
+
+  auto createElementPattern = [&]() -> EnumElementPattern * {
+    // .<elt>
+    EnumElementPattern *eltPattern = new (C) EnumElementPattern(
+        TypeExpr::createImplicit(enumType, C), SourceLoc(), DeclNameLoc(),
+        DeclNameRef(elt->getBaseIdentifier()), elt, nullptr, /*DC*/ parentDC);
+    eltPattern->setImplicit();
+    return eltPattern;
+  };
+
+  Pattern *labelItemPattern;
+  if (subPatternCount > 1) {
+    SmallVector<TuplePatternElt, 2> tuplePatternElts;
+    for (unsigned i = 0; i < subPatternCount; i++) {
+      tuplePatternElts.push_back(TuplePatternElt(createElementPattern()));
+    }
+
+    // (.<elt>, ..., .<elt>)
+    auto caseTuplePattern = TuplePattern::createImplicit(C, tuplePatternElts);
+    caseTuplePattern->setImplicit();
+    labelItemPattern = caseTuplePattern;
+  } else {
+    labelItemPattern = createElementPattern();
+  }
+
+  auto labelItem = CaseLabelItem(labelItemPattern);
+  auto *callExpr =
+      DerivedConformance::createDiagnoseUnavailableCodeReachedCallExpr(C);
+  auto body = BraceStmt::create(C, SourceLoc(), {callExpr}, SourceLoc());
+  return CaseStmt::create(C, CaseParentKind::Switch, SourceLoc(), labelItem,
+                          SourceLoc(), SourceLoc(), body, {},
+                          /*implicit*/ true);
+}
+
 /// Creates a named variable based on a prefix character and a numeric index.
 /// \p prefixChar The prefix character for the variable's name.
 /// \p index The numeric index to append to the variable's name.

lib/Sema/DerivedConformances.h

@@ -29,6 +29,7 @@ class AssociatedTypeDecl;
 class ASTContext;
 struct ASTNode;
 class CallExpr;
+class CaseStmt;
 class Decl;
 class DeclContext;
 class DeclRefExpr;
@@ -389,6 +390,11 @@ class DerivedConformance {
                               ArrayRef<ProtocolConformanceRef> conformances,
                                      ArrayRef<Expr *> args);
 
+  /// Build a call to the stdlib function that should be called when unavailable
+  /// code is reached unexpectedly.
+  static CallExpr *
+  createDiagnoseUnavailableCodeReachedCallExpr(ASTContext &ctx);
+
   /// Returns true if this derivation is trying to use a context that isn't
   /// appropriate for deriving.
   ///
@@ -453,6 +459,16 @@ class DerivedConformance {
       EnumElementDecl *enumElementDecl, char varPrefix, DeclContext *varContext,
       SmallVectorImpl<VarDecl *> &boundVars, bool useLabels = false);
 
+  /// Creates a synthesized case statement that has the following structure:
+  ///
+  ///     case .<elt>, ..., .<elt>:
+  ///       _diagnoseUnavailableCodeReached()
+  ///
+  /// The number of \c .<elt> matches is equal to \p subPatternCount.
+  static CaseStmt *unavailableEnumElementCaseStmt(
+      Type enumType, EnumElementDecl *enumElementDecl, DeclContext *parentDC,
+      unsigned subPatternCount = 1);
+
   static VarDecl *indexedVarDecl(char prefixChar, int index, Type type,
                                  DeclContext *varContext);
 };

lib/Sema/TypeCheckSwitchStmt.cpp

@@ -828,8 +828,8 @@ namespace {
           auto children = E->getAllElements();
           llvm::transform(
               children, std::back_inserter(arr), [&](EnumElementDecl *eed) {
-                // Don't force people to match unavailable cases; they can't
-                // even write them.
+                // Don't force people to match unavailable cases since they
+                // should not be instantiated at run time.
                 if (AvailableAttr::isUnavailable(eed)) {
                   return Space();
                 }

test/Interpreter/enum_equatable_hashable_correctness.swift

@@ -1,4 +1,5 @@
 // RUN: %target-run-simple-swift
+// RUN: %target-run-simple-swift(-Xfrontend -unavailable-decl-optimization=complete)
 // REQUIRES: executable_test
 
 import StdlibUnittest
@@ -17,6 +18,20 @@ enum Combo<T: Hashable, U: Hashable>: Hashable {
   case both(T, U)
 }
 
+@available(*, unavailable)
+struct UnavailableStruct: Hashable {}
+
+enum HasUnavailableCases: Hashable {
+  case available
+  case availablePayload(Int)
+
+  @available(*, unavailable)
+  case unavailable
+
+  @available(*, unavailable)
+  case unavailablePayload(UnavailableStruct)
+}
+
 var EnumSynthesisTests = TestSuite("EnumSynthesis")
 
 EnumSynthesisTests.test("BasicEquatability/Hashability") {
@@ -52,6 +67,13 @@ EnumSynthesisTests.test("CloseGenericValuesDoNotCollide") {
   expectNotEqual(Combo<String, Int>.both("foo", 3).hashValue, Combo<String, Int>.both("goo", 4).hashValue)
 }
 
+EnumSynthesisTests.test("HasUnavailableCasesEquatability/Hashability") {
+  checkHashable([
+    HasUnavailableCases.available,
+    HasUnavailableCases.availablePayload(2),
+  ], equalityOracle: { $0 == $1 })
+}
+
 func hashEncode(_ body: (inout Hasher) -> ()) -> Int {
   var hasher = Hasher()
   body(&hasher)