Add diagnostic for invalid whitespace in @ Attribute

Author: TiagoMaiaL

Sources/SwiftParser/Attributes.swift

@@ -166,8 +166,28 @@ extension Parser {
   }
 
   mutating func parseAttribute(argumentMode: AttributeArgumentMode, parseArguments: (inout Parser) -> RawAttributeSyntax.Argument) -> RawAttributeListSyntax.Element {
-    let (unexpectedBeforeAtSign, atSign) = self.expect(.atSign)
+    var (unexpectedBeforeAtSign, atSign) = self.expect(.atSign)
     let attributeName = self.parseType()
+    
+    if !atSign.trailingTriviaPieces.isEmpty {
+      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 let attributeTokenView = attributeName.raw.tokenView, attributeTokenView.leadingRawTriviaPieces.isEmpty {
+      unexpectedBeforeAtSign = RawUnexpectedNodesSyntax(
+        combining: unexpectedBeforeAtSign, attributeName.as(RawTokenSyntax.self),
+        arena: self.arena
+      )
+    }
+    
     let shouldParseArgument: Bool
     switch argumentMode {
     case .required:

Sources/SwiftParserDiagnostics/ParseDiagnosticsGenerator.swift

@@ -456,7 +456,31 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
         ],
         handledNodes: [argument.id]
       )
-      return .visitChildren
+    }
+    if node.hasError,
+       let unexpectedBeforeAtSignTokens = node.unexpectedBeforeAtSignToken,
+       let unexpectedBeforeAtSign = unexpectedBeforeAtSignTokens.lastToken(viewMode: .sourceAccurate),
+       unexpectedBeforeAtSign.tokenKind == .atSign,
+       !unexpectedBeforeAtSign.trailingTrivia.isEmpty {
+      let tokensToDiscard = unexpectedBeforeAtSignTokens.compactMap { $0.as(TokenSyntax.self) }
+      addDiagnostic(
+        unexpectedBeforeAtSign,
+        StaticParserError.invalidWhitespaceBetweenAttributeAtSignAndIdentifier,
+        fixIts: [
+          FixIt(
+            message: StaticParserFixIt.removeExtraneousWhitespace,
+            changes: [
+              .makeMissing(tokensToDiscard, transferTrivia: false),
+              .makePresent(node.atSignToken),
+            ]
+          )
+        ],
+        handledNodes: [
+          node.id,
+          unexpectedBeforeAtSignTokens.id,
+          node.atSignToken.id
+        ]
+      )
     }
     return .visitChildren
   }

Sources/SwiftParserDiagnostics/ParserDiagnosticMessages.swift

@@ -164,6 +164,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")
   }

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(
+      """
+      @
+      MyAttribute
+      func foo() {}
+      """,
+      diagnostics: [
+        DiagnosticSpec(
+          message: "extraneous whitespace after '@' is not permitted",
+          fixIts: ["remove whitespace"]
+        )
+      ],
+      fixedSource: """
+      @MyAttribute
+      func foo() {}
+      """
+    )
+  }
 }