Split notion of rawNodeArena and allocationArena

The `arena` parameter to `replacingSelf` and `replacingChild` served two purposes: To keep the arena of `raw`/`newChild` alive and to specify which arena new nodes should be allocated in. In almost all cases, these areneas coincided but for `SyntaxRewriter.rewrite` they diverged, causing an arena to be added as a child arena after it itself has children, causing a precondition failure.

Split these different purposes of arenas into two different parameters.

Author: ahoppen

CodeGeneration/Sources/generate-swiftsyntax/templates/swiftparser/ParserEntryFile.swift

@@ -72,7 +72,7 @@ let parserEntryFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
             defer { withExtendedLifetime(parser) {} }
             let node = parser.\(raw: parserFunction)()
             let raw = RawSyntax(parser.parseRemainder(into: node))
-            return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
+            return Syntax(raw: raw, rawNodeArena: raw.arena).cast(Self.self)
           }
         }
         """

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

@@ -171,7 +171,7 @@ let syntaxCollectionsFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
           let arena = SyntaxArena()
           let raw = RawSyntax.makeLayout(kind: SyntaxKind.\(node.varOrCaseName),
                                          from: children.map { $0.raw }, arena: arena)
-          self.init(SyntaxData.forRoot(raw, arena: arena))
+          self.init(SyntaxData.forRoot(raw, rawNodeArena: arena))
         }
         """
       )
@@ -194,7 +194,7 @@ let syntaxCollectionsFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
         internal func replacingLayout(_ layout: [RawSyntax?]) -> \(node.kind.syntaxType) {
           let arena = SyntaxArena()
           let newRaw = layoutView.replacingLayout(with: layout, arena: arena)
-          let newData = data.replacingSelf(newRaw, arena: arena)
+          let newData = data.replacingSelf(newRaw, rawNodeArena: arena, allocationArena: arena)
           return \(node.kind.syntaxType)(newData)
         }
         """

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

@@ -129,7 +129,7 @@ func syntaxNode(emitKind: SyntaxNodeKind) -> SourceFileSyntax {
                   """
                 )
               }
-              StmtSyntax("return SyntaxData.forRoot(raw, arena: arena)")
+              StmtSyntax("return SyntaxData.forRoot(raw, rawNodeArena: arena)")
             }
           )
 
@@ -207,7 +207,7 @@ func syntaxNode(emitKind: SyntaxNodeKind) -> SourceFileSyntax {
                   collection = RawSyntax.makeLayout(kind: SyntaxKind.\(raw: childNode.varOrCaseName),
                                                     from: [element.raw], arena: arena)
                 }
-                let newData = data.replacingChild(at: \(raw: index), with: collection, arena: arena)
+                let newData = data.replacingChild(at: \(raw: index), with: collection, rawNodeArena: arena, allocationArena: arena)
                 return \(node.kind.syntaxType)(newData)
               }
               """

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

@@ -326,7 +326,7 @@ let syntaxRewriterFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
           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, arena: arena).cast(SyntaxType.self)
+            Syntax(raw: newRaw, rawNodeArena: arena).cast(SyntaxType.self)
           }
         } else {
           // No child node was rewritten. So no need to change this node as well.
@@ -343,7 +343,7 @@ let syntaxRewriterFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
       public func rewrite(_ node: Syntax) -> Syntax {
         let rewritten = self.visit(node)
         return withExtendedLifetime(rewritten) {
-          return Syntax(node.data.replacingSelf(rewritten.raw, arena: rewritten.raw.arena))
+          return Syntax(node.data.replacingSelf(rewritten.raw, rawNodeArena: rewritten.raw.arena, allocationArena: SyntaxArena()))
         }
       }
       """

Sources/SwiftParser/generated/Parser+Entry.swift

@@ -55,7 +55,7 @@ extension AccessorDeclSyntax: SyntaxParseable {
     }
     let node = parser.parseAccessorDecl()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
+    return Syntax(raw: raw, rawNodeArena: raw.arena).cast(Self.self)
   }
 }
 
@@ -73,7 +73,7 @@ extension AttributeSyntax: SyntaxParseable {
     }
     let node = parser.parseAttribute()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
+    return Syntax(raw: raw, rawNodeArena: raw.arena).cast(Self.self)
   }
 }
 
@@ -91,7 +91,7 @@ extension CatchClauseSyntax: SyntaxParseable {
     }
     let node = parser.parseCatchClause()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
+    return Syntax(raw: raw, rawNodeArena: raw.arena).cast(Self.self)
   }
 }
 
@@ -109,7 +109,7 @@ extension ClosureParameterSyntax: SyntaxParseable {
     }
     let node = parser.parseClosureParameter()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
+    return Syntax(raw: raw, rawNodeArena: raw.arena).cast(Self.self)
   }
 }
 
@@ -127,7 +127,7 @@ extension CodeBlockItemSyntax: SyntaxParseable {
     }
     let node = parser.parseNonOptionalCodeBlockItem()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
+    return Syntax(raw: raw, rawNodeArena: raw.arena).cast(Self.self)
   }
 }
 
@@ -145,7 +145,7 @@ extension DeclSyntax: SyntaxParseable {
     }
     let node = parser.parseDeclaration()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
+    return Syntax(raw: raw, rawNodeArena: raw.arena).cast(Self.self)
   }
 }
 
@@ -163,7 +163,7 @@ extension EnumCaseParameterSyntax: SyntaxParseable {
     }
     let node = parser.parseEnumCaseParameter()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
+    return Syntax(raw: raw, rawNodeArena: raw.arena).cast(Self.self)
   }
 }
 
@@ -181,7 +181,7 @@ extension ExprSyntax: SyntaxParseable {
     }
     let node = parser.parseExpression()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
+    return Syntax(raw: raw, rawNodeArena: raw.arena).cast(Self.self)
   }
 }
 
@@ -199,7 +199,7 @@ extension FunctionParameterSyntax: SyntaxParseable {
     }
     let node = parser.parseFunctionParameter()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
+    return Syntax(raw: raw, rawNodeArena: raw.arena).cast(Self.self)
   }
 }
 
@@ -217,7 +217,7 @@ extension GenericParameterClauseSyntax: SyntaxParseable {
     }
     let node = parser.parseGenericParameters()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
+    return Syntax(raw: raw, rawNodeArena: raw.arena).cast(Self.self)
   }
 }
 
@@ -235,7 +235,7 @@ extension MemberDeclBlockSyntax: SyntaxParseable {
     }
     let node = parser.parseMemberDeclList()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
+    return Syntax(raw: raw, rawNodeArena: raw.arena).cast(Self.self)
   }
 }
 
@@ -253,7 +253,7 @@ extension PatternSyntax: SyntaxParseable {
     }
     let node = parser.parsePattern()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
+    return Syntax(raw: raw, rawNodeArena: raw.arena).cast(Self.self)
   }
 }
 
@@ -271,7 +271,7 @@ extension SourceFileSyntax: SyntaxParseable {
     }
     let node = parser.parseSourceFile()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
+    return Syntax(raw: raw, rawNodeArena: raw.arena).cast(Self.self)
   }
 }
 
@@ -289,7 +289,7 @@ extension StmtSyntax: SyntaxParseable {
     }
     let node = parser.parseStatement()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
+    return Syntax(raw: raw, rawNodeArena: raw.arena).cast(Self.self)
   }
 }
 
@@ -307,7 +307,7 @@ extension SwitchCaseSyntax: SyntaxParseable {
     }
     let node = parser.parseSwitchCase()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
+    return Syntax(raw: raw, rawNodeArena: raw.arena).cast(Self.self)
   }
 }
 
@@ -325,7 +325,7 @@ extension TypeSyntax: SyntaxParseable {
     }
     let node = parser.parseType()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
+    return Syntax(raw: raw, rawNodeArena: raw.arena).cast(Self.self)
   }
 }
 

Sources/SwiftSyntax/Raw/RawSyntaxLayoutView.swift

@@ -54,6 +54,8 @@ public struct RawSyntaxLayoutView {
   }
 
   /// Creates a new node of the same kind but with children replaced by `elements`.
+  ///
+  /// The newly created syntax node is allocated in `arena`.
   @_spi(RawSyntax)
   public func replacingLayout(
     with elements: some Collection<RawSyntax?>,

Sources/SwiftSyntax/Syntax.swift

@@ -69,8 +69,8 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
   }
 
   @_spi(RawSyntax)
-  public init(raw: RawSyntax, arena: __shared SyntaxArena) {
-    self.init(.forRoot(raw, arena: arena))
+  public init(raw: RawSyntax, rawNodeArena: __shared SyntaxArena) {
+    self.init(.forRoot(raw, rawNodeArena: rawNodeArena))
   }
 
   /// Create a ``Syntax`` node from a specialized syntax node.
@@ -229,7 +229,7 @@ extension SyntaxProtocol {
     // Make sure `self` (and thus the arena of `self.raw`) can’t get deallocated
     // before the detached node can be created.
     return withExtendedLifetime(self) {
-      return Syntax(raw: self.raw, arena: self.raw.arena).cast(Self.self)
+      return Syntax(raw: self.raw, rawNodeArena: self.raw.arena).cast(Self.self)
     }
   }
 

Sources/SwiftSyntax/SyntaxData.swift

@@ -278,9 +278,9 @@ struct SyntaxData {
   /// `arena` must be the arena in which `raw` is allocated. It is passed to
   /// make sure the arena doesn’t get de-allocated before the ``SyntaxData`` has
   /// a chance to retain it.
-  static func forRoot(_ raw: RawSyntax, arena: SyntaxArena) -> SyntaxData {
-    precondition(arena === raw.arena)
-    return SyntaxData(raw, info: .root(.init(arena: arena)))
+  static func forRoot(_ raw: RawSyntax, rawNodeArena: SyntaxArena) -> SyntaxData {
+    precondition(rawNodeArena === raw.arena)
+    return SyntaxData(raw, info: .root(.init(arena: rawNodeArena)))
   }
 
   /// Returns the child data at the provided index in this data's layout.
@@ -304,19 +304,23 @@ struct SyntaxData {
   /// Creates a copy of `self` and recursively creates `SyntaxData` nodes up to
   /// the root.
   ///
-  /// `arena` is the arena in which `newRaw` resides. The new nodes will be
-  /// allocated in this arena.
-  func replacingSelf(_ newRaw: RawSyntax, arena: SyntaxArena) -> SyntaxData {
-    precondition(newRaw.arena === arena)
+  /// - Parameters:
+  ///   - newRaw: The node that should replace `self`
+  ///   - rawNodeArena: The arena in which `newRaw` resides
+  ///   - arena: The arena in which  new nodes should be allocated
+  /// - Returns: A syntax tree with all parents where this node has been
+  ///            replaced by `newRaw`
+  func replacingSelf(_ newRaw: RawSyntax, rawNodeArena: SyntaxArena, allocationArena: SyntaxArena) -> SyntaxData {
+    precondition(newRaw.arena === rawNodeArena)
     // If we have a parent already, then ask our current parent to copy itself
     // recursively up to the root.
     if let parent {
-      let parentData = parent.replacingChild(at: indexInParent, with: newRaw, arena: arena)
+      let parentData = parent.replacingChild(at: indexInParent, with: newRaw, rawNodeArena: rawNodeArena, allocationArena: allocationArena)
       let newParent = Syntax(parentData)
       return SyntaxData(absoluteRaw.replacingSelf(newRaw, newRootId: parentData.nodeId.rootId), parent: newParent)
     } else {
       // Otherwise, we're already the root, so return the new root data.
-      return .forRoot(newRaw, arena: arena)
+      return .forRoot(newRaw, rawNodeArena: rawNodeArena)
     }
   }
 
@@ -327,43 +331,50 @@ struct SyntaxData {
   ///   - index: The index pointing to where in the raw layout to place this
   ///            child.
   ///   - newChild: The raw syntax for the new child to replace.
+  ///   - newChildArena: The arena in which `newChild` resides.
   ///   - arena: The arena in which the new node will be allocated.
   /// - Returns: The new root node created by this operation, and the new child
   ///            syntax data.
   /// - SeeAlso: replacingSelf(_:)
-  func replacingChild(at index: Int, with newChild: RawSyntax?, arena: SyntaxArena) -> SyntaxData {
-    let newRaw = raw.layoutView!.replacingChild(at: index, with: newChild, arena: arena)
-    return replacingSelf(newRaw, arena: arena)
+  func replacingChild(at index: Int, with newChild: RawSyntax?, rawNodeArena: SyntaxArena?, allocationArena: SyntaxArena) -> SyntaxData {
+    precondition(newChild?.arena === rawNodeArena || newChild == nil)
+    // After newRaw has been allocated in `allocationArena`, `rawNodeArena` will
+    // be a child arena of `allocationArena` and thus, `allocationArena` will
+    // keep `newChild` alive.
+    let newRaw = withExtendedLifetime(rawNodeArena) {
+      raw.layoutView!.replacingChild(at: index, with: newChild, arena: allocationArena)
+    }
+    return replacingSelf(newRaw, rawNodeArena: allocationArena, allocationArena: allocationArena)
   }
 
   /// Identical to `replacingChild(at: Int, with: RawSyntax?, arena: SyntaxArena)`
   /// that ensures that the arena of`newChild` doesn’t get de-allocated before
   /// `newChild` has been addded to the result.
   func replacingChild(at index: Int, with newChild: SyntaxData?, arena: SyntaxArena) -> SyntaxData {
     return withExtendedLifetime(newChild) {
-      return replacingChild(at: index, with: newChild?.raw, arena: arena)
+      return replacingChild(at: index, with: newChild?.raw, rawNodeArena: newChild?.raw.arena, allocationArena: arena)
     }
   }
 
   func withLeadingTrivia(_ leadingTrivia: Trivia, arena: SyntaxArena) -> SyntaxData {
     if let raw = raw.withLeadingTrivia(leadingTrivia, arena: arena) {
-      return replacingSelf(raw, arena: arena)
+      return replacingSelf(raw, rawNodeArena: arena, allocationArena: arena)
     } else {
       return self
     }
   }
 
   func withTrailingTrivia(_ trailingTrivia: Trivia, arena: SyntaxArena) -> SyntaxData {
     if let raw = raw.withTrailingTrivia(trailingTrivia, arena: arena) {
-      return replacingSelf(raw, arena: arena)
+      return replacingSelf(raw, rawNodeArena: arena, allocationArena: arena)
     } else {
       return self
     }
   }
 
   func withPresence(_ presence: SourcePresence, arena: SyntaxArena) -> SyntaxData {
     if let raw = raw.tokenView?.withPresence(presence, arena: arena) {
-      return replacingSelf(raw, arena: arena)
+      return replacingSelf(raw, rawNodeArena: arena, allocationArena: arena)
     } else {
       return self
     }

Sources/SwiftSyntax/TokenSyntax.swift

@@ -58,7 +58,7 @@ public struct TokenSyntax: SyntaxProtocol, SyntaxHashable {
       tokenDiagnostic: nil,
       arena: arena
     )
-    self.init(SyntaxData.forRoot(raw, arena: arena))
+    self.init(SyntaxData.forRoot(raw, rawNodeArena: arena))
   }
 
   /// Whether the token is present or missing.
@@ -107,7 +107,7 @@ public struct TokenSyntax: SyntaxProtocol, SyntaxHashable {
       }
       let arena = SyntaxArena()
       let newRaw = tokenView.withKind(newValue, arena: arena)
-      let newData = data.replacingSelf(newRaw, arena: arena)
+      let newData = data.replacingSelf(newRaw, rawNodeArena: arena, allocationArena: arena)
       self = TokenSyntax(newData)
     }
   }