[Macros] Diagnose macro expansions containing invalid declarations, such as

new macro decls or value decls whose names are not covered by the macro.

Author: hborla

include/swift/AST/DiagnosticsSema.def

@@ -6909,6 +6909,15 @@ 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 %1",
+      (DescriptiveDeclKind))
+ERROR(invalid_main_type_in_macro_expansion,none,
+      "macro expansion cannot introduce '@main' type'",
+      ())
+ERROR(invalid_macro_introduced_name,none,
+      "declaration name %0 is not covered by macro %1",
+      (DeclName, DeclName))
 
 //------------------------------------------------------------------------------
 // MARK: Move Only Errors

lib/Sema/TypeCheckMacros.cpp

@@ -476,6 +476,76 @@ ExpandPeerMacroRequest::evaluate(Evaluator &evaluator, Decl *decl) const {
   return decl->getASTContext().AllocateCopy(bufferIDs);
 }
 
+/// 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());
+    }
+
+    // Diagnose `@main` types.
+    if (auto *typeDecl = dyn_cast<TypeDecl>(decl)) {
+      if (typeDecl->getAttrs().hasAttribute<MainTypeAttr>()) {
+        typeDecl->diagnose(diag::invalid_main_type_in_macro_expansion);
+      }
+    }
+
+    // Diagnose default literal type overrides.
+    if (auto *typeAlias = dyn_cast<TypeAliasDecl>(decl)) {
+      // TODO
+    }
+
+    // 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)) {
+        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,
@@ -1102,6 +1172,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;
 }
 

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/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_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 {