[Perf] Optimize TokenSequence

i.e. 'previousToken', 'nextToken', 'firstToken', and 'lastToken'

Iterate over 'layoutBuffer' directly to avoid 'SyntaxChildren' which
implies retaining 'SyntaxDataArena' for each element.

This way, it only retains the found 'TokenSyntax', nothing else.

Author: rintaro

Diff for: Sources/SwiftSyntax/Syntax.swift

@@ -36,7 +36,7 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
     self.dataRef = dataRef
   }
 
-  private var data: SyntaxData {
+  var data: SyntaxData {
     @_transparent unsafeAddress { dataRef.pointer }
   }
 

Diff for: Sources/SwiftSyntax/SyntaxProtocol.swift

@@ -253,20 +253,7 @@ extension SyntaxProtocol {
   /// Recursively walks through the tree to find the token semantically before
   /// this node.
   public func previousToken(viewMode: SyntaxTreeViewMode) -> TokenSyntax? {
-    guard let parent = self.parent else {
-      return nil
-    }
-    let siblings = parent.children(viewMode: viewMode)
-    // `self` could be a missing node at index 0 and `viewMode` be `.sourceAccurate`.
-    // In that case `siblings` skips over the missing `self` node and has a `startIndex > 0`.
-    if siblings.startIndex < self.indexInParent {
-      for child in siblings[..<self.indexInParent].reversed() {
-        if let token = child.lastToken(viewMode: viewMode) {
-          return token
-        }
-      }
-    }
-    return parent.previousToken(viewMode: viewMode)
+    return self._syntaxNode.previousToken(viewMode: viewMode)
   }
 
   @available(*, deprecated, message: "Use nextToken(viewMode:) instead")
@@ -277,16 +264,7 @@ extension SyntaxProtocol {
   /// Recursively walks through the tree to find the next token semantically
   /// after this node.
   public func nextToken(viewMode: SyntaxTreeViewMode) -> TokenSyntax? {
-    guard let parent = self.parent else {
-      return nil
-    }
-    let siblings = parent.children(viewMode: viewMode)
-    for child in siblings[siblings.index(after: self.indexInParent)...] {
-      if let token = child.firstToken(viewMode: viewMode) {
-        return token
-      }
-    }
-    return parent.nextToken(viewMode: viewMode)
+    return self._syntaxNode.nextToken(viewMode: viewMode)
   }
 
   @available(*, deprecated, message: "Use firstToken(viewMode: .sourceAccurate) instead")
@@ -296,17 +274,7 @@ extension SyntaxProtocol {
 
   /// Returns the first token node that is part of this syntax node.
   public func firstToken(viewMode: SyntaxTreeViewMode) -> TokenSyntax? {
-    guard viewMode.shouldTraverse(node: raw) else { return nil }
-    if let token = _syntaxNode.as(TokenSyntax.self) {
-      return token
-    }
-
-    for child in children(viewMode: viewMode) {
-      if let token = child.firstToken(viewMode: viewMode) {
-        return token
-      }
-    }
-    return nil
+    return self._syntaxNode.firstToken(viewMode: viewMode)
   }
 
   @available(*, deprecated, message: "Use lastToken(viewMode: .sourceAccurate) instead")
@@ -316,17 +284,7 @@ extension SyntaxProtocol {
 
   /// Returns the last token node that is part of this syntax node.
   public func lastToken(viewMode: SyntaxTreeViewMode) -> TokenSyntax? {
-    guard viewMode.shouldTraverse(node: raw) else { return nil }
-    if let token = _syntaxNode.as(TokenSyntax.self) {
-      return token
-    }
-
-    for child in children(viewMode: viewMode).reversed() {
-      if let tok = child.lastToken(viewMode: viewMode) {
-        return tok
-      }
-    }
-    return nil
+    return self._syntaxNode.lastToken(viewMode: viewMode)
   }
 
   /// Sequence of tokens that are part of this Syntax node.

Diff for: Sources/SwiftSyntax/TokenSequence.swift

@@ -10,6 +10,70 @@
 //
 //===----------------------------------------------------------------------===//
 
+extension Syntax {
+  /// Implementation of 'SyntaxProtocol.previousToken(viewMode:)'
+  func previousToken(viewMode: SyntaxTreeViewMode) -> TokenSyntax? {
+    guard let parentDataRef = data.parent else {
+      return nil
+    }
+    for case let childDataRef? in arena.layout(for: parentDataRef)[..<layoutIndexInParent].reversed() {
+      if let token = Syntax(arena: arena, dataRef: childDataRef).lastToken(viewMode: viewMode) {
+        return token
+      }
+    }
+    return Syntax(arena: arena, dataRef: parentDataRef).previousToken(viewMode: viewMode)
+  }
+
+  /// Implementation of 'SyntaxProtocol.nextToken(viewMode:)'
+  func nextToken(viewMode: SyntaxTreeViewMode) -> TokenSyntax? {
+    guard let parentDataRef = data.parent else {
+      return nil
+    }
+    for case let childDataRef? in arena.layout(for: parentDataRef)[(layoutIndexInParent &+ 1)...] {
+      if let token = Syntax(arena: arena, dataRef: childDataRef).firstToken(viewMode: viewMode) {
+        return token
+      }
+    }
+    return Syntax(arena: arena, dataRef: parentDataRef).nextToken(viewMode: viewMode)
+  }
+
+  /// Implementation of 'SyntaxProtocol.firstToken(viewMode:)'
+  ///
+  /// - Note: Can't use 'RawSyntax.firstToken(viewMode:)' because it loses absolute info.
+  func firstToken(viewMode: SyntaxTreeViewMode) -> TokenSyntax? {
+    guard viewMode.shouldTraverse(node: raw) else {
+      return nil
+    }
+    if raw.isToken {
+      return TokenSyntax(self)!
+    }
+    for case let childDataRef? in layoutBuffer {
+      if let token = Syntax(arena: arena, dataRef: childDataRef).firstToken(viewMode: viewMode) {
+        return token
+      }
+    }
+    return nil
+  }
+
+  /// Implementation of 'SyntaxProtocol.lastToken(viewMode:)'
+  ///
+  /// - Note: Can't use 'RawSyntax.lastToken(viewMode:)' because it loses absolute info.
+  func lastToken(viewMode: SyntaxTreeViewMode) -> TokenSyntax? {
+    guard viewMode.shouldTraverse(node: raw) else {
+      return nil
+    }
+    if raw.isToken {
+      return TokenSyntax(self)!
+    }
+    for case let childDataRef? in layoutBuffer.reversed() {
+      if let token = Syntax(arena: arena, dataRef: childDataRef).lastToken(viewMode: viewMode) {
+        return token
+      }
+    }
+    return nil
+  }
+}
+
 /// Sequence of tokens that are part of the provided Syntax node.
 public struct TokenSequence: Sequence, Sendable {
   /// Iterates over a ``TokenSequence``.