[AST] Avoid walking macro expanded bodies in non-expanded mode

include/swift/AST/Decl.h

@@ -8065,6 +8065,11 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// type-checked.
   BraceStmt *getMacroExpandedBody() const;
 
+  /// Whether the body of the function has been expanded from a body macro.
+  bool isBodyMacroExpanded() const {
+    return getBodyExpandedStatus() == BodyExpandedStatus::Expanded;
+  }
+
   /// Retrieve the type-checked body of the given function, or \c nullptr if
   /// there's no body available.
   BraceStmt *getTypecheckedBody() const;

lib/AST/ASTWalker.cpp

@@ -558,7 +558,12 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
     if (WalkGenerics && visitTrailingRequirements(AFD))
       return true;
 
-    if (AFD->getBody(/*canSynthesize=*/false)) {
+    // If we're not walking macro expansions, avoid walking into a body if it
+    // was expanded from a macro.
+    auto SkipBody = !Walker.shouldWalkMacroArgumentsAndExpansion().second &&
+                    AFD->isBodyMacroExpanded();
+
+    if (!SkipBody && AFD->getBody(/*canSynthesize=*/false)) {
       AbstractFunctionDecl::BodyKind PreservedKind = AFD->getBodyKind();
       if (BraceStmt *S = cast_or_null<BraceStmt>(doIt(AFD->getBody())))
         AFD->setBody(S, PreservedKind);

test/SourceKit/SemanticTokens/issue-85705.swift

@@ -0,0 +1,40 @@
+// RUN: %empty-directory(%t)
+// RUN: split-file %s %t
+
+// RUN: %host-build-swift -emit-library -o %t/%target-library-name(MacroDefinition) -module-name=MacroDefinition %t/macro.swift -g -no-toolchain-stdlib-rpath
+
+// RUN: %sourcekitd-test -req=semantic-tokens %t/main.swift -- %t/main.swift -load-plugin-library %t/%target-library-name(MacroDefinition) > %t/result.response
+// RUN: diff -u %t/expected.response %t/result.response
+
+// REQUIRES: swift_swift_parser
+
+//--- macro.swift
+import SwiftSyntax
+import SwiftSyntaxMacros
+
+public struct LogMacro: BodyMacro {
+    public static func expansion(
+        of node: AttributeSyntax,
+        providingBodyFor declaration: some DeclSyntaxProtocol & WithOptionalCodeBlockSyntax,
+        in context: some MacroExpansionContext,
+    ) throws -> [CodeBlockItemSyntax] {
+        return ["print()"]
+    }
+}
+//--- main.swift
+@attached(body)
+public macro log() = #externalMacro(module: "MacroDefinition", type: "LogMacro")
+
+// Make sure we don't walk into the macro expanded body here.
+@log func foo() {}
+
+//--- expected.response
+{
+  key.semantic_tokens: [
+    {
+      key.kind: source.lang.swift.ref.macro,
+      key.offset: 161,
+      key.length: 3
+    }
+  ]
+}

tools/SourceKit/lib/SwiftLang/SwiftEditor.cpp

@@ -1027,12 +1027,11 @@ class SemanticAnnotator : public SourceEntityWalker {
     if (!Range.isValid())
       return;
 
-    // If we are walking into macro expansions, make sure we only report ranges
-    // from the requested buffer, not any buffers of child macro expansions.
-    if (IsWalkingMacroExpansionBuffer &&
-        SM.findBufferContainingLoc(Range.getStart()) != BufferID) {
+    // Make sure we only report from the requested buffer, not any buffers of
+    // child macro expansions.
+    if (SM.findBufferContainingLoc(Range.getStart()) != BufferID)
       return;
-    }
+
     unsigned ByteOffset = SM.getLocOffsetInBuffer(Range.getStart(), BufferID);
     unsigned Length = Range.getByteLength();
     auto Kind = ContextFreeCodeCompletionResult::getCodeCompletionDeclKind(D);