Add diagnostic for invalid whitespace in @ Attribute

Author: TiagoMaiaL

Sources/SwiftParser/Attributes.swift

@@ -166,8 +166,29 @@ 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.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 {
+      unexpectedBeforeAtSign = RawUnexpectedNodesSyntax(
+        combining: unexpectedBeforeAtSign, attributeName.as(RawTokenSyntax.self),
+        arena: self.arena
+      )
+      // TODO: Synthesize fixed attributeName (RawTypeSyntax)
+    }
+    
     let shouldParseArgument: Bool
     switch argumentMode {
     case .required:

Sources/SwiftParserDiagnostics/ParseDiagnosticsGenerator.swift

@@ -456,7 +456,50 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
         ],
         handledNodes: [argument.id]
       )
-      return .visitChildren
+    }
+    if node.hasError,
+       let unexpectedBeforeAtSignTokens = node.unexpectedBeforeAtSignToken,
+       let unexpectedBeforeAtSign = unexpectedBeforeAtSignTokens.lastToken(viewMode: .sourceAccurate) {
+      let tokensToDiscard = unexpectedBeforeAtSignTokens.compactMap { $0.as(TokenSyntax.self) }
+      if unexpectedBeforeAtSign.tokenKind == .atSign, !unexpectedBeforeAtSign.trailingTrivia.isEmpty {
+        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
+          ]
+        )
+      } else if unexpectedBeforeAtSign.tokenKind == .identifier(unexpectedBeforeAtSign.text) {
+        addDiagnostic(
+          unexpectedBeforeAtSign,
+          StaticParserError.invalidWhitespaceBetweenAttributeAtSignAndIdentifier,
+          fixIts: [
+            FixIt(
+              message: StaticParserFixIt.removeExtraneousWhitespace,
+              changes: [
+                .makeMissing(tokensToDiscard, transferTrivia: false),
+                .makePresent(node.attributeName), // TODO: Provide synthesized attributeName as a replacement.
+              ]
+            )
+          ],
+          handledNodes: [
+            node.id,
+            unexpectedBeforeAtSignTokens.id,
+            node.attributeName.id
+          ]
+        )
+      }
     }
     return .visitChildren
   }

Sources/SwiftParserDiagnostics/ParserDiagnosticMessages.swift

@@ -170,6 +170,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() {}
+      """
+    )
+  }
 }