[Macros] Apply generic arguments from attached macro custom attributes when

hborla · hborla · commit 09e34f274d33 · 2023-02-06T18:45:26.000-08:00
resolving the macro reference.
diff --git a/lib/AST/TypeCheckRequests.cpp b/lib/AST/TypeCheckRequests.cpp
@@ -1705,8 +1705,16 @@ SourceRange UnresolvedMacroReference::getGenericArgsRange() const {
     return med->getGenericArgsRange();
   if (auto *mee = pointer.dyn_cast<MacroExpansionExpr *>())
     return mee->getGenericArgsRange();
-  if (auto *attr = pointer.dyn_cast<CustomAttr *>())
-    return SourceRange();
+
+  if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
+    auto *typeRepr = attr->getTypeRepr();
+    auto *genericTypeRepr = dyn_cast_or_null<GenericIdentTypeRepr>(typeRepr);
+    if (!genericTypeRepr)
+      return SourceRange();
+
+    return genericTypeRepr->getAngleBrackets();
+  }
+
   llvm_unreachable("Unhandled case");
 }
 
@@ -1715,8 +1723,16 @@ ArrayRef<TypeRepr *> UnresolvedMacroReference::getGenericArgs() const {
     return med->getGenericArgs();
   if (auto *mee = pointer.dyn_cast<MacroExpansionExpr *>())
     return mee->getGenericArgs();
-  if (auto *attr = pointer.dyn_cast<CustomAttr *>())
-    return {};
+
+  if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
+    auto *typeRepr = attr->getTypeRepr();
+    auto *genericTypeRepr = dyn_cast_or_null<GenericIdentTypeRepr>(typeRepr);
+    if (!genericTypeRepr)
+      return {};
+
+    return genericTypeRepr->getGenericArgs();
+  }
+
   llvm_unreachable("Unhandled case");
 }
 
diff --git a/test/Macros/Inputs/syntax_macro_definitions.swift b/test/Macros/Inputs/syntax_macro_definitions.swift
@@ -717,3 +717,48 @@ public struct ObservablePropertyMacro: AccessorMacro {
     return [getAccessor, setAccessor]
   }
 }
+
+extension DeclModifierSyntax {
+  fileprivate var isNeededAccessLevelModifier: Bool {
+    switch self.name.tokenKind {
+    case .keyword(.public): return true
+    default: return false
+    }
+  }
+}
+
+extension SyntaxStringInterpolation {
+  fileprivate mutating func appendInterpolation<Node: SyntaxProtocol>(_ node: Node?) {
+    if let node {
+      appendInterpolation(node)
+    }
+  }
+}
+
+public struct NewTypeMacro: MemberMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingMembersOf declaration: some DeclGroupSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    guard let type = node.attributeName.as(SimpleTypeIdentifierSyntax.self),
+          let genericArguments = type.genericArgumentClause?.arguments,
+          genericArguments.count == 1,
+          let rawType = genericArguments.first
+    else {
+      throw CustomError.message(#"@NewType requires the raw type as an argument, in the form "<RawType>"."#)
+    }
+
+    guard let declaration = declaration.as(StructDeclSyntax.self) else {
+      throw CustomError.message("@NewType can only be applied to a struct declarations.")
+    }
+
+    let access = declaration.modifiers?.first(where: \.isNeededAccessLevelModifier)
+
+    return [
+      "\(access)typealias RawValue = \(rawType)",
+      "\(access)var rawValue: RawValue",
+      "\(access)init(_ rawValue: RawValue) { self.rawValue = rawValue }",
+    ]
+  }
+}
diff --git a/test/Macros/macro_expand_synthesized_members.swift b/test/Macros/macro_expand_synthesized_members.swift
@@ -24,3 +24,18 @@ let s = S()
 // CHECK: synthesized method
 // CHECK: Storage
 s.useSynthesized()
+
+@attached(
+  member,
+  names: named(RawValue), named(rawValue), named(`init`)
+)
+public macro NewType<T>() = #externalMacro(module: "MacroDefinition", type: "NewTypeMacro")
+
+@NewType<String>
+public struct MyString {}
+
+// CHECK: String
+// CHECK: hello
+let myString = MyString("hello")
+print(MyString.RawValue.self)
+print(myString.rawValue)
