Run ParseDiagnosticGenerator on tree with flipped token presence

This uncovered a couple of implicit assumptions, mostly around the fact that tokens insided `UnexpectedNodesSyntax` are present, which isn’t true in general for manually generated trees.

It also uncovered an issue where we weren’t able to retrieve the trivia pieces of a token after it had been modified using `with` because after the modification the token was a `parsedToken` that no longer resided in a `ParsingSyntaxArena`.

Author: ahoppen

CodeGeneration/Sources/SyntaxSupport/DeclNodes.swift

@@ -464,7 +464,7 @@ public let DECL_NODES: [Node] = [
       An editor placeholder, e.g. `<#declaration#>` that is used in a position that expects a declaration.
       """,
     kind: "Decl",
-    traits:[
+    traits: [
       "WithAttributes",
       "WithModifiers",
     ],

Sources/SwiftParserDiagnostics/MissingTokenError.swift

@@ -17,6 +17,7 @@ import SwiftDiagnostics
 extension ParseDiagnosticsGenerator {
   func handleMissingToken(_ missingToken: TokenSyntax) {
     guard let invalidToken = missingToken.previousToken(viewMode: .all),
+      invalidToken.presence == .present,
       let invalidTokenContainer = invalidToken.parent?.as(UnexpectedNodesSyntax.self),
       invalidTokenContainer.count == 1
     else {

Sources/SwiftParserDiagnostics/ParseDiagnosticsGenerator.swift

@@ -133,7 +133,7 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
     removeRedundantFixIt: (_ misplacedTokens: [TokenSyntax]) -> FixItMessage? = { _ in nil }
   ) {
     guard let incorrectContainer = unexpected,
-      let misplacedTokens = incorrectContainer.onlyTokens(satisfying: unexpectedTokenCondition)
+      let misplacedTokens = incorrectContainer.onlyPresentTokens(satisfying: unexpectedTokenCondition)
     else {
       // If there are no unexpected nodes or the unexpected contain multiple tokens, don't emit a diagnostic.
       return
@@ -180,7 +180,7 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
     message: (TokenSyntax) -> Message
   ) {
     guard let unexpected = unexpected,
-      let misplacedToken = unexpected.onlyToken(where: predicate)
+      let misplacedToken = unexpected.onlyPresentToken(where: predicate)
     else {
       // If there is no unexpected node or the unexpected doesn't have the
       // expected token, don't emit a diagnostic.
@@ -262,7 +262,7 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
       }
       if specifier.presence == .present {
         for case .some(let unexpected) in unexpectedNodes {
-          for duplicateSpecifier in unexpected.tokens(satisfying: isOfSameKind) {
+          for duplicateSpecifier in unexpected.presentTokens(satisfying: isOfSameKind) {
             addDiagnostic(
               duplicateSpecifier,
               DuplicateEffectSpecifiers(correctSpecifier: specifier, unexpectedSpecifier: duplicateSpecifier),
@@ -298,29 +298,36 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
       suppressRemainingDiagnostics = true
       return .skipChildren
     }
-    if let tryKeyword = node.onlyToken(where: { $0.tokenKind == .keyword(.try) }),
+    if let tryKeyword = node.onlyPresentToken(where: { $0.tokenKind == .keyword(.try) }),
       let nextToken = tryKeyword.nextToken(viewMode: .sourceAccurate),
       nextToken.tokenKind.isLexerClassifiedKeyword,
       !(node.parent?.is(TypeEffectSpecifiersSyntax.self) ?? false)
     {
       addDiagnostic(node, TryCannotBeUsed(nextToken: nextToken))
-    } else if let semicolons = node.onlyTokens(satisfying: { $0.tokenKind == .semicolon }) {
+    } else if let semicolons = node.onlyPresentTokens(satisfying: { $0.tokenKind == .semicolon }) {
       addDiagnostic(
         node,
         .unexpectedSemicolon,
         fixIts: [
           FixIt(message: RemoveNodesFixIt(semicolons), changes: semicolons.map { FixIt.MultiNodeChange.makeMissing($0) })
         ]
       )
-    } else if node.first?.as(TokenSyntax.self)?.tokenKind.isIdentifier == true,
+    } else if let firstToken = node.first?.as(TokenSyntax.self),
+      firstToken.tokenKind.isIdentifier == true,
+      firstToken.presence == .present,
       let previousToken = node.previousToken(viewMode: .sourceAccurate),
       previousToken.tokenKind.isIdentifier,
       previousToken.parent?.is(DeclSyntax.self) == true || previousToken.parent?.is(IdentifierPatternSyntax.self) == true
     {
       // If multiple identifiers are used for a declaration name, offer to join them together.
       let tokens =
         node
-        .prefix(while: { $0.as(TokenSyntax.self)?.tokenKind.isIdentifier == true })
+        .prefix(while: {
+          guard let token = $0.as(TokenSyntax.self) else {
+            return false
+          }
+          return token.tokenKind.isIdentifier == true && token.presence == .present
+        })
         .map({ $0.as(TokenSyntax.self)! })
       let joined = previousToken.text + tokens.map(\.text).joined()
       var fixIts: [FixIt] = [
@@ -449,7 +456,7 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
       return .skipChildren
     }
     if let unexpected = node.unexpectedBetweenPlatformAndVersion,
-      unexpected.onlyToken(where: { $0.tokenKind == .binaryOperator(">=") }) != nil
+      unexpected.onlyPresentToken(where: { $0.tokenKind == .binaryOperator(">=") }) != nil
     {
       addDiagnostic(
         unexpected,
@@ -552,7 +559,7 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
     }
 
     if let unexpected = node.unexpectedBetweenBodyAndTrailingComma,
-      let token = unexpected.tokens(satisfying: { $0.tokenKind == .binaryOperator("&&") }).first,
+      let token = unexpected.presentTokens(satisfying: { $0.tokenKind == .binaryOperator("&&") }).first,
       let trailingComma = node.trailingComma,
       trailingComma.presence == .missing,
       let previous = node.unexpectedBetweenBodyAndTrailingComma?.previousToken(viewMode: .sourceAccurate)
@@ -583,7 +590,7 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
       return .skipChildren
     }
     if let unexpected = node.unexpectedBetweenDeinitKeywordAndBody,
-      let name = unexpected.filter({ $0.as(TokenSyntax.self)?.tokenKind.isIdentifier == true }).only?.as(TokenSyntax.self)
+      let name = unexpected.presentTokens(satisfying: { $0.tokenKind.isIdentifier == true }).only?.as(TokenSyntax.self)
     {
       addDiagnostic(
         name,
@@ -617,7 +624,7 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
     // Detect C-style for loops based on two semicolons which could not be parsed between the 'for' keyword and the '{'
     // This is mostly a proof-of-concept implementation to produce more complex diagnostics.
     if let unexpectedCondition = node.body.unexpectedBeforeLeftBrace,
-      unexpectedCondition.tokens(withKind: .semicolon).count == 2
+      unexpectedCondition.presentTokens(withKind: .semicolon).count == 2
     {
       // FIXME: This is aweful. We should have a way to either get all children between two cursors in a syntax node or highlight a range from one node to another.
       addDiagnostic(
@@ -690,7 +697,7 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
         message: { _ in .typeParameterPackEllipsis }
       )
     } else if let unexpected = node.unexpectedBetweenNameAndColon,
-      let unexpectedEllipsis = unexpected.onlyToken(where: { $0.tokenKind == .ellipsis }),
+      let unexpectedEllipsis = unexpected.onlyPresentToken(where: { $0.tokenKind == .ellipsis }),
       let each = node.each
     {
       addDiagnostic(
@@ -929,7 +936,7 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
     if shouldSkip(node) {
       return .skipChildren
     }
-    if let token = node.unexpectedBetweenModuleLabelAndColon?.onlyToken(where: { $0.tokenKind.isIdentifier }),
+    if let token = node.unexpectedBetweenModuleLabelAndColon?.onlyPresentToken(where: { $0.tokenKind.isIdentifier }),
       node.moduleLabel.presence == .missing
     {
       addDiagnostic(
@@ -976,9 +983,9 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
         return
       }
       let message: DiagnosticMessage?
-      if let identifier = unexpected.onlyToken(where: { $0.tokenKind.isIdentifier }) {
+      if let identifier = unexpected.onlyPresentToken(where: { $0.tokenKind.isIdentifier }) {
         message = IdentifierNotAllowedInOperatorName(identifier: identifier)
-      } else if let tokens = unexpected.onlyTokens(satisfying: { _ in true }) {
+      } else if let tokens = unexpected.onlyPresentTokens(satisfying: { _ in true }) {
         message = TokensNotAllowedInOperatorName(tokens: tokens)
       } else {
         message = nil
@@ -1068,7 +1075,7 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
     if shouldSkip(node) {
       return .skipChildren
     }
-    if let singleQuote = node.unexpectedBetweenOpenDelimiterAndOpenQuote?.onlyToken(where: { $0.tokenKind == .singleQuote }) {
+    if let singleQuote = node.unexpectedBetweenOpenDelimiterAndOpenQuote?.onlyPresentToken(where: { $0.tokenKind == .singleQuote }) {
       let fixIt = FixIt(
         message: ReplaceTokensFixIt(replaceTokens: [singleQuote], replacement: node.openQuote),
         changes: [
@@ -1094,7 +1101,7 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
       addDiagnostic(diagnostic, handledNodes: handledNodes)
     }
     if case .stringSegment(let segment) = node.segments.last {
-      if let invalidContent = segment.unexpectedBeforeContent?.onlyToken(where: { $0.trailingTrivia.contains(where: { $0.isBackslash }) }) {
+      if let invalidContent = segment.unexpectedBeforeContent?.onlyPresentToken(where: { $0.trailingTrivia.contains(where: { $0.isBackslash }) }) {
         let fixIt = FixIt(
           message: .removeBackslash,
           changes: [
@@ -1119,7 +1126,7 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
       return .skipChildren
     }
     if let unexpected = node.unexpectedBetweenSubscriptKeywordAndGenericParameterClause,
-      let nameTokens = unexpected.onlyTokens(satisfying: { !$0.tokenKind.isLexerClassifiedKeyword })
+      let nameTokens = unexpected.onlyPresentTokens(satisfying: { !$0.tokenKind.isLexerClassifiedKeyword })
     {
       addDiagnostic(
         unexpected,
@@ -1131,7 +1138,7 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
       )
     }
     if let unexpected = node.indices.unexpectedBeforeLeftParen,
-      let nameTokens = unexpected.onlyTokens(satisfying: { !$0.tokenKind.isLexerClassifiedKeyword })
+      let nameTokens = unexpected.onlyPresentTokens(satisfying: { !$0.tokenKind.isLexerClassifiedKeyword })
     {
       addDiagnostic(
         unexpected,
@@ -1255,7 +1262,8 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
       return .skipChildren
     }
 
-    if let token = node.unexpectedBetweenMessageLabelAndColon?.onlyToken(where: { $0.tokenKind.isIdentifier }),
+    if let token = node.unexpectedBetweenMessageLabelAndColon?.onlyPresentToken(where: { $0.tokenKind.isIdentifier }),
+      token.presence == .present,
       node.messageLabel.presence == .missing
     {
       addDiagnostic(

Sources/SwiftParserDiagnostics/SyntaxExtensions.swift

@@ -14,26 +14,30 @@
 @_spi(Diagnostics) import SwiftParser
 
 extension UnexpectedNodesSyntax {
-  func tokens(satisfying isIncluded: (TokenSyntax) -> Bool) -> [TokenSyntax] {
+  func presentTokens(satisfying isIncluded: (TokenSyntax) -> Bool) -> [TokenSyntax] {
     return self.children(viewMode: .sourceAccurate).compactMap({ $0.as(TokenSyntax.self) }).filter(isIncluded)
   }
 
-  func tokens(withKind kind: TokenKind) -> [TokenSyntax] {
-    return self.tokens(satisfying: { $0.tokenKind == kind })
+  func presentTokens(withKind kind: TokenKind) -> [TokenSyntax] {
+    return self.presentTokens(satisfying: { $0.tokenKind == kind })
   }
 
-  /// If this only contains a single item, which is a token satisfying `condition`, return that token, otherwise return `nil`.
-  func onlyToken(where condition: (TokenSyntax) -> Bool) -> TokenSyntax? {
-    if self.count == 1, let token = self.first?.as(TokenSyntax.self), condition(token) {
+  /// If this only contains a single item, which is a present token satisfying `condition`, return that token, otherwise return `nil`.
+  func onlyPresentToken(where condition: (TokenSyntax) -> Bool) -> TokenSyntax? {
+    if self.count == 1,
+      let token = self.first?.as(TokenSyntax.self),
+      condition(token),
+      token.presence == .present
+    {
       return token
     } else {
       return nil
     }
   }
 
-  /// If this only contains tokens satisfying `condition`, return an array containing those tokens, otherwise return `nil`.
-  func onlyTokens(satisfying condition: (TokenSyntax) -> Bool) -> [TokenSyntax]? {
-    let tokens = tokens(satisfying: condition)
+  /// If this only contains present tokens satisfying `condition`, return an array containing those tokens, otherwise return `nil`.
+  func onlyPresentTokens(satisfying condition: (TokenSyntax) -> Bool) -> [TokenSyntax]? {
+    let tokens = presentTokens(satisfying: condition)
     if tokens.count == self.count {
       return tokens
     } else {
@@ -79,7 +83,7 @@ extension SyntaxProtocol {
       } else {
         return "braces"
       }
-    } else if let token = Syntax(self).as(UnexpectedNodesSyntax.self)?.onlyTokens(satisfying: { $0.tokenKind.isLexerClassifiedKeyword })?.only {
+    } else if let token = Syntax(self).as(UnexpectedNodesSyntax.self)?.onlyPresentTokens(satisfying: { $0.tokenKind.isLexerClassifiedKeyword })?.only {
       return "'\(token.text)' keyword"
     } else if let token = Syntax(self).as(TokenSyntax.self) {
       return "'\(token.text)' keyword"

Sources/SwiftSyntax/Raw/RawSyntaxTokenView.swift

@@ -95,7 +95,7 @@ public struct RawSyntaxTokenView {
   public var leadingRawTriviaPieces: [RawTriviaPiece] {
     switch raw.rawData.payload {
     case .parsedToken(let dat):
-      let arena = raw.arena as! ParsingSyntaxArena
+      let arena = raw.arena.parsingArena!
       return arena.parseTrivia(source: dat.leadingTriviaText, position: .leading)
     case .materializedToken(let dat):
       return Array(dat.leadingTrivia)
@@ -108,7 +108,7 @@ public struct RawSyntaxTokenView {
   public var trailingRawTriviaPieces: [RawTriviaPiece] {
     switch raw.rawData.payload {
     case .parsedToken(let dat):
-      let arena = raw.arena as! ParsingSyntaxArena
+      let arena = raw.arena.parsingArena!
       return arena.parseTrivia(source: dat.trailingTriviaText, position: .trailing)
     case .materializedToken(let dat):
       return Array(dat.trailingTrivia)

Sources/SwiftSyntax/SyntaxArena.swift

@@ -69,6 +69,33 @@ public class SyntaxArena {
     }
   }
 
+  /// If this arena or any of its child arenas is a ``ParsingSyntaxArena``
+  /// return one of these arenas, otherwise return `nil`.
+  ///
+  /// If the arena has multiple child nodes that are ``ParsingSyntaxArena``s, it
+  /// is undefined which one will be returned.
+  ///
+  /// The use case for this is to get the trivia parsing function of the arena.
+  /// Parsed tokens created by `SwiftParser` automatically reside in a
+  /// ``ParsingSyntaxArena`` but if they are modified (e.g. using the `with`
+  /// functions), they might reside in a new arena. But we still want to be able
+  /// to retrieve trivia from those modified tokens, which requires calling into
+  /// the `parseTrivia` function of the ``ParsingSyntaxArena`` that created the
+  /// token. Since the modified syntax arena needs to keep the original
+  /// ``ParsingSyntaxArena`` alive, we can search this arena’s `childRefs` for
+  /// the ``ParsingSyntaxArena`` that created the token.
+  var parsingArena: ParsingSyntaxArena? {
+    if let parsingArena = self as? ParsingSyntaxArena {
+      return parsingArena
+    }
+    for child in childRefs {
+      if let parsingArena = child.value.parsingArena {
+        return parsingArena
+      }
+    }
+    return nil
+  }
+
   /// Allocates a buffer of `RawSyntax?` with the given count, then returns the
   /// uninitlialized memory range as a `UnsafeMutableBufferPointer<RawSyntax?>`.
   func allocateRawSyntaxBuffer(count: Int) -> UnsafeMutableBufferPointer<RawSyntax?> {

Sources/SwiftSyntax/generated/SyntaxRewriter.swift

@@ -6997,7 +6997,7 @@ open class SyntaxRewriter {
   public func rewrite(_ node: Syntax) -> Syntax {
     let rewritten = self.visit(node)
     return withExtendedLifetime(rewritten) {
-      return Syntax(node.data.replacingSelf(rewritten.raw, arena: rewritten.raw.arena))
+      return Syntax(node.data.replacingSelf(rewritten.raw, arena: SyntaxArena()))
     }
   }
 }

Tests/SwiftParserTest/Assertions.swift

@@ -681,6 +681,7 @@ func assertParse<S: SyntaxProtocol>(
   if enableLongTests {
     DispatchQueue.concurrentPerform(iterations: Array(tree.tokens(viewMode: .all)).count) { tokenIndex in
       let flippedTokenTree = TokenPresenceFlipper(flipTokenAtIndex: tokenIndex).rewrite(Syntax(tree))
+      _ = ParseDiagnosticsGenerator.diagnostics(for: flippedTokenTree)
       assertRoundTrip(source: flippedTokenTree.syntaxTextBytes, parse, file: file, line: line)
     }