2 files changed, 34 insertions, 17 deletions

diff --git a/clang-tools-extra/clangd/refactor/tweaks/AnnotateHighlightings.cpp b/clang-tools-extra/clangd/refactor/tweaks/AnnotateHighlightings.cpp
index 1a3f022c1d9..10437f29030 100644
--- a/clang-tools-extra/clangd/refactor/tweaks/AnnotateHighlightings.cpp
+++ b/clang-tools-extra/clangd/refactor/tweaks/AnnotateHighlightings.cpp
@@ -23,33 +23,39 @@ class AnnotateHighlightings : public Tweak {
 public:
   const char *id() const override final;
 
-  bool prepare(const Selection &Inputs) override {
-    for (auto N = Inputs.ASTSelection.commonAncestor(); N && !InterestedDecl;
-         N = N->Parent)
-      InterestedDecl = N->ASTNode.get<Decl>();
-    return InterestedDecl;
-  }
+  bool prepare(const Selection &Inputs) override { return true; }
   Expected<Effect> apply(const Selection &Inputs) override;
 
   std::string title() const override { return "Annotate highlighting tokens"; }
   Intent intent() const override { return Refactor; }
   bool hidden() const override { return true; }
-
-private:
-  const Decl *InterestedDecl = nullptr;
 };
 REGISTER_TWEAK(AnnotateHighlightings)
 
 Expected<Tweak::Effect> AnnotateHighlightings::apply(const Selection &Inputs) {
-  // Store the existing scopes.
-  const auto &BackupScopes = Inputs.AST.getASTContext().getTraversalScope();
-  // Narrow the traversal scope to the selected node.
-  Inputs.AST.getASTContext().setTraversalScope(
-      {const_cast<Decl *>(InterestedDecl)});
-  auto HighlightingTokens = getSemanticHighlightings(Inputs.AST);
-  // Restore the traversal scope.
-  Inputs.AST.getASTContext().setTraversalScope(BackupScopes);
+  // TUDecl is always the root ancestor.
+  const Decl *CommonDecl =
+      Inputs.ASTSelection.root().ASTNode.get<TranslationUnitDecl>();
+  for (auto N = Inputs.ASTSelection.commonAncestor(); N && !CommonDecl;
+       N = N->Parent)
+    CommonDecl = N->ASTNode.get<Decl>();
 
+  std::vector<HighlightingToken> HighlightingTokens;
+  if (llvm::isa<TranslationUnitDecl>(CommonDecl)) {
+    // We only annotate tokens in the main file, if CommonDecl is a TUDecl,
+    // we use the default traversal scope (which is the top level decls of the
+    // main file).
+    HighlightingTokens = getSemanticHighlightings(Inputs.AST);
+  } else {
+    // Store the existing scopes.
+    const auto &BackupScopes = Inputs.AST.getASTContext().getTraversalScope();
+    // Narrow the traversal scope to the selected node.
+    Inputs.AST.getASTContext().setTraversalScope(
+        {const_cast<Decl *>(CommonDecl)});
+    HighlightingTokens = getSemanticHighlightings(Inputs.AST);
+    // Restore the traversal scope.
+    Inputs.AST.getASTContext().setTraversalScope(BackupScopes);
+  }
   auto &SM = Inputs.AST.getSourceManager();
   tooling::Replacements Result;
   for (const auto &Token : HighlightingTokens) {
diff --git a/clang-tools-extra/clangd/unittests/TweakTests.cpp b/clang-tools-extra/clangd/unittests/TweakTests.cpp
index e06b28d986b..748e96f3cc6 100644
--- a/clang-tools-extra/clangd/unittests/TweakTests.cpp
+++ b/clang-tools-extra/clangd/unittests/TweakTests.cpp
@@ -487,9 +487,20 @@ TEST(TweakTest, ExtractVariable) {
 TEST(TweakTest, AnnotateHighlightings) {
   llvm::StringLiteral ID = "AnnotateHighlightings";
   checkAvailable(ID, "^vo^id^ ^f(^) {^}^"); // available everywhere.
+  checkAvailable(ID, "[[int a; int b;]]");
   const char *Input = "void ^f() {}";
   const char *Output = "void /* entity.name.function.cpp */f() {}";
   checkTransform(ID, Input, Output);
+
+  checkTransform(ID,
+  R"cpp(
+[[void f1();
+void f2();]]
+)cpp",
+  R"cpp(
+void /* entity.name.function.cpp */f1();
+void /* entity.name.function.cpp */f2();
+)cpp");
 }
 
 TEST(TweakTest, ExpandMacro) {