Separate SyntaxArena for parsing

Sources/SwiftParser/Parser.swift

@@ -118,7 +118,7 @@ extension Parser {
 /// operates on a copy of the lexical stream, no input tokens are lost..
 public struct Parser: TokenConsumer {
   @_spi(RawSyntax)
-  public var arena: SyntaxArena
+  public var arena: ParsingSyntaxArena
   /// A view of the sequence of lexemes in the input.
   var lexemes: Lexer.LexemeSequence
   /// The current token. If there was no input, this token will have a kind of `.eof`.
@@ -153,7 +153,7 @@ public struct Parser: TokenConsumer {
   ///            arena is created automatically, and `input` copied into the
   ///            arena. If non-`nil`, `input` must be the registered source
   ///            buffer of `arena` or a slice of the source buffer.
-  public init(_ input: UnsafeBufferPointer<UInt8>, maximumNestingLevel: Int? = nil, arena: SyntaxArena? = nil) {
+  public init(_ input: UnsafeBufferPointer<UInt8>, maximumNestingLevel: Int? = nil, arena: ParsingSyntaxArena? = nil) {
     self.maximumNestingLevel = maximumNestingLevel ?? Self.defaultMaximumNestingLevel
 
     var sourceBuffer: UnsafeBufferPointer<UInt8>
@@ -162,7 +162,7 @@ public struct Parser: TokenConsumer {
       sourceBuffer = input
       assert(arena.contains(text: SyntaxText(baseAddress: input.baseAddress, count: input.count)))
    } else {
-     self.arena = SyntaxArena(
+     self.arena = ParsingSyntaxArena(
       parseTriviaFunction: TriviaParser.parseTrivia(_:position:))
      sourceBuffer = self.arena.internSourceBuffer(input)
     }

Sources/SwiftSyntax/BumpPtrAllocator.swift

@@ -15,10 +15,12 @@
 public class BumpPtrAllocator {
   typealias Slab = UnsafeMutableRawBufferPointer
 
-  static private var SLAB_SIZE: Int = 4096
   static private var GLOWTH_DELAY: Int = 128
   static private var SLAB_ALIGNMENT: Int = 8
 
+  /// Initial slab size.
+  private var slabSize: Int
+
   private var slabs: [Slab]
   /// Pair of pointers in the current slab.
   /// - pointer: Points to the next unused address in `slabs.last`.
@@ -30,7 +32,8 @@ public class BumpPtrAllocator {
   private var customSizeSlabs: [Slab]
   private var _totalBytesAllocated: Int
 
-  public init() {
+  public init(slabSize: Int) {
+    self.slabSize = slabSize
     slabs = []
     current = nil
     customSizeSlabs = []
@@ -50,13 +53,13 @@ public class BumpPtrAllocator {
   }
 
   /// Calculate the size of the slab at the index.
-  private static func slabSize(at index: Int) -> Int {
+  private func slabSize(at index: Int) -> Int {
     // Double the slab size every 'GLOWTH_DELAY' slabs.
-    return SLAB_SIZE * (1 << min(30, index / GLOWTH_DELAY))
+    return self.slabSize * (1 << min(30, index / Self.GLOWTH_DELAY))
   }
 
   private func startNewSlab() {
-    let newSlabSize = Self.slabSize(at: slabs.count)
+    let newSlabSize = self.slabSize(at: slabs.count)
     let newSlab = Slab.allocate(
       byteCount: newSlabSize, alignment: Self.SLAB_ALIGNMENT)
     let pointer = newSlab.baseAddress!
@@ -103,7 +106,7 @@ public class BumpPtrAllocator {
     }
 
     // If the size is too big, allocate a dedicated slab for it.
-    if byteCount >= Self.SLAB_SIZE {
+    if byteCount >= self.slabSize {
       let customSlab = Slab.allocate(
         byteCount: byteCount, alignment: alignment)
       customSizeSlabs.append(customSlab)

Sources/SwiftSyntax/Raw/RawSyntax.swift

@@ -459,6 +459,8 @@ extension RawSyntax {
   ) -> 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/RawSyntaxTokenView.swift

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

Sources/SwiftSyntax/SyntaxArena.swift

@@ -11,74 +11,38 @@
 //===----------------------------------------------------------------------===//
 
 public class SyntaxArena {
-
-  @_spi(RawSyntax)
-  public typealias ParseTriviaFunction = (_ source: SyntaxText, _ position: TriviaPosition) -> [RawTriviaPiece]
-
   /// Bump-pointer allocator for all "intern" methods.
-  private let allocator: BumpPtrAllocator
-  /// Source file buffer the Syntax tree represents.
-  private var sourceBuffer: UnsafeBufferPointer<UInt8>
+  fileprivate let allocator: BumpPtrAllocator
 
   /// If the syntax tree that’s allocated in this arena references nodes from
   /// other arenas, `childRefs` contains references to the arenas. Child arenas
   /// are retained in `addChild()` and are released in `deinit`.
   private var childRefs: Set<SyntaxArenaRef>
-  private var parseTriviaFunction: ParseTriviaFunction
 
 #if DEBUG
   /// 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
 #endif
 
-  @_spi(RawSyntax)
-  public init(parseTriviaFunction: @escaping ParseTriviaFunction) {
-    self.allocator = BumpPtrAllocator()
+  public convenience init() {
+    self.init(slabSize: 128)
+  }
+
+  fileprivate init(slabSize: Int) {
+    self.allocator = BumpPtrAllocator(slabSize: slabSize)
     self.childRefs = []
-    self.sourceBuffer = .init(start: nil, count: 0)
-    self.parseTriviaFunction = parseTriviaFunction
 #if DEBUG
     self.hasParent = false
 #endif
   }
 
-  public convenience init() {
-    self.init(parseTriviaFunction: _defaultParseTriviaFunction(_:_:))
-  }
-
   deinit {
     for child in childRefs {
       child.release()
     }
   }
 
-  /// Copies a source buffer in to the memory this arena manages, and returns
-  /// the interned buffer.
-  ///
-  /// The interned buffer is guaranteed to be null-terminated.
-  /// `contains(address _:)` is faster if the address is inside the memory
-  /// range this function returned.
-  public func internSourceBuffer(_ buffer: UnsafeBufferPointer<UInt8>) -> UnsafeBufferPointer<UInt8> {
-    let allocated = allocator.allocate(
-      UInt8.self, count: buffer.count + /* for NULL */1)
-    precondition(sourceBuffer.baseAddress == nil, "SourceBuffer should only be set once.")
-    _ = allocated.initialize(from: buffer)
-
-    // NULL terminate.
-    allocated.baseAddress!.advanced(by: buffer.count).initialize(to: 0)
-
-    sourceBuffer = UnsafeBufferPointer(start: allocated.baseAddress!, count: buffer.count)
-    return sourceBuffer
-  }
-
-  /// Checks if the given memory address is inside the memory range returned
-  /// from `internSourceBuffer(_:)` method.
-  func sourceBufferContains(_ address: UnsafePointer<UInt8>) -> Bool {
-    guard let sourceStart = sourceBuffer.baseAddress else { return false }
-    return sourceStart <= address && address < sourceStart.advanced(by: sourceBuffer.count)
-  }
-
   /// 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?> {
@@ -173,10 +137,63 @@ public class SyntaxArena {
   @_spi(RawSyntax)
   public func contains(text: SyntaxText) -> Bool {
     return (text.isEmpty ||
-            sourceBufferContains(text.baseAddress!) ||
             allocator.contains(address: text.baseAddress!))
   }
+}
+
+/// SyntaxArena for parsing.
+public class ParsingSyntaxArena: SyntaxArena {
+  @_spi(RawSyntax)
+  public typealias ParseTriviaFunction = (_ source: SyntaxText, _ position: TriviaPosition) -> [RawTriviaPiece]
+
+  /// Source file buffer the Syntax tree represents.
+  private var sourceBuffer: UnsafeBufferPointer<UInt8>
+
+  /// Function to parse trivia.
+  private var parseTriviaFunction: ParseTriviaFunction
+
+  @_spi(RawSyntax)
+  public init(parseTriviaFunction: @escaping ParseTriviaFunction) {
+    self.sourceBuffer = .init(start: nil, count: 0)
+    self.parseTriviaFunction = parseTriviaFunction
+    super.init(slabSize: 4096)
+  }
 
+  /// Copies a source buffer in to the memory this arena manages, and returns
+  /// the interned buffer.
+  ///
+  /// The interned buffer is guaranteed to be null-terminated.
+  /// `contains(address _:)` is faster if the address is inside the memory
+  /// range this function returned.
+  public func internSourceBuffer(_ buffer: UnsafeBufferPointer<UInt8>) -> UnsafeBufferPointer<UInt8> {
+    let allocated = allocator.allocate(
+      UInt8.self, count: buffer.count + /* for NULL */1)
+    precondition(sourceBuffer.baseAddress == nil, "SourceBuffer should only be set once.")
+    _ = allocated.initialize(from: buffer)
+
+    // NULL terminate.
+    allocated.baseAddress!.advanced(by: buffer.count).initialize(to: 0)
+
+    sourceBuffer = UnsafeBufferPointer(start: allocated.baseAddress!, count: buffer.count)
+    return sourceBuffer
+  }
+
+  @_spi(RawSyntax)
+  public override func contains(text: SyntaxText) -> Bool {
+    if let addr = text.baseAddress, self.sourceBufferContains(addr) {
+      return true
+    }
+    return super.contains(text: text)
+  }
+
+  /// Checks if the given memory address is inside the memory range returned
+  /// from `internSourceBuffer(_:)` method.
+  func sourceBufferContains(_ address: UnsafePointer<UInt8>) -> Bool {
+    guard let sourceStart = sourceBuffer.baseAddress else { return false }
+    return sourceStart <= address && address < sourceStart.advanced(by: sourceBuffer.count)
+  }
+
+  /// Parse `source` into a list of `RawTriviaPiece` using `parseTriviaFunction`.
   @_spi(RawSyntax)
   public func parseTrivia(source: SyntaxText, position: TriviaPosition) -> [RawTriviaPiece] {
     return self.parseTriviaFunction(source, position)
@@ -217,7 +234,3 @@ struct SyntaxArenaRef: Hashable {
     return lhs._value.toOpaque() == rhs._value.toOpaque()
   }
 }
-
-private func _defaultParseTriviaFunction(_ source: SyntaxText, _ position: TriviaPosition) -> [RawTriviaPiece] {
-  preconditionFailure("Trivia parsing not supported")
-}

Tests/SwiftSyntaxTest/BumpPtrAllocatorTests.swift

@@ -16,7 +16,7 @@ import XCTest
 final class BumpPtrAllocatorTests: XCTestCase {
 
     func testBasic() {
-      let allocator = BumpPtrAllocator()
+      let allocator = BumpPtrAllocator(slabSize: 4096)
 
       let byteBuffer = allocator.allocate(byteCount: 42, alignment: 4)
       XCTAssertNotNil(byteBuffer.baseAddress)

Tests/SwiftSyntaxTest/RawSyntaxTests.swift

@@ -104,7 +104,7 @@ final class RawSyntaxTests: XCTestCase {
       return [.unexpectedText(source)]
     }
 
-    withExtendedLifetime(SyntaxArena(parseTriviaFunction: dummyParseToken)) { arena in
+    withExtendedLifetime(ParsingSyntaxArena(parseTriviaFunction: dummyParseToken)) { arena in
       let ident = RawTokenSyntax(
         kind: .identifier, wholeText: arena.intern("\nfoo "), textRange: 1..<4,
         presence: .present,