Implement tracking of syntax nodes if the syntax tree is edited

Author: ahoppen

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

@@ -152,6 +152,23 @@ func syntaxNode(nodesStartingWith: [Character]) -> SourceFileSyntax {
             initializer: InitializerClauseSyntax(value: initializer)
           )
 
+          VariableDeclSyntax(
+            .let,
+            name: "nodes",
+            initializer: InitializerClauseSyntax(
+              value: ArrayExprSyntax {
+                for child in node.children {
+                  if child.isOptional {
+                    ArrayElementSyntax(expression: ExprSyntax("\(child.varOrCaseName.backtickedIfNeeded)?.data"))
+                  } else {
+                    ArrayElementSyntax(expression: ExprSyntax("\(child.varOrCaseName.backtickedIfNeeded).data"))
+                  }
+                }
+              }
+            )
+          )
+          ExprSyntax("data.setSyntaxTrackingOfTree(SyntaxTracking(tracking: nodes))")
+
           ExprSyntax("self.init(data)")
         }
 

Sources/SwiftSyntax/Syntax.swift

@@ -235,7 +235,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).cast(Self.self)
+      result.data.setSyntaxTrackingOfTree(SyntaxTracking(trackingRoot: self.data))
+      return result
     }
   }
 }

Sources/SwiftSyntax/SyntaxCollection.swift

@@ -56,7 +56,9 @@ extension SyntaxCollection {
         arena: arena
       )
     }
-    self.init(SyntaxData.forRoot(raw, rawNodeArena: arena))
+    let data = SyntaxData.forRoot(raw, rawNodeArena: arena)
+    data.setSyntaxTrackingOfTree(SyntaxTracking(tracking: children.map { $0.data }))
+    self.init(data)
   }
 
   public init(arrayLiteral elements: Element...) {
@@ -269,7 +271,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 = self.data.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: element.data
+          )
+        }
+        self = replacingLayout(newLayout)
+        self.data.setSyntaxTrackingOfTree(newSyntaxTracking)
+      } else {
+        self = replacingLayout(newLayout)
+      }
     }
   }
 

Sources/SwiftSyntax/SyntaxData.swift

@@ -194,9 +194,11 @@ struct SyntaxData {
     // For root node.
     class Root {
       var arena: SyntaxArena
+      var syntaxTracking: SyntaxTracking?
 
-      init(arena: SyntaxArena) {
+      init(arena: SyntaxArena, syntaxTracking: SyntaxTracking?) {
         self.arena = arena
+        self.syntaxTracking = syntaxTracking
       }
     }
 
@@ -258,6 +260,18 @@ struct SyntaxData {
     absoluteInfo.nodeId
   }
 
+  var syntaxTracking: SyntaxTracking? {
+    rootInfo.syntaxTracking
+  }
+
+  /// Set the translation ranges of the entire tree.
+  ///
+  /// Must only be called once for every tree.
+  func setSyntaxTrackingOfTree(_ syntaxTracking: SyntaxTracking?) {
+    precondition(rootInfo.syntaxTracking == nil)
+    rootInfo.syntaxTracking = syntaxTracking
+  }
+
   /// The position of the start of this node's leading trivia
   var position: AbsolutePosition {
     AbsolutePosition(utf8Offset: Int(absoluteInfo.offset))
@@ -305,7 +319,7 @@ struct SyntaxData {
   ///     has a chance to retain it.
   static func forRoot(_ raw: RawSyntax, rawNodeArena: SyntaxArena) -> SyntaxData {
     precondition(rawNodeArena === raw.arena)
-    return SyntaxData(raw, info: .root(.init(arena: rawNodeArena)))
+    return SyntaxData(raw, info: .root(.init(arena: rawNodeArena, syntaxTracking: nil)))
   }
 
   /// Returns the child data at the provided index in this data's layout.
@@ -377,7 +391,14 @@ struct SyntaxData {
   /// `newChild` has been addded to the result.
   func replacingChild(at index: Int, with newChild: SyntaxData?, arena: SyntaxArena) -> SyntaxData {
     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.syntaxTracking = syntaxTracking?.replacing(oldIndexInTree: info.nodeId.indexInTree, oldTotalNodes: raw?.totalNodes ?? 0, by: newChild)
+      }
+      return result
     }
   }