[Macros] Add attached extension macros.

Author: hborla

Sources/SwiftParser/Attributes.swift

@@ -332,14 +332,44 @@ extension Parser {
         )
       )
     case nil:
+      let isAttached = self.peek().isAttachedKeyword
       return parseAttribute(argumentMode: .customAttribute) { parser in
-        let arguments = parser.parseArgumentListElements(pattern: .none)
+        let arguments: [RawTupleExprElementSyntax]
+        if isAttached {
+          arguments = parser.parseAttachedArguments()
+        } else {
+          arguments = parser.parseArgumentListElements(pattern: .none)
+        }
+
         return .argumentList(RawTupleExprElementListSyntax(elements: arguments, arena: parser.arena))
       }
     }
   }
 }
 
+extension Parser {
+  mutating func parseAttachedArguments() -> [RawTupleExprElementSyntax] {
+    let (unexpectedBeforeRole, role) = self.expect(.identifier, TokenSpec(.extension, remapping: .identifier), default: .identifier)
+    let roleTrailingComma = self.consume(if: .comma)
+    let roleElement = RawTupleExprElementSyntax(
+      label: nil,
+      colon: nil,
+      expression: RawExprSyntax(
+        RawIdentifierExprSyntax(
+          unexpectedBeforeRole,
+          identifier: role,
+          declNameArguments: nil,
+          arena: self.arena
+        )
+      ),
+      trailingComma: roleTrailingComma,
+      arena: self.arena
+    )
+    let additionalArgs = self.parseArgumentListElements(pattern: .none)
+    return [roleElement] + additionalArgs
+  }
+}
+
 extension Parser {
   mutating func parseDifferentiableAttribute() -> RawAttributeSyntax {
     let (unexpectedBeforeAtSign, atSign) = self.expect(.atSign)

Sources/SwiftParser/Types.swift

@@ -1076,6 +1076,10 @@ extension Lexer.Lexeme {
       || self.rawTokenKind == .prefixOperator
   }
 
+  var isAttachedKeyword: Bool {
+    return self.rawTokenKind == .identifier && self.tokenText == "attached"
+  }
+
   var isEllipsis: Bool {
     return self.isAnyOperator && self.tokenText == "..."
   }

Sources/SwiftSyntaxMacroExpansion/MacroExpansion.swift

@@ -10,6 +10,7 @@ public enum MacroRole {
   case peer
   case conformance
   case codeItem
+  case `extension`
 }
 
 extension MacroRole {
@@ -23,6 +24,7 @@ extension MacroRole {
     case .peer: return "PeerMacro"
     case .conformance: return "ConformanceMacro"
     case .codeItem: return "CodeItemMacro"
+    case .extension: return "ExtensionMacro"
     }
   }
 }
@@ -101,7 +103,7 @@ public func expandFreestandingMacro(
         let rewritten = try codeItemMacroDef.expansion(of: node, in: context)
         expandedSyntax = Syntax(CodeBlockItemListSyntax(rewritten))
 
-      case (.accessor, _), (.memberAttribute, _), (.member, _), (.peer, _), (.conformance, _), (.expression, _), (.declaration, _),
+      case (.accessor, _), (.memberAttribute, _), (.member, _), (.peer, _), (.conformance, _), (.extension, _), (.expression, _), (.declaration, _),
         (.codeItem, _):
         throw MacroExpansionError.unmatchedMacroRole(definition, macroRole)
       }
@@ -283,6 +285,42 @@ public func expandAttachedMacro<Context: MacroExpansionContext>(
         return "extension \(typeName) : \(protocolName) \(whereClause) {}"
       }
 
+    case (let attachedMacro as ExtensionMacro.Type, .extension):
+      guard let declGroup = declarationNode.asProtocol(DeclGroupSyntax.self) else {
+        // Compiler error: type mismatch.
+        throw MacroExpansionError.declarationNotDeclGroup
+      }
+      guard let identified = declarationNode.asProtocol(IdentifiedDeclSyntax.self)
+      else {
+        // Compiler error: type mismatch.
+        throw MacroExpansionError.declarationNotIdentified
+      }
+
+      let type: TypeSyntax = "\(identified.identifier)"
+
+      // Local function to expand a extension macro once we've opened up
+      // the existential.
+      func expandExtensionMacro(
+        _ node: some DeclGroupSyntax
+      ) throws -> [ExtensionDeclSyntax] {
+        return try attachedMacro.expansion(
+          of: attributeNode,
+          attachedTo: node,
+          providingExtensionsOf: type,
+          in: context
+        )
+      }
+
+      let extensions = try _openExistential(
+        declGroup,
+        do: expandExtensionMacro
+      )
+
+      // Form a buffer of peer declarations to return to the caller.
+      return extensions.map {
+        $0.formattedExpansion(definition.formatMode)
+      }
+
     default:
       throw MacroExpansionError.unmatchedMacroRole(definition, macroRole)
     }

Sources/SwiftSyntaxMacros/CMakeLists.txt

@@ -13,6 +13,7 @@ add_swift_host_library(SwiftSyntaxMacros
   MacroProtocols/ConformanceMacro.swift
   MacroProtocols/DeclarationMacro.swift
   MacroProtocols/ExpressionMacro.swift
+  MacroProtocols/ExtensionMacro.swift
   MacroProtocols/FreestandingMacro.swift
   MacroProtocols/Macro.swift
   MacroProtocols/Macro+Format.swift

Sources/SwiftSyntaxMacros/MacroProtocols/ExtensionMacro.swift

@@ -0,0 +1,34 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SwiftSyntax
+
+/// Describes a macro that can add conformances to the declaration it's
+/// attached to.
+public protocol ExtensionMacro: AttachedMacro {
+  /// Expand an attached conformance macro to produce a set of conformances.
+  ///
+  /// - Parameters:
+  ///   - node: The custom attribute describing the attached macro.
+  ///   - declaration: The declaration the macro attribute is attached to.
+  ///   - context: The context in which to perform the macro expansion.
+  ///
+  /// - Returns: the set of `(type, where-clause?)` pairs that each provide the
+  ///   protocol type to which the declared type conforms, along with
+  ///   an optional where clause.
+  static func expansion(
+    of node: AttributeSyntax,
+    attachedTo decl: some DeclGroupSyntax,
+    providingExtensionsOf type: some TypeSyntaxProtocol,
+    in context: some MacroExpansionContext
+  ) throws -> [ExtensionDeclSyntax]
+}

Tests/SwiftParserTest/AttributeTests.swift

@@ -629,4 +629,20 @@ final class AttributeTests: XCTestCase {
       """
     )
   }
+
+  func testAttachedExtensionAttribute() {
+    assertParse(
+      """
+      @attached(extension)
+      macro m()
+      """
+    )
+
+    assertParse(
+      """
+      @attached(extension, names: named(test))
+      macro m()
+      """
+    )
+  }
 }