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,43 @@ extension SyntaxNodeId: ByteTreeScalarDecodable {
     return SyntaxNodeId(rawValue: UInt(rawValue))
   }
 }
+
+extension RawSyntax {
+  func accept(_ visitor: RawSyntaxVisitor) {
+    guard isPresent else { return }
+    switch data {
+    case .node(let kind,let layout):
+      let shouldVisit = visitor.shouldVisit(kind)
+      var visitedNode: Syntax?
+      var forwardVisitor = visitor
+      if shouldVisit {
+
+        // Visit this node realizes a syntax node, so we can re-create a forward
+        // visitor rooted on this node.
+        visitedNode = visitor.visitCurrent()
+        forwardVisitor = RawSyntaxVisitor(visitor, visitedNode!)
+      }
+
+      // Create a slot for child index
+      forwardVisitor.initializeChildIdx()
+      layout.forEach { l in
+
+        // Increment child index for next child to visit
+        forwardVisitor.incrementLastChildIdx()
+        l?.accept(forwardVisitor)
+      }
+
+      // Remove the slot for child index
+      forwardVisitor.removeLastChildIdx()
+      if shouldVisit {
+
+        // We finished visiting the current node.
+        forwardVisitor.visitPost()
+      }
+    case .token(let kind, _, _):
+      if visitor.shouldVisit(kind) {
+        _ = visitor.visitCurrent()
+      }
+    }
+  }
+}

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,80 @@ 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 let root: Syntax
+  private var childIdxChain: [Int]
+
+  required init(_ visitor: SyntaxVisitor, _ root: Syntax) {
+    self.visitor = visitor
+    self.root = 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 initializeChildIdx() {
+    childIdxChain.append(-1)
+  }
+
+  func incrementLastChildIdx() {
+    childIdxChain[childIdxChain.count - 1] += 1
+  }
+
+  func removeLastChildIdx() {
+    childIdxChain.removeLast()
+  }
+
+  // The current raw syntax node is interesting for the user, so realize a
+  // correponding syntax node and feed it into the visitor.
+  func visitCurrent() -> Syntax {
+    var node = root
+    for id in childIdxChain {
+      node = node.child(at: id)!
     }
+    visitor.visitPre(node)
+    visitor.visit(node)
+    return node
   }
 
-  func visitChildren(_ node: Syntax) {
-    node.children.forEach { visit($0) }
+  // Called when we finished traversing the tree.
+  func visitPost() {
+    visitor.visitPost(root)
+  }
+}
+
+extension Syntax {
+  public func walk(_ visitor: SyntaxVisitor) {
+
+    // Traverse the raw syntax tree by using the current syntax node as root.
+    data.raw.accept(RawSyntaxVisitor(visitor, self))
   }
 }