Implement tracking of syntax nodes if the syntax tree is edited

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

@@ -60,53 +60,44 @@ func syntaxNode(nodesStartingWith: [Character]) -> SourceFileSyntax {
           \(node.generateInitializerDeclHeader())
           """
         ) {
-          let parameters = ClosureParameterListSyntax {
-            for child in node.children {
-              ClosureParameterSyntax(firstName: child.varOrCaseName.backtickedIfNeeded)
-            }
-          }
-
           let closureSignature = ClosureSignatureSyntax(
             parameterClause: .parameterClause(
               ClosureParameterClauseSyntax(
                 parameters: ClosureParameterListSyntax {
-                  ClosureParameterSyntax(firstName: .identifier("arena"))
-                  ClosureParameterSyntax(firstName: .wildcardToken())
+                  ClosureParameterSyntax(firstName: "arena")
                 }
               )
             )
           )
-          let layoutList = ArrayExprSyntax {
-            for child in node.children {
-              ArrayElementSyntax(
-                expression: MemberAccessExprSyntax(
-                  base: child.buildableType.optionalChained(
-                    expr: ExprSyntax("\(child.varOrCaseName.backtickedIfNeeded)")
-                  ),
-                  period: .periodToken(),
-                  name: "raw"
-                )
-              )
-            }
-          }
+          VariableDeclSyntax(
+            .let,
+            name: "nodes",
+            type: TypeAnnotationSyntax(type: TypeSyntax("[Syntax?]")),
+            initializer: InitializerClauseSyntax(
+              value: ArrayExprSyntax {
+                for child in node.children {
+                  ArrayElementSyntax(expression: ExprSyntax("Syntax(\(child.varOrCaseName.backtickedIfNeeded))"))
+                }
+              }
+            )
+          )
 
           let initializer = FunctionCallExprSyntax(
             calledExpression: ExprSyntax("withExtendedLifetime"),
             leftParen: .leftParenToken(),
             arguments: LabeledExprListSyntax {
-              LabeledExprSyntax(expression: ExprSyntax("(SyntaxArena(), (\(parameters)))"))
+              LabeledExprSyntax(expression: ExprSyntax("SyntaxArena()"))
             },
             rightParen: .rightParenToken(),
             trailingClosure: ClosureExprSyntax(signature: closureSignature) {
               if node.children.isEmpty {
                 DeclSyntax("let raw = RawSyntax.makeEmptyLayout(kind: SyntaxKind.\(node.varOrCaseName), arena: arena)")
               } else {
-                DeclSyntax("let layout: [RawSyntax?] = \(layoutList)")
                 DeclSyntax(
                   """
                   let raw = RawSyntax.makeLayout(
                     kind: SyntaxKind.\(node.varOrCaseName),
-                    from: layout,
+                    from: nodes,
                     arena: arena,
                     leadingTrivia: leadingTrivia,
                     trailingTrivia: trailingTrivia
@@ -128,6 +119,8 @@ func syntaxNode(nodesStartingWith: [Character]) -> SourceFileSyntax {
             operator: ExprSyntax(AssignmentExprSyntax()),
             rightOperand: initializer
           )
+
+          ExprSyntax("Syntax(self).setSyntaxTrackingOfTree(SyntaxTracking(tracking: nodes))")
         }
 
         for (index, child) in node.children.enumerated() {
@@ -188,7 +181,7 @@ func syntaxNode(nodesStartingWith: [Character]) -> SourceFileSyntax {
                   collection = col.layoutView!.appending(element.raw, arena: arena)
                 } else {
                   collection = RawSyntax.makeLayout(kind: SyntaxKind.\(childNode.varOrCaseName),
-                                                    from: [element.raw], arena: arena)
+                                                    from: [Syntax(element)], arena: arena)
                 }
                 return Syntax(self)
                   .replacingChild(at: \(raw: index), with: collection, rawNodeArena: arena, allocationArena: arena)

Release Notes/601.md

@@ -6,7 +6,7 @@
   - Description: Allows retrieving the represented literal value when valid.
   - Issue: https://github.com/apple/swift-syntax/issues/405
   - Pull Request: https://github.com/apple/swift-syntax/pull/2605
-    
+
 - `SyntaxProtocol` now has a method `ancestorOrSelf`.
   - Description: Returns the node or the first ancestor that satisfies `condition`.
   - Pull Request: https://github.com/swiftlang/swift-syntax/pull/2696
@@ -19,6 +19,11 @@
   - Description: This new library provides facilities for evaluating `#if` conditions and determining which regions of a syntax tree are active according to a given build configuration.
   - Pull Request: https://github.com/swiftlang/swift-syntax/pull/1816
 
+- `SyntaxProtocol.tracked` / `SyntaxProtocol.originalNode(in:)`
+  - Description: `tracked` enables node tracking of a tree. For every tree derived from a tracked tree, `originalNode(in:)` returns the original node in the tracked tree. This allows clients to e.g. get the original location of a node in a source file after a tree has been modified.
+  - Issue: rdar://112679655
+  - Pull Request: https://github.com/apple/swift-syntax/pull/2118
+
 ## API Behavior Changes
 
 - `SyntaxProtocol.trimmed` detaches the node
@@ -38,7 +43,7 @@
   - Description: Allows retrieving the radix value from the `literal.text`.
   - Issue: https://github.com/apple/swift-syntax/issues/405
   - Pull Request: https://github.com/apple/swift-syntax/pull/2605
-  
+
 - `FixIt.Change` gained a new case `replaceChild(data:)`.
   - Description: The new case covers the replacement of a child node with another node.
   - Issue: https://github.com/swiftlang/swift-syntax/issues/2205

Sources/SwiftSyntax/AbsoluteRawSyntax.swift

@@ -39,7 +39,7 @@ struct AbsoluteRawSyntax: Sendable {
     return nil
   }
 
-  func replacingSelf(_ newRaw: RawSyntax, newRootId: UInt) -> AbsoluteRawSyntax {
+  func replacingSelf(_ newRaw: RawSyntax, newRootId: RootID) -> AbsoluteRawSyntax {
     let nodeId = SyntaxIdentifier(rootId: newRootId, indexInTree: info.nodeId.indexInTree)
     let newInfo = AbsoluteSyntaxInfo(position: info.position, nodeId: nodeId)
     return .init(raw: newRaw, info: newInfo)

Sources/SwiftSyntax/CMakeLists.txt

@@ -34,6 +34,7 @@ add_swift_syntax_library(SwiftSyntax
   SyntaxNodeStructure.swift
   SyntaxProtocol.swift
   SyntaxText.swift
+  SyntaxTracking.swift
   SyntaxTreeViewMode.swift
   TokenDiagnostic.swift
   TokenSequence.swift

Sources/SwiftSyntax/Raw/RawSyntax.swift

@@ -817,13 +817,13 @@ extension RawSyntax {
 
   static func makeLayout(
     kind: SyntaxKind,
-    from collection: some Collection<RawSyntax?>,
+    from collection: some Collection<Syntax?>,
     arena: __shared SyntaxArena,
     leadingTrivia: Trivia? = nil,
     trailingTrivia: Trivia? = nil
   ) -> RawSyntax {
     if leadingTrivia != nil || trailingTrivia != nil {
-      var layout = Array(collection)
+      var layout = collection.map { $0?.raw }
       if let leadingTrivia = leadingTrivia,
         // Find the index of the first non-empty node so we can attach the trivia to it.
         let idx = layout.firstIndex(where: { $0 != nil && ($0!.isToken || $0!.totalNodes > 1) })
@@ -842,11 +842,13 @@ extension RawSyntax {
           arena: arena
         )
       }
-      return .makeLayout(kind: kind, from: layout, arena: arena)
+      return .makeLayout(kind: kind, uninitializedCount: collection.count, arena: arena) {
+        _ = $0.initialize(from: layout)
+      }
     }
 
     return .makeLayout(kind: kind, uninitializedCount: collection.count, arena: arena) {
-      _ = $0.initialize(from: collection)
+      _ = $0.initialize(from: collection.lazy.map { $0?.raw })
     }
   }
 }

Sources/SwiftSyntax/Syntax.swift

@@ -23,21 +23,71 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
   final class Info: @unchecked Sendable {
     // For root node.
     struct Root: Sendable {
-      private var arena: RetainedSyntaxArena
+      private let arena: RetainedSyntaxArena
+      var syntaxTracking: SyntaxTracking?
 
-      init(arena: RetainedSyntaxArena) {
+      init(arena: RetainedSyntaxArena, syntaxTracking: SyntaxTracking?) {
         self.arena = arena
+        self.syntaxTracking = syntaxTracking
+      }
+    }
+
+    /// Class that owns a `Root`, is responsible for keeping it alive and also offers an unsafe pointer way to
+    /// access it.
+    ///
+    /// This way, the root node of a tree can own the root info and all other nodes in the tree can have an unsafe
+    /// pointer reference to the root info, which doesn't involve ref counting.
+    final class RefCountedRoot: Sendable {
+      /// `nonisolated(unsafe)` if fine because there are only two ways this gets accessed:
+      ///  - `pointer`: Here we reference `value` via inout to get a pointer to `Root` but the pointer is not mutable
+      ///    so no mutation happens here
+      #if swift(>=6)
+      private nonisolated(unsafe) var value: Root
+      #else
+      private var value: Root
+      #endif
+
+      fileprivate init(_ value: Root) {
+        self.value = value
+      }
+
+      fileprivate var pointer: UnsafePointer<Root> {
+        return withUnsafePointer(to: &value) { $0 }
+      }
+
+      /// Get a reference to the root info which can be mutated.
+      ///
+      ///  - Warning: This must only be used if the caller is guaranteed to have exclusive access to the tree and no
+      ///    concurrent accesses can happen. Effectively, this can only be guaranteed if the tree has just been created,
+      ///    and it hasn't been returned to any function which might concurrently access it.
+      ///    In practice, this should only be used to set the `SyntaxTracking` on the root after a new tree has been
+      ///    created (eg. by replacing a child in an existing tree) but before that new tree is returned to the client.
+      fileprivate var mutablePointer: UnsafeMutablePointer<Root> {
+        return withUnsafeMutablePointer(to: &value) { $0 }
       }
     }
 
     // For non-root nodes.
     struct NonRoot: Sendable {
       var parent: Syntax
       var absoluteInfo: AbsoluteSyntaxInfo
+      // `nonisolated(unsafe)` is fine because `Root` is owned by `RefCountedRoot` and `RefCountedRoot` guarantees that
+      // `Root` is not changing after the tree has been created.
+      #if swift(>=6)
+      nonisolated(unsafe) var rootInfo: UnsafePointer<Root>
+      #else
+      var rootInfo: UnsafePointer<Root>
+      #endif
+
+      init(parent: Syntax, absoluteInfo: AbsoluteSyntaxInfo, rootInfo: UnsafePointer<Root>) {
+        self.parent = parent
+        self.absoluteInfo = absoluteInfo
+        self.rootInfo = rootInfo
+      }
     }
 
     enum InfoImpl: Sendable {
-      case root(Root)
+      case root(RefCountedRoot)
       case nonRoot(NonRoot)
     }
 
@@ -67,13 +117,27 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
   var info: Info!
   let raw: RawSyntax
 
-  private var rootInfo: Info.Root {
+  var rootInfo: UnsafePointer<Info.Root> {
     switch info.info! {
-    case .root(let info): return info
+    case .root(let info): return info.pointer
     case .nonRoot(let info): return info.parent.rootInfo
     }
   }
 
+  /// Get a reference to the root info which can be mutated.
+  ///
+  ///  - Warning: This must only be used if the caller is guaranteed to have exclusive access to the tree and no
+  ///    concurrent accesses can happen. Effectively, this can only be guaranteed if the tree has just been created,
+  ///    and it hasn't been returned to any function which might concurrently access it.
+  ///    In practice, this should only be used to set the `SyntaxTracking` on the root after a new tree has been
+  ///    created (eg. by replacing a child in an existing tree) but before that new tree is returned to the client.
+  var mutableRootInfo: UnsafeMutablePointer<Info.Root> {
+    switch info.info! {
+    case .root(let info): return info.mutablePointer
+    case .nonRoot(let info): return info.parent.mutableRootInfo
+    }
+  }
+
   private var nonRootInfo: Info.NonRoot? {
     switch info.info! {
     case .root(_): return nil
@@ -108,6 +172,20 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
     absoluteInfo.nodeId
   }
 
+  var syntaxTracking: SyntaxTracking? {
+    rootInfo.pointee.syntaxTracking
+  }
+
+  /// Set the translation ranges of the entire tree.
+  ///
+  ///  - Warning: This must only be used if the caller is guaranteed to have exclusive access to the tree and no
+  ///    concurrent accesses can happen. Effectively, this can only be guaranteed if the tree has just been created,
+  ///    and it hasn't been returned to any function which might concurrently access it.
+  func setSyntaxTrackingOfTree(_ syntaxTracking: SyntaxTracking?) {
+    precondition(rootInfo.pointee.syntaxTracking == nil)
+    mutableRootInfo.pointee.syntaxTracking = syntaxTracking
+  }
+
   /// The position of the start of this node's leading trivia
   public var position: AbsolutePosition {
     AbsolutePosition(utf8Offset: Int(absoluteInfo.offset))
@@ -135,7 +213,7 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
   }
 
   init(_ raw: RawSyntax, parent: Syntax, absoluteInfo: AbsoluteSyntaxInfo) {
-    self.init(raw, info: Info(.nonRoot(.init(parent: parent, absoluteInfo: absoluteInfo))))
+    self.init(raw, info: Info(.nonRoot(.init(parent: parent, absoluteInfo: absoluteInfo, rootInfo: parent.rootInfo))))
   }
 
   /// Creates a `Syntax` with the provided raw syntax and parent.
@@ -155,12 +233,22 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
   ///     has a chance to retain it.
   static func forRoot(_ raw: RawSyntax, rawNodeArena: RetainedSyntaxArena) -> Syntax {
     precondition(rawNodeArena == raw.arenaReference)
-    return Syntax(raw, info: Info(.root(.init(arena: rawNodeArena))))
+    return Syntax(
+      raw,
+      info: Info(.root(Syntax.Info.RefCountedRoot(Syntax.Info.Root(arena: rawNodeArena, syntaxTracking: nil))))
+    )
   }
 
   static func forRoot(_ raw: RawSyntax, rawNodeArena: SyntaxArena) -> Syntax {
     precondition(rawNodeArena == raw.arenaReference)
-    return Syntax(raw, info: Info(.root(.init(arena: RetainedSyntaxArena(rawNodeArena)))))
+    return Syntax(
+      raw,
+      info: Info(
+        .root(
+          Syntax.Info.RefCountedRoot(Syntax.Info.Root(arena: RetainedSyntaxArena(rawNodeArena), syntaxTracking: nil))
+        )
+      )
+    )
   }
 
   /// Returns the child data at the provided index in this data's layout.
@@ -256,12 +344,23 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
   /// `newChild` has been addded to the result.
   func replacingChild(at index: Int, with newChild: Syntax?, arena: SyntaxArena) -> Syntax {
     return withExtendedLifetime(newChild) {
-      return replacingChild(
+      let result = replacingChild(
         at: index,
         with: newChild?.raw,
         rawNodeArena: newChild?.raw.arenaReference.retained,
         allocationArena: arena
       )
+      if trackedTree != nil {
+        var iter = RawSyntaxChildren(absoluteRaw).makeIterator()
+        for _ in 0..<index { _ = iter.next() }
+        let (raw, info) = iter.next()!
+        result.mutableRootInfo.pointee.syntaxTracking = syntaxTracking?.replacing(
+          oldIndexInTree: info.nodeId.indexInTree,
+          oldTotalNodes: raw?.totalNodes ?? 0,
+          by: newChild
+        )
+      }
+      return result
     }
   }
 

Sources/SwiftSyntax/SyntaxChildren.swift

@@ -102,7 +102,7 @@ public struct SyntaxChildrenIndex: Hashable, Comparable, ExpressibleByNilLiteral
 
 fileprivate extension AbsoluteSyntaxInfo {
   /// Construct `AbsoluteSyntaxInfo` from the given index data and a `rootId`.
-  init(index: SyntaxChildrenIndexData, rootId: UInt) {
+  init(index: SyntaxChildrenIndexData, rootId: RootID) {
     let position = AbsoluteSyntaxPosition(
       offset: index.offset,
       indexInParent: index.indexInParent
@@ -152,7 +152,7 @@ struct RawSyntaxChildren: BidirectionalCollection, Sendable {
   }
 
   /// The rootId of the tree the child nodes belong to
-  private let rootId: UInt
+  private let rootId: RootID
 
   /// The number of children in `parent`. Cached to avoid reaching into `parent` for every index
   /// advancement
@@ -222,7 +222,7 @@ struct RawSyntaxChildren: BidirectionalCollection, Sendable {
       let offset = startIndex.offset + UInt32(parent.totalLength.utf8Length)
       let indexInParent = startIndex.indexInParent + UInt32(parentLayoutView.children.count)
       let indexInTree = startIndex.indexInTree.indexInTree + UInt32(parent.totalNodes) - 1
-      let syntaxIndexInTree = SyntaxIdentifier.SyntaxIndexInTree(indexInTree: indexInTree)
+      let syntaxIndexInTree = SyntaxIdentifier.SyntaxIndexInTree(indexInTree)
       let materialized = SyntaxChildrenIndex(
         offset: offset,
         indexInParent: indexInParent,