1 files changed, 13 insertions, 6 deletions
diff --git a/clang-tools-extra/clangd/SemanticHighlighting.cpp b/clang-tools-extra/clangd/SemanticHighlighting.cpp
index e71c0a733a6..eaa57f8ab66 100644
--- a/clang-tools-extra/clangd/SemanticHighlighting.cpp
+++ b/clang-tools-extra/clangd/SemanticHighlighting.cpp
@@ -24,16 +24,20 @@ namespace {
 class HighlightingTokenCollector
     : public RecursiveASTVisitor<HighlightingTokenCollector> {
   std::vector<HighlightingToken> Tokens;
-  ASTContext &Ctx;
-  const SourceManager &SM;
+  ParsedAST &AST;
 
 public:
-  HighlightingTokenCollector(ParsedAST &AST)
-      : Ctx(AST.getASTContext()), SM(AST.getSourceManager()) {}
+  HighlightingTokenCollector(ParsedAST &AST) : AST(AST) {}
 
   std::vector<HighlightingToken> collectTokens() {
     Tokens.clear();
-    TraverseAST(Ctx);
+    TraverseAST(AST.getASTContext());
+    // Add highlightings for macro expansions as they are not traversed by the
+    // visitor.
+    // FIXME: Should add highlighting to the macro definitions as well. But this
+    // information is not collected in ParsedAST right now.
+    for (const SourceLocation &L : AST.getMainFileExpansions())
+      addToken(L, HighlightingKind::Macro);
     // Initializer lists can give duplicates of tokens, therefore all tokens
     // must be deduplicated.
     llvm::sort(Tokens);
@@ -264,6 +268,7 @@ private:
   }
 
   void addToken(SourceLocation Loc, HighlightingKind Kind) {
+    const auto &SM = AST.getSourceManager();
     if (Loc.isMacroID()) {
       // Only intereseted in highlighting arguments in macros (DEF_X(arg)).
       if (!SM.isMacroArgExpansion(Loc))
@@ -279,7 +284,7 @@ private:
     if (!isInsideMainFile(Loc, SM))
       return;
 
-    auto R = getTokenRange(SM, Ctx.getLangOpts(), Loc);
+    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");
@@ -466,6 +471,8 @@ llvm::StringRef toTextMateScope(HighlightingKind Kind) {
     return "entity.name.type.template.cpp";
   case HighlightingKind::Primitive:
     return "storage.type.primitive.cpp";
+  case HighlightingKind::Macro:
+    return "entity.name.function.preprocessor.cpp";
   case HighlightingKind::NumKinds:
     llvm_unreachable("must not pass NumKinds to the function");
   }