[6.0] Performance improvements for SyntaxRewriter

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

@@ -183,7 +183,7 @@ let syntaxBaseNodesFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
       DeclSyntax(
         """
         /// Create a \(raw: node.kind.doccLink) node from a specialized syntax node.
-        public init(_ syntax: some \(node.kind.protocolType)) {
+        public init(_ syntax: __shared some \(node.kind.protocolType)) {
           // We know this cast is going to succeed. Go through init(_: SyntaxData)
           // to do a sanity check and verify the kind matches in debug builds and get
           // maximum performance in release builds.
@@ -195,7 +195,7 @@ let syntaxBaseNodesFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
       DeclSyntax(
         """
         /// Create a \(raw: node.kind.doccLink) node from a specialized optional syntax node.
-        public init?(_ syntax: (some \(node.kind.protocolType))?) {
+        public init?(_ syntax: __shared (some \(node.kind.protocolType))?) {
           guard let syntax = syntax else { return nil }
           self.init(syntax)
         }
@@ -204,7 +204,7 @@ let syntaxBaseNodesFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
 
       DeclSyntax(
         """
-        public init(fromProtocol syntax: \(node.kind.protocolType)) {
+        public init(fromProtocol syntax: __shared \(node.kind.protocolType)) {
           // We know this cast is going to succeed. Go through init(_: SyntaxData)
           // to do a sanity check and verify the kind matches in debug builds and get
           // maximum performance in release builds.
@@ -216,14 +216,14 @@ let syntaxBaseNodesFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
       DeclSyntax(
         """
         /// Create a \(raw: node.kind.doccLink) node from a specialized optional syntax node.
-        public init?(fromProtocol syntax: \(node.kind.protocolType)?) {
+        public init?(fromProtocol syntax: __shared \(node.kind.protocolType)?) {
           guard let syntax = syntax else { return nil }
           self.init(fromProtocol: syntax)
         }
         """
       )
 
-      try InitializerDeclSyntax("public init?(_ node: some SyntaxProtocol)") {
+      try InitializerDeclSyntax("public init?(_ node: __shared some SyntaxProtocol)") {
         try SwitchExprSyntax("switch node.raw.kind") {
           SwitchCaseListSyntax {
             SwitchCaseSyntax(

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

@@ -47,7 +47,7 @@ func syntaxNode(nodesStartingWith: [Character]) -> SourceFileSyntax {
 
         DeclSyntax(
           """
-          public init?(_ node: some SyntaxProtocol) {
+          public init?(_ node: __shared some SyntaxProtocol) {
             guard node.raw.kind == .\(node.varOrCaseName) else { return nil }
             self._syntaxNode = node._syntaxNode
           }
@@ -199,17 +199,14 @@ func syntaxNode(nodesStartingWith: [Character]) -> SourceFileSyntax {
           }
         }
 
-        try! VariableDeclSyntax("public static var structure: SyntaxNodeStructure") {
-          let layout = ArrayExprSyntax {
-            for child in node.children {
-              ArrayElementSyntax(
-                expression: ExprSyntax(#"\Self.\#(child.varOrCaseName)"#)
-              )
-            }
+        let layout = ArrayExprSyntax {
+          for child in node.children {
+            ArrayElementSyntax(
+              expression: ExprSyntax(#"\Self.\#(child.varOrCaseName)"#)
+            )
           }
-
-          StmtSyntax("return .layout(\(layout))")
         }
+        "public static let structure: SyntaxNodeStructure = .layout(\(layout))"
       }
     }
   }
@@ -256,7 +253,7 @@ private func generateSyntaxChildChoices(for child: Child) throws -> EnumDeclSynt
       }
     }
 
-    try! InitializerDeclSyntax("public init?(_ node: some SyntaxProtocol)") {
+    try! InitializerDeclSyntax("public init?(_ node: __shared some SyntaxProtocol)") {
       for choice in choices {
         StmtSyntax(
           """

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

@@ -32,6 +32,13 @@ let syntaxRewriterFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
   ) {
     DeclSyntax("public let viewMode: SyntaxTreeViewMode")
 
+    DeclSyntax(
+      """
+      /// 'Syntax' object factory recycling 'Syntax.Info' instances.
+      private let nodeFactory: SyntaxNodeFactory = SyntaxNodeFactory()
+      """
+    )
+
     DeclSyntax(
       """
       public init(viewMode: SyntaxTreeViewMode = .sourceAccurate) {
@@ -44,7 +51,8 @@ let syntaxRewriterFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
       """
       /// Rewrite `node`, keeping its parent unless `detach` is `true`.
       public func rewrite(_ node: some SyntaxProtocol, detach: Bool = false) -> Syntax {
-        let rewritten = self.dispatchVisit(Syntax(node))
+        var rewritten = Syntax(node)
+        self.dispatchVisit(&rewritten)
         if detach {
           return rewritten
         }
@@ -105,15 +113,19 @@ let syntaxRewriterFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
       ///   - Returns: the rewritten node
       @available(*, deprecated, renamed: "rewrite(_:detach:)")
       public func visit(_ node: Syntax) -> Syntax {
-        return dispatchVisit(node)
+        var rewritten = node
+        dispatchVisit(&rewritten)
+        return rewritten
       }
       """
     )
 
     DeclSyntax(
       """
       public func visit<T: SyntaxChildChoices>(_ node: T) -> T {
-        return dispatchVisit(Syntax(node)).cast(T.self)
+        var rewritten = Syntax(node)
+        dispatchVisit(&rewritten)
+        return rewritten.cast(T.self)
       }
       """
     )
@@ -127,7 +139,7 @@ let syntaxRewriterFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
           ///   - Returns: the rewritten node
           \(node.apiAttributes())\
           open func visit(_ node: \(node.kind.syntaxType)) -> \(node.kind.syntaxType) {
-            return visitChildren(node)
+            return visitChildren(node._syntaxNode).cast(\(node.kind.syntaxType).self)
           }
           """
         )
@@ -139,7 +151,7 @@ let syntaxRewriterFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
           ///   - Returns: the rewritten node
           \(node.apiAttributes())\
           open func visit(_ node: \(node.kind.syntaxType)) -> \(node.baseType.syntaxBaseName) {
-            return \(node.baseType.syntaxBaseName)(visitChildren(node))
+            return \(node.baseType.syntaxBaseName)(visitChildren(node._syntaxNode).cast(\(node.kind.syntaxType).self))
           }
           """
         )
@@ -156,7 +168,9 @@ let syntaxRewriterFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
         ///   - Returns: the rewritten node
         \(baseNode.apiAttributes())\
         public func visit(_ node: \(baseKind.syntaxType)) -> \(baseKind.syntaxType) {
-          return dispatchVisit(Syntax(node)).cast(\(baseKind.syntaxType).self)
+          var node: Syntax = Syntax(node)
+          dispatchVisit(&node)
+          return node.cast(\(baseKind.syntaxType).self)
         }
         """
       )
@@ -166,21 +180,16 @@ let syntaxRewriterFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
       """
       /// Interpret `node` as a node of type `nodeType`, visit it, calling
       /// the `visit` to transform the node.
+      @inline(__always)
       private func visitImpl<NodeType: SyntaxProtocol>(
-        _ node: Syntax,
+        _ node: inout Syntax,
         _ nodeType: NodeType.Type,
         _ visit: (NodeType) -> some SyntaxProtocol
-      ) -> Syntax {
-        let castedNode = node.cast(NodeType.self)
-        // Accessing _syntaxNode directly is faster than calling Syntax(node)
-        visitPre(node)
-        defer {
-          visitPost(node)
-        }
-        if let newNode = visitAny(node) {
-          return newNode
-        }
-        return Syntax(visit(castedNode))
+      ) {
+        let origNode = node
+        visitPre(origNode)
+        node = visitAny(origNode) ?? Syntax(visit(origNode.cast(NodeType.self)))
+        visitPost(origNode)
       }
       """
     )
@@ -221,26 +230,26 @@ let syntaxRewriterFile = 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) -> Syntax)
+                private func visitationFunc(for node: Syntax) -> ((inout Syntax) -> Void)
                 """
               ) {
                 try SwitchExprSyntax("switch node.raw.kind") {
                   SwitchCaseSyntax("case .token:") {
-                    StmtSyntax("return { self.visitImpl($0, TokenSyntax.self, self.visit) }")
+                    StmtSyntax("return { self.visitImpl(&$0, TokenSyntax.self, self.visit) }")
                   }
 
                   for node in NON_BASE_SYNTAX_NODES {
                     SwitchCaseSyntax("case .\(node.varOrCaseName):") {
-                      StmtSyntax("return { self.visitImpl($0, \(node.kind.syntaxType).self, self.visit) }")
+                      StmtSyntax("return { self.visitImpl(&$0, \(node.kind.syntaxType).self, self.visit) }")
                     }
                   }
                 }
               }
 
               DeclSyntax(
                 """
-                private func dispatchVisit(_ node: Syntax) -> Syntax {
-                  return visitationFunc(for: node)(node)
+                private func dispatchVisit(_ node: inout Syntax) {
+                  visitationFunc(for: node)(&node)
                 }
                 """
               )
@@ -251,15 +260,15 @@ let syntaxRewriterFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
           poundKeyword: .poundElseToken(),
           elements: .statements(
             CodeBlockItemListSyntax {
-              try! FunctionDeclSyntax("private func dispatchVisit(_ node: Syntax) -> Syntax") {
+              try! FunctionDeclSyntax("private func dispatchVisit(_ node: inout Syntax)") {
                 try SwitchExprSyntax("switch node.raw.kind") {
                   SwitchCaseSyntax("case .token:") {
-                    StmtSyntax("return visitImpl(node, TokenSyntax.self, visit)")
+                    StmtSyntax("return visitImpl(&node, TokenSyntax.self, visit)")
                   }
 
                   for node in NON_BASE_SYNTAX_NODES {
                     SwitchCaseSyntax("case .\(node.varOrCaseName):") {
-                      StmtSyntax("return visitImpl(node, \(node.kind.syntaxType).self, visit)")
+                      StmtSyntax("return visitImpl(&node, \(node.kind.syntaxType).self, visit)")
                     }
                   }
                 }
@@ -272,9 +281,7 @@ let syntaxRewriterFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
 
     DeclSyntax(
       """
-      private func visitChildren<SyntaxType: SyntaxProtocol>(
-        _ node: SyntaxType
-      ) -> SyntaxType {
+      private func visitChildren(_ node: Syntax) -> Syntax {
         // Walk over all children of this node and rewrite them. Don't store any
         // rewritten nodes until the first non-`nil` value is encountered. When this
         // happens, retrieve all previous syntax nodes from the parent node to
@@ -284,73 +291,48 @@ let syntaxRewriterFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
 
         // newLayout is nil until the first child node is rewritten and rewritten
         // nodes are being collected.
-        var newLayout: ContiguousArray<RawSyntax?>?
-
-        // Rewritten children just to keep their 'SyntaxArena' alive until they are
-        // wrapped with 'Syntax'
-        var rewrittens: ContiguousArray<Syntax> = []
+        var newLayout: UnsafeMutableBufferPointer<RawSyntax?> = .init(start: nil, count: 0)
 
-        let syntaxNode = node._syntaxNode
+        // Keep 'SyntaxArena' of rewritten nodes alive until they are wrapped
+        // with 'Syntax'
+        var rewrittens: ContiguousArray<RetainedSyntaxArena> = []
 
-        // Incrementing i manually is faster than using .enumerated()
-        var childIndex = 0
-        for (raw, info) in RawSyntaxChildren(syntaxNode) {
-          defer { childIndex += 1 }
-
-          guard let child = raw, viewMode.shouldTraverse(node: child) else {
-            // Node does not exist or should not be visited. If we are collecting
-            // rewritten nodes, we need to collect this one as well, otherwise we
-            // can ignore it.
-            if newLayout != nil {
-              newLayout!.append(raw)
-            }
-            continue
-          }
+        for case let (child?, info) in RawSyntaxChildren(node) where viewMode.shouldTraverse(node: child) {
 
           // Build the Syntax node to rewrite
-          let absoluteRaw = AbsoluteRawSyntax(raw: child, info: info)
+          var childNode = nodeFactory.create(parent: node, raw: child, absoluteInfo: info)
 
-          let rewritten = dispatchVisit(Syntax(absoluteRaw, parent: syntaxNode))
-          if rewritten.id != info.nodeId {
+          dispatchVisit(&childNode)
+          if childNode.raw.id != child.id {
             // The node was rewritten, let's handle it
-            if newLayout == nil {
+
+            if newLayout.baseAddress == nil {
               // We have not yet collected any previous rewritten nodes. Initialize
-              // the new layout with the previous nodes of the parent. This is
-              // possible, since we know they were not rewritten.
-
-              // The below implementation is based on Collection.map but directly
-              // reserves enough capacity for the entire layout.
-              newLayout = ContiguousArray<RawSyntax?>()
-              newLayout!.reserveCapacity(node.raw.layoutView!.children.count)
-              for j in 0..<childIndex {
-                newLayout!.append(node.raw.layoutView!.children[j])
-              }
+              // the new layout with the previous nodes of the parent.
+              newLayout = .allocate(capacity: node.raw.layoutView!.children.count)
+              _ = newLayout.initialize(fromContentsOf: node.raw.layoutView!.children)
             }
 
-            // Now that we know we have a new layout in which we collect rewritten
-            // nodes, add it.
-            rewrittens.append(rewritten)
-            newLayout!.append(rewritten.raw)
-          } else {
-            // The node was not changed by the rewriter. Only store it if a previous
-            // node has been rewritten and we are collecting a rewritten layout.
-            if newLayout != nil {
-              newLayout!.append(raw)
-            }
+            // Update the rewritten child.
+            newLayout[Int(info.indexInParent)] = childNode.raw
+            // Retain the syntax arena of the new node until it's wrapped with Syntax node.
+            rewrittens.append(childNode.raw.arenaReference.retained)
           }
+
+          // Recycle 'childNode.info'
+          nodeFactory.dispose(&childNode)
         }
 
-        if let newLayout {
+        if newLayout.baseAddress != nil {
           // A child node was rewritten. Build the updated node.
 
-          // Sanity check, ensure the new children are the same length.
-          precondition(newLayout.count == node.raw.layoutView!.children.count)
-
           let arena = SyntaxArena()
-          let newRaw = node.raw.layoutView!.replacingLayout(with: Array(newLayout), arena: arena)
+          let newRaw = node.raw.layoutView!.replacingLayout(with: newLayout, arena: arena)
+          newLayout.deinitialize()
+          newLayout.deallocate()
           // 'withExtendedLifetime' to keep 'SyntaxArena's of them alive until here.
           return withExtendedLifetime(rewrittens) {
-            Syntax(raw: newRaw, rawNodeArena: arena).cast(SyntaxType.self)
+            Syntax(raw: newRaw, rawNodeArena: arena)
           }
         } else {
           // No child node was rewritten. So no need to change this node as well.