Significantly improve performance of SyntaxVisitor

CodeGeneration/Sources/generate-swift-syntax/templates/swiftsyntax/SyntaxVisitorFile.swift

@@ -32,6 +32,32 @@ let syntaxVisitorFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
   try! ClassDeclSyntax("open class SyntaxVisitor") {
     DeclSyntax("public let viewMode: SyntaxTreeViewMode")
 
+    DeclSyntax(
+      """
+      /// `Syntax.Info` objects created in `visitChildren` but whose `Syntax` nodes were not retained by the `visit`
+      /// functions implemented by a subclass of `SyntaxVisitor`.
+      ///
+      /// Instead of deallocating them and allocating memory for new syntax nodes, store the allocated memory in an array.
+      /// We can then re-use them to create new syntax nodes.
+      ///
+      /// The array's size should be a typical nesting depth of a Swift file. That way we can store all allocated syntax
+      /// nodes when unwinding the visitation stack.
+      ///
+      /// The actual `info` stored in the `Syntax.Info` objects is garbage. It needs to be set when any of the `Syntax.Info`
+      /// objects get re-used.
+      private var recyclableNodeInfos: ContiguousArray<Syntax.Info?> =  ContiguousArray(repeating: nil, count: 64)
+      """
+    )
+
+    DeclSyntax(
+      """
+      /// A bit is set to 1 if the corresponding index in `recyclableNodeInfos` is occupied and ready to be reused.
+      ///
+      /// The last bit in this UInt64 corresponds to index 0 in `recyclableNodeInfos`.
+      private var recyclableNodeInfosUsageBitmap: UInt64 = 0
+      """
+    )
+
     DeclSyntax(
       """
       public init(viewMode: SyntaxTreeViewMode) {
@@ -45,7 +71,8 @@ let syntaxVisitorFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
       /// Walk all nodes of the given syntax tree, calling the corresponding `visit`
       /// function for every node that is being visited.
       public func walk(_ node: some SyntaxProtocol) {
-        visit(Syntax(node))
+        var syntaxNode = Syntax(node)
+        visit(&syntaxNode)
       }
       """
     )
@@ -94,21 +121,30 @@ let syntaxVisitorFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
 
     DeclSyntax(
       """
-      /// Interpret `data` as a node of type `nodeType`, visit it, calling
+      /// Cast `node` to a node of type `nodeType`, visit it, calling
       /// the `visit` and `visitPost` functions during visitation.
+      ///
+      /// - Note: node is an `inout` parameter so that callers don't have to retain it before passing it to `visitImpl`.
+      ///   With it being an `inout` parameter, the caller and `visitImpl` can work on the same reference of `node` without
+      ///   any reference counting.
+      /// - Note: Inline so that the optimizer can look through the calles to `visit` and `visitPost`, which means it
+      ///   doesn't need to retain `self` when forming closures to the unapplied function references on `self`.
+      @inline(__always)
       private func visitImpl<NodeType: SyntaxProtocol>(
-        _ node: Syntax,
+        _ node: inout Syntax,
         _ nodeType: NodeType.Type,
         _ visit: (NodeType) -> SyntaxVisitorContinueKind,
         _ visitPost: (NodeType) -> Void
       ) {
-        let node = node.cast(NodeType.self)
-        let needsChildren = (visit(node) == .visitChildren)
+        let castedNode = node.cast(NodeType.self)
+        // We retain castedNode.info here before passing it to visit.
+        // I don't think that's necessary because castedNode is already retained but don't know how to prevent it.
+        let needsChildren = (visit(castedNode) == .visitChildren)
         // Avoid calling into visitChildren if possible.
         if needsChildren && !node.raw.layoutView!.children.isEmpty {
-          visitChildren(node)
+          visitChildren(&node)
         }
-        visitPost(node)
+        visitPost(castedNode)
       }
       """
     )
@@ -149,7 +185,7 @@ let syntaxVisitorFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
                 /// that determines the correct visitation function will be popped of the
                 /// stack before the function is being called, making the switch's stack
                 /// space transient instead of having it linger in the call stack.
-                private func visitationFunc(for node: Syntax) -> ((Syntax) -> Void)
+                private func visitationFunc(for node: Syntax) -> ((inout Syntax) -> Void)
                 """
               ) {
                 try SwitchExprSyntax("switch node.raw.kind") {
@@ -168,16 +204,16 @@ let syntaxVisitorFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
 
                   for node in NON_BASE_SYNTAX_NODES {
                     SwitchCaseSyntax("case .\(node.varOrCaseName):") {
-                      StmtSyntax("return { self.visitImpl($0, \(node.kind.syntaxType).self, self.visit, self.visitPost) }")
+                      StmtSyntax("return { self.visitImpl(&$0, \(node.kind.syntaxType).self, self.visit, self.visitPost) }")
                     }
                   }
                 }
               }
 
               DeclSyntax(
                 """
-                private func visit(_ node: Syntax) {
-                  return visitationFunc(for: node)(node)
+                private func visit(_ node: inout Syntax) {
+                  return visitationFunc(for: node)(&node)
                 }
                 """
               )
@@ -188,7 +224,12 @@ let syntaxVisitorFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
           poundKeyword: .poundElseToken(),
           elements: .statements(
             CodeBlockItemListSyntax {
-              try! FunctionDeclSyntax("private func visit(_ node: Syntax)") {
+              try! FunctionDeclSyntax(
+                """
+                /// - Note: `node` is `inout` to avoid ref-counting. See comment in `visitImpl`
+                private func visit(_ node: inout Syntax)
+                """
+              ) {
                 try SwitchExprSyntax("switch node.raw.kind") {
                   SwitchCaseSyntax("case .token:") {
                     DeclSyntax("let node = node.cast(TokenSyntax.self)")
@@ -203,7 +244,7 @@ let syntaxVisitorFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
 
                   for node in NON_BASE_SYNTAX_NODES {
                     SwitchCaseSyntax("case .\(node.varOrCaseName):") {
-                      ExprSyntax("visitImpl(node, \(node.kind.syntaxType).self, visit, visitPost)")
+                      ExprSyntax("visitImpl(&node, \(node.kind.syntaxType).self, visit, visitPost)")
                     }
                   }
                 }
@@ -217,13 +258,66 @@ let syntaxVisitorFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
 
     DeclSyntax(
       """
-      private func visitChildren(_ node: some SyntaxProtocol) {
-        let syntaxNode = Syntax(node)
+      /// - Note: `node` is `inout` to avoid reference counting. See comment in `visitImpl`.
+      private func visitChildren(_ syntaxNode: inout Syntax) {
         for childRaw in NonNilRawSyntaxChildren(syntaxNode, viewMode: viewMode) {
-          visit(Syntax(childRaw, parent: syntaxNode))
+          // syntaxNode gets retained here. That seems unnecessary but I don't know how to remove it.
+          var childNode: Syntax
+          if let recycledInfoIndex = recyclableNodeInfosUsageBitmap.indexOfRightmostOne {
+            var recycledInfo: Syntax.Info? = nil
+            // Use `swap` to extract the recyclable syntax node without incurring ref-counting.
+            swap(&recycledInfo, &recyclableNodeInfos[recycledInfoIndex])
+            assert(recycledInfo != nil, "Slot indicated by the bitmap did not contain a value")
+            recyclableNodeInfosUsageBitmap.setBitToZero(at: recycledInfoIndex)
+            // syntaxNode.info gets retained here. This is necessary because we build up the parent tree.
+            recycledInfo!.info = .nonRoot(.init(parent: syntaxNode, absoluteInfo: childRaw.info))
+            childNode = Syntax(childRaw.raw, info: recycledInfo!)
+          } else {
+            childNode = Syntax(childRaw, parent: syntaxNode)
+          }
+          visit(&childNode)
+          if isKnownUniquelyReferenced(&childNode.info) {
+            // The node didn't get stored by the subclass's visit method. We can re-use the memory of its `Syntax.Info`
+            // for future syntax nodes.
+            childNode.info.info = nil
+            if let emptySlot = recyclableNodeInfosUsageBitmap.indexOfRightmostZero {
+              // Use `swap` to store the recyclable syntax node without incurring ref-counting.
+              swap(&recyclableNodeInfos[emptySlot], &childNode.info)
+              assert(childNode.info == nil, "Slot should not have contained a value")
+              recyclableNodeInfosUsageBitmap.setBitToOne(at: emptySlot)
+            }
+          }
         }
       }
       """
     )
   }
+
+  DeclSyntax(
+    """
+    fileprivate extension UInt64 {
+      var indexOfRightmostZero: Int? {
+        return (~self).indexOfRightmostOne
+      }
+
+      var indexOfRightmostOne: Int? {
+        let trailingZeroCount = self.trailingZeroBitCount
+        if trailingZeroCount == Self.bitWidth {
+          // All indicies are 0
+          return nil
+        }
+        return trailingZeroCount
+      }
+
+      mutating func setBitToZero(at index: Int) {
+        self &= ~(1 << index)
+      }
+
+      mutating func setBitToOne(at index: Int) {
+        self |= 1 << index
+      }
+    }
+
+    """
+  )
 }

Sources/SwiftSyntax/Syntax.swift

@@ -14,10 +14,9 @@
 /// Each node has accessors for its known children, and allows efficient
 /// iteration over the children through its `children` property.
 public struct Syntax: SyntaxProtocol, SyntaxHashable {
-  fileprivate enum Info: Sendable {
-    case root(Root)
-    indirect case nonRoot(NonRoot)
-
+  /// We need a heap indirection to store a syntax node's parent. We could use an indirect enum here but explicitly
+  /// modelling it using a class allows us to re-use these heap-allocated objects in `SyntaxVisitor`.
+  final class Info: Sendable {
     // For root node.
     struct Root: Sendable {
       private var arena: RetainedSyntaxArena
@@ -32,27 +31,50 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
       var parent: Syntax
       var absoluteInfo: AbsoluteSyntaxInfo
     }
+
+    enum InfoImpl: Sendable {
+      case root(Root)
+      case nonRoot(NonRoot)
+    }
+
+    init(_ info: InfoImpl) {
+      self.info = info
+    }
+
+    /// The actual stored information that references the parent or the tree's root.
+    ///
+    /// - Important: Must only be set to `nil` when `Syntax.Info` is used in a memory recycling pool
+    ///    (eg. in `SyntaxVisitor`). In that case the `Syntax.Info` is considered garbage memory that can be re-used
+    ///    later. `info` needs to be set to a real value when `Syntax.Info` is recycled from the memory recycling pool.
+    var info: InfoImpl!
   }
 
-  private let info: Info
+  /// Reference to the node's parent or, if this node is the root of a tree, a reference to the `SyntaxArena` to keep
+  /// the syntax tree alive.
+  ///
+  /// - Important: In almost all use cases you should not access this directly. Prefer accessors like `parent`.
+  /// - Important: Must only be set to `nil` when this `Syntax` node is known to get destroyed and the `Info` should be
+  ///   stored in a memory recycling pool (eg. in `SyntaxVisitor`). After setting `info` to `nil`, this `Syntax` node
+  ///   is considered garbage and should not be accessed anymore in any way.
+  var info: Info!
   let raw: RawSyntax
 
   private var rootInfo: Info.Root {
-    switch info {
+    switch info.info! {
     case .root(let info): return info
     case .nonRoot(let info): return info.parent.rootInfo
     }
   }
 
   private var nonRootInfo: Info.NonRoot? {
-    switch info {
+    switch info.info! {
     case .root(_): return nil
     case .nonRoot(let info): return info
     }
   }
 
   private var root: Syntax {
-    switch info {
+    switch info.info! {
     case .root(_): return self
     case .nonRoot(let info): return info.parent.root
     }
@@ -99,13 +121,13 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
   }
 
   /// "designated" memberwise initializer of `Syntax`.
-  private init(_ raw: RawSyntax, info: Info) {
+  init(_ raw: RawSyntax, info: Info) {
     self.raw = raw
     self.info = info
   }
 
   init(_ raw: RawSyntax, parent: Syntax, absoluteInfo: AbsoluteSyntaxInfo) {
-    self.init(raw, info: .nonRoot(.init(parent: parent, absoluteInfo: absoluteInfo)))
+    self.init(raw, info: Info(.nonRoot(.init(parent: parent, absoluteInfo: absoluteInfo))))
   }
 
   /// Creates a `Syntax` with the provided raw syntax and parent.
@@ -125,12 +147,12 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
   ///     has a chance to retain it.
   static func forRoot(_ raw: RawSyntax, rawNodeArena: RetainedSyntaxArena) -> Syntax {
     precondition(rawNodeArena == raw.arenaReference)
-    return Syntax(raw, info: .root(.init(arena: rawNodeArena)))
+    return Syntax(raw, info: Info(.root(.init(arena: rawNodeArena))))
   }
 
   static func forRoot(_ raw: RawSyntax, rawNodeArena: SyntaxArena) -> Syntax {
     precondition(rawNodeArena == raw.arenaReference)
-    return Syntax(raw, info: .root(.init(arena: RetainedSyntaxArena(rawNodeArena))))
+    return Syntax(raw, info: Info(.root(.init(arena: RetainedSyntaxArena(rawNodeArena)))))
   }
 
   /// Returns the child data at the provided index in this data's layout.
@@ -252,6 +274,9 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
   }
 
   /// Create a ``Syntax`` node from a specialized syntax node.
+  // Inline always so the optimizer can optimize this to a member access on `syntax` without having to go through
+  // generics.
+  @inline(__always)
   public init(_ syntax: some SyntaxProtocol) {
     self = syntax._syntaxNode
   }

Sources/SwiftSyntax/SyntaxChildren.swift

@@ -411,6 +411,8 @@ struct NonNilRawSyntaxChildren: BidirectionalCollection, Sendable {
     self.viewMode = viewMode
   }
 
+  /// - Note: Inline so we don't retain `Syntax.Info` when creating `NonNilRawSyntaxChildren` from a `Syntax`.
+  @inline(__always)
   init(_ node: Syntax, viewMode: SyntaxTreeViewMode) {
     self.init(node.absoluteRaw, viewMode: viewMode)
   }