Merge pull request #72425 from apple/ApolloZhu/6.0/const-extract/extension-macro

[6.0][Compile Time Constant Extraction] Extract macro added extensions

Author: ApolloZhu

lib/ConstExtract/ConstExtract.cpp

@@ -560,6 +560,11 @@ gatherConstValuesForPrimary(const std::unordered_set<std::string> &Protocols,
                                                        ConformanceDecls);
   for (auto D : SF->getTopLevelDecls())
     D->walk(ConformanceCollector);
+  // Visit macro expanded extensions
+  if (auto *synthesizedSF = SF->getSynthesizedFile())
+    for (auto D : synthesizedSF->getTopLevelDecls())
+      if (isa<ExtensionDecl>(D))
+        D->walk(ConformanceCollector);
 
   for (auto *CD : ConformanceDecls)
     Result.emplace_back(evaluateOrDefault(

test/ConstExtraction/ExtractFromMacroExpansion.swift

@@ -15,7 +15,7 @@ macro AddMacroAddedStruct() = #externalMacro(module: "MacroDefinition", type: "A
 @freestanding(declaration, names: named(macroAddedVar))
 macro AddMacroAddedVar() = #externalMacro(module: "MacroDefinition", type: "AddVarDeclMacro")
 
-@attached(extension, conformances: MyProto, names: prefixed(_extension_))
+@attached(extension, conformances: MyProto, names: prefixed(_extension_), named(_Extension_MyProto))
 macro AddExtension() = #externalMacro(module: "MacroDefinition", type: "AddExtensionMacro")
 
 @attached(peer, names: prefixed(_peer_))
@@ -43,6 +43,9 @@ struct MyStruct {
   #AddMacroAddedVar
 
   @AddPeerVar
+  @AddExtension
+  @AddMemberVar
+  @AddPeerStruct
   struct Inner { }
 }
 
@@ -131,3 +134,83 @@ extension MyStruct {
 // CHECK:   "type": "Swift.Int",
 // CHECK:   "valueKind": "RawLiteral",
 // CHECK:   "value": "3"
+
+
+// CHECK: "typeName": "ExtractFromMacroExpansion.MyStruct.Inner",
+// CHECK: "properties": [
+// CHECK:   "label": "_member_Inner",
+// CHECK:   "type": "Swift.Int",
+// CHECK:   "valueKind": "RawLiteral",
+// CHECK:   "value": "5"
+
+// CHECK:   "label": "_extension_Inner",
+// CHECK:   "type": "Swift.Int",
+// CHECK:   "valueKind": "RawLiteral",
+// CHECK:   "value": "3"
+
+
+// CHECK: "typeName": "ExtractFromMacroExpansion.MyStruct._Peer_Inner",
+// CHECK: "properties": [
+// CHECK:   "label": "peerMacroVar",
+// CHECK:   "type": "Swift.Int",
+// CHECK:   "valueKind": "RawLiteral",
+// CHECK:   "value": "7"
+
+// CHECK:   "label": "macroAddedVar",
+// CHECK:   "type": "Swift.Int",
+// CHECK:   "valueKind": "RawLiteral",
+// CHECK:   "value": "2"
+
+// CHECK:   "label": "_peer_peerMacroVar",
+// CHECK:   "type": "Swift.Int",
+// CHECK:   "valueKind": "RawLiteral",
+// CHECK:   "value": "4"
+
+// CHECK:   "label": "_member__Peer_Inner",
+// CHECK:   "type": "Swift.Int",
+// CHECK:   "valueKind": "RawLiteral",
+// CHECK:   "value": "5"
+
+// CHECK:   "label": "_extension__Peer_Inner",
+// CHECK:   "type": "Swift.Int",
+// CHECK:   "valueKind": "RawLiteral",
+// CHECK:   "value": "3"
+
+
+// CHECK: "typeName": "ExtractFromMacroExpansion.MacroAddedStruct._Extension_MyProto",
+// CHECK: "properties": [
+// CHECK:   "label": "nested",
+// CHECK:   "type": "Swift.Int",
+// CHECK:   "valueKind": "RawLiteral",
+// CHECK:   "value": "8"
+
+
+// CHECK: "typeName": "ExtractFromMacroExpansion._Peer_MyStruct._Extension_MyProto",
+// CHECK: "properties": [
+// CHECK:   "label": "nested",
+// CHECK:   "type": "Swift.Int",
+// CHECK:   "valueKind": "RawLiteral",
+// CHECK:   "value": "8"
+
+// CHECK: "typeName": "ExtractFromMacroExpansion.MyStruct._Extension_MyProto",
+// CHECK: "properties": [
+// CHECK:   "label": "nested",
+// CHECK:   "type": "Swift.Int",
+// CHECK:   "valueKind": "RawLiteral",
+// CHECK:   "value": "8"
+
+
+// CHECK: "typeName": "ExtractFromMacroExpansion.MyStruct._Peer_Inner._Extension_MyProto",
+// CHECK: "properties": [
+// CHECK:   "label": "nested",
+// CHECK:   "type": "Swift.Int",
+// CHECK:   "valueKind": "RawLiteral",
+// CHECK:   "value": "8"
+
+
+// CHECK: "typeName": "ExtractFromMacroExpansion.MyStruct.Inner._Extension_MyProto",
+// CHECK: "properties": [
+// CHECK:   "label": "nested",
+// CHECK:   "type": "Swift.Int",
+// CHECK:   "valueKind": "RawLiteral",
+// CHECK:   "value": "8"

test/ConstExtraction/Inputs/Macros.swift

@@ -43,8 +43,14 @@ public struct AddExtensionMacro: ExtensionMacro {
   ) throws -> [ExtensionDeclSyntax] {
     let typeName = declaration.declGroupName
     return protocols.map {
-      ("extension \(typeName): \($0) { }" as DeclSyntax)
-        .cast(ExtensionDeclSyntax.self)
+      ("""
+      extension \(typeName): \($0) {
+        struct _Extension_\($0): \($0) {
+          var nested = 8
+        }
+      }
+      """ as DeclSyntax)
+      .cast(ExtensionDeclSyntax.self)
     } + [
     ("""
     extension \(typeName) {