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/CMakeLists.txt

@@ -25,6 +25,7 @@ add_swift_host_library(SwiftSyntax
   SyntaxCollection.swift
   SyntaxData.swift
   SyntaxText.swift
+  SyntaxTracking.swift
   SyntaxTreeViewMode.swift
   TokenDiagnostic.swift
   TokenSyntax.swift

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

@@ -206,9 +206,11 @@ struct SyntaxData {
     // 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
       }
     }
 
@@ -270,6 +272,18 @@ struct SyntaxData {
     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
   var position: AbsolutePosition {
     AbsolutePosition(utf8Offset: Int(absoluteInfo.offset))
@@ -317,7 +331,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(Wrapper(Info.Root(arena: rawNodeArena))))
+    return SyntaxData(raw, info: .root(Wrapper(Info.Root(arena: rawNodeArena, syntaxTracking: nil))))
   }
 
   /// Returns the child data at the provided index in this data's layout.
@@ -387,7 +401,18 @@ 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.pointee.syntaxTracking = syntaxTracking?.replacing(
+          oldIndexInTree: info.nodeId.indexInTree,
+          oldTotalNodes: raw?.totalNodes ?? 0,
+          by: newChild
+        )
+      }
+      return result
     }
   }