Merge pull request #64097 from hborla/macro-expansion-validation

[Macros] Diagnose macro expansions containing invalid declarations.

Author: hborla

include/swift/AST/Decl.h

@@ -8463,6 +8463,11 @@ class MacroDecl : public GenericContext, public ValueDecl {
   void getIntroducedNames(MacroRole role, ValueDecl *attachedTo,
                           SmallVectorImpl<DeclName> &names) const;
 
+  /// Returns a DeclName that represents arbitrary names.
+  static DeclName getArbitraryName() {
+    return DeclName();
+  }
+
   /// Retrieve the definition of this macro.
   MacroDefinition getDefinition() const;
 

include/swift/AST/DiagnosticsSema.def

@@ -6909,6 +6909,18 @@ ERROR(external_macro_arg_not_type_name,none,
       "the external macro's %select{module|type}0", (unsigned))
 ERROR(attached_declaration_macro_not_supported,none,
       "attached declaration macros are not yet supported", ())
+ERROR(invalid_decl_in_macro_expansion,none,
+      "macro expansion cannot introduce %0",
+      (DescriptiveDeclKind))
+ERROR(invalid_main_type_in_macro_expansion,none,
+      "macro expansion cannot introduce '@main' type",
+      ())
+ERROR(literal_type_in_macro_expansion,none,
+      "macro expansion cannot introduce default literal type %0",
+      (Identifier))
+ERROR(invalid_macro_introduced_name,none,
+      "declaration name %0 is not covered by macro %1",
+      (DeclName, DeclName))
 
 //------------------------------------------------------------------------------
 // MARK: Move Only Errors

lib/AST/Decl.cpp

@@ -10122,7 +10122,7 @@ void MacroDecl::getIntroducedNames(MacroRole role, ValueDecl *attachedTo,
     }
 
     case MacroIntroducedDeclNameKind::Arbitrary:
-      // FIXME: Indicate that the macro covers arbitrary names.
+      names.push_back(MacroDecl::getArbitraryName());
       break;
     }
   }

lib/Sema/TypeCheckMacros.cpp

@@ -476,6 +476,104 @@ ExpandPeerMacroRequest::evaluate(Evaluator &evaluator, Decl *decl) const {
   return decl->getASTContext().AllocateCopy(bufferIDs);
 }
 
+static Identifier makeIdentifier(ASTContext &ctx, StringRef name) {
+  return ctx.getIdentifier(name);
+}
+
+static Identifier makeIdentifier(ASTContext &ctx, std::nullptr_t) {
+  return Identifier();
+}
+
+/// Diagnose macro expansions that produce any of the following declarations:
+///   - Import declarations
+///   - Operator and precedence group declarations
+///   - Macro declarations
+///   - Extensions
+///   - Types with `@main` attributes
+///   - Top-level default literal type overrides
+///   - Value decls with names not covered by the macro declaration.
+static void validateMacroExpansion(SourceFile *expansionBuffer,
+                                   MacroDecl *macro,
+                                   ValueDecl *attachedTo,
+                                   MacroRole role) {
+  // Gather macro-introduced names
+  llvm::SmallVector<DeclName, 2> introducedNames;
+  macro->getIntroducedNames(role, attachedTo, introducedNames);
+
+  llvm::SmallDenseSet<DeclName, 2> coversName(introducedNames.begin(),
+                                              introducedNames.end());
+
+  for (auto *decl : expansionBuffer->getTopLevelDecls()) {
+    auto &ctx = decl->getASTContext();
+
+    // Certain macro roles can generate special declarations.
+    if ((isa<AccessorDecl>(decl) && role == MacroRole::Accessor) ||
+        (isa<ExtensionDecl>(decl) && role == MacroRole::Conformance)) {
+      continue;
+    }
+
+    // Diagnose invalid declaration kinds.
+    if (isa<ImportDecl>(decl) ||
+        isa<OperatorDecl>(decl) ||
+        isa<PrecedenceGroupDecl>(decl) ||
+        isa<MacroDecl>(decl) ||
+        isa<ExtensionDecl>(decl)) {
+      decl->diagnose(diag::invalid_decl_in_macro_expansion,
+                     decl->getDescriptiveKind());
+      decl->setInvalid();
+
+      if (auto *extension = dyn_cast<ExtensionDecl>(decl)) {
+        extension->setExtendedNominal(nullptr);
+      }
+
+      continue;
+    }
+
+    // Diagnose `@main` types.
+    if (auto *mainAttr = decl->getAttrs().getAttribute<MainTypeAttr>()) {
+      ctx.Diags.diagnose(mainAttr->getLocation(),
+                         diag::invalid_main_type_in_macro_expansion);
+      mainAttr->setInvalid();
+    }
+
+    // Diagnose default literal type overrides.
+    if (auto *typeAlias = dyn_cast<TypeAliasDecl>(decl)) {
+      auto name = typeAlias->getBaseIdentifier();
+#define EXPRESSIBLE_BY_LITERAL_PROTOCOL_WITH_NAME(_, __, typeName,     \
+                                                  supportsOverride)    \
+      if (supportsOverride && name == makeIdentifier(ctx, typeName)) { \
+        typeAlias->diagnose(diag::literal_type_in_macro_expansion,     \
+                            makeIdentifier(ctx, typeName));            \
+        typeAlias->setInvalid();                                       \
+        continue;                                                      \
+      }
+#include "swift/AST/KnownProtocols.def"
+#undef EXPRESSIBLE_BY_LITERAL_PROTOCOL_WITH_NAME
+    }
+
+    // Diagnose value decls with names not covered by the macro
+    if (auto *value = dyn_cast<ValueDecl>(decl)) {
+      auto baseName = value->getBaseName();
+      if (baseName.isSpecial()) {
+        baseName = ctx.getIdentifier(baseName.userFacingName());
+      }
+
+      // $-prefixed names are unique names. These are always allowed.
+      if (baseName.getIdentifier().hasDollarPrefix()) {
+        continue;
+      }
+
+      if (coversName.count(baseName) ||
+          coversName.count(MacroDecl::getArbitraryName())) {
+        continue;
+      }
+
+      value->diagnose(diag::invalid_macro_introduced_name,
+                      baseName, macro->getBaseName());
+    }
+  }
+}
+
 /// Determine whether the given source file is from an expansion of the given
 /// macro.
 static bool isFromExpansionOfMacro(SourceFile *sourceFile, MacroDecl *macro,
@@ -1101,6 +1199,8 @@ evaluateAttachedMacro(MacroDecl *macro, Decl *attachedTo, CustomAttr *attr,
       /*parsingOpts=*/{}, /*isPrimary=*/false);
   macroSourceFile->setImports(declSourceFile->getImports());
 
+  validateMacroExpansion(macroSourceFile, macro,
+                         dyn_cast<ValueDecl>(attachedTo), role);
   return macroSourceFile;
 }
 

stdlib/public/core/Macros.swift

@@ -86,7 +86,7 @@ public macro dsohandle() -> UnsafeRawPointer = Builtin.DSOHandleMacro
 ///         case standard
 ///       }
 ///     }
-@attached(member)
+@attached(member, names: named(RawValue), named(rawValue), named(`init`), arbitrary)
 @attached(conformance)
 public macro OptionSet<RawType>() =
   #externalMacro(module: "SwiftMacros", type: "OptionSetMacro")

test/Macros/Inputs/macro_library.swift

@@ -23,7 +23,10 @@ public struct ObservationRegistrar<Subject: Observable> {
   public func register<Value>(observable: Subject, didSet: KeyPath<Subject, Value>) {}
 }
 
-@attached(member)
+@attached(
+  member,
+  names: named(_registrar), named(addObserver), named(removeObserver), named(withTransaction), named(Storage), named(_storage)
+)
 @attached(memberAttribute)
 public macro Observable() = #externalMacro(module: "MacroDefinition", type: "ObservableMacro")
 

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -494,6 +494,25 @@ public struct AddMembers: MemberMacro {
   }
 }
 
+public struct AddArbitraryMembers: MemberMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingMembersOf decl: some DeclGroupSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    guard let identified = decl.asProtocol(IdentifiedDeclSyntax.self) else {
+      return []
+    }
+
+    let parentName = identified.identifier.trimmed
+    return [
+      "struct \(parentName)1 {}",
+      "struct \(parentName)2 {}",
+      "struct \(parentName)3 {}",
+    ]
+  }
+}
+
 /// Implementation of the `wrapStoredProperties` macro, which can be
 /// used to apply an attribute to all of the stored properties of a type.
 ///
@@ -745,6 +764,38 @@ public struct AddCompletionHandler: PeerMacro {
   }
 }
 
+public struct InvalidMacro: PeerMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingPeersOf declaration: some DeclSyntaxProtocol,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    return [
+      "import Swift",
+      "precedencegroup MyPrecedence {}",
+      "@attached(member) macro myMacro()",
+      "extension Int {}",
+      """
+      @main
+      struct MyMain {
+        static func main() {}
+      }
+      """,
+      "typealias Array = Void",
+      "typealias Dictionary = Void",
+      "typealias BooleanLiteralType = Void",
+      "typealias ExtendedGraphemeClusterType = Void",
+      "typealias FloatLiteralType = Void",
+      "typealias IntegerLiteralType = Void",
+      "typealias StringLiteralType = Void",
+      "typealias UnicodeScalarType = Void",
+      "typealias _ColorLiteralType = Void",
+      "typealias _ImageLiteralType = Void",
+      "typealias _FileReferenceLiteralType = Void",
+    ]
+  }
+}
+
 public struct WrapInType: PeerMacro {
   public static func expansion(
     of node: AttributeSyntax,

test/Macros/composed_macro.swift

@@ -10,7 +10,7 @@
 // REQUIRES: OS=macosx
 
 @attached(memberAttribute)
-@attached(member)
+@attached(member, names: named(_storage))
 macro myTypeWrapper() = #externalMacro(module: "MacroDefinition", type: "TypeWrapperMacro")
 @attached(accessor) macro accessViaStorage() = #externalMacro(module: "MacroDefinition", type: "AccessViaStorageMacro")
 

test/Macros/macro_expand.swift

@@ -26,6 +26,55 @@
 // FIXME: Swift parser is not enabled on Linux CI yet.
 // REQUIRES: OS=macosx
 
+#if TEST_DIAGNOSTICS
+@attached(peer)
+macro Invalid() = #externalMacro(module: "MacroDefinition", type: "InvalidMacro")
+
+@Invalid
+struct Bad {}
+// expected-note@-1 18 {{in expansion of macro 'Invalid' here}}
+
+// CHECK-DIAGS: error: macro expansion cannot introduce import
+// CHECK-DIAGS: error: macro expansion cannot introduce precedence group
+// CHECK-DIAGS: error: macro expansion cannot introduce macro
+// CHECK-DIAGS: error: macro expansion cannot introduce extension
+// CHECK-DIAGS: error: macro expansion cannot introduce '@main' type
+// CHECK-DIAGS: error: declaration name 'MyMain' is not covered by macro 'Invalid'
+// CHECK-DIAGS: error: declaration name 'Array' is not covered by macro 'Invalid'
+// CHECK-DIAGS: error: declaration name 'Dictionary' is not covered by macro 'Invalid'
+// CHECK-DIAGS: error: macro expansion cannot introduce default literal type 'BooleanLiteralType'
+// CHECK-DIAGS: error: macro expansion cannot introduce default literal type 'ExtendedGraphemeClusterType'
+// CHECK-DIAGS: error: macro expansion cannot introduce default literal type 'FloatLiteralType'
+// CHECK-DIAGS: error: macro expansion cannot introduce default literal type 'IntegerLiteralType'
+// CHECK-DIAGS: error: macro expansion cannot introduce default literal type 'StringLiteralType'
+// CHECK-DIAGS: error: macro expansion cannot introduce default literal type 'UnicodeScalarType'
+// CHECK-DIAGS: error: macro expansion cannot introduce default literal type '_ColorLiteralType'
+// CHECK-DIAGS: error: macro expansion cannot introduce default literal type '_ImageLiteralType'
+// CHECK-DIAGS: error: macro expansion cannot introduce default literal type '_FileReferenceLiteralType'
+
+// CHECK-DIAGS: CONTENTS OF FILE @__swiftmacro_9MacroUser3BadV7InvalidfMp_.swift
+// CHECK-DIAGS: import Swift
+// CHECK-DIAGS: precedencegroup MyPrecedence {}
+// CHECK-DIAGS: @attached(member) macro myMacro()
+// CHECK-DIAGS: extension Int {}
+// CHECK-DIAGS: @main
+// CHECK-DIAGS: struct MyMain {
+// CHECK-DIAGS:   static func main() {}
+// CHECK-DIAGS: }
+// CHECK-DIAGS: typealias Array = Void
+// CHECK-DIAGS: typealias Dictionary = Void
+// CHECK-DIAGS: typealias BooleanLiteralType = Void
+// CHECK-DIAGS: typealias ExtendedGraphemeClusterType = Void
+// CHECK-DIAGS: typealias FloatLiteralType = Void
+// CHECK-DIAGS: typealias IntegerLiteralType = Void
+// CHECK-DIAGS: typealias StringLiteralType = Void
+// CHECK-DIAGS: typealias UnicodeScalarType = Void
+// CHECK-DIAGS: typealias _ColorLiteralType = Void
+// CHECK-DIAGS: typealias _ImageLiteralType = Void
+// CHECK-DIAGS: typealias _FileReferenceLiteralType = Void
+// CHECK-DIAGS: END CONTENTS OF FILE
+#endif
+
 @freestanding(expression) macro customFileID() -> String = #externalMacro(module: "MacroDefinition", type: "FileIDMacro")
 @freestanding(expression) macro stringify<T>(_ value: T) -> (T, String) = #externalMacro(module: "MacroDefinition", type: "StringifyMacro")
 @freestanding(expression) macro fileID<T: ExpressibleByStringLiteral>() -> T = #externalMacro(module: "MacroDefinition", type: "FileIDMacro")

test/Macros/macro_expand_conformances.swift

@@ -40,7 +40,7 @@ requireEquatable(S2())
 requireHashable(S2())
 
 @attached(conformance)
-@attached(member)
+@attached(member, names: named(requirement))
 macro DelegatedConformance() = #externalMacro(module: "MacroDefinition", type: "DelegatedConformanceMacro")
 
 protocol P {

test/Macros/macro_expand_synthesized_members.swift

@@ -9,7 +9,11 @@
 // FIXME: Swift parser is not enabled on Linux CI yet.
 // REQUIRES: OS=macosx
 
-@attached(member) macro addMembers() = #externalMacro(module: "MacroDefinition", type: "AddMembers")
+@attached(
+  member,
+  names: named(Storage), named(storage), named(getStorage), named(method), named(`init`)
+)
+macro addMembers() = #externalMacro(module: "MacroDefinition", type: "AddMembers")
 
 @addMembers
 struct S {
@@ -25,6 +29,19 @@ let s = S()
 // CHECK: Storage
 s.useSynthesized()
 
+@attached(member, names: arbitrary)
+macro addArbitraryMembers() = #externalMacro(module: "MacroDefinition", type: "AddArbitraryMembers")
+
+@addArbitraryMembers
+struct MyType {}
+
+// CHECK: MyType1
+// CHECK: MyType2
+// CHECK: MyType3
+print(MyType.MyType1.self)
+print(MyType.MyType2.self)
+print(MyType.MyType3.self)
+
 @attached(
   member,
   names: named(RawValue), named(rawValue), named(`init`)