Merge pull request swiftlang#70249 from apple/ApolloZhu/const-extract/attribute-on-all-nominals

[Compile Time Constant Extraction] Allow attribute on all nominal decls

Author: ApolloZhu

include/swift/AST/Attr.def

@@ -534,7 +534,7 @@ SIMPLE_DECL_ATTR(_staticExclusiveOnly, StaticExclusiveOnly,
   OnStruct | UserInaccessible | ABIStableToAdd | ABIStableToRemove | APIBreakingToAdd | APIStableToRemove,
   151)
 SIMPLE_DECL_ATTR(extractConstantsFromMembers, ExtractConstantsFromMembers,
-  OnProtocol | UserInaccessible | ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove,
+  OnClass | OnEnum | OnProtocol | OnStruct | UserInaccessible | ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove,
   152)
 SIMPLE_DECL_ATTR(_noRuntime, NoRuntime,
   OnAbstractFunction | OnSubscript | UserInaccessible | ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove,

lib/ConstExtract/ConstExtract.cpp

@@ -58,11 +58,17 @@ class NominalTypeConformanceCollector : public ASTWalker {
       if (auto *ETD = dyn_cast<ExtensionDecl>(D))
         NTD = ETD->getExtendedNominal();
     if (NTD)
-      if (!isa<ProtocolDecl>(NTD) && CheckedDecls.insert(NTD).second)
+      if (!isa<ProtocolDecl>(NTD) && CheckedDecls.insert(NTD).second) {
+        if (NTD->getAttrs().hasAttribute<ExtractConstantsFromMembersAttr>()) {
+          ConformanceTypeDecls.push_back(NTD);
+          return Action::Continue();
+        }
+
         for (auto &Protocol : NTD->getAllProtocols())
           if (Protocol->getAttrs().hasAttribute<ExtractConstantsFromMembersAttr>() ||
               Protocols.count(Protocol->getName().str().str()) != 0)
             ConformanceTypeDecls.push_back(NTD);
+      }
     return Action::Continue();
   }
 

lib/Serialization/ModuleFormat.h

@@ -58,7 +58,7 @@ const uint16_t SWIFTMODULE_VERSION_MAJOR = 0;
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
 /// Don't worry about adhering to the 80-column limit for this line.
-const uint16_t SWIFTMODULE_VERSION_MINOR = 821; // extractConstantsFromMembers
+const uint16_t SWIFTMODULE_VERSION_MINOR = 822; // extractConstantsFromMembers on all nominals
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///

test/ConstExtraction/ExtractConstantsFromMembersAttribute.swift

@@ -12,9 +12,56 @@ public struct TestStruct : MyProto {
     let cane: [String] = ["bar", "baz"]
 }
 
+@extractConstantsFromMembers
+struct DirectOnStruct {
+  var really = true
+}
+
+@extractConstantsFromMembers
+class DirectOnClass {
+  static let answer = 42
+}
+
+@extractConstantsFromMembers
+enum DirectOnEnum {
+  case yes
+}
+
+// CHECK:         "typeName": "ExtractConstantsFromMembersAttribute.TestStruct",
+// CHECK:         "kind": "struct",
+// CHECK:         "conformances": [
+// CHECK-NEXT:         "ExtractConstantsFromMembersAttribute.MyProto"
+// CHECK-NEXT:    ],
+// CHECK:         "properties": [
 // CHECK:             "label": "foo",
+// CHECK-NEXT:        "type": "Swift.String",
 // CHECK:             "valueKind": "RawLiteral",
 // CHECK:             "value": "foo"
 
 // CHECK:             "label": "cane",
 // CHECK-NEXT:        "type": "Swift.Array<Swift.String>",
+// CHECK:             "valueKind": "Array",
+
+
+// CHECK:         "typeName": "ExtractConstantsFromMembersAttribute.DirectOnStruct",
+// CHECK:         "kind": "struct",
+// CHECK:         "properties": [
+// CHECK:             "label": "really",
+// CHECK-NEXT:        "type": "Swift.Bool",
+// CHECK:             "valueKind": "RawLiteral",
+// CHECK:             "value": "true"
+
+
+// CHECK:         "ExtractConstantsFromMembersAttribute.DirectOnClass",
+// CHECK:         "kind": "class",
+// CHECK:         "properties": [
+// CHECK:             "label": "answer",
+// CHECK-NEXT:        "type": "Swift.Int",
+// CHECK:             "valueKind": "RawLiteral",
+// CHECK:             "value": "42"
+
+
+// CHECK:         "ExtractConstantsFromMembersAttribute.DirectOnEnum",
+// CHECK:         "kind": "enum",
+// CHECK:         "cases": [
+// CHECK:             "name": "yes"