[Macros] Diagnose attached and freestanding declaration macros that produce

something other than a declaration.

The validation code already diagnosed all sorts of invalid declarations, but
it was ignoring AST nodes that aren't declarations at all.

Author: hborla

include/swift/AST/DiagnosticsSema.def

@@ -7284,7 +7284,7 @@ ERROR(experimental_no_metadata_feature_can_only_be_used_when_enabled,
 ERROR(expected_macro_expansion_expr,PointsToFirstBadToken,
       "expected macro expansion to produce an expression", ())
 ERROR(expected_macro_expansion_decls,PointsToFirstBadToken,
-      "expected macro expansion to produce declarations", ())
+      "expected macro expansion to produce a declaration", ())
 ERROR(macro_undefined,PointsToFirstBadToken,
       "no macro named %0", (Identifier))
 ERROR(external_macro_not_found,none,

lib/AST/Decl.cpp

@@ -11281,8 +11281,24 @@ void MacroExpansionDecl::forEachExpandedNode(
     return;
   auto startLoc = sourceMgr.getLocForBufferStart(*bufferID);
   auto *sourceFile = moduleDecl->getSourceFileContainingLocation(startLoc);
-  for (auto node : sourceFile->getTopLevelItems())
+
+  auto *macro = dyn_cast<MacroDecl>(getMacroRef().getDecl());
+  auto roles = macro->getMacroRoles();
+
+  for (auto node : sourceFile->getTopLevelItems()) {
+    // The assumption here is that macros can only have a single
+    // freestanding macro role. Expression macros can only produce
+    // expressions, declaration macros can only produce declarations,
+    // and code item macros can produce expressions, declarations, and
+    // statements.
+    if (roles.contains(MacroRole::Expression) && !node.is<Expr *>())
+      continue;
+
+    if (roles.contains(MacroRole::Declaration) && !node.is<Decl *>())
+      continue;
+
     callback(node);
+  }
 }
 
 /// Adjust the declaration context to find a point in the context hierarchy

lib/Sema/TypeCheckMacros.cpp

@@ -688,7 +688,18 @@ static void validateMacroExpansion(SourceFile *expansionBuffer,
             introducedNameSet.count(MacroDecl::getArbitraryName()));
   };
 
-  for (auto *decl : expansionBuffer->getTopLevelDecls()) {
+  for (auto item : expansionBuffer->getTopLevelItems()) {
+    auto *decl = item.dyn_cast<Decl *>();
+    if (!decl) {
+      if (role != MacroRole::CodeItem) {
+        auto &ctx = expansionBuffer->getASTContext();
+        ctx.Diags.diagnose(item.getStartLoc(),
+                           diag::expected_macro_expansion_decls);
+      }
+
+      continue;
+    }
+
     // Certain macro roles can generate special declarations.
     if ((isa<AccessorDecl>(decl) && role == MacroRole::Accessor) ||
         (isa<ExtensionDecl>(decl) && role == MacroRole::Conformance)) {
@@ -1156,11 +1167,16 @@ swift::expandFreestandingMacro(MacroExpansionDecl *med) {
     return llvm::None;
 
   MacroDecl *macro = cast<MacroDecl>(med->getMacroRef().getDecl());
+  auto macroRoles = macro->getMacroRoles();
+  assert(macroRoles.contains(MacroRole::Declaration) ||
+         macroRoles.contains(MacroRole::CodeItem));
   DeclContext *dc = med->getDeclContext();
 
   validateMacroExpansion(macroSourceFile, macro,
                          /*attachedTo*/nullptr,
-                         MacroRole::Declaration);
+                         macroRoles.contains(MacroRole::Declaration) ?
+                             MacroRole::Declaration :
+                             MacroRole::CodeItem);
 
   PrettyStackTraceDecl debugStack(
       "type checking expanded declaration macro", med);

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -2054,3 +2054,20 @@ extension RequiredDefaultInitMacro: MemberMacro {
     return [ decl ]
   }
 }
+
+public struct FakeCodeItemMacro: DeclarationMacro, PeerMacro {
+  public static func expansion(
+    of node: some FreestandingMacroExpansionSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    return ["guard true else { return }"]
+  }
+
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingPeersOf declaration: some DeclSyntaxProtocol,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    return ["if true { return }"]
+  }
+}

test/Macros/macro_expand.swift

@@ -13,7 +13,7 @@
 // RUN: %target-typecheck-verify-swift -swift-version 5 -load-plugin-library %t/%target-library-name(MacroDefinition) -module-name MacroUser -DTEST_DIAGNOSTICS -I %t -DIMPORT_MACRO_LIBRARY
 
 // RUN: not %target-swift-frontend -swift-version 5 -typecheck -load-plugin-library %t/%target-library-name(MacroDefinition) -module-name MacroUser -DTEST_DIAGNOSTICS -serialize-diagnostics-path %t/macro_expand.dia %s -emit-macro-expansion-files no-diagnostics -Rmacro-loading > %t/macro-printing.txt
-// RUN: c-index-test -read-diagnostics %t/macro_expand.dia 2>&1 | %FileCheck -check-prefix CHECK-DIAGS %s
+// RUN: c-index-test -read-diagnostics %t/macro_expand.dia 2>&1 | %FileCheck -check-prefix CHECK-DIAGS -dump-input=always %s
 
 // RUN: %FileCheck %s  --check-prefix CHECK-MACRO-PRINTED < %t/macro-printing.txt
 
@@ -127,6 +127,24 @@ struct Bad {}
 // CHECK-DIAGS: END CONTENTS OF FILE
 #endif
 
+@freestanding(declaration)
+macro accidentalCodeItem() = #externalMacro(module: "MacroDefinition", type: "FakeCodeItemMacro")
+
+@attached(peer)
+macro AccidentalCodeItem() = #externalMacro(module: "MacroDefinition", type: "FakeCodeItemMacro")
+
+#if TEST_DIAGNOSTICS
+func invalidDeclarationMacro() {
+  #accidentalCodeItem
+  // expected-note@-1 {{in expansion of macro 'accidentalCodeItem' here}}
+  // CHECK-DIAGS: @__swiftmacro_9MacroUser018invalidDeclarationA0yyF18accidentalCodeItemfMf0_.swift:1:1: error: expected macro expansion to produce a declaration
+
+  @AccidentalCodeItem struct S {}
+  // expected-note@-1 {{in expansion of macro 'AccidentalCodeItem' on struct 'S' here}}
+  // CHECK-DIAGS: @__swiftmacro_9MacroUser018invalidDeclarationA0yyF1SL_18AccidentalCodeItemfMp_.swift:1:1: error: expected macro expansion to produce a declaration
+}
+#endif
+
 @freestanding(expression) macro customFileID() -> String = #externalMacro(module: "MacroDefinition", type: "FileIDMacro")
 @freestanding(expression) macro fileID<T: ExpressibleByStringLiteral>() -> T = #externalMacro(module: "MacroDefinition", type: "FileIDMacro")
 @freestanding(expression) macro recurse(_: Bool) = #externalMacro(module: "MacroDefinition", type: "RecursiveMacro")