Add an arena parameter to SyntaxData.forRoot

The `arena` parameter must be the same one that `raw` is allocated in and ensures that the arena of `raw` doesn’t get deallocated before the `SyntaxData` can retain it.

This removes the need for a few `withExtendedLifetime` calls that we had earlier.

Author: ahoppen

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

@@ -63,9 +63,16 @@ let parserEntryFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
         """
         extension \(node.kind.syntaxType): SyntaxParseable {
           public static func parse(from parser: inout Parser) -> Self {
+            // Keep the parser alive so that the arena in which `raw` is allocated
+            // doesn’t get deallocated before we have a chance to create a syntax node
+            // from it. We can’t use `parser.arena` as the parameter to
+            // `Syntax(raw:arena:)` because the node might have been re-used during an
+            // incremental parse and would then live in a different arena than
+            // `parser.arena`.
+            defer { withExtendedLifetime(parser) {} }
             let node = parser.\(raw: parserFunction)()
             let raw = RawSyntax(parser.parseRemainder(into: node))
-            return Syntax(raw: raw).cast(Self.self)
+            return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
           }
         }
         """

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

@@ -168,12 +168,10 @@ let syntaxCollectionsFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
       DeclSyntax(
         """
         public init(_ children: [Element]) {
-          let data: SyntaxData = withExtendedLifetime(SyntaxArena()) { arena in
-            let raw = RawSyntax.makeLayout(kind: SyntaxKind.\(node.varOrCaseName),
-                                           from: children.map { $0.raw }, arena: arena)
-            return SyntaxData.forRoot(raw)
-          }
-          self.init(data)
+          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))
         }
         """
       )

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

@@ -129,7 +129,7 @@ func syntaxNode(emitKind: SyntaxNodeKind) -> SourceFileSyntax {
                   """
                 )
               }
-              StmtSyntax("return SyntaxData.forRoot(raw)")
+              StmtSyntax("return SyntaxData.forRoot(raw, arena: arena)")
             }
           )
 

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

Sources/SwiftParser/generated/Parser+Entry.swift

@@ -43,129 +43,289 @@ public protocol SyntaxParseable: SyntaxProtocol {
 
 extension AccessorDeclSyntax: SyntaxParseable {
   public static func parse(from parser: inout Parser) -> Self {
+    // Keep the parser alive so that the arena in which `raw` is allocated
+    // doesn’t get deallocated before we have a chance to create a syntax node
+    // from it. We can’t use `parser.arena` as the parameter to
+    // `Syntax(raw:arena:)` because the node might have been re-used during an
+    // incremental parse and would then live in a different arena than
+    // `parser.arena`.
+    defer {
+      withExtendedLifetime(parser) {
+      }
+    }
     let node = parser.parseAccessorDecl()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw).cast(Self.self)
+    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
   }
 }
 
 extension AttributeSyntax: SyntaxParseable {
   public static func parse(from parser: inout Parser) -> Self {
+    // Keep the parser alive so that the arena in which `raw` is allocated
+    // doesn’t get deallocated before we have a chance to create a syntax node
+    // from it. We can’t use `parser.arena` as the parameter to
+    // `Syntax(raw:arena:)` because the node might have been re-used during an
+    // incremental parse and would then live in a different arena than
+    // `parser.arena`.
+    defer {
+      withExtendedLifetime(parser) {
+      }
+    }
     let node = parser.parseAttribute()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw).cast(Self.self)
+    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
   }
 }
 
 extension CatchClauseSyntax: SyntaxParseable {
   public static func parse(from parser: inout Parser) -> Self {
+    // Keep the parser alive so that the arena in which `raw` is allocated
+    // doesn’t get deallocated before we have a chance to create a syntax node
+    // from it. We can’t use `parser.arena` as the parameter to
+    // `Syntax(raw:arena:)` because the node might have been re-used during an
+    // incremental parse and would then live in a different arena than
+    // `parser.arena`.
+    defer {
+      withExtendedLifetime(parser) {
+      }
+    }
     let node = parser.parseCatchClause()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw).cast(Self.self)
+    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
   }
 }
 
 extension ClosureParameterSyntax: SyntaxParseable {
   public static func parse(from parser: inout Parser) -> Self {
+    // Keep the parser alive so that the arena in which `raw` is allocated
+    // doesn’t get deallocated before we have a chance to create a syntax node
+    // from it. We can’t use `parser.arena` as the parameter to
+    // `Syntax(raw:arena:)` because the node might have been re-used during an
+    // incremental parse and would then live in a different arena than
+    // `parser.arena`.
+    defer {
+      withExtendedLifetime(parser) {
+      }
+    }
     let node = parser.parseClosureParameter()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw).cast(Self.self)
+    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
   }
 }
 
 extension CodeBlockItemSyntax: SyntaxParseable {
   public static func parse(from parser: inout Parser) -> Self {
+    // Keep the parser alive so that the arena in which `raw` is allocated
+    // doesn’t get deallocated before we have a chance to create a syntax node
+    // from it. We can’t use `parser.arena` as the parameter to
+    // `Syntax(raw:arena:)` because the node might have been re-used during an
+    // incremental parse and would then live in a different arena than
+    // `parser.arena`.
+    defer {
+      withExtendedLifetime(parser) {
+      }
+    }
     let node = parser.parseNonOptionalCodeBlockItem()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw).cast(Self.self)
+    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
   }
 }
 
 extension DeclSyntax: SyntaxParseable {
   public static func parse(from parser: inout Parser) -> Self {
+    // Keep the parser alive so that the arena in which `raw` is allocated
+    // doesn’t get deallocated before we have a chance to create a syntax node
+    // from it. We can’t use `parser.arena` as the parameter to
+    // `Syntax(raw:arena:)` because the node might have been re-used during an
+    // incremental parse and would then live in a different arena than
+    // `parser.arena`.
+    defer {
+      withExtendedLifetime(parser) {
+      }
+    }
     let node = parser.parseDeclaration()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw).cast(Self.self)
+    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
   }
 }
 
 extension EnumCaseParameterSyntax: SyntaxParseable {
   public static func parse(from parser: inout Parser) -> Self {
+    // Keep the parser alive so that the arena in which `raw` is allocated
+    // doesn’t get deallocated before we have a chance to create a syntax node
+    // from it. We can’t use `parser.arena` as the parameter to
+    // `Syntax(raw:arena:)` because the node might have been re-used during an
+    // incremental parse and would then live in a different arena than
+    // `parser.arena`.
+    defer {
+      withExtendedLifetime(parser) {
+      }
+    }
     let node = parser.parseEnumCaseParameter()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw).cast(Self.self)
+    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
   }
 }
 
 extension ExprSyntax: SyntaxParseable {
   public static func parse(from parser: inout Parser) -> Self {
+    // Keep the parser alive so that the arena in which `raw` is allocated
+    // doesn’t get deallocated before we have a chance to create a syntax node
+    // from it. We can’t use `parser.arena` as the parameter to
+    // `Syntax(raw:arena:)` because the node might have been re-used during an
+    // incremental parse and would then live in a different arena than
+    // `parser.arena`.
+    defer {
+      withExtendedLifetime(parser) {
+      }
+    }
     let node = parser.parseExpression()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw).cast(Self.self)
+    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
   }
 }
 
 extension FunctionParameterSyntax: SyntaxParseable {
   public static func parse(from parser: inout Parser) -> Self {
+    // Keep the parser alive so that the arena in which `raw` is allocated
+    // doesn’t get deallocated before we have a chance to create a syntax node
+    // from it. We can’t use `parser.arena` as the parameter to
+    // `Syntax(raw:arena:)` because the node might have been re-used during an
+    // incremental parse and would then live in a different arena than
+    // `parser.arena`.
+    defer {
+      withExtendedLifetime(parser) {
+      }
+    }
     let node = parser.parseFunctionParameter()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw).cast(Self.self)
+    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
   }
 }
 
 extension GenericParameterClauseSyntax: SyntaxParseable {
   public static func parse(from parser: inout Parser) -> Self {
+    // Keep the parser alive so that the arena in which `raw` is allocated
+    // doesn’t get deallocated before we have a chance to create a syntax node
+    // from it. We can’t use `parser.arena` as the parameter to
+    // `Syntax(raw:arena:)` because the node might have been re-used during an
+    // incremental parse and would then live in a different arena than
+    // `parser.arena`.
+    defer {
+      withExtendedLifetime(parser) {
+      }
+    }
     let node = parser.parseGenericParameters()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw).cast(Self.self)
+    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
   }
 }
 
 extension MemberDeclBlockSyntax: SyntaxParseable {
   public static func parse(from parser: inout Parser) -> Self {
+    // Keep the parser alive so that the arena in which `raw` is allocated
+    // doesn’t get deallocated before we have a chance to create a syntax node
+    // from it. We can’t use `parser.arena` as the parameter to
+    // `Syntax(raw:arena:)` because the node might have been re-used during an
+    // incremental parse and would then live in a different arena than
+    // `parser.arena`.
+    defer {
+      withExtendedLifetime(parser) {
+      }
+    }
     let node = parser.parseMemberDeclList()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw).cast(Self.self)
+    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
   }
 }
 
 extension PatternSyntax: SyntaxParseable {
   public static func parse(from parser: inout Parser) -> Self {
+    // Keep the parser alive so that the arena in which `raw` is allocated
+    // doesn’t get deallocated before we have a chance to create a syntax node
+    // from it. We can’t use `parser.arena` as the parameter to
+    // `Syntax(raw:arena:)` because the node might have been re-used during an
+    // incremental parse and would then live in a different arena than
+    // `parser.arena`.
+    defer {
+      withExtendedLifetime(parser) {
+      }
+    }
     let node = parser.parsePattern()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw).cast(Self.self)
+    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
   }
 }
 
 extension SourceFileSyntax: SyntaxParseable {
   public static func parse(from parser: inout Parser) -> Self {
+    // Keep the parser alive so that the arena in which `raw` is allocated
+    // doesn’t get deallocated before we have a chance to create a syntax node
+    // from it. We can’t use `parser.arena` as the parameter to
+    // `Syntax(raw:arena:)` because the node might have been re-used during an
+    // incremental parse and would then live in a different arena than
+    // `parser.arena`.
+    defer {
+      withExtendedLifetime(parser) {
+      }
+    }
     let node = parser.parseSourceFile()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw).cast(Self.self)
+    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
   }
 }
 
 extension StmtSyntax: SyntaxParseable {
   public static func parse(from parser: inout Parser) -> Self {
+    // Keep the parser alive so that the arena in which `raw` is allocated
+    // doesn’t get deallocated before we have a chance to create a syntax node
+    // from it. We can’t use `parser.arena` as the parameter to
+    // `Syntax(raw:arena:)` because the node might have been re-used during an
+    // incremental parse and would then live in a different arena than
+    // `parser.arena`.
+    defer {
+      withExtendedLifetime(parser) {
+      }
+    }
     let node = parser.parseStatement()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw).cast(Self.self)
+    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
   }
 }
 
 extension SwitchCaseSyntax: SyntaxParseable {
   public static func parse(from parser: inout Parser) -> Self {
+    // Keep the parser alive so that the arena in which `raw` is allocated
+    // doesn’t get deallocated before we have a chance to create a syntax node
+    // from it. We can’t use `parser.arena` as the parameter to
+    // `Syntax(raw:arena:)` because the node might have been re-used during an
+    // incremental parse and would then live in a different arena than
+    // `parser.arena`.
+    defer {
+      withExtendedLifetime(parser) {
+      }
+    }
     let node = parser.parseSwitchCase()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw).cast(Self.self)
+    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
   }
 }
 
 extension TypeSyntax: SyntaxParseable {
   public static func parse(from parser: inout Parser) -> Self {
+    // Keep the parser alive so that the arena in which `raw` is allocated
+    // doesn’t get deallocated before we have a chance to create a syntax node
+    // from it. We can’t use `parser.arena` as the parameter to
+    // `Syntax(raw:arena:)` because the node might have been re-used during an
+    // incremental parse and would then live in a different arena than
+    // `parser.arena`.
+    defer {
+      withExtendedLifetime(parser) {
+      }
+    }
     let node = parser.parseType()
     let raw = RawSyntax(parser.parseRemainder(into: node))
-    return Syntax(raw: raw).cast(Self.self)
+    return Syntax(raw: raw, arena: raw.arena).cast(Self.self)
   }
 }
 

Sources/SwiftSyntax/Raw/RawSyntax.swift

@@ -209,7 +209,8 @@ public struct RawSyntax {
     rawData.arenaReference
   }
 
-  internal var arena: SyntaxArena {
+  @_spi(RawSyntax)
+  public var arena: SyntaxArena {
     rawData.arenaReference.value
   }