[SourceKit] Support getting semantic tokens for macro expansion buffers

Author: ahoppen

lib/IDE/SourceEntityWalker.cpp

@@ -609,14 +609,19 @@ ASTWalker::PreWalkResult<Expr *> SemaAnnotator::walkToExprPre(Expr *E) {
     // We already visited the children.
     return doSkipChildren();
   } else if (auto ME = dyn_cast<MacroExpansionExpr>(E)) {
-    // Add a reference to the macro
-    auto macroRef = ME->getMacroRef();
-    if (auto *macroDecl = dyn_cast_or_null<MacroDecl>(macroRef.getDecl())) {
-      auto macroRefType = macroDecl->getDeclaredInterfaceType();
-      if (!passReference(
-              macroDecl, macroRefType, ME->getMacroNameLoc(),
-              ReferenceMetaData(SemaReferenceKind::DeclRef, llvm::None)))
-        return Action::Stop();
+    // Add a reference to the macro if this is a true macro expansion *expression*.
+    // If this is a `MacroExpansionExpr` that expands a declaration macro, the
+    // substitute decl will be visited by ASTWalker and we would be passing its
+    // reference if we didn't have this check.
+    if (!ME->getSubstituteDecl()) {
+      auto macroRef = ME->getMacroRef();
+      if (auto *macroDecl = dyn_cast_or_null<MacroDecl>(macroRef.getDecl())) {
+        auto macroRefType = macroDecl->getDeclaredInterfaceType();
+        if (!passReference(
+                           macroDecl, macroRefType, ME->getMacroNameLoc(),
+                           ReferenceMetaData(SemaReferenceKind::DeclRef, llvm::None)))
+          return Action::Stop();
+      }
     }
   }
 

test/SourceKit/Macros/macro_semantic_token.swift

@@ -0,0 +1,28 @@
+
+@freestanding(declaration)
+macro anonymousTypes(_: () -> String) = #externalMacro(module: "MacroDefinition", type: "DefineAnonymousTypesMacro")
+
+#anonymousTypes { "hello" }
+
+// REQUIRES: swift_swift_parser, executable_test
+
+// RUN: %empty-directory(%t)
+
+//##-- Prepare the macro plugin.
+// RUN: %host-build-swift -swift-version 5 -emit-library -o %t/%target-library-name(MacroDefinition) -module-name=MacroDefinition %S/../../Macros/Inputs/syntax_macro_definitions.swift -g -no-toolchain-stdlib-rpath
+
+// Check the output of the the `#anonymousTypes` macro
+// RUN: %sourcekitd-test -req=semantic-tokens @__swiftmacro_9MacroUser33_8C2BB8A10AE555140C0EDFDEB4A9572DLl14anonymousTypesfMf0_.swift -primary-file %s -- -swift-version 5 -load-plugin-library %t/%target-library-name(MacroDefinition) -module-name MacroUser %s | %FileCheck %s --check-prefix IN_BUFFER
+
+// Check that we get some semantic tokens. Checking exact offsets is brittle.
+// IN_BUFFER: source.lang.swift.ref.struct
+// IN_BUFFER: key.kind: source.lang.swift.ref.class
+
+// Check that we don't get semantic tokens inside the buffer when requesting semantic tokens for the outer file
+// RUN: %sourcekitd-test -req=semantic-tokens %s -- -swift-version 5 -load-plugin-library %t/%target-library-name(MacroDefinition) -module-name MacroUser %s | %FileCheck %s --check-prefix PRIMARY_FILE
+// Reference to String in line 2
+// PRIMARY_FILE: key.kind: source.lang.swift.ref.struct
+// Reference to macro in line 4
+// PRIMARY_FILE: key.kind: source.lang.swift.ref.macro
+// There should be no other references
+// PRIMARY_FILE-NOT: key.kind

tools/SourceKit/lib/SwiftLang/SwiftEditor.cpp

@@ -952,9 +952,39 @@ class SemanticAnnotator : public SourceEntityWalker {
 public:
 
   std::vector<SwiftSemanticToken> SemaToks;
+  bool IsWalkingMacroExpansionBuffer = false;
 
   SemanticAnnotator(SourceManager &SM, unsigned BufferID)
-    : SM(SM), BufferID(BufferID) {}
+      : SM(SM), BufferID(BufferID) {
+    if (auto GeneratedSourceInfo = SM.getGeneratedSourceInfo(BufferID)) {
+      switch (GeneratedSourceInfo->kind) {
+      case GeneratedSourceInfo::ExpressionMacroExpansion:
+      case GeneratedSourceInfo::FreestandingDeclMacroExpansion:
+      case GeneratedSourceInfo::AccessorMacroExpansion:
+      case GeneratedSourceInfo::MemberAttributeMacroExpansion:
+      case GeneratedSourceInfo::MemberMacroExpansion:
+      case GeneratedSourceInfo::PeerMacroExpansion:
+      case GeneratedSourceInfo::ConformanceMacroExpansion:
+      case GeneratedSourceInfo::ExtensionMacroExpansion:
+        IsWalkingMacroExpansionBuffer = true;
+        break;
+      case GeneratedSourceInfo::ReplacedFunctionBody:
+      case GeneratedSourceInfo::PrettyPrinted:
+        break;
+      }
+    }
+  }
+
+  MacroWalking getMacroWalkingBehavior() const override {
+    if (IsWalkingMacroExpansionBuffer) {
+      // When we are walking a macro expansion buffer, we need to set the macro
+      // walking behavior to walk the expansion, otherwise we skip over all the
+      // declarations in the buffer.
+      return MacroWalking::ArgumentsAndExpansion;
+    } else {
+      return SourceEntityWalker::getMacroWalkingBehavior();
+    }
+  }
 
   bool visitDeclReference(ValueDecl *D, CharSourceRange Range,
                           TypeDecl *CtorTyRef, ExtensionDecl *ExtTyRef, Type T,
@@ -996,6 +1026,12 @@ 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) {
+      return;
+    }
     unsigned ByteOffset = SM.getLocOffsetInBuffer(Range.getStart(), BufferID);
     unsigned Length = Range.getByteLength();
     auto Kind = ContextFreeCodeCompletionResult::getCodeCompletionDeclKind(D);