Merge pull request #2939 from rintaro/perf-token-sequence

[Perf] Optimize TokenSequence

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``.