Reimplement SyntaxVisitor to improve visiting performance.

This implementation allows users to specify interested SyntaxNodeKind
and TokenKind; thus we can traverse the RawSyntax Tree instead of
SyntaxTree and realize only the interested RawSyntax nodes into
SyntaxNodes.

Author: nkcsgexi

Sources/SwiftSyntax/RawSyntax.swift

@@ -466,3 +466,35 @@ extension SyntaxNodeId: ByteTreeScalarDecodable {
     return SyntaxNodeId(rawValue: UInt(rawValue))
   }
 }
+
+extension RawSyntax {
+  func accept(_ visitor: RawSyntaxVisitor) {
+    defer { visitor.moveUp() }
+    guard isPresent else { return }
+    switch data {
+    case .node(let kind,let layout):
+      let shouldVisit = visitor.shouldVisit(kind)
+      if shouldVisit {
+
+        // Visit this node realizes a syntax node.
+        visitor.visitCurrent()
+      }
+
+      for (offset, element) in layout.enumerated() {
+        guard let element = element else { continue }
+
+        // Teach the visitor to navigate to this child.
+        visitor.addChildIdx(offset)
+        element.accept(visitor)
+      }
+      if shouldVisit {
+        visitor.visitPost()
+      }
+    case .token(let kind, _, _):
+      if visitor.shouldVisit(kind) {
+        visitor.visitCurrent()
+        visitor.visitPost()
+      }
+    }
+  }
+}

Sources/SwiftSyntax/SyntaxKind.swift.gyb

@@ -24,7 +24,7 @@
 import Foundation
 
 /// Enumerates the known kinds of Syntax represented in the Syntax tree.
-internal enum SyntaxKind: String, Codable {
+public enum SyntaxKind: String, Codable {
   case token = "Token"
   case unknown = "Unknown"
 % for node in SYNTAX_NODES:

Sources/SwiftSyntax/SyntaxRewriter.swift.gyb

@@ -98,12 +98,17 @@ open class SyntaxVisitor {
     public init() {}
 % for node in SYNTAX_NODES:
 %   if is_visitable(node):
-  open func visit(_ node: ${node.name}) {
-    visitChildren(node)
-  }
+  open func visit(_ node: ${node.name}) {}
 %   end
 % end
 
+  open func shouldVisit(_ kind: SyntaxKind) -> Bool {
+    return true
+  }
+  open func shouldVisit(_ kind: TokenKind) -> Bool {
+    return true
+  }
+
   open func visit(_ token: TokenSyntax) {}
 
   /// The function called before visiting the node and its descendents.
@@ -115,20 +120,79 @@ open class SyntaxVisitor {
   open func visitPost(_ node: Syntax) {}
 
   public func visit(_ node: Syntax) {
-    visitPre(node)
-    defer { visitPost(node) }
     switch node.raw.kind {
     case .token: visit(node as! TokenSyntax)
 % for node in SYNTAX_NODES:
 %   if is_visitable(node):
     case .${node.swift_syntax_kind}: visit(node as! ${node.name})
 %   end
 % end
-    default: visitChildren(node)
+    default: return
     }
   }
+}
+
+/// The raw syntax walker traverses the raw syntax tree to find
+/// node kinds the SyntaxVisitor is interested and feed these syntax nodes to
+/// SyntaxVisitor.
+/// By traversing the raw syntax tree, we avoid realizing syntax nodes that're
+/// not interesting to users' SyntaxVisitor.
+class RawSyntaxVisitor {
+  private let visitor: SyntaxVisitor
+  private var childIdxChain: [Int]
+  private var nodeChain: [Syntax]
+
+  required init(_ visitor: SyntaxVisitor, _ root: Syntax) {
+    self.visitor = visitor
+    self.nodeChain = [root]
+    self.childIdxChain = []
+  }
+
+  convenience init(_ another: RawSyntaxVisitor, _ root: Syntax) {
+    self.init(another.visitor, root)
+  }
+
+  func shouldVisit(_ kind: SyntaxKind) -> Bool {
+    return visitor.shouldVisit(kind)
+  }
+
+  func shouldVisit(_ kind: TokenKind) -> Bool {
+    return visitor.shouldVisit(kind)
+  }
+
+  func addChildIdx(_ idx: Int) {
+    childIdxChain.append(idx)
+  }
+
+  func moveUp() {
+    if childIdxChain.isEmpty {
+      nodeChain.removeLast()
+    } else {
+      childIdxChain.removeLast()
+    }
+  }
+
+  func visitPost() {
+    visitor.visitPost(nodeChain.last!)
+  }
+
+  // The current raw syntax node is interesting for the user, so realize a
+  // correponding syntax node and feed it into the visitor.
+  func visitCurrent() {
+    for id in childIdxChain {
+      nodeChain.append(nodeChain.last!.child(at: id)!)
+    }
+    childIdxChain.removeAll(keepingCapacity: true)
+    let node = nodeChain.last!
+    visitor.visitPre(node)
+    visitor.visit(node)
+  }
+}
+
+extension Syntax {
+  public func walk(_ visitor: SyntaxVisitor) {
 
-  func visitChildren(_ node: Syntax) {
-    node.children.forEach { visit($0) }
+    // Traverse the raw syntax tree by using the current syntax node as root.
+    data.raw.accept(RawSyntaxVisitor(visitor, self))
   }
 }