92 changes: 52 additions & 40 deletions clang-tools-extra/clangd/XRefs.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
#include "clang/AST/ExprCXX.h"
#include "clang/AST/Type.h"
#include "clang/Basic/LLVM.h"
#include "clang/Basic/LangOptions.h"
#include "clang/Basic/SourceLocation.h"
#include "clang/Basic/SourceManager.h"
#include "clang/Index/IndexDataConsumer.h"
Expand Down Expand Up @@ -149,25 +150,27 @@ std::vector<const NamedDecl *> getDeclAtPosition(ParsedAST &AST,
return Result;
}

llvm::Optional<Location> makeLocation(ASTContext &AST, SourceLocation TokLoc,
// Expects Loc to be a SpellingLocation, will bail out otherwise as it can't
// figure out a filename.
llvm::Optional<Location> makeLocation(const ASTContext &AST, SourceLocation Loc,
llvm::StringRef TUPath) {
const SourceManager &SourceMgr = AST.getSourceManager();
const FileEntry *F = SourceMgr.getFileEntryForID(SourceMgr.getFileID(TokLoc));
const auto &SM = AST.getSourceManager();
const FileEntry *F = SM.getFileEntryForID(SM.getFileID(Loc));
if (!F)
return None;
auto FilePath = getCanonicalPath(F, SourceMgr);
auto FilePath = getCanonicalPath(F, SM);
if (!FilePath) {
log("failed to get path!");
return None;
}
if (auto Range =
getTokenRange(AST.getSourceManager(), AST.getLangOpts(), TokLoc)) {
Location L;
L.uri = URIForFile::canonicalize(*FilePath, TUPath);
L.range = *Range;
return L;
}
return None;
Location L;
L.uri = URIForFile::canonicalize(*FilePath, TUPath);
// We call MeasureTokenLength here as TokenBuffer doesn't store spelled tokens
// outside the main file.
auto TokLen = Lexer::MeasureTokenLength(Loc, SM, AST.getLangOpts());
L.range = halfOpenToRange(
SM, CharSourceRange::getCharRange(Loc, Loc.getLocWithOffset(TokLen)));
return L;
}

} // namespace
Expand Down Expand Up @@ -373,11 +376,15 @@ namespace {
class ReferenceFinder : public index::IndexDataConsumer {
public:
struct Reference {
SourceLocation Loc;
syntax::Token SpelledTok;
index::SymbolRoleSet Role;

Range range(const SourceManager &SM) const {
return halfOpenToRange(SM, SpelledTok.range(SM).toCharRange(SM));
}
};

ReferenceFinder(ASTContext &AST, Preprocessor &PP,
ReferenceFinder(const ParsedAST &AST,
const std::vector<const NamedDecl *> &TargetDecls)
: AST(AST) {
for (const NamedDecl *D : TargetDecls)
Expand All @@ -386,13 +393,17 @@ class ReferenceFinder : public index::IndexDataConsumer {

std::vector<Reference> take() && {
llvm::sort(References, [](const Reference &L, const Reference &R) {
return std::tie(L.Loc, L.Role) < std::tie(R.Loc, R.Role);
auto LTok = L.SpelledTok.location();
auto RTok = R.SpelledTok.location();
return std::tie(LTok, L.Role) < std::tie(RTok, R.Role);
});
// We sometimes see duplicates when parts of the AST get traversed twice.
References.erase(std::unique(References.begin(), References.end(),
[](const Reference &L, const Reference &R) {
return std::tie(L.Loc, L.Role) ==
std::tie(R.Loc, R.Role);
auto LTok = L.SpelledTok.location();
auto RTok = R.SpelledTok.location();
return std::tie(LTok, L.Role) ==
std::tie(RTok, R.Role);
}),
References.end());
return std::move(References);
Expand All @@ -404,22 +415,27 @@ class ReferenceFinder : public index::IndexDataConsumer {
SourceLocation Loc,
index::IndexDataConsumer::ASTNodeInfo ASTNode) override {
assert(D->isCanonicalDecl() && "expect D to be a canonical declaration");
if (!CanonicalTargets.count(D))
return true;
const auto &TB = AST.getTokens();
const SourceManager &SM = AST.getSourceManager();
Loc = SM.getFileLoc(Loc);
if (isInsideMainFile(Loc, SM) && CanonicalTargets.count(D))
References.push_back({Loc, Roles});
// We are only traversing decls *inside* the main file, so this should hold.
assert(isInsideMainFile(Loc, SM));
if (const auto *Tok = TB.spelledTokenAt(Loc))
References.push_back({*Tok, Roles});
return true;
}

private:
llvm::SmallSet<const Decl *, 4> CanonicalTargets;
std::vector<Reference> References;
const ASTContext &AST;
const ParsedAST &AST;
};

std::vector<ReferenceFinder::Reference>
findRefs(const std::vector<const NamedDecl *> &Decls, ParsedAST &AST) {
ReferenceFinder RefFinder(AST.getASTContext(), AST.getPreprocessor(), Decls);
ReferenceFinder RefFinder(AST, Decls);
index::IndexingOptions IndexOpts;
IndexOpts.SystemSymbolFilter =
index::IndexingOptions::SystemSymbolFilterKind::All;
Expand Down Expand Up @@ -450,18 +466,15 @@ std::vector<DocumentHighlight> findDocumentHighlights(ParsedAST &AST,
// different kinds, deduplicate them.
std::vector<DocumentHighlight> Result;
for (const auto &Ref : References) {
if (auto Range =
getTokenRange(AST.getSourceManager(), AST.getLangOpts(), Ref.Loc)) {
DocumentHighlight DH;
DH.range = *Range;
if (Ref.Role & index::SymbolRoleSet(index::SymbolRole::Write))
DH.kind = DocumentHighlightKind::Write;
else if (Ref.Role & index::SymbolRoleSet(index::SymbolRole::Read))
DH.kind = DocumentHighlightKind::Read;
else
DH.kind = DocumentHighlightKind::Text;
Result.push_back(std::move(DH));
}
DocumentHighlight DH;
DH.range = Ref.range(SM);
if (Ref.Role & index::SymbolRoleSet(index::SymbolRole::Write))
DH.kind = DocumentHighlightKind::Write;
else if (Ref.Role & index::SymbolRoleSet(index::SymbolRole::Read))
DH.kind = DocumentHighlightKind::Read;
else
DH.kind = DocumentHighlightKind::Text;
Result.push_back(std::move(DH));
}
return Result;
}
Expand Down Expand Up @@ -524,16 +537,15 @@ ReferencesResult findReferences(ParsedAST &AST, Position Pos, uint32_t Limit,
MainFileRefs.erase(std::unique(MainFileRefs.begin(), MainFileRefs.end(),
[](const ReferenceFinder::Reference &L,
const ReferenceFinder::Reference &R) {
return L.Loc == R.Loc;
return L.SpelledTok.location() ==
R.SpelledTok.location();
}),
MainFileRefs.end());
for (const auto &Ref : MainFileRefs) {
if (auto Range = getTokenRange(SM, AST.getLangOpts(), Ref.Loc)) {
Location Result;
Result.range = *Range;
Result.uri = URIMainFile;
Results.References.push_back(std::move(Result));
}
Location Result;
Result.range = Ref.range(SM);
Result.uri = URIMainFile;
Results.References.push_back(std::move(Result));
}
if (Index && Results.References.size() <= Limit) {
for (const Decl *D : Decls) {
Expand Down
4 changes: 4 additions & 0 deletions clang/include/clang/Tooling/Syntax/Tokens.h
Original file line number Diff line number Diff line change
Expand Up @@ -245,6 +245,10 @@ class TokenBuffer {
/// "DECL", "(", "a", ")", ";"}
llvm::ArrayRef<syntax::Token> spelledTokens(FileID FID) const;

/// Returns the spelled Token starting at Loc, if there are no such tokens
/// returns nullptr.
const syntax::Token *spelledTokenAt(SourceLocation Loc) const;

/// Get all tokens that expand a macro in \p FID. For the following input
/// #define FOO B
/// #define FOO2(X) int X
Expand Down
10 changes: 10 additions & 0 deletions clang/lib/Tooling/Syntax/Tokens.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -183,6 +183,16 @@ llvm::ArrayRef<syntax::Token> TokenBuffer::spelledTokens(FileID FID) const {
return It->second.SpelledTokens;
}

const syntax::Token *TokenBuffer::spelledTokenAt(SourceLocation Loc) const {
assert(Loc.isFileID());
const auto *Tok = llvm::partition_point(
spelledTokens(SourceMgr->getFileID(Loc)),
[&](const syntax::Token &Tok) { return Tok.location() < Loc; });
if (!Tok || Tok->location() != Loc)
return nullptr;
return Tok;
}

std::string TokenBuffer::Mapping::str() const {
return std::string(
llvm::formatv("spelled tokens: [{0},{1}), expanded tokens: [{2},{3})",
Expand Down
7 changes: 7 additions & 0 deletions clang/unittests/Tooling/Syntax/TokensTest.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ using ::testing::ElementsAre;
using ::testing::Field;
using ::testing::Matcher;
using ::testing::Not;
using ::testing::Pointee;
using ::testing::StartsWith;

namespace {
Expand Down Expand Up @@ -363,6 +364,12 @@ TEST_F(TokenCollectorTest, Locations) {
AllOf(Kind(tok::equal), RangeIs(Code.range("r3"))),
AllOf(Kind(tok::string_literal), RangeIs(Code.range("r4"))),
AllOf(Kind(tok::semi), RangeIs(Code.range("r5")))));

auto StartLoc = SourceMgr->getLocForStartOfFile(SourceMgr->getMainFileID());
for (auto &R : Code.ranges()) {
EXPECT_THAT(Buffer.spelledTokenAt(StartLoc.getLocWithOffset(R.Begin)),
Pointee(RangeIs(R)));
}
}

TEST_F(TokenCollectorTest, MacroDirectives) {
Expand Down