Ensure all parsed tokens are allocated in a ParsingSyntaxArena

Author: ahoppen

Sources/SwiftSyntax/Raw/RawSyntax.swift

@@ -550,16 +550,12 @@ extension RawSyntax {
     textRange: Range<SyntaxText.Index>,
     presence: SourcePresence,
     tokenDiagnostic: TokenDiagnostic?,
-    arena: __shared SyntaxArena
+    arena: __shared ParsingSyntaxArena
   ) -> RawSyntax {
     assert(
       arena.contains(text: wholeText),
       "token text must be managed by the arena"
     )
-    assert(
-      arena is ParsingSyntaxArena || textRange == wholeText.indices,
-      "arena must be able to parse trivia"
-    )
     let payload = RawSyntaxData.ParsedToken(
       tokenKind: kind,
       wholeText: wholeText,

Sources/SwiftSyntax/Raw/RawSyntaxNodeProtocol.swift

@@ -183,7 +183,7 @@ public struct RawTokenSyntax: RawSyntaxNodeProtocol {
     textRange: Range<SyntaxText.Index>,
     presence: SourcePresence,
     tokenDiagnostic: TokenDiagnostic?,
-    arena: __shared SyntaxArena
+    arena: __shared ParsingSyntaxArena
   ) {
     let raw = RawSyntax.parsedToken(
       kind: kind,
@@ -205,7 +205,7 @@ public struct RawTokenSyntax: RawSyntaxNodeProtocol {
     trailingTriviaPieces: [RawTriviaPiece] = [],
     presence: SourcePresence,
     tokenDiagnostic: TokenDiagnostic? = nil,
-    arena: __shared SyntaxArena
+    arena: __shared ParsingSyntaxArena
   ) {
     if leadingTriviaPieces.isEmpty && trailingTriviaPieces.isEmpty {
       // Create it via `RawSyntax.parsedToken()`.

Sources/SwiftSyntax/Raw/RawSyntaxTokenView.swift

@@ -95,7 +95,7 @@ public struct RawSyntaxTokenView: Sendable {
   public var leadingRawTriviaPieces: [RawTriviaPiece] {
     switch raw.rawData.payload {
     case .parsedToken(let dat):
-      let arena = raw.arena.parsingArena!
+      let arena = raw.arena as! ParsingSyntaxArena
       return arena.parseTrivia(source: dat.leadingTriviaText, position: .leading)
     case .materializedToken(let dat):
       return Array(dat.leadingTrivia)
@@ -108,7 +108,7 @@ public struct RawSyntaxTokenView: Sendable {
   public var trailingRawTriviaPieces: [RawTriviaPiece] {
     switch raw.rawData.payload {
     case .parsedToken(let dat):
-      let arena = raw.arena.parsingArena!
+      let arena = raw.arena as! ParsingSyntaxArena
       return arena.parseTrivia(source: dat.trailingTriviaText, position: .trailing)
     case .materializedToken(let dat):
       return Array(dat.trailingTrivia)
@@ -204,8 +204,21 @@ public struct RawSyntaxTokenView: Sendable {
     arena.addChild(self.raw.arenaReference)
     switch raw.rawData.payload {
     case .parsedToken(var payload):
-      payload.presence = newValue
-      return RawSyntax(arena: arena, payload: .parsedToken(payload))
+      if arena === self.raw.arena {
+        payload.presence = newValue
+        return RawSyntax(arena: arena, payload: .parsedToken(payload))
+      }
+      // If the modified token is allocated in a different arena, it might have
+      // a different or no `parseTrivia` function. We thus cannot use a
+      // `parsedToken` anymore.
+      return .makeMaterializedToken(
+        kind: formKind(),
+        leadingTrivia: formLeadingTrivia(),
+        trailingTrivia: formTrailingTrivia(),
+        presence: newValue,
+        tokenDiagnostic: tokenDiagnostic,
+        arena: arena
+      )
     case .materializedToken(var payload):
       payload.presence = newValue
       return RawSyntax(arena: arena, payload: .materializedToken(payload))
@@ -274,8 +287,21 @@ public struct RawSyntaxTokenView: Sendable {
     arena.addChild(self.raw.arenaReference)
     switch raw.rawData.payload {
     case .parsedToken(var dat):
-      dat.tokenDiagnostic = tokenDiagnostic
-      return RawSyntax(arena: arena, payload: .parsedToken(dat))
+      if arena === self.raw.arena {
+        dat.tokenDiagnostic = tokenDiagnostic
+        return RawSyntax(arena: arena, payload: .parsedToken(dat))
+      }
+      // If the modified token is allocated in a different arena, it might have
+      // a different or no `parseTrivia` function. We thus cannot use a
+      // `parsedToken` anymore.
+      return .makeMaterializedToken(
+        kind: formKind(),
+        leadingTrivia: formLeadingTrivia(),
+        trailingTrivia: formTrailingTrivia(),
+        presence: presence,
+        tokenDiagnostic: tokenDiagnostic,
+        arena: arena
+      )
     case .materializedToken(var dat):
       dat.tokenDiagnostic = tokenDiagnostic
       return RawSyntax(arena: arena, payload: .materializedToken(dat))

Sources/SwiftSyntax/SyntaxArena.swift

@@ -71,33 +71,6 @@ 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
   /// uninitialized memory range as a `UnsafeMutableBufferPointer<RawSyntax?>`.
   func allocateRawSyntaxBuffer(count: Int) -> UnsafeMutableBufferPointer<RawSyntax?> {

Tests/SwiftSyntaxTest/DummyParseToken.swift

@@ -0,0 +1,19 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+@_spi(RawSyntax) import SwiftSyntax
+
+/// Dummy trivia parsing function.
+func dummyParseToken(source: SyntaxText, position: TriviaPosition) -> [RawTriviaPiece] {
+  // Emit a single `unexpectedText` trivia of the whole trivia text.
+  return [.unexpectedText(source)]
+}

Tests/SwiftSyntaxTest/RawSyntaxTests.swift

@@ -13,7 +13,7 @@
 @_spi(RawSyntax) import SwiftSyntax
 import XCTest
 
-fileprivate func cannedStructDecl(arena: SyntaxArena) -> RawStructDeclSyntax {
+fileprivate func cannedStructDecl(arena: ParsingSyntaxArena) -> RawStructDeclSyntax {
   let structKW = RawTokenSyntax(
     kind: .keyword,
     text: arena.intern("struct"),
@@ -68,7 +68,7 @@ fileprivate func cannedStructDecl(arena: SyntaxArena) -> RawStructDeclSyntax {
 final class RawSyntaxTests: XCTestCase {
 
   func testFactory() {
-    withExtendedLifetime(SyntaxArena()) { arena in
+    withExtendedLifetime(ParsingSyntaxArena(parseTriviaFunction: dummyParseToken)) { arena in
       let structDecl = cannedStructDecl(arena: arena)
       XCTAssertEqual(
         "\(structDecl.raw)",
@@ -81,7 +81,7 @@ final class RawSyntaxTests: XCTestCase {
   }
 
   func testAccessor() {
-    withExtendedLifetime(SyntaxArena()) { arena in
+    withExtendedLifetime(ParsingSyntaxArena(parseTriviaFunction: dummyParseToken)) { arena in
       let structDecl = cannedStructDecl(arena: arena)
       XCTAssertEqual(structDecl.name.tokenKind, .identifier)
       XCTAssertEqual(structDecl.structKeyword.tokenText, "struct")
@@ -96,7 +96,7 @@ final class RawSyntaxTests: XCTestCase {
   }
 
   func testMaterializedToken() {
-    withExtendedLifetime(SyntaxArena()) { arena in
+    withExtendedLifetime(ParsingSyntaxArena(parseTriviaFunction: dummyParseToken)) { arena in
       let ident = RawTokenSyntax(
         kind: .identifier,
         text: arena.intern("foo"),
@@ -118,12 +118,6 @@ final class RawSyntaxTests: XCTestCase {
   }
 
   func testParsedToken() {
-    // Dummy trivia parsing function.
-    func dummyParseToken(source: SyntaxText, position: TriviaPosition) -> [RawTriviaPiece] {
-      // Emit a single `unexpectedText` trivia of the whole trivia text.
-      return [.unexpectedText(source)]
-    }
-
     withExtendedLifetime(ParsingSyntaxArena(parseTriviaFunction: dummyParseToken)) { arena in
       let ident = RawTokenSyntax(
         kind: .identifier,

Tests/SwiftSyntaxTest/SourceLocationConverterTests.swift

@@ -36,7 +36,7 @@ fileprivate func assertPresumedSourceLocation(
 
 final class SourceLocationConverterTests: XCTestCase {
   func testInvalidUtf8() {
-    let eofToken = withExtendedLifetime(SyntaxArena()) { arena in
+    let eofToken = withExtendedLifetime(ParsingSyntaxArena(parseTriviaFunction: dummyParseToken)) { arena in
       let leadingTriviaText = [UInt8(0xfd)].withUnsafeBufferPointer { buf in
         arena.intern(SyntaxText(buffer: SyntaxArenaAllocatedBufferPointer(buf)))
       }