Reimplement SyntaxVisitor to improve visiting performance.

Sources/SwiftSyntax/RawSyntax.swift

@@ -466,3 +466,37 @@ 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)
+      var visitChildren = true
+      if shouldVisit {
+        // Visit this node realizes a syntax node.
+        visitor.visitPre()
+        visitChildren = visitor.visit()
+      }
+      if visitChildren {
+        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.visitPre()
+        _ = visitor.visit()
+        visitor.visitPost()
+      }
+    }
+  }
+}

Sources/SwiftSyntax/SyntaxClassifier.swift.gyb

@@ -50,7 +50,7 @@ fileprivate class _SyntaxClassifier: SyntaxVisitor {
     force: Bool = false
   ) {
     contextStack.append((classification: classification, force: force))
-    visit(node)
+    node.walk(self)
     contextStack.removeLast()
   }
 
@@ -71,6 +71,7 @@ fileprivate class _SyntaxClassifier: SyntaxVisitor {
   }
 
   override func visit(_ token: TokenSyntax) {
+    assert(token.isPresent)
     // FIXME: We need to come up with some way in which the SyntaxClassifier can
     // classify trivia (i.e. comments). In particular we need to be able to
     // look into the comments to find things like URLs or keywords like MARK.
@@ -94,12 +95,12 @@ fileprivate class _SyntaxClassifier: SyntaxVisitor {
 
 % for node in SYNTAX_NODES:
 %   if is_visitable(node):
-  override func visit(_ node: ${node.name}) {
+  override func visit(_ node: ${node.name}) -> Bool {
     if skipNodeIds.contains(node.raw.id) {
-      return
+      return false
     }
 %     if node.is_unknown() or node.is_syntax_collection():
-    super.visit(node)
+    return true
 %     else:
 %       for child in node.children:
 %         if child.is_optional:
@@ -109,7 +110,7 @@ fileprivate class _SyntaxClassifier: SyntaxVisitor {
             classification: .${child.classification.swift_name}, 
             force: ${"true" if child.force_classification else "false"})
 %           else:
-      visit(${child.swift_name})
+      ${child.swift_name}.walk(self)
 %           end
     }
 %         else:
@@ -118,12 +119,12 @@ fileprivate class _SyntaxClassifier: SyntaxVisitor {
           classification: .${child.classification.swift_name}, 
           force: ${"true" if child.force_classification else "false"})
 %           else:
-    visit(node.${child.swift_name})
+    node.${child.swift_name}.walk(self)
 %           end
 %         end
 %       end
+    return false
 %     end
-
   }
 %   end
 % end
@@ -139,7 +140,7 @@ public enum SyntaxClassifier {
   ) -> [TokenSyntax: SyntaxClassification] {
     let classifier = _SyntaxClassifier()
     classifier.skipNodeIds = skipNodes
-    classifier.visit(syntaxTree)
+    syntaxTree.walk(classifier)
     return classifier.classifications
   }
 }

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,23 @@ 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}) -> Bool {
+    return true
   }
 %   end
 % end
 
+  open func visit(_ node: UnknownSyntax) -> Bool {
+    return true
+  }
+
+  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.
@@ -114,21 +125,109 @@ open class SyntaxVisitor {
   ///   - node: the node we just finished visiting.
   open func visitPost(_ node: Syntax) {}
 
-  public func visit(_ node: Syntax) {
-    visitPre(node)
-    defer { visitPost(node) }
+  public func visit(_ node: Syntax) -> Bool {
     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})
+    case .${node.swift_syntax_kind}: return visit(node as! ${node.name})
 %   end
 % end
-    default: visitChildren(node)
+    case .unknown: return visit(node as! UnknownSyntax)
+    default: break
     }
+    return false
   }
+}
+
+
+/// A wrapper over Syntax. A syntax node is only realized when explicitly asked;
+/// otherwise the node is represented as a child index list from a realized
+/// ancestor.
+class PendingSyntaxNode {
+  let parent: PendingSyntaxNode?
+  private var kind: PendingSyntaxNodeKind
+
+  private enum PendingSyntaxNodeKind {
+    /// We already have a `Syntax` node realised for this node
+    case realized(node: Syntax)
+    /// This node does not have a `Syntax` node instantiated yet. If needed, we
+    /// need to compute it from its parent RawSyntax node
+    case virtual(index: Int)
+  }
+
+  var node: Syntax {
+    switch kind {
+    case .realized(let node):
+      return node
+    case .virtual(let index):
+      let _node = parent!.node.child(at: index)!
+      kind = .realized(node: _node)
+      return _node
+    }
+  }
+
+  init(_ root: Syntax) {
+    self.parent = nil
+    self.kind = .realized(node: root)
+  }
+
+  init(_ parent: PendingSyntaxNode, _ idx: Int) {
+    self.parent = parent
+    self.kind = .virtual(index: idx)
+  }
+}
+
+
+/// 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 currentNode: PendingSyntaxNode?
+
+  required init(_ visitor: SyntaxVisitor, _ root: Syntax) {
+    self.visitor = visitor
+    self.currentNode = PendingSyntaxNode(root)
+  }
+
+  func shouldVisit(_ kind: SyntaxKind) -> Bool {
+    return visitor.shouldVisit(kind)
+  }
+
+  func shouldVisit(_ kind: TokenKind) -> Bool {
+    return visitor.shouldVisit(kind)
+  }
+
+  func addChildIdx(_ idx: Int) {
+    currentNode = PendingSyntaxNode(currentNode!, idx)
+  }
+
+  func moveUp() {
+    currentNode = currentNode!.parent
+  }
+
+  func visitPre() {
+    visitor.visitPre(currentNode!.node)
+  }
+
+  func visitPost() {
+    visitor.visitPost(currentNode!.node)
+  }
+
+  // The current raw syntax node is interesting for the user, so realize a
+  // correponding syntax node and feed it into the visitor.
+  func visit() -> Bool {
+    return visitor.visit(currentNode!.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))
   }
 }

Sources/lit-test-helper/ClassifiedSyntaxTreePrinter.swift

@@ -52,7 +52,7 @@ class ClassifiedSyntaxTreePrinter: SyntaxVisitor {
 
   func print(tree: SourceFileSyntax) -> String {
     result = ""
-    visit(tree)
+    tree.walk(self)
     // Emit the last closing tag
     recordCurrentClassification(.none)
     return result

Tests/SwiftSyntaxTest/AbsolutePosition.swift

@@ -69,7 +69,7 @@ public class AbsolutePositionTestCase: XCTestCase {
             node.position.utf8Offset + node.leadingTrivia.byteSize)
         }
       }
-      Visitor().visit(parsed)
+      parsed.walk(Visitor())
     }())
   }
 

Tests/SwiftSyntaxTest/DiagnosticTest.swift

@@ -76,19 +76,20 @@ public class DiagnosticTestCase: XCTestCase {
         self.url = url
         self.engine = engine
       }
-      override func visit(_ function: FunctionDeclSyntax) {
+      override func visit(_ function: FunctionDeclSyntax) -> Bool {
         let startLoc = function.identifier.startLocation(in: url)
         let endLoc = function.endLocation(in: url)
         engine.diagnose(.badFunction(function.identifier), location: startLoc) {
           $0.highlight(function.identifier.sourceRange(in: self.url))
         }
         engine.diagnose(.endOfFunction(function.identifier), location: endLoc)
+        return true
       }
     }
 
     XCTAssertNoThrow(try {
       let file = try SyntaxTreeParser.parse(url)
-       Visitor(url: url, engine: engine).visit(file)
+       file.walk(Visitor(url: url, engine: engine))
     }())
 
      XCTAssertEqual(6, engine.diagnostics.count)

Tests/SwiftSyntaxTest/VisitorTest.swift

@@ -13,16 +13,16 @@ public class SyntaxVisitorTestCase: XCTestCase {
   public func testBasic() {
     class FuncCounter: SyntaxVisitor {
       var funcCount = 0
-      override func visit(_ node: FunctionDeclSyntax) {
+      override func visit(_ node: FunctionDeclSyntax) -> Bool {
         funcCount += 1
-        super.visit(node)
+        return true
       }
     }
     XCTAssertNoThrow(try {
       let parsed = try SyntaxTreeParser.parse(getInput("visitor.swift"))
       let counter = FuncCounter()
       let hashBefore = parsed.hashValue
-      counter.visit(parsed)
+      parsed.walk(counter)
       XCTAssertEqual(counter.funcCount, 3)
       XCTAssertEqual(hashBefore, parsed.hashValue)
     }())
@@ -70,16 +70,16 @@ public class SyntaxVisitorTestCase: XCTestCase {
     class VisitCollections: SyntaxVisitor {
       var numberOfCodeBlockItems = 0
 
-      override func visit(_ items: CodeBlockItemListSyntax) {
+      override func visit(_ items: CodeBlockItemListSyntax) -> Bool {
         numberOfCodeBlockItems += items.count
-        super.visit(items)
+        return true
       }
     }
 
     XCTAssertNoThrow(try {
       let parsed = try SyntaxTreeParser.parse(getInput("nested-blocks.swift"))
       let visitor = VisitCollections()
-      visitor.visit(parsed)
+      parsed.walk(visitor)
       XCTAssertEqual(4, visitor.numberOfCodeBlockItems)
     }())
   }

update-toolchain.py

@@ -43,8 +43,6 @@ def install(swiftc_exec, build_dir, dylib_dir, module_dir):
         '--build-dir', build_dir, '-v'])
 
 def main():
-    if not os.geteuid() == 0:
-        fatal_error("\n!!!!!!!!Run this script with root access!!!!!!!!\n")
     parser = argparse.ArgumentParser(
         formatter_class=argparse.RawDescriptionHelpFormatter,
         description='''
@@ -61,6 +59,11 @@ def main():
         ''')
 
     args = parser.parse_args(sys.argv[1:])
+    if not args.toolchain_dir:
+        fatal_error("Need to specify --toolchain-dir")
+    if not os.geteuid() == 0:
+        fatal_error("\n!!!!!!!!Run this script with root access!!!!!!!!\n")
+
     build_dir = tempfile.mkdtemp()
     swiftc_exec = args.toolchain_dir + '/usr/bin/swiftc'
     dylib_dir = args.toolchain_dir + '/usr/lib/swift/macosx'