Fallback to root node lookup for detached node

Author: mateusrodriguesxyz

Sources/SwiftSyntaxMacroExpansion/BasicMacroExpansionContext.swift

@@ -210,7 +210,9 @@ 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[Syntax(node)]
+      ?? sharedState.detachedNodes[Syntax(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)

Tests/SwiftSyntaxMacroExpansionTest/BodyMacroTests.swift

@@ -125,4 +125,51 @@ final class BodyMacroTests: XCTestCase {
       indentationWidth: indentationWidth
     )
   }
+
+  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
+    )
+  }
 }