Add diagnostic for invalid whitespace in @ Attribute

Author: TiagoMaiaL

Sources/SwiftParser/Attributes.swift

@@ -166,8 +166,34 @@ extension Parser {
   }
 
   mutating func parseAttribute(argumentMode: AttributeArgumentMode, parseArguments: (inout Parser) -> RawAttributeSyntax.Argument) -> RawAttributeListSyntax.Element {
-    let (unexpectedBeforeAtSign, atSign) = self.expect(.atSign)
-    let attributeName = self.parseType()
+    var (unexpectedBeforeAtSign, atSign) = self.expect(.atSign)
+    var attributeName = self.parseType()
+    var unexpectedBetweenAtSignAndAttributeName: RawUnexpectedNodesSyntax?
+
+    if atSign.trailingTriviaByteLength != 0 {
+      unexpectedBeforeAtSign = RawUnexpectedNodesSyntax(
+        combining: unexpectedBeforeAtSign,
+        atSign,
+        arena: self.arena
+      )
+      atSign = RawTokenSyntax(
+        kind: .atSign,
+        text: atSign.tokenText,
+        leadingTriviaPieces: atSign.leadingTriviaPieces,
+        presence: .missing,
+        arena: self.arena
+      )
+    } else if attributeName.raw.leadingTriviaByteLength != 0 {
+      unexpectedBetweenAtSignAndAttributeName = RawUnexpectedNodesSyntax(attributeName)
+      // `withLeadingTrivia` only returns `nil` if there is no token in `attributeName`.
+      // But since `attributeName` has leadingTriviaLength != 0 there must be trivia and thus a token.
+      // So we can safely force-unwrap here.
+      attributeName = attributeName
+        .raw
+        .withLeadingTrivia([], arena: self.arena)!
+        .as(RawTypeSyntax.self)!
+    }
+
     let shouldParseArgument: Bool
     switch argumentMode {
     case .required:
@@ -185,6 +211,7 @@ extension Parser {
         RawAttributeSyntax(
           unexpectedBeforeAtSign,
           atSignToken: atSign,
+          unexpectedBetweenAtSignAndAttributeName,
           attributeName: attributeName,
           unexpectedBeforeLeftParen,
           leftParen: leftParen,
@@ -199,6 +226,7 @@ extension Parser {
         RawAttributeSyntax(
           unexpectedBeforeAtSign,
           atSignToken: atSign,
+          unexpectedBetweenAtSignAndAttributeName,
           attributeName: attributeName,
           leftParen: nil,
           argument: nil,

Sources/SwiftParserDiagnostics/ParseDiagnosticsGenerator.swift

@@ -457,7 +457,51 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
         ],
         handledNodes: [argument.id]
       )
-      return .visitChildren
+    }
+    if let unexpectedAtSign = node.unexpectedBeforeAtSignToken?.onlyToken(where: { $0.tokenKind == .atSign && !$0.trailingTrivia.isEmpty }),
+      node.atSignToken.presence == .missing
+    {
+      addDiagnostic(
+        unexpectedAtSign,
+        position: unexpectedAtSign.endPositionBeforeTrailingTrivia,
+        StaticParserError.invalidWhitespaceBetweenAttributeAtSignAndIdentifier,
+        fixIts: [
+          FixIt(
+            message: StaticParserFixIt.removeExtraneousWhitespace,
+            changes: [
+              .makeMissing(unexpectedAtSign, transferTrivia: false),
+              .makePresent(node.atSignToken),
+            ]
+          )
+        ],
+        handledNodes: [
+          node.id,
+          unexpectedAtSign.id,
+          node.atSignToken.id,
+        ]
+      )
+    } else if node.attributeName.isMissingAllTokens,
+      let unexpectedBetweenAtSignTokenAndAttributeName = node.unexpectedBetweenAtSignTokenAndAttributeName,
+      unexpectedBetweenAtSignTokenAndAttributeName.trailingTriviaLength.utf8Length != 0
+    {
+      addDiagnostic(
+        unexpectedBetweenAtSignTokenAndAttributeName,
+        StaticParserError.invalidWhitespaceBetweenAttributeAtSignAndIdentifier,
+        fixIts: [
+          FixIt(
+            message: StaticParserFixIt.removeExtraneousWhitespace,
+            changes: [
+              .makeMissing(unexpectedBetweenAtSignTokenAndAttributeName, transferTrivia: false),
+              .makePresent(node.attributeName),
+            ]
+          )
+        ],
+        handledNodes: [
+          node.id,
+          unexpectedBetweenAtSignTokenAndAttributeName.id,
+          node.attributeName.id,
+        ]
+      )
     }
     return .visitChildren
   }

Sources/SwiftParserDiagnostics/ParserDiagnosticMessages.swift

@@ -171,6 +171,9 @@ extension DiagnosticMessage where Self == StaticParserError {
   public static var invalidWhitespaceAfterPeriod: Self {
     .init("extraneous whitespace after '.' is not permitted")
   }
+  public static var invalidWhitespaceBetweenAttributeAtSignAndIdentifier: Self {
+    .init("extraneous whitespace after '@' is not permitted")
+  }
   public static var joinConditionsUsingComma: Self {
     .init("expected ',' joining parts of a multi-clause condition")
   }

Sources/SwiftSyntax/Raw/RawSyntax.swift

@@ -293,7 +293,8 @@ extension RawSyntax {
   /// - Parameters:
   ///   - leadingTrivia: The trivia to attach.
   ///   - arena: SyntaxArena to the result node data resides.
-  func withLeadingTrivia(_ leadingTrivia: Trivia, arena: SyntaxArena) -> RawSyntax? {
+  @_spi(RawSyntax)
+  public func withLeadingTrivia(_ leadingTrivia: Trivia, arena: SyntaxArena) -> RawSyntax? {
     switch view {
     case .token(let tokenView):
       return .makeMaterializedToken(
@@ -319,7 +320,8 @@ extension RawSyntax {
   /// - Parameters:
   ///   - trailingTrivia: The trivia to attach.
   ///   - arena: SyntaxArena to the result node data resides.
-  func withTrailingTrivia(_ trailingTrivia: Trivia, arena: SyntaxArena) -> RawSyntax? {
+  @_spi(RawSyntax)
+  public func withTrailingTrivia(_ trailingTrivia: Trivia, arena: SyntaxArena) -> RawSyntax? {
     switch view {
     case .token(let tokenView):
       return .makeMaterializedToken(

Tests/SwiftParserTest/AttributeTests.swift

@@ -631,4 +631,43 @@ final class AttributeTests: XCTestCase {
       """
     )
   }
+
+  func testInvalidWhitespaceBetweenAtSignAndIdenfifierIsDiagnosed() {
+    assertParse(
+      """
+      @1️⃣ MyAttribute
+      func foo() {}
+      """,
+      diagnostics: [
+        DiagnosticSpec(
+          message: "extraneous whitespace after '@' is not permitted",
+          fixIts: ["remove whitespace"]
+        )
+      ],
+      fixedSource: """
+        @MyAttribute
+        func foo() {}
+        """
+    )
+  }
+
+  func testInvalidNewlineBetweenAtSignAndIdenfifierIsDiagnosed() {
+    assertParse(
+      """
+      @1️⃣
+      MyAttribute
+      func foo() {}
+      """,
+      diagnostics: [
+        DiagnosticSpec(
+          message: "extraneous whitespace after '@' is not permitted",
+          fixIts: ["remove whitespace"]
+        )
+      ],
+      fixedSource: """
+        @MyAttribute
+        func foo() {}
+        """
+    )
+  }
 }