[AST/Sema] SE-0487: Expand @nonexhaustive attribute to support warn argument

The spelling `@nonexhaustive(warn)` replaces `@preEnumExtensibility`
attriubte.

Author: xedin

include/swift/AST/Attr.h

@@ -136,6 +136,11 @@ enum : unsigned {
       InheritActorContextModifier::Last_InheritActorContextKind))
 };
 
+enum : unsigned {
+  NumNonexhaustiveModeBits = countBitsUsed(
+      static_cast<unsigned>(NonexhaustiveMode::Last_NonexhaustiveMode))
+};
+
 enum : unsigned { NumDeclAttrKindBits = countBitsUsed(NumDeclAttrKinds - 1) };
 
 enum : unsigned { NumTypeAttrKindBits = countBitsUsed(NumTypeAttrKinds - 1) };
@@ -277,6 +282,10 @@ class DeclAttribute : public AttributeBase {
     SWIFT_INLINE_BITFIELD(LifetimeAttr, DeclAttribute, 1,
       isUnderscored : 1
     );
+
+    SWIFT_INLINE_BITFIELD(NonexhaustiveAttr, DeclAttribute, NumNonexhaustiveModeBits,
+      mode : NumNonexhaustiveModeBits
+    );
   } Bits;
   // clang-format on
 
@@ -3500,6 +3509,36 @@ class ABIAttr : public DeclAttribute {
   }
 };
 
+/// Defines a @nonexhaustive attribute.
+class NonexhaustiveAttr : public DeclAttribute {
+public:
+  NonexhaustiveAttr(SourceLoc atLoc, SourceRange range, NonexhaustiveMode mode,
+                    bool implicit = false)
+      : DeclAttribute(DeclAttrKind::Nonexhaustive, atLoc, range, implicit) {
+    Bits.NonexhaustiveAttr.mode = unsigned(mode);
+  }
+
+  NonexhaustiveAttr(NonexhaustiveMode mode)
+    : NonexhaustiveAttr(SourceLoc(), SourceRange(), mode) {}
+
+  NonexhaustiveMode getMode() const {
+    return NonexhaustiveMode(Bits.NonexhaustiveAttr.mode);
+  }
+
+  static bool classof(const DeclAttribute *DA) {
+    return DA->getKind() == DeclAttrKind::Nonexhaustive;
+  }
+
+  NonexhaustiveAttr *clone(ASTContext &ctx) const {
+    return new (ctx) NonexhaustiveAttr(AtLoc, Range, getMode(), isImplicit());
+  }
+
+  bool isEquivalent(const NonexhaustiveAttr *other, Decl *attachedTo) const {
+    return getMode() == other->getMode();
+  }
+};
+
+
 /// The kind of unary operator, if any.
 enum class UnaryOperatorKind : uint8_t { None, Prefix, Postfix };
 

include/swift/AST/AttrKind.h

@@ -135,6 +135,12 @@ enum class ENUM_EXTENSIBILITY_ATTR(closed)
       Last_InheritActorContextKind = Always
     };
 
+enum class ENUM_EXTENSIBILITY_ATTR(closed) NonexhaustiveMode : uint8_t {
+  Error SWIFT_NAME("error") = 0,
+  Warning SWIFT_NAME("warning") = 1,
+  Last_NonexhaustiveMode = Warning
+};
+
 enum class ENUM_EXTENSIBILITY_ATTR(closed) DeclAttrKind : unsigned {
 #define DECL_ATTR(_, CLASS, ...) CLASS,
 #define LAST_DECL_ATTR(CLASS) Last_DeclAttr = CLASS,

include/swift/AST/DeclAttr.def

@@ -877,7 +877,7 @@ SIMPLE_DECL_ATTR(constInitialized, ConstInitialized,
   168)
 DECL_ATTR_FEATURE_REQUIREMENT(ConstInitialized, CompileTimeValues)
 
-SIMPLE_DECL_ATTR(nonexhaustive, Nonexhaustive,
+DECL_ATTR(nonexhaustive, Nonexhaustive,
   OnEnum,
   ABIStableToAdd | ABIStableToRemove | APIBreakingToAdd | APIStableToRemove | ForbiddenInABIAttr,
   169)

include/swift/AST/KnownIdentifiers.def

@@ -169,6 +169,7 @@ IDENTIFIER(value)
 IDENTIFIER_WITH_NAME(value_, "_value")
 IDENTIFIER(Void)
 IDENTIFIER(WinSDK)
+IDENTIFIER(warn)
 IDENTIFIER(with)
 IDENTIFIER(withArguments)
 IDENTIFIER(withKeywordArguments)

include/swift/Parse/IDEInspectionCallbacks.h

@@ -40,7 +40,8 @@ enum class ParameterizedDeclAttributeKind {
   FreestandingMacro,
   AttachedMacro,
   StorageRestrictions,
-  InheritActorContext
+  InheritActorContext,
+  Nonexhaustive,
 };
 
 /// A bit of a hack. When completing inside the '@storageRestrictions'

lib/AST/Attr.cpp

@@ -1732,6 +1732,19 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
     break;
   }
 
+  case DeclAttrKind::Nonexhaustive: {
+    auto *attr = cast<NonexhaustiveAttr>(this);
+    Printer << "@nonexhaustive";
+    switch (attr->getMode()) {
+    case NonexhaustiveMode::Error:
+      break;
+    case NonexhaustiveMode::Warning:
+      Printer << "(warn)";
+      break;
+    }
+    break;
+  }
+
 #define SIMPLE_DECL_ATTR(X, CLASS, ...) case DeclAttrKind::CLASS:
 #include "swift/AST/DeclAttr.def"
     llvm_unreachable("handled above");
@@ -1967,6 +1980,8 @@ StringRef DeclAttribute::getAttrName() const {
     }
   case DeclAttrKind::Lifetime:
     return cast<LifetimeAttr>(this)->isUnderscored() ? "_lifetime" : "lifetime";
+  case DeclAttrKind::Nonexhaustive:
+    return "nonexhaustive";
   }
   llvm_unreachable("bad DeclAttrKind");
 }

lib/IDE/CompletionLookup.cpp

@@ -3278,6 +3278,11 @@ void CompletionLookup::getAttributeDeclParamCompletions(
         getStoredPropertyCompletions(NT);
       }
     }
+    break;
+  }
+  case ParameterizedDeclAttributeKind::Nonexhaustive: {
+    addDeclAttrParamKeyword("warn", /*Parameters=*/{}, "", false);
+    break;
   }
   }
 }

lib/Parse/ParseDecl.cpp

@@ -4033,6 +4033,21 @@ ParserStatus Parser::parseNewDeclAttribute(DeclAttributes &Attributes,
       Attributes.add(Attr.get());
     break;
   }
+
+  case DeclAttrKind::Nonexhaustive: {
+    auto mode = parseSingleAttrOption<NonexhaustiveMode>(
+        *this, Loc, AttrRange, AttrName, DK,
+        {{Context.Id_warn, NonexhaustiveMode::Warning}},
+        NonexhaustiveMode::Error,
+        ParameterizedDeclAttributeKind::Nonexhaustive);
+    if (!mode)
+      return makeParserSuccess();
+
+    if (!DiscardAttribute)
+      Attributes.add(new (Context) NonexhaustiveAttr(AtLoc, AttrRange, *mode));
+
+    break;
+  }
   }
 
   if (DuplicateAttribute) {

lib/Sema/TypeCheckSwitchStmt.cpp

@@ -1167,14 +1167,20 @@ namespace {
             DE.diagnose(startLoc, diag::non_exhaustive_switch_unknown_only,
                         subjectType, shouldIncludeFutureVersionComment);
 
-        // Presence of `@preEnumExtensibility` pushed the warning farther
-        // into the future.
-        if (theEnum &&
-            theEnum->getAttrs().hasAttribute<PreEnumExtensibilityAttr>()) {
-          diag.warnUntilFutureSwiftVersion();
-        } else {
-          diag.warnUntilSwiftVersion(6);
-        }
+        auto shouldWarnUntilVersion = [&theEnum]() -> unsigned {
+          if (theEnum) {
+            // Presence of `@nonexhaustive(warn)` pushes the warning farther,
+            // into the future.
+            if (auto *nonexhaustive =
+                    theEnum->getAttrs().getAttribute<NonexhaustiveAttr>()) {
+              if (nonexhaustive->getMode() == NonexhaustiveMode::Warning)
+                return swift::version::Version::getFutureMajorLanguageVersion();
+            }
+          }
+          return 6;
+        };
+
+        diag.warnUntilSwiftVersion(shouldWarnUntilVersion());
 
         mainDiagType = std::nullopt;
       }

lib/Serialization/Deserialization.cpp

@@ -6676,6 +6676,14 @@ llvm::Error DeclDeserializer::deserializeDeclCommon() {
         break;
       }
 
+      case decls_block::Nonexhaustive_DECL_ATTR: {
+        unsigned mode;
+        serialization::decls_block::NonexhaustiveDeclAttrLayout::readRecord(
+            scratch, mode);
+        Attr = new (ctx) NonexhaustiveAttr((NonexhaustiveMode)mode);
+        break;
+      }
+
 #define SIMPLE_DECL_ATTR(NAME, CLASS, ...) \
       case decls_block::CLASS##_DECL_ATTR: { \
         bool isImplicit; \