Merge pull request #80379 from hamishknight/extendable-expansion

Author: hamishknight

lib/Sema/TypeCheckMacros.cpp

@@ -957,8 +957,12 @@ static CharSourceRange getExpansionInsertionRange(MacroRole role,
   }
 
   case MacroRole::Extension: {
+    // Extensions are expanded at the top-level.
+    auto *NTD = cast<NominalTypeDecl>(target.get<Decl *>());
+    auto *topLevelDecl = NTD->getTopmostDeclarationDeclContext();
+
     SourceLoc afterDeclLoc =
-        Lexer::getLocForEndOfToken(sourceMgr, target.getEndLoc());
+        Lexer::getLocForEndOfToken(sourceMgr, topLevelDecl->getEndLoc());
     return CharSourceRange(afterDeclLoc, 0);
   }
 
@@ -2095,7 +2099,7 @@ std::optional<unsigned> swift::expandExtensions(CustomAttr *attr,
                                                 MacroDecl *macro,
                                                 MacroRole role,
                                                 NominalTypeDecl *nominal) {
-  if (nominal->getDeclContext()->isLocalContext()) {
+  if (nominal->getLocalContext()) {
     nominal->diagnose(diag::local_extension_macro);
     return std::nullopt;
   }

test/ConstExtraction/Inputs/Macros.swift

@@ -41,10 +41,9 @@ public struct AddExtensionMacro: ExtensionMacro {
     conformingTo protocols: [TypeSyntax],
     in context: some MacroExpansionContext
   ) throws -> [ExtensionDeclSyntax] {
-    let typeName = declaration.declGroupName
     return protocols.map {
       ("""
-      extension \(typeName): \($0) {
+      extension \(type.trimmed): \($0) {
         struct _Extension_\($0): \($0) {
           var nested = 8
         }
@@ -53,8 +52,8 @@ public struct AddExtensionMacro: ExtensionMacro {
       .cast(ExtensionDeclSyntax.self)
     } + [
     ("""
-    extension \(typeName) {
-      static let _extension_\(typeName) = 3
+    extension \(type.trimmed) {
+      static let _extension_\(declaration.declGroupName) = 3
     }
     """ as DeclSyntax).cast(ExtensionDeclSyntax.self)
     ]

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -1628,6 +1628,20 @@ public struct FooExtensionMacro: ExtensionMacro {
   }
 }
 
+public struct BadExtensionMacro: ExtensionMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    attachedTo declaration: some DeclGroupSyntax,
+    providingExtensionsOf type: some TypeSyntaxProtocol,
+    conformingTo protocols: [TypeSyntax],
+    in context: some MacroExpansionContext
+  ) throws -> [ExtensionDeclSyntax] {
+    // Note this is purposefully not using `providingExtensionsOf`.
+    let unqualifiedName = declaration.as(StructDeclSyntax.self)!.name.trimmed
+    return [try ExtensionDeclSyntax("extension \(unqualifiedName) {}")]
+  }
+}
+
 public struct ConformanceViaExtensionMacro: ExtensionMacro {
   public static func expansion(
     of node: AttributeSyntax,

test/Macros/macro_expand_extensions.swift

@@ -115,6 +115,12 @@ func testLocal() {
   @DelegatedConformance
   struct Local<Element> {}
   // expected-error@-1{{local type cannot have attached extension macro}}
+
+  struct S {
+    @DelegatedConformance
+    struct Local<Element> {}
+    // expected-error@-1{{local type cannot have attached extension macro}}
+  }
 }
 
 @DelegatedConformance
@@ -161,6 +167,17 @@ struct TestUndocumentedEncodable {}
 
 // CHECK-DIAGS: error: conformance to 'Codable' (aka 'Decodable & Encodable') is not covered by macro 'UndocumentedEncodable'
 
+@attached(extension)
+macro BadExtension() = #externalMacro(module: "MacroDefinition", type: "BadExtensionMacro")
+
+// Make sure 'extension Foo' is rejected here as it needs to
+// be a qualified reference.
+struct HasSomeNestedType {
+  @BadExtension // expected-note {{in expansion of macro 'BadExtension' on struct 'SomeNestedType' here}}
+  struct SomeNestedType {}
+}
+// CHECK-DIAGS: error: cannot find type 'SomeNestedType' in scope
+
 #endif
 
 @attached(extension, conformances: Equatable)

test/Reflection/Inputs/Macros.swift

@@ -40,10 +40,9 @@ public struct AddExtensionMacro: ExtensionMacro {
     conformingTo protocols: [TypeSyntax],
     in context: some MacroExpansionContext
   ) throws -> [ExtensionDeclSyntax] {
-    let typeName = declaration.declGroupName
     return protocols.map {
       ("""
-      extension \(typeName): \($0) {
+      extension \(type.trimmed): \($0) {
         struct _Extension_\($0): \($0) {}
       }
       """ as DeclSyntax)
@@ -66,12 +65,3 @@ extension DeclSyntaxProtocol {
         fatalError("Not implemented")
     }
 }
-
-extension DeclGroupSyntax {
-    var declGroupName: TokenSyntax {
-        if let structDecl = self.as(StructDeclSyntax.self) {
-            return structDecl.name.trimmed
-        }
-        fatalError("Not implemented")
-    }
-}

test/SourceKit/Macros/macro_basic.swift

@@ -121,6 +121,14 @@ func remoteCall<Result: ConjureRemoteValue>(function: String, arguments: [String
   return Result.conjureValue()
 }
 
+@attached(extension, conformances: Equatable)
+macro AddEquatable() = #externalMacro(module: "MacroDefinition", type: "EquatableMacro")
+
+struct HasNestedType {
+  @AddEquatable
+  struct Inner {}
+}
+
 // REQUIRES: swift_swift_parser, executable_test, shell, asserts
 // REQUIRES: swift_feature_PreambleMacros
 
@@ -379,3 +387,10 @@ func remoteCall<Result: ConjureRemoteValue>(function: String, arguments: [String
 // BODY_EXPAND-NEXT: return try await remoteCall(function: "f", arguments: ["a": a, "b": b])
 // BODY_EXPAND-NEXT: }"
 // BODY_EXPAND-NEXT: source.edit.kind.active:
+
+// Make sure the extension is added at the top level.
+// RUN: %sourcekitd-test -req=refactoring.expand.macro -pos=128:4 %s -- ${COMPILER_ARGS[@]} | %FileCheck -check-prefix=ADD_EQUATABLE_EXPAND %s
+// ADD_EQUATABLE_EXPAND: source.edit.kind.active:
+// ADD_EQUATABLE_EXPAND-NEXT: 130:2-130:2 (@__swiftmacro_9MacroUser13HasNestedTypeV5Inner12AddEquatablefMe_.swift) "extension HasNestedType.Inner: Equatable {
+// ADD_EQUATABLE_EXPAND-NEXT: }"
+// ADD_EQUATABLE_EXPAND-NEXT: source.edit.kind.active: