[PATCH] D65210: [clangd] Fix the annotate tweak after rL366893

[Haojian Wu via Phabricator via cfe-commits]

hokein created this revision.
hokein added a reviewer: sammccall.
Herald added subscribers: kadircet, arphaman, jkorous, MaskRay, ilya-biryukov.
Herald added a project: clang.

After rL366893 <https://reviews.llvm.org/rL366893>, the annoate tweak is not activated when we select the
whole file (the commonAncestor is TUDecl but we intend to return null).

This patch fixes this, and also avoid traversing the TUDecl.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D65210

Files:
  clang-tools-extra/clangd/refactor/tweaks/AnnotateHighlightings.cpp
  clang-tools-extra/clangd/unittests/TweakTests.cpp


Index: clang-tools-extra/clangd/unittests/TweakTests.cpp
===================================================================
--- clang-tools-extra/clangd/unittests/TweakTests.cpp
+++ clang-tools-extra/clangd/unittests/TweakTests.cpp
@@ -487,6 +487,7 @@
 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);
Index: clang-tools-extra/clangd/refactor/tweaks/AnnotateHighlightings.cpp
===================================================================
--- clang-tools-extra/clangd/refactor/tweaks/AnnotateHighlightings.cpp
+++ clang-tools-extra/clangd/refactor/tweaks/AnnotateHighlightings.cpp
@@ -24,6 +24,7 @@
   const char *id() const override final;
 
   bool prepare(const Selection &Inputs) override {
+    InterestedDecl = Inputs.ASTSelection.root().ASTNode.get<Decl>();
     for (auto N = Inputs.ASTSelection.commonAncestor(); N && !InterestedDecl;
          N = N->Parent)
       InterestedDecl = N->ASTNode.get<Decl>();
@@ -41,15 +42,22 @@
 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);
-
+  std::vector<HighlightingToken> HighlightingTokens;
+  if (llvm::isa<TranslationUnitDecl>(InterestedDecl)) {
+    // We only annotate tokens in the main file, if InterestedDecl 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 *>(InterestedDecl)});
+    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) {


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D65210.211497.patch
Type: text/x-patch
Size: 2736 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/cfe-commits/attachments/20190724/96a3c54b/attachment-0001.bin>