From 75177e0555cda4101305831e8e68001f8e6909e7 Mon Sep 17 00:00:00 2001
From: Mateus Rodrigues <mateus.fr26@gmail.com>
Date: Fri, 9 Aug 2024 09:47:42 -0300
Subject: [PATCH] Use root node for detached node lookup

---
 .../BasicMacroExpansionContext.swift          |  2 +-
 .../BodyMacroTests.swift                      | 46 +++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/Sources/SwiftSyntaxMacroExpansion/BasicMacroExpansionContext.swift b/Sources/SwiftSyntaxMacroExpansion/BasicMacroExpansionContext.swift
index 4701332d1fb..87e044cc18c 100644
--- a/Sources/SwiftSyntaxMacroExpansion/BasicMacroExpansionContext.swift
+++ b/Sources/SwiftSyntaxMacroExpansion/BasicMacroExpansionContext.swift
@@ -210,7 +210,7 @@ extension BasicMacroExpansionContext: MacroExpansionContext {
       // The syntax node came from the source file itself.
       rootSourceFile = directRootSourceFile
       offsetAdjustment = .zero
-    } else if let nodeInOriginalTree = sharedState.detachedNodes[Syntax(node)] {
+    } else if let nodeInOriginalTree = sharedState.detachedNodes[node.root] {
       // The syntax node came from a disconnected root, so adjust for that.
       rootSourceFile = nodeInOriginalTree.root.as(SourceFileSyntax.self)
       offsetAdjustment = SourceLength(utf8Length: nodeInOriginalTree.position.utf8Offset)
diff --git a/Tests/SwiftSyntaxMacroExpansionTest/BodyMacroTests.swift b/Tests/SwiftSyntaxMacroExpansionTest/BodyMacroTests.swift
index 14cc32ff164..af483830071 100644
--- a/Tests/SwiftSyntaxMacroExpansionTest/BodyMacroTests.swift
+++ b/Tests/SwiftSyntaxMacroExpansionTest/BodyMacroTests.swift
@@ -155,4 +155,50 @@ final class BodyMacroTests: XCTestCase {
       macros: ["EmptyBody": EmptyBodyMacro.self]
     )
   }
+
+  func testBodyNodeLocationFromContext() {
+    struct SourceLocationMacro: BodyMacro {
+      public static var formatMode: FormatMode { .disabled }
+
+      public static func expansion(
+        of node: AttributeSyntax,
+        providingBodyFor declaration: some DeclSyntaxProtocol & WithOptionalCodeBlockSyntax,
+        in context: some MacroExpansionContext
+      ) throws -> [CodeBlockItemSyntax] {
+        guard let statements = declaration.body?.statements else {
+          return []
+        }
+        let body =
+          if let location = context.location(of: statements, at: .afterLeadingTrivia, filePathMode: .filePath) {
+            CodeBlockItemListSyntax {
+              "#sourceLocation(file: \(location.file), line: \(location.line))"
+              statements
+              "#sourceLocation()"
+            }
+          } else {
+            statements
+          }
+        return body.map(\.self)
+      }
+    }
+
+    assertMacroExpansion(
+      """
+      @SourceLocationMacro
+      func f() {
+        let x: Int = 1
+      }
+      """,
+      expandedSource:
+        """
+        func f() {
+          #sourceLocation(file: "test.swift", line: 3)
+          let x: Int = 1
+          #sourceLocation()
+        }
+        """,
+      macros: ["SourceLocationMacro": SourceLocationMacro.self],
+      indentationWidth: indentationWidth
+    )
+  }
 }