Implement tracking of syntax nodes if the syntax tree is edited

Author: ahoppen

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

@@ -128,6 +128,20 @@ func syntaxNode(nodesStartingWith: [Character]) -> SourceFileSyntax {
             operator: ExprSyntax(AssignmentExprSyntax()),
             rightOperand: initializer
           )
+
+          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))"))
+                }
+              }
+            )
+          )
+          ExprSyntax("Syntax(self).setSyntaxTrackingOfTree(SyntaxTracking(tracking: nodes))")
         }
 
         for (index, child) in node.children.enumerated() {

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/Syntax.swift

@@ -24,9 +24,11 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
     // For root node.
     struct Root: Sendable {
       private let arena: RetainedSyntaxArena
+      var syntaxTracking: SyntaxTracking?
 
-      init(arena: RetainedSyntaxArena) {
+      init(arena: RetainedSyntaxArena, syntaxTracking: SyntaxTracking?) {
         self.arena = arena
+        self.syntaxTracking = syntaxTracking
       }
     }
 
@@ -52,6 +54,17 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
       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.
@@ -111,6 +124,20 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
     }
   }
 
+  /// 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
@@ -145,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))
@@ -192,14 +233,21 @@ 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(Syntax.Info.RefCountedRoot(Syntax.Info.Root(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(Syntax.Info.RefCountedRoot(Syntax.Info.Root(arena: RetainedSyntaxArena(rawNodeArena)))))
+      info: Info(
+        .root(
+          Syntax.Info.RefCountedRoot(Syntax.Info.Root(arena: RetainedSyntaxArena(rawNodeArena), syntaxTracking: nil))
+        )
+      )
     )
   }
 
@@ -296,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/SyntaxCollection.swift

@@ -50,6 +50,7 @@ extension SyntaxCollection {
       )
     }
     self = Syntax.forRoot(raw, rawNodeArena: RetainedSyntaxArena(arena)).cast(Self.self)
+    Syntax(self).setSyntaxTrackingOfTree(SyntaxTracking(tracking: children.map { Syntax($0) }))
   }
 
   public init(arrayLiteral elements: Element...) {
@@ -266,7 +267,29 @@ extension SyntaxCollection {
       let layoutRangeLowerBound = (subrange.lowerBound.data?.indexInParent).map(Int.init) ?? newLayout.endIndex
       let layoutRangeUpperBound = (subrange.upperBound.data?.indexInParent).map(Int.init) ?? newLayout.endIndex
       newLayout.replaceSubrange(layoutRangeLowerBound..<layoutRangeUpperBound, with: newElements.map { $0.raw })
-      self = replacingLayout(newLayout)
+      if var newSyntaxTracking = Syntax(self).syntaxTracking {
+        // Perform the tracking transformation in two steps.
+        // 1. Remove all old elements
+        newSyntaxTracking = newSyntaxTracking.replacing(
+          oldIndexInTree: subrange.lowerBound.data!.indexInTree,
+          oldTotalNodes: Int(
+            subrange.upperBound.data!.indexInTree.indexInTree - subrange.lowerBound.data!.indexInTree.indexInTree
+          ),
+          by: nil
+        )
+        // 2. Now insert all the new elements.
+        for element in newElements.reversed() {
+          newSyntaxTracking = newSyntaxTracking.replacing(
+            oldIndexInTree: subrange.lowerBound.data!.indexInTree,
+            oldTotalNodes: 0,
+            by: Syntax(element)
+          )
+        }
+        self = replacingLayout(newLayout)
+        Syntax(self).setSyntaxTrackingOfTree(newSyntaxTracking)
+      } else {
+        self = replacingLayout(newLayout)
+      }
     }
   }
 

Sources/SwiftSyntax/SyntaxIdentifier.swift

@@ -39,7 +39,7 @@ public struct SyntaxIdentifier: Comparable, Hashable, Sendable {
   /// This is similar to ``SyntaxIdentifier`` but does not store the root ID of the tree.
   /// It can thus be transferred across trees that are structurally equivalent, for example two copies of the same tree
   /// that live in different processes.
-  public struct SyntaxIndexInTree: Hashable, Sendable {
+  public struct SyntaxIndexInTree: Comparable, Hashable, Sendable {
     /// When traversing the syntax tree using a depth-first traversal, the index at which the node will be visited.
     let indexInTree: UInt32
 
@@ -78,6 +78,11 @@ public struct SyntaxIdentifier: Comparable, Hashable, Sendable {
     public init(fromOpaque opaque: UInt64) {
       self.indexInTree = UInt32(opaque)
     }
+
+    /// Returns `true` if `lhs` occurs before `rhs` in the tree.
+    public static func < (lhs: SyntaxIndexInTree, rhs: SyntaxIndexInTree) -> Bool {
+      return lhs.indexInTree < rhs.indexInTree
+    }
   }
 
   /// Unique value for the root node.

Sources/SwiftSyntax/SyntaxProtocol.swift

@@ -159,7 +159,9 @@ 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.arenaReference.retained).cast(Self.self)
+      let result = Syntax(raw: self.raw, rawNodeArena: self.raw.arenaReference.retained)
+      result.setSyntaxTrackingOfTree(SyntaxTracking(trackingRoot: Syntax(self)))
+      return result.cast(Self.self)
     }
   }
 }