[clangd] Make use of token buffers in semantic highlighting

[kadircet]

• Add forgetten macroid check, before looking for an macroargexpansion.

[sammccall]

clang-tools-extra/clangd/SemanticHighlighting.cpp

149–151

nit: this can be const

158–159

I think using TokenBuffer to translate expanded -> spelled isn't actually an improvement here as the API exposes several possibilities that can't actually happen in our case.

The specialized translation from expanded to spelled location is *policy*, and it's pretty easy to implement I think:

// For a macro usage `DUMP(foo)`, we want:
//  - DUMP --> "macro"
//  - foo --> "variable".
SourceLocation getHighlightableSpellingToken(SourceLocation L) {
  if (L.isFileID()) return SM.isWrittenInMainFile(L) ? L : {};
  // Tokens expanded from the macro body contribute no highlightings.
  if (!SM.isMacroArgExpansion(L)) return {};
  // Tokens expanded from macro args are potentially highlightable.
  return getHighlightableSpellingToken(SM.getImmediateSpellingLocation(L));
}

once you have the spelling location, getting the range from tokenbuffer is easy :-)

164

only if this is a macro!

[kadircet]

• Address comments

[sammccall]

clang-tools-extra/clangd/SemanticHighlighting.cpp

162

&& Tok->location() == Loc, right?

(I think this would make a nice method TokenBuffer::spelledTokenAt(SourceLocation) -> const Token*)

[kadircet]

• Use spelledTokenAt

clang-tools-extra/clangd/SemanticHighlighting.cpp

 #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 "clang/Tooling/Syntax/Tokens.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Base64.h"
 #include "llvm/Support/Casting.h"
 #include <algorithm>
 
 namespace clang {
for (const NamedDecl *Decl : R.Targets) {
     auto Kind = kindForDecl(Decl);
     if (!Kind || (Result && Kind != Result))
       return llvm::None;
     Result = Kind;
   }
   return Result;
 }
 
+// For a macro usage `DUMP(foo)`, we want:
+//  - DUMP --> "macro"
+//  - foo --> "variable".
+SourceLocation getHighlightableSpellingToken(SourceLocation L,
+                                             const SourceManager &SM) {
+  if (L.isFileID())
+    return SM.isWrittenInMainFile(L) ? L : SourceLocation{};
+  // Tokens expanded from the macro body contribute no highlightings.
+  if (!SM.isMacroArgExpansion(L))
+    return {};
+  // Tokens expanded from macro args are potentially highlightable.
+  return getHighlightableSpellingToken(SM.getImmediateSpellingLoc(L), SM);
+}
+
 /// Consumes source locations and maps them to text ranges for highlightings.
 class HighlightingsBuilder {
 public:
-  HighlightingsBuilder(const SourceManager &SourceMgr,
-                       const LangOptions &LangOpts)
-      : SourceMgr(SourceMgr), LangOpts(LangOpts) {}
+  HighlightingsBuilder(const ParsedAST &AST)
+      : TB(AST.getTokens()), SourceMgr(AST.getSourceManager()),
+        LangOpts(AST.getLangOpts()) {}
 
   void addToken(HighlightingToken T) { Tokens.push_back(T); }
 
   void addToken(SourceLocation Loc, HighlightingKind Kind) {
+    Loc = getHighlightableSpellingToken(Loc, SourceMgr);
     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});
+    const auto *Tok = llvm::partition_point(
+        TB.spelledTokens(SourceMgr.getMainFileID()),
+        [&](const syntax::Token &Tok) { return Tok.location() < Loc; });
+    assert(Tok);
+
+    auto Range = halfOpenToRange(SourceMgr,
+                                 Tok->range(SourceMgr).toCharRange(SourceMgr));
+    Tokens.push_back(HighlightingToken{Kind, std::move(Range)});
   }
 
   std::vector<HighlightingToken> collect(ParsedAST &AST) && {
     // 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());
for (const Range &R : AST.getMacros().SkippedRanges) {
       }
     }
     // Re-sort the tokens because that's what the diffing expects.
     llvm::sort(NonConflicting);
     return NonConflicting;
   }
 
 private:
+  const syntax::TokenBuffer &TB;
   const SourceManager &SourceMgr;
   const LangOptions &LangOpts;
   std::vector<HighlightingToken> Tokens;
 };
 
 /// Produces highlightings, which are not captured by findExplicitReferences,
 /// e.g. highlights dependent names and 'auto' as the underlying type.
 class CollectExtraHighlightings
return ArrayRef<HighlightingToken>(StartLineIt, AllTokens.end())
         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());
+  HighlightingsBuilder Builder(AST);
   // 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 references.