Merge pull request #2415 from bnbarham/qualify-extensions

MacroSystem should qualify nested types when expanding extension macros

Author: bnbarham

Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift

@@ -319,12 +319,11 @@ private func expandExtensionMacro(
   in context: some MacroExpansionContext,
   indentationWidth: Trivia
 ) throws -> CodeBlockItemListSyntax? {
-  let extendedType: TypeSyntax
-  if let identified = attachedTo.asProtocol(NamedDeclSyntax.self) {
-    extendedType = "\(identified.name.trimmed)"
-  } else if let ext = attachedTo.as(ExtensionDeclSyntax.self) {
-    extendedType = "\(ext.extendedType.trimmed)"
-  } else {
+  guard attachedTo.isProtocol(DeclGroupSyntax.self) else {
+    return nil
+  }
+
+  guard let extendedType = attachedTo.syntacticQualifiedTypeContext else {
     return nil
   }
 
@@ -1319,3 +1318,25 @@ private extension AccessorDeclSyntax {
     return accessorSpecifier.tokenKind == .keyword(.get) || accessorSpecifier.tokenKind == .keyword(.set)
   }
 }
+
+private extension SyntaxProtocol {
+  /// Retrieve the qualified type for the nearest extension or name decl.
+  ///
+  /// For example, for `struct Foo { struct Bar {} }`, calling this on the
+  /// inner struct (`Bar`) will return `Foo.Bar`.
+  var syntacticQualifiedTypeContext: TypeSyntax? {
+    if let ext = self.as(ExtensionDeclSyntax.self) {
+      // Don't handle nested 'extension' cases - they are invalid anyway.
+      return ext.extendedType.trimmed
+    }
+
+    let base = self.parent?.syntacticQualifiedTypeContext
+    if let named = self.asProtocol(NamedDeclSyntax.self) {
+      if let base = base {
+        return TypeSyntax(MemberTypeSyntax(baseType: base, name: named.name.trimmed))
+      }
+      return TypeSyntax(IdentifierTypeSyntax(name: named.name.trimmed))
+    }
+    return base
+  }
+}

Tests/SwiftSyntaxMacroExpansionTest/ExtensionMacroTests.swift

@@ -25,90 +25,130 @@ import SwiftSyntaxMacros
 import SwiftSyntaxMacrosTestSupport
 import XCTest
 
-fileprivate struct DeclsFromStringsMacro: DeclarationMacro {
-  static func expansion(
-    of node: some FreestandingMacroExpansionSyntax,
-    in context: some MacroExpansionContext
-  ) throws -> [DeclSyntax] {
-    var strings: [String] = []
-    for arg in node.arguments {
-      guard let value = arg.expression.as(StringLiteralExprSyntax.self)?.representedLiteralValue else {
-        continue
-      }
-      strings.append(value)
-    }
-
-    return strings.map { "\(raw: $0)" }
-  }
-}
-
 final class ExtensionMacroTests: XCTestCase {
   private let indentationWidth: Trivia = .spaces(2)
 
-  func testExtensionExpansion() {
-    struct SendableExtensionMacro: ExtensionMacro {
-      static func expansion(
-        of node: AttributeSyntax,
-        attachedTo: some DeclGroupSyntax,
-        providingExtensionsOf type: some TypeSyntaxProtocol,
-        conformingTo protocols: [TypeSyntax],
-        in context: some MacroExpansionContext
-      ) throws -> [ExtensionDeclSyntax] {
-        let sendableExtension: DeclSyntax =
-          """
-          extension \(type.trimmed): Sendable {}
-          """
+  func testSimpleExpansion() {
+    assertMacroExpansion(
+      """
+      @AddSendableExtension
+      struct MyType {}
+      """,
+      expandedSource: """
+        struct MyType {}
 
-        guard let extensionDecl = sendableExtension.as(ExtensionDeclSyntax.self) else {
-          return []
+        extension MyType: Sendable {
         }
+        """,
+      macros: ["AddSendableExtension": SendableExtensionMacro.self],
+      indentationWidth: indentationWidth
+    )
+  }
 
-        return [extensionDecl]
-      }
-    }
-
+  func testNestedExpansion() {
     assertMacroExpansion(
       """
-      @AddSendableExtension
-      struct MyType {
+      struct Wrapper {
+        @AddSendableExtension
+        struct MyType {}
       }
       """,
       expandedSource: """
-
-        struct MyType {
+        struct Wrapper {
+          struct MyType {}
         }
 
-        extension MyType: Sendable {
+        extension Wrapper.MyType: Sendable {
         }
         """,
       macros: ["AddSendableExtension": SendableExtensionMacro.self],
       indentationWidth: indentationWidth
     )
+  }
 
+  func testNestedInExtensionExpansion() {
     assertMacroExpansion(
       """
       struct Wrapper {
+        struct AnotherWrapper {}
+      }
+      extension Wrapper.AnotherWrapper {
         @AddSendableExtension
-        struct MyType {
-        }
+        struct MyType {}
       }
       """,
       expandedSource: """
         struct Wrapper {
-          struct MyType {
+          struct AnotherWrapper {}
+        }
+        extension Wrapper.AnotherWrapper {
+          struct MyType {}
+        }
+
+        extension Wrapper.AnotherWrapper.MyType: Sendable {
+        }
+        """,
+      macros: ["AddSendableExtension": SendableExtensionMacro.self],
+      indentationWidth: indentationWidth
+    )
+  }
+
+  func testComplexNestedExpansion() {
+    assertMacroExpansion(
+      """
+      struct Wrapper {}
+      extension Wrapper {
+        struct AnotherWrapper {
+          @AddSendableExtension
+          struct MyType {}
+        }
+      }
+      """,
+      expandedSource: """
+        struct Wrapper {}
+        extension Wrapper {
+          struct AnotherWrapper {
+            struct MyType {}
           }
         }
 
-        extension MyType: Sendable {
+        extension Wrapper.AnotherWrapper.MyType: Sendable {
         }
         """,
       macros: ["AddSendableExtension": SendableExtensionMacro.self],
       indentationWidth: indentationWidth
     )
   }
 
+  func testAttachedToInvalid() {
+    assertMacroExpansion(
+      "@AddSendableExtension var foo: Int",
+      expandedSource: "var foo: Int",
+      macros: [
+        "AddSendableExtension": SendableExtensionMacro.self
+      ]
+    )
+
+    assertMacroExpansion(
+      """
+      struct Foo {
+        @AddSendableExtension var foo: Int
+      }
+      """,
+      expandedSource:
+        """
+        struct Foo {
+          var foo: Int
+        }
+        """,
+      macros: [
+        "AddSendableExtension": SendableExtensionMacro.self
+      ]
+    )
+  }
+
   func testEmpty() {
-    struct TestMacro: ExtensionMacro {
+    struct EmptyExtensionMacro: ExtensionMacro {
       static func expansion(
         of node: AttributeSyntax,
         attachedTo declaration: some DeclGroupSyntax,
@@ -121,11 +161,32 @@ final class ExtensionMacroTests: XCTestCase {
     }
 
     assertMacroExpansion(
-      "@Test struct Foo {}",
+      "@Empty struct Foo {}",
       expandedSource: "struct Foo {}",
       macros: [
-        "Test": TestMacro.self
+        "Empty": EmptyExtensionMacro.self
       ]
     )
   }
 }
+
+fileprivate struct SendableExtensionMacro: ExtensionMacro {
+  static func expansion(
+    of node: AttributeSyntax,
+    attachedTo: some DeclGroupSyntax,
+    providingExtensionsOf type: some TypeSyntaxProtocol,
+    conformingTo protocols: [TypeSyntax],
+    in context: some MacroExpansionContext
+  ) throws -> [ExtensionDeclSyntax] {
+    let sendableExtension: DeclSyntax =
+      """
+      extension \(type.trimmed): Sendable {}
+      """
+
+    guard let extensionDecl = sendableExtension.as(ExtensionDeclSyntax.self) else {
+      return []
+    }
+
+    return [extensionDecl]
+  }
+}