[Parse] Improve macro expansion parsing

* Unify macro expansion parsing logic between MacroExpansionExpr and
  MacroExpansionDecl
* Diagnose whitespace between '#' and the macro name
* Diagnose keyword as a macro name

Author: rintaro

include/swift/AST/DiagnosticsParse.def

@@ -2057,6 +2057,8 @@ ERROR(macro_expansion_expr_expected_macro_identifier,PointsToFirstBadToken,
       "expected a macro identifier for a pound literal expression", ())
 ERROR(macro_expansion_decl_expected_macro_identifier,PointsToFirstBadToken,
       "expected a macro identifier for a pound literal declaration", ())
+ERROR(extra_whitespace_macro_expansion_identifier,PointsToFirstBadToken,
+      "extraneous whitespace between '#' and macro name is not permitted", ())
 
 ERROR(macro_role_attr_expected_kind,PointsToFirstBadToken,
       "expected %select{a freestanding|an attached}0 macro role such as "

include/swift/Parse/Parser.h

@@ -1731,6 +1731,9 @@ class Parser {
     /// cases this doesn't actually make sense but we need to accept them for
     /// backwards compatibility.
     AllowLowercaseAndUppercaseSelf = 1 << 6,
+
+    /// Parse keyword as identifier but emit fix-it to escape it.
+    AllowKeywordWithBacktickEscape = AllowKeywords | 1 << 7
   };
   using DeclNameOptions = OptionSet<DeclNameFlag>;
 
@@ -1751,6 +1754,16 @@ class Parser {
   DeclNameRef parseDeclNameRef(DeclNameLoc &loc, const Diagnostic &diag,
                                DeclNameOptions flags);
 
+  /// Parse macro expansion.
+  ///
+  ///   macro-expansion:
+  ///     '#' identifier generic-arguments? expr-paren? trailing-closure?
+  ParserStatus parseFreestandingMacroExpansion(
+      SourceLoc &poundLoc, DeclNameLoc &macroNameLoc, DeclNameRef &macroNameRef,
+      SourceLoc &leftAngleLoc, SmallVectorImpl<TypeRepr *> &genericArgs,
+      SourceLoc &rightAngleLoc, ArgumentList *&argList, bool isExprBasic,
+      const Diagnostic &diag);
+
   ParserResult<Expr> parseExprIdentifier();
   Expr *parseExprEditorPlaceholder(Token PlaceholderTok,
                                    Identifier PlaceholderId);

lib/Parse/ParseDecl.cpp

@@ -4945,8 +4945,9 @@ bool Parser::isStartOfSwiftDecl(bool allowPoundIfAttributes,
   }
 
   // 'macro' name
-  if (Tok.isContextualKeyword("macro") && Tok2.is(tok::identifier))
-    return true;
+  if (Tok.isContextualKeyword("macro")) {
+    return Tok2.is(tok::identifier);
+  }
 
   if (Tok.isContextualKeyword("package")) {
     // If `case` is the next token after `return package` statement,
@@ -5031,16 +5032,20 @@ bool Parser::isStartOfSILDecl() {
 }
 
 bool Parser::isStartOfFreestandingMacroExpansion() {
-  // Check if "'#' <identifier>" without any whitespace between them.
+  // Check if "'#' <identifier>" where the identifier is on the sameline.
   if (!Tok.is(tok::pound))
     return false;
-  if (Tok.getRange().getEnd() != peekToken().getLoc())
-    return false;
-  if (!peekToken().isAny(tok::identifier, tok::code_complete) &&
-      // allow keywords right after '#' so we can diagnose it when parsing.
-      !peekToken().isKeyword())
+  const Token &Tok2 = peekToken();
+  if (Tok2.isAtStartOfLine())
     return false;
-  return true;
+
+  if (Tok2.isAny(tok::identifier, tok::code_complete))
+    return true;
+  if (Tok2.isKeyword()) {
+    // allow keywords right after '#' so we can diagnose it when parsing.
+    return Tok.getRange().getEnd() == Tok2.getLoc();
+  }
+  return false;
 }
 
 void Parser::consumeDecl(ParserPosition BeginParserPosition,
@@ -9814,47 +9819,18 @@ ParserResult<MacroDecl> Parser::parseDeclMacro(DeclAttributes &attributes) {
 ParserResult<MacroExpansionDecl>
 Parser::parseDeclMacroExpansion(ParseDeclOptions flags,
                                 DeclAttributes &attributes) {
-  SourceLoc poundLoc = consumeToken(tok::pound);
+  SourceLoc poundLoc;
   DeclNameLoc macroNameLoc;
-  DeclNameRef macroNameRef = parseDeclNameRef(
-      macroNameLoc, diag::macro_expansion_decl_expected_macro_identifier,
-      DeclNameOptions());
-  if (!macroNameRef)
-    return makeParserError();
-
-  ParserStatus status;
+  DeclNameRef macroNameRef;
   SourceLoc leftAngleLoc, rightAngleLoc;
-  SmallVector<TypeRepr *, 8> genericArgs;
-  if (canParseAsGenericArgumentList()) {
-    auto genericArgsStatus = parseGenericArguments(
-        genericArgs, leftAngleLoc, rightAngleLoc);
-    status |= genericArgsStatus;
-    if (genericArgsStatus.isErrorOrHasCompletion())
-      diagnose(leftAngleLoc, diag::while_parsing_as_left_angle_bracket);
-  }
-
+  SmallVector<TypeRepr *, 4> genericArgs;
   ArgumentList *argList = nullptr;
-  if (Tok.isFollowingLParen()) {
-    auto result = parseArgumentList(tok::l_paren, tok::r_paren,
-                                    /*isExprBasic*/ false,
-                                    /*allowTrailingClosure*/ true);
-    status |= result;
-    if (result.hasCodeCompletion())
-      return makeParserCodeCompletionResult<MacroExpansionDecl>();
-    argList = result.getPtrOrNull();
-  } else if (Tok.is(tok::l_brace)) {
-    SmallVector<Argument, 2> trailingClosures;
-    auto closuresStatus = parseTrailingClosures(/*isExprBasic*/ false,
-                                                macroNameLoc.getSourceRange(),
-                                                trailingClosures);
-    status |= closuresStatus;
-
-    if (!trailingClosures.empty()) {
-      argList = ArgumentList::createParsed(Context, SourceLoc(),
-                                           trailingClosures, SourceLoc(),
-                                           /*trailingClosureIdx*/ 0);
-    }
-  }
+  ParserStatus status = parseFreestandingMacroExpansion(
+      poundLoc, macroNameLoc, macroNameRef, leftAngleLoc, genericArgs,
+      rightAngleLoc, argList, /*isExprBasic=*/false,
+      diag::macro_expansion_decl_expected_macro_identifier);
+  if (!macroNameRef)
+    return status;
 
   auto *med = new (Context) MacroExpansionDecl(
       CurDeclContext, poundLoc, macroNameRef, macroNameLoc, leftAngleLoc,

lib/Parse/ParseExpr.cpp

@@ -2269,6 +2269,71 @@ DeclNameRef Parser::parseDeclNameRef(DeclNameLoc &loc,
   return DeclNameRef({ Context, baseName, argumentLabels });
 }
 
+ParserStatus Parser::parseFreestandingMacroExpansion(
+    SourceLoc &poundLoc, DeclNameLoc &macroNameLoc, DeclNameRef &macroNameRef,
+    SourceLoc &leftAngleLoc, SmallVectorImpl<TypeRepr *> &genericArgs,
+    SourceLoc &rightAngleLoc, ArgumentList *&argList, bool isExprBasic,
+    const Diagnostic &diag) {
+  SourceLoc poundEndLoc = Tok.getRange().getEnd();
+  poundLoc = consumeToken(tok::pound);
+
+  if (Tok.isAtStartOfLine()) {
+    // Diagnose and bail out if there's no macro name on the same line.
+    diagnose(poundEndLoc, diag);
+    return makeParserError();
+  }
+
+  bool hasWhitespaceBeforeName = poundEndLoc != Tok.getLoc();
+
+  // Diagnose and parse keyword right after `#`.
+  if (Tok.isKeyword() && !hasWhitespaceBeforeName) {
+    diagnose(Tok, diag::keyword_cant_be_identifier, Tok.getText());
+    diagnose(Tok, diag::backticks_to_escape)
+        .fixItReplace(Tok.getLoc(), "`" + Tok.getText().str() + "`");
+
+    // Let 'parseDeclNameRef' to parse this as an identifier.
+    Tok.setKind(tok::identifier);
+  }
+  macroNameRef = parseDeclNameRef(macroNameLoc, diag, DeclNameOptions());
+  if (!macroNameRef)
+    return makeParserError();
+
+  // Diagnose whitespace between '#' and the macro name, and continue parsing.
+  if (hasWhitespaceBeforeName) {
+    diagnose(poundEndLoc, diag::extra_whitespace_macro_expansion_identifier)
+        .fixItRemoveChars(poundEndLoc, macroNameLoc.getStartLoc());
+  }
+
+  ParserStatus status;
+  if (canParseAsGenericArgumentList()) {
+    auto genericArgsStatus =
+        parseGenericArguments(genericArgs, leftAngleLoc, rightAngleLoc);
+    status |= genericArgsStatus;
+    if (genericArgsStatus.isErrorOrHasCompletion())
+      diagnose(leftAngleLoc, diag::while_parsing_as_left_angle_bracket);
+  }
+
+  if (Tok.isFollowingLParen()) {
+    auto result = parseArgumentList(tok::l_paren, tok::r_paren, isExprBasic,
+                                    /*allowTrailingClosure*/ true);
+    status |= result;
+    argList = result.getPtrOrNull();
+  } else if (Tok.is(tok::l_brace)) {
+    SmallVector<Argument, 2> trailingClosures;
+    auto closuresStatus = parseTrailingClosures(
+        /*isExprBasic*/ false, macroNameLoc.getSourceRange(), trailingClosures);
+    status |= closuresStatus;
+
+    if (!trailingClosures.empty()) {
+      argList = ArgumentList::createParsed(Context, SourceLoc(),
+                                           trailingClosures, SourceLoc(),
+                                           /*trailingClosureIdx*/ 0);
+    }
+  }
+
+  return status;
+}
+
 ///   expr-identifier:
 ///     unqualified-decl-name generic-args?
 ParserResult<Expr> Parser::parseExprIdentifier() {
@@ -3376,45 +3441,18 @@ Parser::parseExprCallSuffix(ParserResult<Expr> fn, bool isExprBasic) {
 }
 
 ParserResult<Expr> Parser::parseExprMacroExpansion(bool isExprBasic) {
-  SourceLoc poundLoc = consumeToken(tok::pound);
+  SourceLoc poundLoc;
   DeclNameLoc macroNameLoc;
-  DeclNameRef macroNameRef = parseDeclNameRef(
-      macroNameLoc, diag::macro_expansion_expr_expected_macro_identifier,
-      DeclNameOptions());
-  if (!macroNameRef)
-    return makeParserError();
-
-  ParserStatus status;
+  DeclNameRef macroNameRef;
   SourceLoc leftAngleLoc, rightAngleLoc;
-  SmallVector<TypeRepr *, 8> genericArgs;
-  if (canParseAsGenericArgumentList()) {
-    auto genericArgsStatus = parseGenericArguments(
-        genericArgs, leftAngleLoc, rightAngleLoc);
-    status |= genericArgsStatus;
-    if (genericArgsStatus.isErrorOrHasCompletion())
-      diagnose(leftAngleLoc, diag::while_parsing_as_left_angle_bracket);
-  }
-
+  SmallVector<TypeRepr *, 4> genericArgs;
   ArgumentList *argList = nullptr;
-  if (Tok.isFollowingLParen()) {
-    auto result = parseArgumentList(tok::l_paren, tok::r_paren, isExprBasic,
-                                /*allowTrailingClosure*/ true);
-    argList = result.getPtrOrNull();
-    status |= result;
-  } else if (Tok.is(tok::l_brace) &&
-             isValidTrailingClosure(isExprBasic, *this)) {
-    SmallVector<Argument, 2> trailingClosures;
-    auto closureStatus = parseTrailingClosures(isExprBasic,
-                                               macroNameLoc.getSourceRange(),
-                                               trailingClosures);
-    status |= closureStatus;
-
-    if (!trailingClosures.empty()) {
-      argList = ArgumentList::createParsed(Context, SourceLoc(),
-                                           trailingClosures, SourceLoc(),
-                                           /*trailingClosureIdx*/ 0);
-    }
-  }
+  ParserStatus status = parseFreestandingMacroExpansion(
+      poundLoc, macroNameLoc, macroNameRef, leftAngleLoc, genericArgs,
+      rightAngleLoc, argList, isExprBasic,
+      diag::macro_expansion_expr_expected_macro_identifier);
+  if (!macroNameRef)
+    return status;
 
   return makeParserResult(
       status,

test/Macros/macro_keywordname.swift

@@ -0,0 +1,61 @@
+// REQUIRES: swift_swift_parser
+
+// RUN: %empty-directory(%t)
+// RUN: mkdir -p %t/src
+// RUN: mkdir -p %t/plugins
+// RUN: mkdir -p %t/lib
+
+// RUN: split-file %s %t/src
+
+//#-- Prepare the macro dylib plugin.
+// RUN: %host-build-swift \
+// RUN:   -swift-version 5 \
+// RUN:   -emit-library -o %t/plugins/%target-library-name(MacroDefinition) \
+// RUN:   -module-name MacroDefinition \
+// RUN:   %t/src/MacroDefinition.swift \
+// RUN:   -g -no-toolchain-stdlib-rpath
+
+//#-- Prepare the macro library.
+// RUN: %target-swift-frontend \
+// RUN:   -swift-version 5 \
+// RUN:   -emit-module -o %t/lib/MacroLib.swiftmodule \
+// RUN:   -module-name MacroLib \
+// RUN:   -plugin-path %t/plugins \
+// RUN:   %t/src/MacroLib.swift
+
+// RUN: %target-swift-frontend \
+// RUN:   -typecheck -verify \
+// RUN:   -I %t/lib \
+// RUN:   -plugin-path %t/plugins \
+// RUN:   %t/src/test.swift
+
+//--- MacroDefinition.swift
+import SwiftSyntax
+import SwiftSyntaxMacros
+
+public struct OneMacro: ExpressionMacro {
+  public static func expansion(
+    of node: some FreestandingMacroExpansionSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> ExprSyntax {
+    return "1"
+  }
+}
+
+
+//--- MacroLib.swift
+@freestanding(expression) public macro `public`() -> Int = #externalMacro(module: "MacroDefinition", type: "OneMacro")
+@freestanding(expression) public macro `escaped`() -> Int = #externalMacro(module: "MacroDefinition", type: "OneMacro")
+@freestanding(expression) public macro normal() -> Int = #externalMacro(module: "MacroDefinition", type: "OneMacro")
+
+//--- test.swift
+import MacroLib
+@freestanding(expression) public macro `class`() -> Int = #externalMacro(module: "MacroDefinition", type: "OneMacro")
+
+func test() {
+  let _: Int = #public() // expected-error {{keyword 'public' cannot be used as an identifier here}} expected-note {{if this name is unavoidable, use backticks to escape it}}
+  let _: Int = #`public`()
+  let _: Int = #escaped()
+  let _: Int = #`class`()
+  let _: Int = #normal()
+}

test/Parse/macro_decl.swift

@@ -46,3 +46,11 @@ func test() {
   @CustomAttr
   isolated #someFunc
 }
+
+public # someFunc // expected-error {{extraneous whitespace between '#' and macro name is not permitted}} {{9-10=}}
+
+struct S {
+  # someFunc // expected-error {{extraneous whitespace between '#' and macro name is not permitted}} {{4-5=}}
+  #class // expected-error {{keyword 'class' cannot be used as an identifier here}} expected-note {{if this name is unavoidable, use backticks to escape it}} {{4-9=`class`}} 
+  # struct Inner {} // expected-error {{expected a macro identifier for a pound literal declaration}} expected-error {{consecutive declarations on a line}}
+}

test/Parse/macro_expr.swift

@@ -35,3 +35,14 @@ _ = #another {
 
 // expected-error @+1 {{expected a macro identifier for a pound literal expression}}
 _ = #()
+
+do {
+  _ = # // expected-error {{expected a macro identifier for a pound literal expression}}
+  name()
+}
+do {
+  _ = # macro() // expected-error {{extraneous whitespace between '#' and macro name is not permitted}} {{8-9=}}
+}
+do {
+  _ = #public() // expected-error {{keyword 'public' cannot be used as an identifier here}} expected-note {{if this name is unavoidable, use backticks to escape it}} {{8-14=`public`}}
+}