Introduce RawSyntaxData.ParsedToken

A new RawSyntax data variation. This is intended to be used for parsed
tokens. Actual usage will be in follow-ups.

In this variation, unlike the existing `MaterializedToken`, leading/trailing
trivia are stored in a SyntaxText form, and are lazily materialized by a
function registered in `SyntaxArena`. Since the parser doesn't need to
materialize them while parsing, it'd gain the "parsing" performance boost,
in exchange for a performance reduction when accessing trivia.

Author: rintaro

Sources/SwiftSyntax/RawSyntax.swift

@@ -16,10 +16,21 @@ typealias RawTriviaPieceBuffer = UnsafeBufferPointer<RawTriviaPiece>
 /// Node data for RawSyntax tree. Tagged union plus common data.
 internal struct RawSyntaxData {
   internal enum Payload {
+    case parsedToken(ParsedToken)
     case materializedToken(MaterializedToken)
     case layout(Layout)
   }
 
+  /// Token with lazy trivia parsing.
+  ///
+  /// The RawSyntax's `arena` must have a valid trivia parsing function to
+  /// lazily materialize the leading/trailing trivia.
+  struct ParsedToken {
+    var tokenKind: RawTokenKind
+    var wholeText: SyntaxText
+    var textRange: Range<SyntaxText.Index>
+  }
+
   /// Token typically created with `TokenSyntax.<someToken>`.
   struct MaterializedToken {
     var tokenKind: RawTokenKind
@@ -42,6 +53,18 @@ internal struct RawSyntaxData {
   fileprivate var arenaReference: SyntaxArenaRef
 }
 
+extension RawSyntaxData.ParsedToken {
+  var tokenText: SyntaxText {
+    SyntaxText(rebasing: wholeText[textRange])
+  }
+  var leadingTriviaText: SyntaxText {
+    SyntaxText(rebasing: wholeText[..<textRange.lowerBound])
+  }
+  var trailingTriviaText: SyntaxText {
+    SyntaxText(rebasing: wholeText[textRange.upperBound...])
+  }
+}
+
 extension RawSyntaxData.MaterializedToken {
   var leadingTrivia: RawTriviaPieceBuffer {
     RawTriviaPieceBuffer(rebasing: triviaPieces[..<Int(numLeadingTrivia)])
@@ -94,6 +117,8 @@ extension RawSyntax {
     switch rawData.payload {
     case .materializedToken(let dat):
       return dat.tokenKind
+    case .parsedToken(let dat):
+      return dat.tokenKind
     case .layout(_):
       preconditionFailure("'tokenKind' is not available for non-token node")
     }
@@ -102,6 +127,8 @@ extension RawSyntax {
   /// Token text of this node assuming this node is a token.
   var rawTokenText: SyntaxText {
     switch rawData.payload {
+    case .parsedToken(let dat):
+      return dat.tokenText
     case .materializedToken(let dat):
       return dat.tokenText
     case .layout(_):
@@ -112,6 +139,8 @@ extension RawSyntax {
   /// The UTF-8 byte length of the leading trivia, assuming this node is a token.
   var tokenLeadingTriviaByteLength: Int {
     switch rawData.payload {
+    case .parsedToken(let dat):
+      return dat.leadingTriviaText.count
     case .materializedToken(let dat):
       return dat.leadingTrivia.reduce(0) { $0 + $1.byteLength }
     case .layout(_):
@@ -122,26 +151,32 @@ extension RawSyntax {
   /// The UTF-8 byte length of the trailing trivia, assuming this node is a token.
   var tokenTrailingTriviaByteLength: Int {
     switch rawData.payload {
+    case .parsedToken(let dat):
+      return dat.trailingTriviaText.count
     case .materializedToken(let dat):
       return dat.trailingTrivia.reduce(0) { $0 + $1.byteLength }
     case .layout(_):
       preconditionFailure("'tokenTrailingTriviaByteLength' is not available for non-token node")
     }
   }
 
-  var tokenLeadingRawTriviaPieces: RawTriviaPieceBuffer {
+  var tokenLeadingRawTriviaPieces: [RawTriviaPiece] {
     switch rawData.payload {
+    case .parsedToken(let dat):
+      return self.arena.parseTrivia(source: dat.leadingTriviaText, isTrailing: false)
     case .materializedToken(let dat):
-      return dat.leadingTrivia
+      return Array(dat.leadingTrivia)
     case .layout(_):
       preconditionFailure("'tokenLeadingRawTriviaPieces' is called on non-token raw syntax")
     }
   }
 
-  var tokenTrailingRawTriviaPieces: RawTriviaPieceBuffer {
+  var tokenTrailingRawTriviaPieces: [RawTriviaPiece] {
     switch rawData.payload {
+    case .parsedToken(let dat):
+      return self.arena.parseTrivia(source: dat.trailingTriviaText, isTrailing: false)
     case .materializedToken(let dat):
-      return dat.trailingTrivia
+      return Array(dat.trailingTrivia)
     case .layout(_):
       preconditionFailure("'tokenTrailingRawTriviaPieces' is called on non-token raw syntax")
     }
@@ -152,6 +187,7 @@ extension RawSyntax {
   /// The syntax kind of this raw syntax.
   var kind: SyntaxKind {
     switch rawData.payload {
+    case .parsedToken(_): return .token
     case .materializedToken(_): return .token
     case .layout(let dat): return dat.kind
     }
@@ -175,8 +211,11 @@ extension RawSyntax {
   /// Child nodes.
   var children: RawSyntaxBuffer {
     switch rawData.payload {
-    case .materializedToken(_): return .init(start: nil, count: 0)
-    case .layout(let dat): return dat.layout
+    case .parsedToken(_),
+         .materializedToken(_):
+      return .init(start: nil, count: 0)
+    case .layout(let dat):
+      return dat.layout
     }
   }
 
@@ -197,8 +236,11 @@ extension RawSyntax {
   /// Total number of nodes in this sub-tree, including `self` node.
   var totalNodes: Int {
     switch rawData.payload {
-    case .materializedToken(_): return 1
-    case .layout(let dat): return dat.descendantCount + 1
+    case .parsedToken(_),
+         .materializedToken(_):
+      return 1
+    case .layout(let dat):
+      return dat.descendantCount + 1
     }
   }
 
@@ -226,6 +268,7 @@ extension RawSyntax {
   /// Sum of text byte lengths of all descendant token nodes.
   var byteLength: Int {
     switch rawData.payload {
+    case .parsedToken(let dat): return dat.wholeText.count
     case .materializedToken(let dat): return Int(dat.byteLength)
     case .layout(let dat): return dat.byteLength
     }
@@ -237,6 +280,8 @@ extension RawSyntax {
 
   func formTokenKind() -> TokenKind? {
     switch rawData.payload {
+    case .parsedToken(let dat):
+      return TokenKind.fromRaw(kind: dat.tokenKind, text: dat.tokenText)
     case .materializedToken(let dat):
       return TokenKind.fromRaw(kind: dat.tokenKind, text: dat.tokenText)
     case .layout(_):
@@ -280,7 +325,13 @@ extension RawSyntax {
   /// Assuming this node is a token, returns a `RawSyntax` node with the same
   /// source text but with the token kind changed to `newValue`.
   func withTokenKind(_ newValue: TokenKind) -> RawSyntax {
-    switch payload {
+    switch rawData.payload {
+    case .parsedToken(_):
+      return .makeMaterializedToken(
+        kind: newValue,
+        leadingTrivia: formLeadingTrivia()!,
+        trailingTrivia: formTokenTrailingTrivia()!,
+        arena: arena)
     case .materializedToken(var payload):
       let decomposed = newValue.decomposeToRaw()
       let rawKind = decomposed.rawKind
@@ -517,6 +568,9 @@ extension RawSyntax: TextOutputStreamable, CustomStringConvertible {
   /// - Parameter stream: The stream on which to output this node.
   public func write<Target: TextOutputStream>(to target: inout Target) {
     switch rawData.payload {
+    case .parsedToken(let dat):
+      String(syntaxText: dat.wholeText).write(to: &target)
+      break
     case .materializedToken(let dat):
       for p in dat.leadingTrivia { p.write(to: &target) }
       String(syntaxText: dat.tokenText).write(to: &target)
@@ -594,6 +648,8 @@ extension RawSyntax {
   /// is a token node.
   var tokenTextByteLength: Int {
     switch rawData.payload {
+    case .parsedToken(let dat):
+      return dat.tokenText.count
     case .materializedToken(let dat):
       return dat.tokenText.count
     case .layout(_):
@@ -640,6 +696,24 @@ private func makeRawTriviaPieces(leadingTrivia: Trivia, trailingTrivia: Trivia,
 }
 
 extension RawSyntax {
+  /// "Designated" factory method to create a parsed token node.
+  ///
+  /// - Parameters:
+  ///   - kind: Token kind.
+  ///   - wholeText: Whole text of this token including trailing/leading trivia.
+  ///   - textRange: Range of the token text in `wholeText`.
+  ///   - arena: SyntaxArea to the result node data resides.
+  internal static func parsedToken(
+    kind: RawTokenKind,
+    wholeText: SyntaxText,
+    textRange: Range<SyntaxText.Index>,
+    arena: SyntaxArena
+  ) -> RawSyntax {
+    let payload = RawSyntaxData.ParsedToken(
+      tokenKind: kind, wholeText: wholeText, textRange: textRange)
+    return RawSyntax(arena: arena, payload: .parsedToken(payload))
+  }
+
   /// "Designated" factory method to create a materialized token node.
   ///
   /// This should not be called directly.
@@ -864,6 +938,11 @@ extension RawSyntax: CustomDebugStringConvertible {
   private func debugWrite<Target: TextOutputStream>(to target: inout Target, indent: Int, withChildren: Bool = false) {
     let childIndent = indent + 2
     switch rawData.payload {
+    case .parsedToken(let dat):
+      target.write(".parsedToken(")
+      target.write(String(describing: dat.tokenKind))
+      target.write(" wholeText=\(dat.tokenText.debugDescription)")
+      target.write(" textRange=\(dat.textRange.description)")
     case .materializedToken(let dat):
       target.write(".materializedToken(")
       target.write(String(describing: dat.tokenKind))

Sources/SwiftSyntax/RawSyntaxNodeProtocol.swift

@@ -86,6 +86,10 @@ public struct RawTokenSyntax: RawSyntaxNodeProtocol {
     return raw.rawTokenText
   }
 
+  public var byteLength: Int {
+    return raw.byteLength
+  }
+
   public var presence: SourcePresence {
     raw.presence
   }
@@ -94,7 +98,30 @@ public struct RawTokenSyntax: RawSyntaxNodeProtocol {
     presence == .missing
   }
 
-  /// Creates a `TokenSyntax`. `text` and trivia must be managed by the same
+  public var leadingTriviaPieces: [RawTriviaPiece] {
+    raw.tokenLeadingRawTriviaPieces
+  }
+
+  public var trailingTriviaPieces: [RawTriviaPiece] {
+    raw.tokenTrailingRawTriviaPieces
+  }
+
+  /// Creates a `RawTokenSyntax`. `wholeText` must be managed by the same
+  /// `arena`. `textRange` is a range of the token text in `wholeText`.
+  public init(
+    kind: RawTokenKind,
+    wholeText: SyntaxText,
+    textRange: Range<SyntaxText.Index>,
+    arena: SyntaxArena
+  ) {
+    assert(arena.contains(text: wholeText),
+           "token text must be managed by the arena")
+    let raw = RawSyntax.parsedToken(
+      kind: kind, wholeText: wholeText, textRange: textRange, arena: arena)
+    self = RawTokenSyntax(raw: raw)
+  }
+
+  /// Creates a `RawTokenSyntax`. `text` and trivia must be managed by the same
   /// `arena`.
   public init(
     kind: RawTokenKind,

Sources/SwiftSyntax/SyntaxArena.swift

@@ -11,6 +11,10 @@
 //===----------------------------------------------------------------------===//
 
 public class SyntaxArena {
+
+  @_spi(RawSyntax)
+  public typealias ParseTriviaFunction = (_ source: SyntaxText, _ isTrailing: Bool) -> [RawTriviaPiece]
+
   /// Bump-pointer allocator for all "intern" methods.
   private var allocator: BumpPtrAllocator
   /// Source file buffer the Syntax tree represents.
@@ -22,12 +26,19 @@ public class SyntaxArena {
   /// Whether or not this arena has been added to other arenas as a child.
   /// Used to make sure we don’t introduce retain cycles between arenas.
   private var hasParent: Bool
+  private var parseTriviaFunction: ParseTriviaFunction
 
-  public init() {
+  @_spi(RawSyntax)
+  public init(parseTriviaFunction: @escaping ParseTriviaFunction) {
     allocator = BumpPtrAllocator()
     children = []
     sourceBuffer = .init(start: nil, count: 0)
     hasParent = false
+    self.parseTriviaFunction = parseTriviaFunction
+  }
+
+  public convenience init() {
+    self.init(parseTriviaFunction: _defaultParseTriviaFunction(_:_:))
   }
 
   /// Copies a source buffer in to the memory this arena manages, and returns
@@ -145,6 +156,11 @@ public class SyntaxArena {
             sourceBufferContains(text.baseAddress!) ||
             allocator.contains(address: text.baseAddress!))
   }
+
+  @_spi(RawSyntax)
+  public func parseTrivia(source: SyntaxText, isTrailing: Bool) -> [RawTriviaPiece] {
+    return parseTriviaFunction(source, isTrailing)
+  }
 }
 
 extension SyntaxArena: Hashable {
@@ -184,3 +200,7 @@ extension SyntaxArena {
   //@available(*, deprecated, message: ".default SyntaxArena is subject to remove soon")
   public static let `default` = SyntaxArena()
 }
+
+private func _defaultParseTriviaFunction(_ source: SyntaxText, _ isTrailnig: Bool) -> [RawTriviaPiece] {
+  preconditionFailure("Trivia parsing not supported")
+}

Sources/SwiftSyntax/SyntaxClassifier.swift

@@ -276,9 +276,9 @@ fileprivate struct FastTokenSequence: Sequence {
 /// Provides a sequence of `SyntaxClassifiedRange`s for a token.
 fileprivate struct TokenClassificationIterator: IteratorProtocol {
   enum State {
-    case atLeadingTrivia(RawTriviaPieceBuffer, Int)
+    case atLeadingTrivia([RawTriviaPiece], Int)
     case atTokenText
-    case atTrailingTrivia(RawTriviaPieceBuffer, Int)
+    case atTrailingTrivia([RawTriviaPiece], Int)
   }
 
   let token: AbsoluteNode

Sources/SwiftSyntax/Trivia.swift.gyb

@@ -70,7 +70,7 @@ extension TriviaPiece: CustomDebugStringConvertible {
       return "${trivia.lower_name}s(\(data))"
 %   else:
     case .${trivia.lower_name}(let name):
-      return "${trivia.lower_name}(\(name))"
+      return "${trivia.lower_name}(\(name.debugDescription))"
 %   end
 % end
     }
@@ -263,6 +263,11 @@ extension RawTriviaPiece: TextOutputStreamable {
     TriviaPiece(raw: self).write(to: &target)
   }
 }
+extension RawTriviaPiece: CustomDebugStringConvertible {
+  public var debugDescription: String {
+    TriviaPiece(raw: self).debugDescription
+  }
+}
 
 extension TriviaPiece {
   init(raw: RawTriviaPiece) {

Sources/SwiftSyntax/gyb_generated/Trivia.swift

@@ -108,17 +108,17 @@ extension TriviaPiece: CustomDebugStringConvertible {
     case .carriageReturnLineFeeds(let data):
       return "carriageReturnLineFeeds(\(data))"
     case .lineComment(let name):
-      return "lineComment(\(name))"
+      return "lineComment(\(name.debugDescription))"
     case .blockComment(let name):
-      return "blockComment(\(name))"
+      return "blockComment(\(name.debugDescription))"
     case .docLineComment(let name):
-      return "docLineComment(\(name))"
+      return "docLineComment(\(name.debugDescription))"
     case .docBlockComment(let name):
-      return "docBlockComment(\(name))"
+      return "docBlockComment(\(name.debugDescription))"
     case .unexpectedText(let name):
-      return "unexpectedText(\(name))"
+      return "unexpectedText(\(name.debugDescription))"
     case .shebang(let name):
-      return "shebang(\(name))"
+      return "shebang(\(name.debugDescription))"
     }
   }
 }
@@ -415,6 +415,11 @@ extension RawTriviaPiece: TextOutputStreamable {
     TriviaPiece(raw: self).write(to: &target)
   }
 }
+extension RawTriviaPiece: CustomDebugStringConvertible {
+  public var debugDescription: String {
+    TriviaPiece(raw: self).debugDescription
+  }
+}
 
 extension TriviaPiece {
   init(raw: RawTriviaPiece) {