Merge pull request #2893 from rintaro/perf-syntaxbytes

[Perf] Optimize SourceLocationConverter

Author: rintaro

Sources/SwiftSyntax/Raw/RawSyntax.swift

@@ -361,16 +361,23 @@ extension RawSyntax {
 }
 
 extension RawTriviaPiece {
-  func withSyntaxText(body: (SyntaxText) throws -> Void) rethrows {
+  /// Call `body` with the syntax text of this trivia piece.
+  ///
+  /// If `isEphemeral` is `true`, the ``SyntaxText`` argument is only guaranteed
+  /// to be valid within the call.
+  func withSyntaxText(body: (SyntaxText, _ isEphemeral: Bool) throws -> Void) rethrows {
     if let syntaxText = storedText {
-      try body(syntaxText)
+      try body(syntaxText, /*isEphemeral*/ false)
       return
     }
 
     var description = ""
     write(to: &description)
     try description.withUTF8 { buffer in
-      try body(SyntaxText(baseAddress: buffer.baseAddress, count: buffer.count))
+      try body(
+        SyntaxText(baseAddress: buffer.baseAddress, count: buffer.count),
+        /*isEphemeral*/ true
+      )
     }
   }
 }
@@ -382,21 +389,21 @@ extension RawSyntax {
   /// Unlike `description`, this provides a source-accurate representation
   /// even in the presence of malformed UTF-8 in the input source.
   ///
-  /// The ``SyntaxText`` arguments passed to the visitor are only guaranteed
-  /// to be valid within that call. It is unsafe to escape the `SyntaxValue`
-  /// values outside of the closure.
-  public func withEachSyntaxText(body: (SyntaxText) throws -> Void) rethrows {
+  /// If `isEphemeral` is `true`, the ``SyntaxText`` arguments passed to the
+  /// visitor are only guaranteed to be valid within that call. Otherwise, they
+  /// are valid as long as the raw syntax is alive.
+  public func withEachSyntaxText(body: (SyntaxText, _ isEphemeral: Bool) throws -> Void) rethrows {
     switch rawData.payload {
     case .parsedToken(let dat):
       if dat.presence == .present {
-        try body(dat.wholeText)
+        try body(dat.wholeText, /*isEphemeral*/ false)
       }
     case .materializedToken(let dat):
       if dat.presence == .present {
         for p in dat.leadingTrivia {
           try p.withSyntaxText(body: body)
         }
-        try body(dat.tokenText)
+        try body(dat.tokenText, /*isEphemeral*/ false)
         for p in dat.trailingTrivia {
           try p.withSyntaxText(body: body)
         }
@@ -412,9 +419,20 @@ extension RawSyntax {
   /// source even in the presence of invalid UTF-8.
   public var syntaxTextBytes: [UInt8] {
     var result: [UInt8] = []
-    withEachSyntaxText { syntaxText in
-      result.append(contentsOf: syntaxText)
+    var buf: SyntaxText = ""
+    withEachSyntaxText { syntaxText, isEphemeral in
+      if isEphemeral {
+        result.append(contentsOf: buf)
+        result.append(contentsOf: syntaxText)
+        buf = ""
+      } else if let base = buf.baseAddress, base + buf.count == syntaxText.baseAddress {
+        buf = SyntaxText(baseAddress: base, count: buf.count + syntaxText.count)
+      } else {
+        result.append(contentsOf: buf)
+        buf = syntaxText
+      }
     }
+    result.append(contentsOf: buf)
     return result
   }
 }

Sources/SwiftSyntax/SourceLocation.swift

@@ -522,38 +522,33 @@ fileprivate extension SyntaxText {
     prefix: SourceLength = .zero,
     body: (SourceLength) -> ()
   ) -> SourceLength {
-    //    let startIndex = utf8.startIndex
-    //    let endIndex = utf8.endIndex
     var curIdx = startIndex
-    var lineLength = prefix
-    let advanceLengthByOne = { () -> () in
-      lineLength += SourceLength(utf8Length: 1)
-      curIdx = self.index(after: curIdx)
-    }
+    let endIdx = endIndex
+    var lineStart = curIdx - prefix.utf8Length
 
-    while curIdx < endIndex {
-      let char = self[curIdx]
-      advanceLengthByOne()
+    while curIdx != endIdx {
+      let chr = self[curIdx]
+      curIdx &+= 1
 
       /// From https://docs.swift.org/swift-book/ReferenceManual/LexicalStructure.html#grammar_line-break
       /// * line-break → U+000A
       /// * line-break → U+000D
       /// * line-break → U+000D followed by U+000A
-      let isNewline = { () -> Bool in
-        if char == 10 { return true }
-        if char == 13 {
-          if curIdx < endIndex && self[curIdx] == 10 { advanceLengthByOne() }
-          return true
+      switch chr {
+      case UInt8(ascii: "\n"):
+        break
+      case UInt8(ascii: "\r"):
+        if curIdx != endIdx && self[curIdx] == UInt8(ascii: "\n") {
+          curIdx &+= 1
         }
-        return false
-      }
-
-      if isNewline() {
-        body(lineLength)
-        lineLength = .zero
+        break
+      default:
+        continue
       }
+      body(SourceLength(utf8Length: curIdx - lineStart))
+      lineStart = curIdx
     }
-    return lineLength
+    return SourceLength(utf8Length: curIdx - lineStart)
   }
 
   func containsSwiftNewline() -> Bool {

Sources/SwiftSyntax/Syntax.swift

@@ -365,20 +365,6 @@ extension Syntax: Identifiable {
 }
 
 extension Syntax {
-  /// Enumerate all of the syntax text present in this node, and all
-  /// of its children, to give a source-accurate view of the bytes.
-  ///
-  /// Unlike `description`, this provides a source-accurate representation
-  /// even in the presence of malformed UTF-8 in the input source.
-  ///
-  /// The ``SyntaxText`` arguments passed to the visitor are only guaranteed
-  /// to be valid within that call. It is unsafe to escape the `SyntaxValue`
-  /// values outside of the closure.
-  @_spi(RawSyntax)
-  public func withEachSyntaxText(body: (SyntaxText) throws -> Void) rethrows {
-    try raw.withEachSyntaxText(body: body)
-  }
-
   /// Retrieve the syntax text as an array of bytes that models the input
   /// source even in the presence of invalid UTF-8.
   public var syntaxTextBytes: [UInt8] {

Sources/SwiftSyntax/SyntaxArenaAllocatedBuffer.swift

@@ -17,7 +17,8 @@
 /// reference its contents, we know that the pointer's contents won't get
 /// deallocated while being accessed and thus we can add an unchecked `Sendable`
 /// conformance.
-@_spi(RawSyntax) public struct SyntaxArenaAllocatedPointer<Element: Sendable>: @unchecked Sendable {
+@_spi(RawSyntax)
+public struct SyntaxArenaAllocatedPointer<Element: Sendable>: @unchecked Sendable {
   private let pointer: UnsafePointer<Element>
 
   /// Create a pointer from an `UnsafePointer` that was allocated inside a
@@ -61,22 +62,21 @@ public struct SyntaxArenaAllocatedBufferPointer<Element: Sendable>: RandomAccess
   /// - Important: The client needs to ensure sure that the buffer is indeed
   ///   allocated by a ``SyntaxArena`` and that the ``SyntaxArena`` will outlive
   ///   any users of this ``SyntaxArenaAllocatedBufferPointer``.
-  @_spi(RawSyntax) public init(_ buffer: UnsafeBufferPointer<Element>) {
+  public init(_ buffer: UnsafeBufferPointer<Element>) {
     self.buffer = buffer
   }
 
-  @_spi(RawSyntax)
   public subscript<RangeType: RangeExpression<Int>>(
     range: RangeType
   ) -> SyntaxArenaAllocatedBufferPointer<Element> {
     return SyntaxArenaAllocatedBufferPointer(UnsafeBufferPointer(rebasing: self.buffer[range]))
   }
 
-  @_spi(RawSyntax) public subscript(_ index: Int) -> Element {
+  public subscript(_ index: Int) -> Element {
     return self.buffer[index]
   }
 
-  @_spi(RawSyntax) public func makeIterator() -> UnsafeBufferPointer<Element>.Iterator {
+  public func makeIterator() -> UnsafeBufferPointer<Element>.Iterator {
     return buffer.makeIterator()
   }
 
@@ -111,4 +111,8 @@ public struct SyntaxArenaAllocatedBufferPointer<Element: Sendable>: RandomAccess
   public func withContiguousStorageIfAvailable<R>(_ body: (UnsafeBufferPointer<Element>) throws -> R) rethrows -> R? {
     try body(buffer)
   }
+
+  public func _copyContents(initializing ptr: UnsafeMutableBufferPointer<Element>) -> (Iterator, Int) {
+    buffer._copyContents(initializing: ptr)
+  }
 }

Sources/SwiftSyntax/SyntaxText.swift

@@ -47,10 +47,11 @@ import Musl
 /// replacement character (`\u{FFFD}`).
 @_spi(RawSyntax)
 public struct SyntaxText: Sendable {
-  var buffer: SyntaxArenaAllocatedBufferPointer<UInt8>
+  public typealias Buffer = SyntaxArenaAllocatedBufferPointer<UInt8>
+  var buffer: Buffer
 
   /// Construct a ``SyntaxText`` whose text is represented by the given `buffer`.
-  public init(buffer: SyntaxArenaAllocatedBufferPointer<UInt8>) {
+  public init(buffer: Buffer) {
     self.buffer = buffer
   }
 
@@ -180,6 +181,18 @@ extension SyntaxText: RandomAccessCollection {
   public subscript(index: Index) -> Element {
     get { return buffer[index] }
   }
+
+  public func makeIterator() -> Buffer.Iterator {
+    buffer.makeIterator()
+  }
+
+  public func withContiguousStorageIfAvailable<R>(_ body: (UnsafeBufferPointer<Element>) throws -> R) rethrows -> R? {
+    try buffer.withContiguousStorageIfAvailable(body)
+  }
+
+  public func _copyContents(initializing ptr: UnsafeMutableBufferPointer<Element>) -> (Iterator, Int) {
+    buffer._copyContents(initializing: ptr)
+  }
 }
 
 extension SyntaxText: Hashable {