diff --git a/clang-tools-extra/clangd/ClangdLSPServer.cpp b/clang-tools-extra/clangd/ClangdLSPServer.cpp --- a/clang-tools-extra/clangd/ClangdLSPServer.cpp +++ b/clang-tools-extra/clangd/ClangdLSPServer.cpp @@ -1152,7 +1152,13 @@ void ClangdLSPServer::onReference(const ReferenceParams &Params, Callback> Reply) { Server->findReferences(Params.textDocument.uri.file(), Params.position, - CCOpts.Limit, std::move(Reply)); + CCOpts.Limit, + [Reply = std::move(Reply)]( + llvm::Expected Refs) mutable { + if (!Refs) + return Reply(Refs.takeError()); + return Reply(std::move(Refs->References)); + }); } void ClangdLSPServer::onSymbolInfo(const TextDocumentPositionParams &Params, diff --git a/clang-tools-extra/clangd/ClangdServer.h b/clang-tools-extra/clangd/ClangdServer.h --- a/clang-tools-extra/clangd/ClangdServer.h +++ b/clang-tools-extra/clangd/ClangdServer.h @@ -225,7 +225,7 @@ /// Retrieve locations for symbol references. void findReferences(PathRef File, Position Pos, uint32_t Limit, - Callback> CB); + Callback CB); /// Run formatting for \p Rng inside \p File with content \p Code. llvm::Expected formatRange(StringRef Code, diff --git a/clang-tools-extra/clangd/ClangdServer.cpp b/clang-tools-extra/clangd/ClangdServer.cpp --- a/clang-tools-extra/clangd/ClangdServer.cpp +++ b/clang-tools-extra/clangd/ClangdServer.cpp @@ -562,7 +562,7 @@ } void ClangdServer::findReferences(PathRef File, Position Pos, uint32_t Limit, - Callback> CB) { + Callback CB) { auto Action = [Pos, Limit, CB = std::move(CB), this](llvm::Expected InpAST) mutable { if (!InpAST) diff --git a/clang-tools-extra/clangd/XRefs.h b/clang-tools-extra/clangd/XRefs.h --- a/clang-tools-extra/clangd/XRefs.h +++ b/clang-tools-extra/clangd/XRefs.h @@ -124,11 +124,14 @@ format::FormatStyle Style, const SymbolIndex *Index); -/// Returns reference locations of the symbol at a specified \p Pos. +struct ReferencesResult { + std::vector References; + bool HasMore = false; +}; +/// Returns references of the symbol at a specified \p Pos. /// \p Limit limits the number of results returned (0 means no limit). -std::vector findReferences(ParsedAST &AST, Position Pos, - uint32_t Limit, - const SymbolIndex *Index = nullptr); +ReferencesResult findReferences(ParsedAST &AST, Position Pos, uint32_t Limit, + const SymbolIndex *Index = nullptr); /// Get info about symbols at \p Pos. std::vector getSymbolInfo(ParsedAST &AST, Position Pos); diff --git a/clang-tools-extra/clangd/XRefs.cpp b/clang-tools-extra/clangd/XRefs.cpp --- a/clang-tools-extra/clangd/XRefs.cpp +++ b/clang-tools-extra/clangd/XRefs.cpp @@ -883,11 +883,11 @@ return HI; } -std::vector findReferences(ParsedAST &AST, Position Pos, - uint32_t Limit, const SymbolIndex *Index) { +ReferencesResult findReferences(ParsedAST &AST, Position Pos, uint32_t Limit, + const SymbolIndex *Index) { if (!Limit) Limit = std::numeric_limits::max(); - std::vector Results; + ReferencesResult Results; const SourceManager &SM = AST.getSourceManager(); auto MainFilePath = getCanonicalPath(SM.getFileEntryForID(SM.getMainFileID()), SM); @@ -922,12 +922,11 @@ Location Result; Result.range = *Range; Result.uri = URIForFile::canonicalize(*MainFilePath, *MainFilePath); - Results.push_back(std::move(Result)); + Results.References.push_back(std::move(Result)); } } - // Now query the index for references from other files. - if (Index && Results.size() < Limit) { + if (Index && Results.References.size() <= Limit) { RefsRequest Req; Req.Limit = Limit; @@ -942,15 +941,22 @@ } if (Req.IDs.empty()) return Results; - Index->refs(Req, [&](const Ref &R) { + Results.HasMore |= Index->refs(Req, [&](const Ref &R) { + // no need to continue process if we reach the limit. + if (Results.References.size() > Limit) + return; auto LSPLoc = toLSPLocation(R.Location, *MainFilePath); // Avoid indexed results for the main file - the AST is authoritative. - if (LSPLoc && LSPLoc->uri.file() != *MainFilePath) - Results.push_back(std::move(*LSPLoc)); + if (!LSPLoc || LSPLoc->uri.file() == *MainFilePath) + return; + + Results.References.push_back(std::move(*LSPLoc)); }); } - if (Results.size() > Limit) - Results.resize(Limit); + if (Results.References.size() > Limit) { + Results.HasMore = true; + Results.References.resize(Limit); + } return Results; } diff --git a/clang-tools-extra/clangd/unittests/XRefsTests.cpp b/clang-tools-extra/clangd/unittests/XRefsTests.cpp --- a/clang-tools-extra/clangd/unittests/XRefsTests.cpp +++ b/clang-tools-extra/clangd/unittests/XRefsTests.cpp @@ -2109,7 +2109,7 @@ std::vector> ExpectedLocations; for (const auto &R : T.ranges()) ExpectedLocations.push_back(RangeIs(R)); - EXPECT_THAT(findReferences(AST, T.point(), 0), + EXPECT_THAT(findReferences(AST, T.point(), 0).References, ElementsAreArray(ExpectedLocations)) << Test; } @@ -2170,7 +2170,7 @@ for (const auto &R : T.ranges()) ExpectedLocations.push_back(RangeIs(R)); ASSERT_THAT(ExpectedLocations, Not(IsEmpty())); - EXPECT_THAT(findReferences(AST, T.point(), 0), + EXPECT_THAT(findReferences(AST, T.point(), 0).References, ElementsAreArray(ExpectedLocations)) << Test; } @@ -2185,8 +2185,9 @@ auto AST = TU.build(); // References in main file are returned without index. - EXPECT_THAT(findReferences(AST, Main.point(), 0, /*Index=*/nullptr), - ElementsAre(RangeIs(Main.range()))); + EXPECT_THAT( + findReferences(AST, Main.point(), 0, /*Index=*/nullptr).References, + ElementsAre(RangeIs(Main.range()))); Annotations IndexedMain(R"cpp( int main() { [[f^oo]](); } )cpp"); @@ -2196,17 +2197,18 @@ IndexedTU.Code = IndexedMain.code(); IndexedTU.Filename = "Indexed.cpp"; IndexedTU.HeaderCode = Header; - EXPECT_THAT(findReferences(AST, Main.point(), 0, IndexedTU.index().get()), - ElementsAre(RangeIs(Main.range()), RangeIs(IndexedMain.range()))); - - EXPECT_EQ(1u, findReferences(AST, Main.point(), /*Limit*/ 1, - IndexedTU.index().get()) - .size()); + EXPECT_THAT( + findReferences(AST, Main.point(), 0, IndexedTU.index().get()).References, + ElementsAre(RangeIs(Main.range()), RangeIs(IndexedMain.range()))); + auto LimitRefs = + findReferences(AST, Main.point(), /*Limit*/ 1, IndexedTU.index().get()); + EXPECT_EQ(1u, LimitRefs.References.size()); + EXPECT_TRUE(LimitRefs.HasMore); // If the main file is in the index, we don't return duplicates. // (even if the references are in a different location) TU.Code = ("\n\n" + Main.code()).str(); - EXPECT_THAT(findReferences(AST, Main.point(), 0, TU.index().get()), + EXPECT_THAT(findReferences(AST, Main.point(), 0, TU.index().get()).References, ElementsAre(RangeIs(Main.range()))); }