Add OptionSet macro example and improve MacroToolkit to work well for it (new attribute, expression, and struct wrapps)

stackotter · stackotter · commit d982a14af8d5 · 2023-06-17T12:06:05.000+10:00
diff --git a/Sources/MacroToolkit/MacroToolkit.swift b/Sources/MacroToolkit/MacroToolkit.swift
@@ -290,6 +290,10 @@ public struct Type: SyntaxExpressibleByStringInterpolation {
         _syntax = syntax
     }
 
+    public init(_ syntax: any TypeSyntaxProtocol) {
+        _syntax = TypeSyntax(syntax)
+    }
+
     public init(stringInterpolation: SyntaxStringInterpolation) {
         _syntax = TypeSyntax(stringInterpolation: stringInterpolation)
     }
@@ -448,6 +452,104 @@ extension Collection where Element == AttributeListElement {
     }
 }
 
+public struct Expression {
+    public var _syntax: ExprSyntax
+
+    public init(_ syntax: ExprSyntax) {
+        _syntax = syntax
+    }
+
+    public init(_ syntax: any ExprSyntaxProtocol) {
+        _syntax = ExprSyntax(syntax)
+    }
+
+    /// Gets the contents of the expression if it's a string literal with no interpolation.
+    public var asSimpleStringLiteral: String? {
+        guard
+            let literal = _syntax.as(StringLiteralExprSyntax.self),
+            literal.segments.count == 1,
+            case let .stringSegment(segment)? = literal.segments.first
+        else {
+            return nil
+        }
+        return segment.content.text
+    }
+}
+
+public struct MacroAttribute {
+    public var _syntax: AttributeSyntax
+    public var _argumentListSyntax: TupleExprElementListSyntax
+
+    public init?(_ syntax: AttributeSyntax) {
+        guard case let .argumentList(arguments) = syntax.argument else {
+            return nil
+        }
+        _syntax = syntax
+        _argumentListSyntax = arguments
+    }
+
+    public func argument(labeled label: String) -> Expression? {
+        (_argumentListSyntax.first { element in
+            return element.label?.text == label
+        }?.expression).map(Expression.init)
+    }
+
+    public var arguments: [Expression] {
+        Array(_argumentListSyntax).map { argument in
+            Expression(argument.expression)
+        }
+    }
+
+    public var name: Type {
+        Type(_syntax.attributeName)
+    }
+}
+
+public struct Struct {
+    public var _syntax: StructDeclSyntax
+
+    public init(_ syntax: StructDeclSyntax) {
+        _syntax = syntax
+    }
+
+    public init?(_ syntax: any DeclGroupSyntax) {
+        guard let syntax = syntax.as(StructDeclSyntax.self) else {
+            return nil
+        }
+        _syntax = syntax
+    }
+
+    // TODO: Add members property to all declgroupsyntax decls through protocol default impl
+    public var members: [Decl] {
+        _syntax.memberBlock.members.map(\.decl).map(Decl.init)
+    }
+
+    public var inheritedTypes: [Type] {
+        _syntax.inheritanceClause?.inheritedTypeCollection.map(\.typeName).map(Type.init) ?? []
+    }
+}
+
+public struct Decl {
+    public var _syntax: DeclSyntax
+
+    public init(_ syntax: DeclSyntax) {
+        _syntax = syntax
+    }
+
+    public init(_ syntax: any DeclSyntaxProtocol) {
+        _syntax = DeclSyntax(syntax)
+    }
+
+    // TODO: Add conversions for all possible member types
+    public var asEnum: Enum? {
+        _syntax.as(EnumDeclSyntax.self).map(Enum.init)
+    }
+
+    public var asStruct: Struct? {
+        _syntax.as(StructDeclSyntax.self).map(Struct.init)
+    }
+}
+
 extension FunctionDeclSyntax {
     /// Gets the signature's effect specifiers, or returns a default effect specifiers
     /// syntax (without any specifiers).
@@ -552,6 +654,12 @@ extension CodeBlockSyntax {
     }
 }
 
+extension DeclGroupSyntax {
+    public var isPublic: Bool {
+        modifiers?.contains { $0.name.tokenKind == .keyword(.public) } == true
+    }
+}
+
 // TODO: Figure out a destructuring implementation that uses variadic generics (tricky without same type requirements)
 public func destructure<Element>(_ elements: some Sequence<Element>) -> ()? {
     let array = Array(elements)
@@ -561,7 +669,9 @@ public func destructure<Element>(_ elements: some Sequence<Element>) -> ()? {
     return ()
 }
 
-public func destructure<Element>(_ elements: some Sequence<Element>) -> (Element)? {
+/// Named differently to allow type inference to still work correctly (single element tuples
+/// are weird in Swift).
+public func destructureSingle<Element>(_ elements: some Sequence<Element>) -> (Element)? {
     let array = Array(elements)
     guard array.count == 1 else {
         return nil
@@ -607,8 +717,10 @@ public func destructure(_ type: NominalType) -> (String, ())? {
     }
 }
 
-public func destructure(_ type: NominalType) -> (String, (Type))? {
-    destructure(type.genericArguments ?? []).map { arguments in
+/// Named differently to allow type inference to still work correctly (single element tuples
+/// are weird in Swift).
+public func destructureSingle(_ type: NominalType) -> (String, (Type))? {
+    destructureSingle(type.genericArguments ?? []).map { arguments in
         (type.name, arguments)
     }
 }
@@ -643,8 +755,10 @@ public func destructure(_ type: FunctionType) -> ((), Type)? {
     }
 }
 
-public func destructure(_ type: FunctionType) -> ((Type), Type)? {
-    destructure(type.parameters).map { parameters in
+/// Named differently to allow type inference to still work correctly (single element tuples
+/// are weird in Swift).
+public func destructureSingle(_ type: FunctionType) -> ((Type), Type)? {
+    destructureSingle(type.parameters).map { parameters in
         (parameters, type.returnType)
     }
 }
@@ -687,13 +801,15 @@ public func destructure(_ type: Type) -> DestructuredType<()>? {
     }
 }
 
-public func destructure(_ type: Type) -> DestructuredType<(Type)>? {
+/// Named differently to allow type inference to still work correctly (single element tuples
+/// are weird in Swift).
+public func destructureSingle(_ type: Type) -> DestructuredType<(Type)>? {
     if let type = type.asNominalType {
-        return destructure(type).map { destructured in
+        return destructureSingle(type).map { destructured in
             .nominal(name: destructured.0, genericArguments: destructured.1)
         }
     } else if let type = type.asFunctionType {
-        return destructure(type).map { destructured in
+        return destructureSingle(type).map { destructured in
             .function(parameterTypes: destructured.0, returnType: destructured.1)
         }
     } else {
diff --git a/Sources/MacroToolkitExample/MacroToolkitExample.swift b/Sources/MacroToolkitExample/MacroToolkitExample.swift
@@ -13,3 +13,8 @@ public macro AddCompletionHandler() =
 @freestanding(expression)
 public macro addBlocker<T>(_ value: T) -> T =
     #externalMacro(module: "MacroToolkitExamplePlugin", type: "AddBlocker")
+
+@attached(member, names: arbitrary)
+@attached(conformance)
+public macro MyOptionSet<RawType>() =
+    #externalMacro(module: "MacroToolkitExamplePlugin", type: "OptionSetMacro")
diff --git a/Sources/MacroToolkitExamplePlugin/AddBlocker.swift b/Sources/MacroToolkitExamplePlugin/AddBlocker.swift
@@ -47,7 +47,7 @@ public struct AddBlocker: ExpressionMacro {
         of node: some FreestandingMacroExpansionSyntax,
         in context: some MacroExpansionContext
     ) throws -> ExprSyntax {
-        guard let (argument): (TupleExprElementSyntax) = destructure(node.argumentList) else {
+        guard let (argument) = destructureSingle(node.argumentList) else {
             throw MacroError("#addBlocker only expects one argument")
         }
 
diff --git a/Sources/MacroToolkitExamplePlugin/OptionSetMacro.swift b/Sources/MacroToolkitExamplePlugin/OptionSetMacro.swift
@@ -0,0 +1,165 @@
+import MacroToolkit
+import SwiftDiagnostics
+import SwiftSyntax
+import SwiftSyntaxBuilder
+import SwiftSyntaxMacros
+
+// Modified from: https://github.com/DougGregor/swift-macro-examples/blob/f61ac7cdca8dc3557e53f86e7e03df1353908d3e/MacroExamplesPlugin/OptionSetMacro.swift
+
+enum OptionSetMacroDiagnostic {
+    case requiresStruct
+    case requiresStringLiteral(String)
+    case requiresOptionsEnum(String)
+    case requiresOptionsEnumRawType
+    case invalidAttribute
+}
+
+extension OptionSetMacroDiagnostic: DiagnosticMessage {
+    func diagnose(at node: some SyntaxProtocol) -> Diagnostic {
+        Diagnostic(node: Syntax(node), message: self)
+    }
+
+    var message: String {
+        switch self {
+            case .requiresStruct:
+                return "'OptionSet' macro can only be applied to a struct"
+
+            case .requiresStringLiteral(let name):
+                return "'OptionSet' macro argument \(name) must be a string literal"
+
+            case .requiresOptionsEnum(let name):
+                return "'OptionSet' macro requires nested options enum '\(name)'"
+
+            case .requiresOptionsEnumRawType:
+                return "'OptionSet' macro requires a raw type"
+
+            case .invalidAttribute:
+                return "'OptionSet' macro attribute is invalid"
+        }
+    }
+
+    var severity: DiagnosticSeverity { .error }
+
+    var diagnosticID: MessageID {
+        MessageID(domain: "Swift", id: "OptionSet.\(self)")
+    }
+}
+
+public struct OptionSetMacro {
+    /// Decodes the arguments to the macro expansion.
+    /// - Returns: the important arguments used by the various roles of this
+    ///   macro inhabits, or nil if an error occurred.
+    static func decodeExpansion(
+        of attribute: AttributeSyntax,
+        attachedTo decl: some DeclGroupSyntax,
+        in context: some MacroExpansionContext
+    ) -> (Struct, Enum, Type)? {
+        guard let attribute = MacroAttribute(attribute) else {
+            context.diagnose(OptionSetMacroDiagnostic.invalidAttribute.diagnose(at: attribute))
+            return nil
+        }
+
+        // Determine the name of the options enum
+        let optionsEnumName: String
+        if let argument = attribute.argument(labeled: "optionsName") {
+            guard let stringLiteral = argument.asSimpleStringLiteral else {
+                context.diagnose(
+                    OptionSetMacroDiagnostic
+                        .requiresStringLiteral("optionsName")
+                        .diagnose(at: argument._syntax)
+                )
+                return nil
+            }
+
+            optionsEnumName = stringLiteral
+        } else {
+            optionsEnumName = "Options"
+        }
+
+        // Only apply to structs
+        guard let structDecl = Struct(decl) else {
+            context.diagnose(OptionSetMacroDiagnostic.requiresStruct.diagnose(at: decl))
+            return nil
+        }
+
+        // Find the option enum within the struct
+        guard
+            let optionsEnum = structDecl.members.compactMap(\.asEnum)
+                .first(where: { $0.identifier == optionsEnumName })
+        else {
+            context.diagnose(
+                OptionSetMacroDiagnostic.requiresOptionsEnum(optionsEnumName).diagnose(at: decl)
+            )
+            return nil
+        }
+
+        // Retrieve the raw type from the attribute
+        // TODO: Improve destructuring on single-element arrays
+        guard case let .nominal(_, (rawType)) = destructureSingle(attribute.name) else {
+            context.diagnose(
+                OptionSetMacroDiagnostic.requiresOptionsEnumRawType.diagnose(at: attribute._syntax))
+            return nil
+        }
+
+        return (structDecl, optionsEnum, rawType)
+    }
+}
+
+extension OptionSetMacro: ConformanceMacro {
+    public static func expansion(
+        of attribute: AttributeSyntax,
+        providingConformancesOf decl: some DeclGroupSyntax,
+        in context: some MacroExpansionContext
+    ) throws -> [(TypeSyntax, GenericWhereClauseSyntax?)] {
+        // Decode the expansion arguments. If there is an explicit conformance to
+        // OptionSet already, don't add one.
+        guard
+            let (structDecl, _, _) = decodeExpansion(of: attribute, attachedTo: decl, in: context),
+            !structDecl.inheritedTypes.contains(where: { type in
+                type.description == "OptionSet"
+            })
+        else {
+            return []
+        }
+
+        return [("OptionSet", nil)]
+    }
+}
+
+extension OptionSetMacro: MemberMacro {
+    public static func expansion(
+        of attribute: AttributeSyntax,
+        providingMembersOf decl: some DeclGroupSyntax,
+        in context: some MacroExpansionContext
+    ) throws -> [DeclSyntax] {
+        // Decode the expansion arguments.
+        guard
+            let (_, optionsEnum, rawType) = decodeExpansion(
+                of: attribute,
+                attachedTo: decl,
+                in: context
+            )
+        else {
+            return []
+        }
+
+        let cases = optionsEnum.cases
+
+        // TODO: This seems wrong, surely other modifiers would also make sense to passthrough?
+        let access = decl.isPublic ? "public " : ""
+
+        let staticVars = cases.map { (case_) -> DeclSyntax in
+            """
+            \(raw: access)static let \(raw: case_.identifier): Self =
+                Self(rawValue: 1 << \(raw: optionsEnum.identifier).\(raw: case_.identifier).rawValue)
+            """
+        }
+
+        return [
+            "\(raw: access)typealias RawValue = \(rawType._syntax)",
+            "\(raw: access)var rawValue: RawValue",
+            "\(raw: access)init() { self.rawValue = 0 }",
+            "\(raw: access)init(rawValue: RawValue) { self.rawValue = rawValue }",
+        ] + staticVars
+    }
+}
diff --git a/Tests/MacroToolkitTests/MacroToolkitTests.swift b/Tests/MacroToolkitTests/MacroToolkitTests.swift

Original file line number	Diff line number	Diff line change
`@@ -47,7 +47,7 @@ public struct AddBlocker: ExpressionMacro {`
`47`	`47`	`of node: some FreestandingMacroExpansionSyntax,`
`48`	`48`	`in context: some MacroExpansionContext`
`49`	`49`	`) throws -> ExprSyntax {`
`50`		`- guard let (argument): (TupleExprElementSyntax) = destructure(node.argumentList) else {`
	`50`	`+ guard let (argument) = destructureSingle(node.argumentList) else {`
`51`	`51`	`throw MacroError("#addBlocker only expects one argument")`
`52`	`52`	`}`
`53`	`53`