Implement tracking of syntax nodes if the syntax tree is edited

Author: ahoppen

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

@@ -137,6 +137,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

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

Sources/SwiftSyntax/Syntax.swift

@@ -33,9 +33,11 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
     // For root node.
     struct Root {
       var arena: SyntaxArena
+      var syntaxTracking: SyntaxTracking?
 
-      init(arena: SyntaxArena) {
+      init(arena: SyntaxArena, syntaxTracking: SyntaxTracking?) {
         self.arena = arena
+        self.syntaxTracking = syntaxTracking
       }
     }
 
@@ -97,6 +99,18 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
     absoluteInfo.nodeId
   }
 
+  var syntaxTracking: SyntaxTracking? {
+    rootInfo.pointee.syntaxTracking
+  }
+
+  /// Set the translation ranges of the entire tree.
+  ///
+  /// Must only be called once for every tree.
+  func setSyntaxTrackingOfTree(_ syntaxTracking: SyntaxTracking?) {
+    precondition(rootInfo.pointee.syntaxTracking == nil)
+    rootInfo.pointee.syntaxTracking = syntaxTracking
+  }
+
   /// The position of the start of this node's leading trivia
   public var position: AbsolutePosition {
     AbsolutePosition(utf8Offset: Int(absoluteInfo.offset))
@@ -144,7 +158,7 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
   ///     has a chance to retain it.
   static func forRoot(_ raw: RawSyntax, rawNodeArena: SyntaxArena) -> Syntax {
     precondition(rawNodeArena === raw.arena)
-    return Syntax(raw, info: .root(Wrapper(Info.Root(arena: rawNodeArena))))
+    return Syntax(raw, info: .root(Wrapper(Info.Root(arena: rawNodeArena, syntaxTracking: nil))))
   }
 
   /// Returns the child data at the provided index in this data's layout.
@@ -215,7 +229,18 @@ 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(at: index, with: newChild?.raw, rawNodeArena: newChild?.raw.arena, allocationArena: arena)
+      let result = replacingChild(at: index, with: newChild?.raw, rawNodeArena: newChild?.raw.arena, allocationArena: arena)
+      if trackedTree != nil {
+        var iter = RawSyntaxChildren(absoluteRaw).makeIterator()
+        for _ in 0..<index { _ = iter.next() }
+        let (raw, info) = iter.next()!
+        result.rootInfo.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: arena).cast(Self.self)
+    Syntax(self).setSyntaxTrackingOfTree(SyntaxTracking(tracking: children.map { Syntax($0) }))
   }
 
   public init(arrayLiteral elements: Element...) {
@@ -260,7 +261,27 @@ 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/SyntaxProtocol.swift

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