Add arena to syntax nodes

CodeGeneration/Package.swift

@@ -21,8 +21,8 @@ let package = Package(
     // checkout of swift-syntax that is unaffected by the newly generated files.
     // Be sure to revert the change before committing your changes.
     //
-    // .package(url: "https://github.com/apple/swift-syntax", branch: "main")
-    .package(path: "..")
+    .package(url: "https://github.com/apple/swift-syntax", branch: "main")
+    // .package(path: "..")
   ],
   targets: [
     .executableTarget(

CodeGeneration/Sources/generate-swiftsyntax/LayoutNode+Extensions.swift

@@ -66,6 +66,9 @@ extension LayoutNode {
 
       FunctionParameterSyntax("trailingTrivia: Trivia? = nil")
         .with(\.leadingTrivia, .newline)
+
+      FunctionParameterSyntax("arena: __shared SyntaxArena = SyntaxArena()")
+        .with(\.leadingTrivia, .newline)
     }
 
     return """

CodeGeneration/Sources/generate-swiftsyntax/templates/swiftsyntax/RenamedChildrenCompatibilityFile.swift

@@ -93,6 +93,7 @@ let renamedChildrenCompatibilityFile = try! SourceFileSyntax(leadingTrivia: copy
             }
           }
           LabeledExprSyntax(label: "trailingTrivia", expression: ExprSyntax("trailingTrivia"))
+          LabeledExprSyntax(label: "arena", expression: ExprSyntax("arena"))
         }
       }
     }

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

@@ -116,7 +116,7 @@ func syntaxNode(emitKind: SyntaxNodeKind) -> SourceFileSyntax {
             calledExpression: ExprSyntax("withExtendedLifetime"),
             leftParen: .leftParenToken(),
             arguments: LabeledExprListSyntax {
-              LabeledExprSyntax(expression: ExprSyntax("(SyntaxArena(), (\(parameters)))"))
+              LabeledExprSyntax(expression: ExprSyntax("(arena, (\(parameters)))"))
             },
             rightParen: .rightParenToken(),
             trailingClosure: ClosureExprSyntax(signature: closureSignature) {

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

@@ -31,11 +31,13 @@ let syntaxRewriterFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
     """
   ) {
     DeclSyntax("public let viewMode: SyntaxTreeViewMode")
+    DeclSyntax("private let arena: SyntaxArena")
 
     DeclSyntax(
       """
-      public init(viewMode: SyntaxTreeViewMode = .sourceAccurate) {
+      public init(viewMode: SyntaxTreeViewMode = .sourceAccurate, arena: __shared SyntaxArena = SyntaxArena()) {
         self.viewMode = viewMode
+        self.arena = arena
       }
       """
     )
@@ -290,11 +292,7 @@ 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?>?
 
         let syntaxNode = node._syntaxNode
 
@@ -308,7 +306,7 @@ let syntaxRewriterFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
             // rewritten nodes, we need to collect this one as well, otherwise we
             // can ignore it.
             if newLayout != nil {
-              newLayout!.append(raw)
+              newLayout!.initializeElement(at: childIndex, to: raw)
             }
             continue
           }
@@ -319,6 +317,9 @@ let syntaxRewriterFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
 
           let rewritten = visit(data)
           if rewritten.data.nodeId != info.nodeId {
+            //        if let sharedArena = self.arena {
+            //          precondition(rewritten.raw.arenaReference == SyntaxArenaRef(sharedArena))
+            //        }
             // The node was rewritten, let's handle it
             if newLayout == nil {
               // We have not yet collected any previous rewritten nodes. Initialize
@@ -327,22 +328,23 @@ let syntaxRewriterFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
 
               // 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])
+              newLayout = arena.allocateRawSyntaxBuffer(count: node.raw.layoutView!.children.count)
+              for j in 0 ..< childIndex {
+                newLayout!.initializeElement(at: j, to: node.raw.layoutView!.children[j])
               }
             }
 
             // Now that we know we have a new layout in which we collect rewritten
             // nodes, add it.
-            rewrittens.append(rewritten)
-            newLayout!.append(rewritten.raw)
+            withExtendedLifetime(rewritten) {
+              arena.addChild(rewritten.raw.arenaReference)
+            }
+            newLayout!.initializeElement(at: childIndex, to: 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)
+              newLayout!.initializeElement(at: childIndex, to: raw)
             }
           }
         }
@@ -353,12 +355,8 @@ let syntaxRewriterFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
           // 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)
-          // 'withExtendedLifetime' to keep 'SyntaxArena's of them alive until here.
-          return withExtendedLifetime(rewrittens) {
-            Syntax(raw: newRaw, rawNodeArena: arena).cast(SyntaxType.self)
-          }
+          let newRaw = RawSyntax.makeLayout(kind: node.raw.kind, layoutBuffer: newLayout, arena: arena)
+          return Syntax(raw: newRaw, rawNodeArena: arena).cast(SyntaxType.self)
         } else {
           // No child node was rewritten. So no need to change this node as well.
           return node

Sources/SwiftSyntax/Raw/RawSyntax.swift

@@ -771,7 +771,27 @@ extension RawSyntax {
     let layoutBuffer = arena.allocateRawSyntaxBuffer(count: count)
     initializer(layoutBuffer)
 
-    // Calculate the "byte width".
+    return self.makeLayout(
+      kind: kind,
+      isMaximumNestingLevelOverflow: isMaximumNestingLevelOverflow,
+      layoutBuffer: layoutBuffer,
+      arena: arena
+    )
+  }
+
+  /// Factory method to create a layout node.
+  ///
+  /// - Parameters:
+  ///   - arena: SyntaxArena to the result node data resides.
+  ///   - kind: Syntax kind.
+  ///   - count: Number of children.
+  ///   - initializer: A closure that initializes elements.
+  public static func makeLayout(
+    kind: SyntaxKind,
+    isMaximumNestingLevelOverflow: Bool = false,
+    layoutBuffer: UnsafeMutableBufferPointer<RawSyntax?>,
+    arena: __shared SyntaxArena
+  ) -> RawSyntax {  // Calculate the "byte width".
     var byteLength = 0
     var descendantCount = 0
     var recursiveFlags = RecursiveRawSyntaxFlags()

Sources/SwiftSyntax/SyntaxArena.swift

@@ -163,12 +163,12 @@ public class SyntaxArena {
   func addChild(_ otherRef: SyntaxArenaRef) {
     if SyntaxArenaRef(self) == otherRef { return }
 
-    #if DEBUG || SWIFTSYNTAX_ENABLE_ASSERTIONS
-    precondition(
-      !self.hasParent,
-      "an arena can't have a new child once it's owned by other arenas"
-    )
-    #endif
+    //    #if DEBUG || SWIFTSYNTAX_ENABLE_ASSERTIONS
+    //    precondition(
+    //      !self.hasParent,
+    //      "an arena can't have a new child once it's owned by other arenas"
+    //    )
+    //    #endif
 
     if childRefs.insert(otherRef).inserted {
       otherRef.retain()