[clang-tools-extra] r366996 - [clangd] Fix the annotate tweak after rL366893

[Haojian Wu via cfe-commits]

Author: hokein
Date: Thu Jul 25 01:48:05 2019
New Revision: 366996

URL: http://llvm.org/viewvc/llvm-project?rev=366996&view=rev
Log:
[clangd] Fix the annotate tweak after rL366893

Summary:
After 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.

Reviewers: sammccall

Subscribers: ilya-biryukov, MaskRay, jkorous, arphaman, kadircet, cfe-commits

Tags: #clang

Differential Revision: https://reviews.llvm.org/D65210

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

Modified: clang-tools-extra/trunk/clangd/refactor/tweaks/AnnotateHighlightings.cpp
URL: http://llvm.org/viewvc/llvm-project/clang-tools-extra/trunk/clangd/refactor/tweaks/AnnotateHighlightings.cpp?rev=366996&r1=366995&r2=366996&view=diff
==============================================================================
--- clang-tools-extra/trunk/clangd/refactor/tweaks/AnnotateHighlightings.cpp (original)
+++ clang-tools-extra/trunk/clangd/refactor/tweaks/AnnotateHighlightings.cpp Thu Jul 25 01:48:05 2019
@@ -23,33 +23,39 @@ class AnnotateHighlightings : public Twe
 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) {

Modified: clang-tools-extra/trunk/clangd/unittests/TweakTests.cpp
URL: http://llvm.org/viewvc/llvm-project/clang-tools-extra/trunk/clangd/unittests/TweakTests.cpp?rev=366996&r1=366995&r2=366996&view=diff
==============================================================================
--- clang-tools-extra/trunk/clangd/unittests/TweakTests.cpp (original)
+++ clang-tools-extra/trunk/clangd/unittests/TweakTests.cpp Thu Jul 25 01:48:05 2019
@@ -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) {