[clangd] Implement semantic highlightings via findExplicitReferences

[ilya-biryukov]

Summary

To keep the logic of finding locations of interesting AST nodes in one
place.

The advantage is better coverage of various AST nodes, both now and in
the future: as new nodes get added to findExplicitReferences, semantic
highlighting will automatically pick them up.

The drawback of this change is that we have to traverse declarations
inside our file twice in order to highlight dependent names, 'auto'
and 'decltype'. Hopefully, this should not affect the actual latency
too much, most time should be spent in building the AST and not
traversing it.

[merge_guards_bot]

Build result: pass - 59818 tests passed, 0 failed and 762 were skipped.
Log files: console-log.txt, CMakeCache.txt

[hokein]

clang-tools-extra/clangd/SemanticHighlighting.cpp

195

nit: I'd call it LangOpts

201

maybe call it CollectExtraHiglightings or CollectAdditionalHighlightings?

222

IIUC, if the decls don't have the same kind, we will not highlight it (previously, we fall back to DependentName).

[ilya-biryukov]

• Rename to LangOpts and SourceMgr
• Rename to CollectExtraHighlightings

[ilya-biryukov]

clang-tools-extra/clangd/SemanticHighlighting.cpp

222

You're right. I don't think highlighting as dependent name is any better than not highlighting at all in that case, though.
Therefore, I would argue it's not worth adding more code to handle this corner case (I'm not even sure it's possible in practice, it'll take quite some time to come up with examples like these)

[merge_guards_bot]

Build result: pass - 59845 tests passed, 0 failed and 763 were skipped.
Log files: console-log.txt, CMakeCache.txt

[hokein]

clang-tools-extra/clangd/SemanticHighlighting.cpp

222

I thought it is better to highlight it rather than not highlighting, but I agree with you, it may be not worth handling this edge case.

clang-tools-extra/clangd/FindTarget.h

 // approach based around AST nodes (DynTypedNode), and can be combined with
 // SelectionTree or other traversals.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_UNITTESTS_CLANGD_FINDTARGET_H
 #define LLVM_CLANG_TOOLS_EXTRA_UNITTESTS_CLANGD_FINDTARGET_H
 
+#include "clang/AST/ASTContext.h"
 #include "clang/AST/ASTTypeTraits.h"
 #include "clang/AST/NestedNameSpecifier.h"
 #include "clang/AST/Stmt.h"
 #include "clang/Basic/SourceLocation.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"

 /// references in some subrange of the file and apply simple edits, e.g. add
 /// qualifiers.
 /// FIXME: currently this does not report references to overloaded operators.
 /// FIXME: extend to report location information about declaration names too.
 void findExplicitReferences(const Stmt *S,
                             llvm::function_ref<void(ReferenceLoc)> Out);
 void findExplicitReferences(const Decl *D,
                             llvm::function_ref<void(ReferenceLoc)> Out);
+void findExplicitReferences(const ASTContext &AST,
+                            llvm::function_ref<void(ReferenceLoc)> Out);
 
 /// Similar to targetDecl(), however instead of applying a filter, all possible
 /// decls are returned along with their DeclRelationSets.
 /// This is suitable for indexing, where everything is recorded and filtering
 /// is applied later.
 llvm::SmallVector<std::pair<const Decl *, DeclRelationSet>, 1>
 allTargetDecls(const ast_type_traits::DynTypedNode &);
 

clang-tools-extra/clangd/FindTarget.cpp

void findExplicitReferences(const Stmt *S,
   assert(S);
   ExplicitReferenceColletor(Out).TraverseStmt(const_cast<Stmt *>(S));
 }
 void findExplicitReferences(const Decl *D,
                             llvm::function_ref<void(ReferenceLoc)> Out) {
   assert(D);
   ExplicitReferenceColletor(Out).TraverseDecl(const_cast<Decl *>(D));
 }
+void findExplicitReferences(const ASTContext &AST,
+                            llvm::function_ref<void(ReferenceLoc)> Out) {
+  ExplicitReferenceColletor(Out).TraverseAST(const_cast<ASTContext &>(AST));
+}
 
 llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, DeclRelation R) {
   switch (R) {
 #define REL_CASE(X)                                                            \
   case DeclRelation::X:                                                        \
     return OS << #X;
     REL_CASE(Alias);
     REL_CASE(Underlying);

clang-tools-extra/clangd/SemanticHighlighting.cpp

 //===--- SemanticHighlighting.cpp - ------------------------- ---*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "SemanticHighlighting.h"
+#include "FindTarget.h"
 #include "Logger.h"
 #include "ParsedAST.h"
 #include "Protocol.h"
 #include "SourceCode.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclCXX.h"
 #include "clang/AST/DeclarationName.h"
+#include "clang/AST/ExprCXX.h"
 #include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/AST/Type.h"
 #include "clang/AST/TypeLoc.h"
+#include "clang/Basic/LangOptions.h"
 #include "clang/Basic/SourceLocation.h"
+#include "clang/Basic/SourceManager.h"
+#include "llvm/ADT/None.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/Support/Casting.h"
 #include <algorithm>
 
 namespace clang {
 namespace clangd {
 namespace {
 
 /// Some names are not written in the source code and cannot be highlighted,
 /// e.g. anonymous classes. This function detects those cases.
llvm::Optional<HighlightingKind> kindForType(const Type *TP) {
   if (TP->isBuiltinType()) // Builtins are special, they do not have decls.
     return HighlightingKind::Primitive;
   if (auto *TD = dyn_cast<TemplateTypeParmType>(TP))
     return kindForDecl(TD->getDecl());
   if (auto *TD = TP->getAsTagDecl())
     return kindForDecl(TD);
   return llvm::None;
 }
-// Given a set of candidate declarations, if the declarations all have the same
-// highlighting kind, return that highlighting kind, otherwise return None.
-template <typename IteratorRange>
-llvm::Optional<HighlightingKind> kindForCandidateDecls(IteratorRange Decls) {
+
+llvm::Optional<HighlightingKind> kindForReference(const ReferenceLoc &R) {
   llvm::Optional<HighlightingKind> Result;
-  for (NamedDecl *Decl : Decls) {
+  for (const NamedDecl *Decl : R.Targets) {
+    if (!canHighlightName(Decl->getDeclName()))
+      return llvm::None;
     auto Kind = kindForDecl(Decl);
     if (!Kind || (Result && Kind != Result))
       return llvm::None;
     Result = Kind;
   }
   return Result;
 }
 
-// Collects all semantic tokens in an ASTContext.
-class HighlightingTokenCollector
-    : public RecursiveASTVisitor<HighlightingTokenCollector> {
-  std::vector<HighlightingToken> Tokens;
-  ParsedAST &AST;
-
+/// Consumes source locations and maps them to text ranges for highlightings.
+class HighlightingsBuilder {
 public:
-  HighlightingTokenCollector(ParsedAST &AST) : AST(AST) {}
+  HighlightingsBuilder(const SourceManager &SourceMgr,
+                       const LangOptions &LangOpts)
+      : SourceMgr(SourceMgr), LangOpts(LangOpts) {}
 
-  std::vector<HighlightingToken> collectTokens() {
-    Tokens.clear();
-    TraverseAST(AST.getASTContext());
-    // Add highlightings for macro expansions as they are not traversed by the
-    // visitor.
-    for (const auto &M : AST.getMacros().Ranges)
-      Tokens.push_back({HighlightingKind::Macro, M});
+  void addToken(HighlightingToken T) { Tokens.push_back(T); }
+
+  void addToken(SourceLocation Loc, HighlightingKind Kind) {
+    if (Loc.isInvalid())
+      return;
+    if (Loc.isMacroID()) {
+      // Only intereseted in highlighting arguments in macros (DEF_X(arg)).
+      if (!SourceMgr.isMacroArgExpansion(Loc))
+        return;
+      Loc = SourceMgr.getSpellingLoc(Loc);
+    }
+
+    // Non top level decls that are included from a header are not filtered by
+    // topLevelDecls. (example: method declarations being included from
+    // another file for a class from another file).
+    // There are also cases with macros where the spelling loc will not be in
+    // the main file and the highlighting would be incorrect.
+    if (!isInsideMainFile(Loc, SourceMgr))
+      return;
+
+    auto Range = getTokenRange(SourceMgr, LangOpts, Loc);
+    if (!Range) {
+      // R should always have a value, if it doesn't something is very wrong.
+      elog("Tried to add semantic token with an invalid range");
+      return;
+    }
+    Tokens.push_back(HighlightingToken{Kind, *Range});
+  }
+
+  std::vector<HighlightingToken> collect() && {
     // Initializer lists can give duplicates of tokens, therefore all tokens
     // must be deduplicated.
     llvm::sort(Tokens);
     auto Last = std::unique(Tokens.begin(), Tokens.end());
     Tokens.erase(Last, Tokens.end());
+
     // Macros can give tokens that have the same source range but conflicting
     // kinds. In this case all tokens sharing this source range should be
     // removed.
     std::vector<HighlightingToken> NonConflicting;
     NonConflicting.reserve(Tokens.size());
     for (ArrayRef<HighlightingToken> TokRef = Tokens; !TokRef.empty();) {
       ArrayRef<HighlightingToken> Conflicting =
           TokRef.take_while([&](const HighlightingToken &T) {
             // TokRef is guaranteed at least one element here because otherwise
             // this predicate would never fire.
             return T.R == TokRef.front().R;
           });
       // If there is exactly one token with this range it's non conflicting and
       // should be in the highlightings.
       if (Conflicting.size() == 1)
         NonConflicting.push_back(TokRef.front());
       // TokRef[Conflicting.size()] is the next token with a different range (or
       // the end of the Tokens).
       TokRef = TokRef.drop_front(Conflicting.size());
     }
     return NonConflicting;
   }
 
-  bool VisitNamespaceAliasDecl(NamespaceAliasDecl *NAD) {
-    // The target namespace of an alias can not be found in any other way.
-    addToken(NAD->getTargetNameLoc(), NAD->getAliasedNamespace());
-    return true;
-  }
-
-  bool VisitMemberExpr(MemberExpr *ME) {
-    if (canHighlightName(ME->getMemberNameInfo().getName()))
-      addToken(ME->getMemberLoc(), ME->getMemberDecl());
-    return true;
-  }
-
-  bool VisitOverloadExpr(OverloadExpr *E) {
-    if (canHighlightName(E->getName()))
-      addToken(E->getNameLoc(),
-               kindForCandidateDecls(E->decls())
-                   .getValueOr(HighlightingKind::DependentName));
-    return true;
-  }
-
-  bool VisitDependentScopeDeclRefExpr(DependentScopeDeclRefExpr *E) {
-    if (canHighlightName(E->getDeclName()))
-      addToken(E->getLocation(), HighlightingKind::DependentName);
-    return true;
-  }
-
-  bool VisitCXXDependentScopeMemberExpr(CXXDependentScopeMemberExpr *E) {
-    if (canHighlightName(E->getMember()))
-      addToken(E->getMemberLoc(), HighlightingKind::DependentName);
-    return true;
-  }
+private:
+  const SourceManager &SourceMgr;
+  const LangOptions &LangOpts;
+  std::vector<HighlightingToken> Tokens;
+};
 
-  bool VisitNamedDecl(NamedDecl *ND) {
-    if (canHighlightName(ND->getDeclName()))
-      addToken(ND->getLocation(), ND);
-    return true;
-  }
+/// Produces highlightings, which are not captured by findExplicitReferences,
+/// e.g. highlights dependent names and 'auto' as the underlying type.
+class CollectExtraHighlightings
+    : public RecursiveASTVisitor<CollectExtraHighlightings> {
+public:
+  CollectExtraHighlightings(HighlightingsBuilder &H) : H(H) {}
 
-  bool VisitUsingDecl(UsingDecl *UD) {
-    if (auto K = kindForCandidateDecls(UD->shadows()))
-      addToken(UD->getLocation(), *K);
+  bool VisitDecltypeTypeLoc(DecltypeTypeLoc L) {
+    if (auto K = kindForType(L.getTypePtr()))
+      H.addToken(L.getBeginLoc(), *K);
     return true;
   }
 
-  bool VisitDeclRefExpr(DeclRefExpr *Ref) {
-    if (canHighlightName(Ref->getNameInfo().getName()))
-      addToken(Ref->getLocation(), Ref->getDecl());
+  bool VisitDeclaratorDecl(DeclaratorDecl *D) {
+    auto *AT = D->getType()->getContainedAutoType();
+    if (!AT)
       return true;
-  }
-
-  bool VisitTypedefTypeLoc(TypedefTypeLoc TL) {
-    addToken(TL.getBeginLoc(), TL.getTypedefNameDecl());
+    if (auto K = kindForType(AT->getDeducedType().getTypePtrOrNull()))
+      H.addToken(D->getTypeSpecStartLoc(), *K);
     return true;
   }
 
-  bool VisitTemplateSpecializationTypeLoc(TemplateSpecializationTypeLoc TL) {
-    if (const TemplateDecl *TD =
-            TL.getTypePtr()->getTemplateName().getAsTemplateDecl())
-      addToken(TL.getBeginLoc(), TD);
+  bool VisitOverloadExpr(OverloadExpr *E) {
+    if (!E->decls().empty())
+      return true; // handled by findExplicitReferences.
+    H.addToken(E->getNameLoc(), HighlightingKind::DependentName);
     return true;
   }
 
-  bool VisitTagTypeLoc(TagTypeLoc L) {
-    if (L.isDefinition())
-      return true; // Definition will be highligthed by VisitNamedDecl.
-    if (auto K = kindForType(L.getTypePtr()))
-      addToken(L.getBeginLoc(), *K);
+  bool VisitCXXDependentScopeMemberExpr(CXXDependentScopeMemberExpr *E) {
+    H.addToken(E->getMemberNameInfo().getLoc(),
+               HighlightingKind::DependentName);
     return true;
   }
 
-  bool VisitDecltypeTypeLoc(DecltypeTypeLoc L) {
-    if (auto K = kindForType(L.getTypePtr()))
-      addToken(L.getBeginLoc(), *K);
+  bool VisitDependentScopeDeclRefExpr(DependentScopeDeclRefExpr *E) {
+    H.addToken(E->getNameInfo().getLoc(), HighlightingKind::DependentName);
     return true;
   }
 
   bool VisitDependentNameTypeLoc(DependentNameTypeLoc L) {
-    addToken(L.getNameLoc(), HighlightingKind::DependentType);
-    return true;
-  }
-
-  bool VisitTemplateTypeParmTypeLoc(TemplateTypeParmTypeLoc TL) {
-    addToken(TL.getBeginLoc(), HighlightingKind::TemplateParameter);
-    return true;
-  }
-
-  bool TraverseNestedNameSpecifierLoc(NestedNameSpecifierLoc NNSLoc) {
-    if (auto *NNS = NNSLoc.getNestedNameSpecifier()) {
-      if (NNS->getKind() == NestedNameSpecifier::Namespace ||
-          NNS->getKind() == NestedNameSpecifier::NamespaceAlias)
-        addToken(NNSLoc.getLocalBeginLoc(), HighlightingKind::Namespace);
-    }
-    return RecursiveASTVisitor<
-        HighlightingTokenCollector>::TraverseNestedNameSpecifierLoc(NNSLoc);
-  }
-
-  bool TraverseConstructorInitializer(CXXCtorInitializer *CI) {
-    if (const FieldDecl *FD = CI->getMember())
-      addToken(CI->getSourceLocation(), FD);
-    return RecursiveASTVisitor<
-        HighlightingTokenCollector>::TraverseConstructorInitializer(CI);
-  }
-
-  bool VisitDeclaratorDecl(DeclaratorDecl *D) {
-    // Highlight 'auto' with its underlying type.
-    auto *AT = D->getType()->getContainedAutoType();
-    if (!AT)
-      return true;
-    auto K = kindForType(AT->getDeducedType().getTypePtrOrNull());
-    if (!K)
-      return true;
-    addToken(D->getTypeSpecStartLoc(), *K);
+    H.addToken(L.getNameLoc(), HighlightingKind::DependentType);
     return true;
   }
 
 private:
-  void addToken(SourceLocation Loc, HighlightingKind Kind) {
-    if (Loc.isInvalid())
-      return;
-    const auto &SM = AST.getSourceManager();
-    if (Loc.isMacroID()) {
-      // Only intereseted in highlighting arguments in macros (DEF_X(arg)).
-      if (!SM.isMacroArgExpansion(Loc))
-        return;
-      Loc = SM.getSpellingLoc(Loc);
-    }
-
-    // Non top level decls that are included from a header are not filtered by
-    // topLevelDecls. (example: method declarations being included from
-    // another file for a class from another file).
-    // There are also cases with macros where the spelling loc will not be in
-    // the main file and the highlighting would be incorrect.
-    if (!isInsideMainFile(Loc, SM))
-      return;
-
-    auto R = getTokenRange(SM, AST.getASTContext().getLangOpts(), Loc);
-    if (!R) {
-      // R should always have a value, if it doesn't something is very wrong.
-      elog("Tried to add semantic token with an invalid range");
-      return;
-    }
-
-    Tokens.push_back({Kind, R.getValue()});
-  }
-
-  void addToken(SourceLocation Loc, const NamedDecl *D) {
-    if (auto K = kindForDecl(D))
-      addToken(Loc, *K);
-  }
+  HighlightingsBuilder &H;
 };
 
 // Encode binary data into base64.
 // This was copied from compiler-rt/lib/fuzzer/FuzzerUtil.cpp.
 // FIXME: Factor this out into llvm/Support?
 std::string encodeBase64(const llvm::SmallVectorImpl<char> &Bytes) {
   static const char Table[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
                               "abcdefghijklmnopqrstuvwxyz"
takeLine(ArrayRef<HighlightingToken> AllTokens,
          ArrayRef<HighlightingToken>::iterator StartLineIt, int Line) {
   return ArrayRef<HighlightingToken>(StartLineIt, AllTokens.end())
       .take_while([Line](const HighlightingToken &Token) {
         return Token.R.start.line == Line;
       });
 }
 } // namespace
 
+std::vector<HighlightingToken> getSemanticHighlightings(ParsedAST &AST) {
+  auto &C = AST.getASTContext();
+  // Add highlightings for AST nodes.
+  HighlightingsBuilder Builder(AST.getSourceManager(), C.getLangOpts());
+  // Highlight 'decltype' and 'auto' as their underlying types.
+  CollectExtraHighlightings(Builder).TraverseAST(C);
+  // Highlight all decls and references coming from the AST.
+  findExplicitReferences(C, [&](ReferenceLoc R) {
+    if (auto Kind = kindForReference(R))
+      Builder.addToken(R.NameLoc, *Kind);
+  });
+  // Add highlightings for macro expansions.
+  for (const auto &M : AST.getMacros().Ranges)
+    Builder.addToken({HighlightingKind::Macro, M});
+  return std::move(Builder).collect();
+}
+
 llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, HighlightingKind K) {
   switch (K) {
   case HighlightingKind::Variable:
     return OS << "Variable";
   case HighlightingKind::LocalVariable:
     return OS << "LocalVariable";
   case HighlightingKind::Parameter:
     return OS << "Parameter";

 }
 bool operator<(const HighlightingToken &L, const HighlightingToken &R) {
   return std::tie(L.R, L.Kind) < std::tie(R.R, R.Kind);
 }
 bool operator==(const LineHighlightings &L, const LineHighlightings &R) {
   return std::tie(L.Line, L.Tokens) == std::tie(R.Line, R.Tokens);
 }
 
-std::vector<HighlightingToken> getSemanticHighlightings(ParsedAST &AST) {
-  return HighlightingTokenCollector(AST).collectTokens();
-}
-
 std::vector<SemanticHighlightingInformation>
 toSemanticHighlightingInformation(llvm::ArrayRef<LineHighlightings> Tokens) {
   if (Tokens.size() == 0)
     return {};
 
   // FIXME: Tokens might be multiple lines long (block comments) in this case
   // this needs to add multiple lines for those tokens.
   std::vector<SemanticHighlightingInformation> Lines;

clang-tools-extra/clangd/unittests/SemanticHighlightingTests.cpp

)cpp",
           return $TemplateParameter[[T]]::$DependentName[[Field]];
         }
       };
     )cpp",
       // Highlighting the using decl as the underlying using shadow decl.
       R"cpp(
       void $Function[[foo]]();
       using ::$Function[[foo]];
+    )cpp",
+      // Highlighting of template template arguments.
+      R"cpp(
+      template <template <class> class $TemplateParameter[[TT]],
+                template <class> class ...$TemplateParameter[[TTs]]>
+      struct $Class[[Foo]] {
+        $Class[[Foo]]<$TemplateParameter[[TT]], $TemplateParameter[[TTs]]...>
+          *$Field[[t]];
+      }
     )cpp"};
   for (const auto &TestCase : TestCases) {
     checkHighlightings(TestCase);
   }
 
   checkHighlightings(R"cpp(
     class $Class[[A]] {
       #include "imp.h"