Check that we are actually tracking the tree passed in as originalTree to originalNode

Author: ahoppen

Sources/SwiftSyntax/SyntaxTracking.swift

@@ -289,6 +289,9 @@ extension SyntaxProtocol {
   /// - SeeAlso: ``SyntaxProtocol/tracked``
   /// - Complexity: O(tracked ranges in this tree) + O(width * depth of original tree)
   public func originalNode(in originalTree: some SyntaxProtocol) -> Self? {
+    guard self.data.syntaxTracking?.trackedTree == originalTree.id.rootId else {
+      return nil
+    }
     guard let originalIndexInTree = self.data.syntaxTracking?.originalIndexInTree(of: self.id.indexInTree) else {
       return nil
     }

Tests/SwiftSyntaxTest/SyntaxTrackingTests.swift

@@ -105,4 +105,36 @@ public class SyntaxTrackingTests: XCTestCase {
 
     XCTAssertEqual(codeBlockItemList.last?.item.originalNode(in: originalFunction)?.id, originalFunction.id)
   }
+
+  func testTwoLayerModification() throws {
+    let originalFunction = try FunctionDeclSyntax("func foo() {}").tracked
+    let modifiedFunction = originalFunction.with(\.name, "bar")
+    let secondModification = modifiedFunction.with(\.body, nil)
+
+    XCTAssertEqual(secondModification.funcKeyword.originalNode(in: originalFunction)?.id, originalFunction.funcKeyword.id)
+    XCTAssertNil(secondModification.name.originalNode(in: originalFunction))
+    XCTAssertNil(secondModification.name.originalNode(in: modifiedFunction))
+  }
+
+  func testEditAndThenCreation() throws {
+    let originalFunction = try FunctionDeclSyntax("func foo() {}").tracked
+    let modifiedFunction = originalFunction.with(\.name, "bar")
+
+    let codeBlockItem = CodeBlockItemSyntax(item: .decl(DeclSyntax(modifiedFunction)))
+
+    XCTAssertEqual(codeBlockItem.item.as(FunctionDeclSyntax.self)?.funcKeyword.originalNode(in: originalFunction)?.id, originalFunction.funcKeyword.id)
+    XCTAssertNil(codeBlockItem.item.as(FunctionDeclSyntax.self)?.name.originalNode(in: originalFunction)?.id)
+    XCTAssertNil(codeBlockItem.item.as(FunctionDeclSyntax.self)?.name.originalNode(in: modifiedFunction)?.id)
+  }
+
+  func testStartTrackingModifiedNode() throws {
+    let originalFunction = try FunctionDeclSyntax("func foo() {}").tracked
+    let modifiedFunction = originalFunction.with(\.name, "bar").tracked
+    let secondModification = modifiedFunction.with(\.body, nil)
+
+    XCTAssertNil(secondModification.funcKeyword.originalNode(in: originalFunction))
+    XCTAssertEqual(secondModification.funcKeyword.originalNode(in: modifiedFunction)?.id, modifiedFunction.funcKeyword.id)
+    XCTAssertNil(secondModification.name.originalNode(in: originalFunction))
+    XCTAssertEqual(secondModification.name.originalNode(in: modifiedFunction)?.id, modifiedFunction.name.id)
+  }
 }