Introduce RetainedSyntaxArena so we don’t need to pass SyntaxArena across actor boundaries

Author: ahoppen

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

@@ -50,7 +50,7 @@ let syntaxRewriterFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
         }
 
         return withExtendedLifetime(rewritten) {
-          return Syntax(node).replacingSelf(rewritten.raw, rawNodeArena: rewritten.raw.arena, allocationArena: SyntaxArena())
+          return Syntax(node).replacingSelf(rewritten.raw, rawNodeArena: rewritten.raw.arenaReference.retained, allocationArena: SyntaxArena())
         }
       }
       """

Sources/SwiftSyntax/Raw/RawSyntax.swift

@@ -216,11 +216,6 @@ public struct RawSyntax: Sendable {
     rawData.arenaReference
   }
 
-  @_spi(RawSyntax)
-  public var arena: SyntaxArena {
-    rawData.arenaReference.value
-  }
-
   internal var payload: RawSyntaxData.Payload {
     get { rawData.payload }
   }

Sources/SwiftSyntax/Syntax.swift

@@ -19,12 +19,10 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
     indirect case nonRoot(NonRoot)
 
     // For root node.
-    struct Root: @unchecked Sendable {
-      // Unchecked conformance to sendable is fine because `arena` is not
-      // accessible. It is just used to keep the arena alive.
-      private var arena: SyntaxArena
+    struct Root: Sendable {
+      private var arena: RetainedSyntaxArena
 
-      init(arena: SyntaxArena) {
+      init(arena: RetainedSyntaxArena) {
         self.arena = arena
       }
     }
@@ -125,11 +123,16 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
   ///   - rawNodeArena: The arena in which `raw` is allocated. It is passed to
   ///     make sure the arena doesn’t get de-allocated before the ``Syntax``
   ///     has a chance to retain it.
-  static func forRoot(_ raw: RawSyntax, rawNodeArena: SyntaxArena) -> Syntax {
+  static func forRoot(_ raw: RawSyntax, rawNodeArena: RetainedSyntaxArena) -> Syntax {
     precondition(rawNodeArena == raw.arenaReference)
     return Syntax(raw, 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))))
+  }
+
   /// Returns the child data at the provided index in this data's layout.
   /// - Note: This has O(n) performance, prefer using a proper Sequence type
   ///         if applicable, instead of this.
@@ -157,7 +160,7 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
   ///   - allocationArena: 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) -> Syntax {
+  func replacingSelf(_ newRaw: RawSyntax, rawNodeArena: RetainedSyntaxArena, allocationArena: SyntaxArena) -> Syntax {
     precondition(newRaw.arenaReference == rawNodeArena)
     // If we have a parent already, then ask our current parent to copy itself
     // recursively up to the root.
@@ -182,45 +185,50 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
   /// - 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?, rawNodeArena: SyntaxArena?, allocationArena: SyntaxArena) -> Syntax {
+  func replacingChild(at index: Int, with newChild: RawSyntax?, rawNodeArena: RetainedSyntaxArena?, allocationArena: SyntaxArena) -> Syntax {
     precondition(newChild == nil || (rawNodeArena != nil && newChild!.arenaReference == rawNodeArena!))
     // 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)
+    return replacingSelf(newRaw, rawNodeArena: RetainedSyntaxArena(allocationArena), allocationArena: allocationArena)
+  }
+
+  /// Same as `replacingChild(at:with:rawNodeArena:allocationArena:)` but takes a `__SyntaxArena` instead of a `RetainedSyntaxArena`.
+  func replacingChild(at index: Int, with newChild: RawSyntax?, rawNodeArena: SyntaxArena?, allocationArena: SyntaxArena) -> Syntax {
+    return self.replacingChild(at: index, with: newChild, rawNodeArena: rawNodeArena.map(RetainedSyntaxArena.init), 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: Syntax?, arena: SyntaxArena) -> Syntax {
     return withExtendedLifetime(newChild) {
-      return replacingChild(at: index, with: newChild?.raw, rawNodeArena: newChild?.raw.arena, allocationArena: arena)
+      return replacingChild(at: index, with: newChild?.raw, rawNodeArena: newChild?.raw.arenaReference.retained, allocationArena: arena)
     }
   }
 
   func withLeadingTrivia(_ leadingTrivia: Trivia, arena: SyntaxArena) -> Syntax {
     if let raw = raw.withLeadingTrivia(leadingTrivia, arena: arena) {
-      return replacingSelf(raw, rawNodeArena: arena, allocationArena: arena)
+      return replacingSelf(raw, rawNodeArena: RetainedSyntaxArena(arena), allocationArena: arena)
     } else {
       return self
     }
   }
 
   func withTrailingTrivia(_ trailingTrivia: Trivia, arena: SyntaxArena) -> Syntax {
     if let raw = raw.withTrailingTrivia(trailingTrivia, arena: arena) {
-      return replacingSelf(raw, rawNodeArena: arena, allocationArena: arena)
+      return replacingSelf(raw, rawNodeArena: RetainedSyntaxArena(arena), allocationArena: arena)
     } else {
       return self
     }
   }
 
   func withPresence(_ presence: SourcePresence, arena: SyntaxArena) -> Syntax {
     if let raw = raw.tokenView?.withPresence(presence, arena: arena) {
-      return replacingSelf(raw, rawNodeArena: arena, allocationArena: arena)
+      return replacingSelf(raw, rawNodeArena: RetainedSyntaxArena(arena), allocationArena: arena)
     } else {
       return self
     }
@@ -234,10 +242,15 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
   }
 
   @_spi(RawSyntax)
-  public init(raw: RawSyntax, rawNodeArena: __shared SyntaxArena) {
+  public init(raw: RawSyntax, rawNodeArena: __shared RetainedSyntaxArena) {
     self = .forRoot(raw, rawNodeArena: rawNodeArena)
   }
 
+  @_spi(RawSyntax)
+  public init(raw: RawSyntax, rawNodeArena: __shared SyntaxArena) {
+    self = .forRoot(raw, rawNodeArena: RetainedSyntaxArena(rawNodeArena))
+  }
+
   /// Create a ``Syntax`` node from a specialized syntax node.
   public init(_ syntax: some SyntaxProtocol) {
     self = syntax._syntaxNode

Sources/SwiftSyntax/SyntaxArena.swift

@@ -234,6 +234,22 @@ public class ParsingSyntaxArena: SyntaxArena {
   }
 }
 
+/// An opaque wrapper around `SyntaxArena` that keeps the arena alive.
+@_spi(RawSyntax)
+public struct RetainedSyntaxArena: @unchecked Sendable {
+  // Unchecked conformance to sendable is fine because `arena` is not
+  // accessible. It is just used to keep the arena alive.
+  private let arena: SyntaxArena
+
+  init(_ arena: SyntaxArena) {
+    self.arena = arena
+  }
+
+  fileprivate func arenaRef() -> SyntaxArenaRef {
+    return SyntaxArenaRef(arena)
+  }
+}
+
 /// Unsafely unowned reference to ``SyntaxArena``. The user is responsible to
 /// maintain the lifetime of the ``SyntaxArena``.
 ///
@@ -267,6 +283,11 @@ struct SyntaxArenaRef: Hashable, @unchecked Sendable {
     self._value.release()
   }
 
+  /// Get an opaque wrapper that keeps the syntax arena alive.
+  var retained: RetainedSyntaxArena {
+    return RetainedSyntaxArena(_value.takeUnretainedValue())
+  }
+
   func hash(into hasher: inout Hasher) {
     hasher.combine(_value.toOpaque())
   }
@@ -282,4 +303,12 @@ struct SyntaxArenaRef: Hashable, @unchecked Sendable {
   static func == (lhs: __shared SyntaxArena, rhs: SyntaxArenaRef) -> Bool {
     return rhs == lhs
   }
+
+  static func == (lhs: SyntaxArenaRef, rhs: RetainedSyntaxArena) -> Bool {
+    return lhs == rhs.arenaRef()
+  }
+
+  static func == (lhs: RetainedSyntaxArena, rhs: SyntaxArenaRef) -> Bool {
+    return rhs == lhs
+  }
 }

Sources/SwiftSyntax/SyntaxCollection.swift

@@ -49,7 +49,7 @@ extension SyntaxCollection {
         arena: arena
       )
     }
-    self = Syntax.forRoot(raw, rawNodeArena: arena).cast(Self.self)
+    self = Syntax.forRoot(raw, rawNodeArena: RetainedSyntaxArena(arena)).cast(Self.self)
   }
 
   public init(arrayLiteral elements: Element...) {
@@ -70,7 +70,7 @@ extension SyntaxCollection {
   internal func replacingLayout(_ layout: [RawSyntax?]) -> Self {
     let arena = SyntaxArena()
     let newRaw = layoutView.replacingLayout(with: layout, arena: arena)
-    return Syntax(self).replacingSelf(newRaw, rawNodeArena: arena, allocationArena: arena).cast(Self.self)
+    return Syntax(self).replacingSelf(newRaw, rawNodeArena: RetainedSyntaxArena(arena), allocationArena: arena).cast(Self.self)
   }
 
   /// Return the index of `node` within this collection.

Sources/SwiftSyntax/SyntaxProtocol.swift

@@ -159,7 +159,7 @@ public 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, rawNodeArena: self.raw.arena).cast(Self.self)
+      return Syntax(raw: self.raw, rawNodeArena: self.raw.arenaReference.retained).cast(Self.self)
     }
   }
 }

Sources/SwiftSyntax/TokenSyntax.swift

@@ -57,7 +57,7 @@ public struct TokenSyntax: SyntaxProtocol, SyntaxHashable {
       tokenDiagnostic: nil,
       arena: arena
     )
-    self = Syntax.forRoot(raw, rawNodeArena: arena).cast(TokenSyntax.self)
+    self = Syntax.forRoot(raw, rawNodeArena: RetainedSyntaxArena(arena)).cast(TokenSyntax.self)
   }
 
   /// Whether the token is present or missing.
@@ -111,7 +111,7 @@ public struct TokenSyntax: SyntaxProtocol, SyntaxHashable {
       }
       let arena = SyntaxArena()
       let newRaw = tokenView.withKind(newValue, arena: arena)
-      self = Syntax(self).replacingSelf(newRaw, rawNodeArena: arena, allocationArena: arena).cast(TokenSyntax.self)
+      self = Syntax(self).replacingSelf(newRaw, rawNodeArena: RetainedSyntaxArena(arena), allocationArena: arena).cast(TokenSyntax.self)
     }
   }
 

Sources/SwiftSyntax/generated/SyntaxRewriter.swift

@@ -36,7 +36,7 @@ open class SyntaxRewriter {
     }
 
     return withExtendedLifetime(rewritten) {
-      return Syntax(node).replacingSelf(rewritten.raw, rawNodeArena: rewritten.raw.arena, allocationArena: SyntaxArena())
+      return Syntax(node).replacingSelf(rewritten.raw, rawNodeArena: rewritten.raw.arenaReference.retained, allocationArena: SyntaxArena())
     }
   }