Index: clangd/ClangdLSPServer.cpp =================================================================== --- clangd/ClangdLSPServer.cpp +++ clangd/ClangdLSPServer.cpp @@ -710,7 +710,7 @@ void ClangdLSPServer::onReference(const ReferenceParams &Params, Callback> Reply) { Server->findReferences(Params.textDocument.uri.file(), Params.position, - std::move(Reply)); + CCOpts.Limit, std::move(Reply)); } void ClangdLSPServer::onSymbolInfo(const TextDocumentPositionParams &Params, Index: clangd/ClangdServer.h =================================================================== --- clangd/ClangdServer.h +++ clangd/ClangdServer.h @@ -181,7 +181,7 @@ Callback> CB); /// Retrieve locations for symbol references. - void findReferences(PathRef File, Position Pos, + void findReferences(PathRef File, Position Pos, uint32_t Limit, Callback> CB); /// Run formatting for \p Rng inside \p File with content \p Code. Index: clangd/ClangdServer.cpp =================================================================== --- clangd/ClangdServer.cpp +++ clangd/ClangdServer.cpp @@ -500,13 +500,13 @@ Bind(Action, std::move(CB))); } -void ClangdServer::findReferences(PathRef File, Position Pos, +void ClangdServer::findReferences(PathRef File, Position Pos, uint32_t Limit, Callback> CB) { - auto Action = [Pos, this](Callback> CB, - llvm::Expected InpAST) { + auto Action = [Pos, Limit, this](Callback> CB, + llvm::Expected InpAST) { if (!InpAST) return CB(InpAST.takeError()); - CB(clangd::findReferences(InpAST->AST, Pos, Index)); + CB(clangd::findReferences(InpAST->AST, Pos, Limit, Index)); }; WorkScheduler.runWithAST("References", File, Bind(Action, std::move(CB))); Index: clangd/XRefs.h =================================================================== --- clangd/XRefs.h +++ clangd/XRefs.h @@ -35,7 +35,9 @@ llvm::Optional getHover(ParsedAST &AST, Position Pos); /// Returns reference locations 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); /// Get info about symbols at \p Pos. Index: clangd/XRefs.cpp =================================================================== --- clangd/XRefs.cpp +++ clangd/XRefs.cpp @@ -704,7 +704,9 @@ } std::vector findReferences(ParsedAST &AST, Position Pos, - const SymbolIndex *Index) { + uint32_t Limit, const SymbolIndex *Index) { + if (!Limit) + Limit = std::numeric_limits::max(); std::vector Results; const SourceManager &SM = AST.getASTContext().getSourceManager(); auto MainFilePath = @@ -733,26 +735,30 @@ } // Now query the index for references from other files. - if (!Index) - return Results; - RefsRequest Req; - for (const Decl *D : TargetDecls) { - // Not all symbols can be referenced from outside (e.g. function-locals). - // TODO: we could skip TU-scoped symbols here (e.g. static functions) if - // we know this file isn't a header. The details might be tricky. - if (D->getParentFunctionOrMethod()) - continue; - if (auto ID = getSymbolID(D)) - Req.IDs.insert(*ID); + if (Index && Results.size() < Limit) { + RefsRequest Req; + Req.Limit = Limit; + + for (const Decl *D : TargetDecls) { + // Not all symbols can be referenced from outside (e.g. function-locals). + // TODO: we could skip TU-scoped symbols here (e.g. static functions) if + // we know this file isn't a header. The details might be tricky. + if (D->getParentFunctionOrMethod()) + continue; + if (auto ID = getSymbolID(D)) + Req.IDs.insert(*ID); + } + if (Req.IDs.empty()) + return Results; + Index->refs(Req, [&](const Ref &R) { + 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 (Req.IDs.empty()) - return Results; - Index->refs(Req, [&](const Ref &R) { - 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 (Results.size() > Limit) + Results.resize(Limit); return Results; } Index: clangd/index/Index.h =================================================================== --- clangd/index/Index.h +++ clangd/index/Index.h @@ -474,6 +474,10 @@ struct RefsRequest { llvm::DenseSet IDs; RefKind Filter = RefKind::All; + /// If set, limit the number of refers returned from the index. The index may + /// choose to return less than this, e.g. it tries to avoid returning stale + /// results. + llvm::Optional Limit; }; /// Interface for symbol indexes that can be used for searching or Index: clangd/index/MemIndex.cpp =================================================================== --- clangd/index/MemIndex.cpp +++ clangd/index/MemIndex.cpp @@ -69,13 +69,18 @@ void MemIndex::refs(const RefsRequest &Req, llvm::function_ref Callback) const { trace::Span Tracer("MemIndex refs"); + uint32_t Remaining = + Req.Limit.getValueOr(std::numeric_limits::max()); for (const auto &ReqID : Req.IDs) { auto SymRefs = Refs.find(ReqID); if (SymRefs == Refs.end()) continue; - for (const auto &O : SymRefs->second) - if (static_cast(Req.Filter & O.Kind)) + for (const auto &O : SymRefs->second) { + if (Remaining > 0 && static_cast(Req.Filter & O.Kind)) { + --Remaining; Callback(O); + } + } } } Index: clangd/index/Merge.cpp =================================================================== --- clangd/index/Merge.cpp +++ clangd/index/Merge.cpp @@ -87,6 +87,8 @@ void MergedIndex::refs(const RefsRequest &Req, llvm::function_ref Callback) const { trace::Span Tracer("MergedIndex refs"); + uint32_t Remaining = + Req.Limit.getValueOr(std::numeric_limits::max()); // We don't want duplicated refs from the static/dynamic indexes, // and we can't reliably duplicate them because offsets may differ slightly. // We consider the dynamic index authoritative and report all its refs, @@ -99,10 +101,18 @@ Dynamic->refs(Req, [&](const Ref &O) { DynamicIndexFileURIs.insert(O.Location.FileURI); Callback(O); + --Remaining; }); + assert(Remaining >= 0); + if (Remaining == 0) + return; + // We return less than Req.Limit if static index returns more refs for dirty + // files. Static->refs(Req, [&](const Ref &O) { - if (!DynamicIndexFileURIs.count(O.Location.FileURI)) + if (Remaining > 0 && !DynamicIndexFileURIs.count(O.Location.FileURI)) { + --Remaining; Callback(O); + } }); } Index: clangd/index/dex/Dex.cpp =================================================================== --- clangd/index/dex/Dex.cpp +++ clangd/index/dex/Dex.cpp @@ -236,10 +236,15 @@ void Dex::refs(const RefsRequest &Req, llvm::function_ref Callback) const { trace::Span Tracer("Dex refs"); + uint32_t Remaining = + Req.Limit.getValueOr(std::numeric_limits::max()); for (const auto &ID : Req.IDs) - for (const auto &Ref : Refs.lookup(ID)) - if (static_cast(Req.Filter & Ref.Kind)) + for (const auto &Ref : Refs.lookup(ID)) { + if (Remaining > 0 && static_cast(Req.Filter & Ref.Kind)) { + --Remaining; Callback(Ref); + } + } } size_t Dex::estimateMemoryUsage() const { Index: unittests/clangd/DexTests.cpp =================================================================== --- unittests/clangd/DexTests.cpp +++ unittests/clangd/DexTests.cpp @@ -667,18 +667,26 @@ auto Foo = symbol("foo"); auto Bar = symbol("bar"); AddRef(Foo, "foo.h", RefKind::Declaration); + AddRef(Foo, "foo.cc", RefKind::Definition); AddRef(Foo, "reffoo.h", RefKind::Reference); AddRef(Bar, "bar.h", RefKind::Declaration); - std::vector Files; RefsRequest Req; Req.IDs.insert(Foo.ID); Req.Filter = RefKind::Declaration | RefKind::Definition; + + std::vector Files; Dex(std::vector{Foo, Bar}, Refs).refs(Req, [&](const Ref &R) { Files.push_back(R.Location.FileURI); }); + EXPECT_THAT(Files, UnorderedElementsAre("foo.h", "foo.cc")); - EXPECT_THAT(Files, ElementsAre("foo.h")); + Req.Limit = 1; + Files.clear(); + Dex(std::vector{Foo, Bar}, Refs).refs(Req, [&](const Ref &R) { + Files.push_back(R.Location.FileURI); + }); + EXPECT_THAT(Files, ElementsAre(AnyOf("foo.h", "foo.cc"))); } } // namespace Index: unittests/clangd/IndexTests.cpp =================================================================== --- unittests/clangd/IndexTests.cpp +++ unittests/clangd/IndexTests.cpp @@ -19,6 +19,7 @@ using testing::_; using testing::AllOf; +using testing::AnyOf; using testing::ElementsAre; using testing::Pair; using testing::Pointee; @@ -292,7 +293,6 @@ Request.IDs = {Foo.ID}; RefSlab::Builder Results; Merge.refs(Request, [&](const Ref &O) { Results.insert(Foo.ID, O); }); - EXPECT_THAT( std::move(Results).build(), ElementsAre(Pair( @@ -300,6 +300,14 @@ FileURI("unittest:///test.cc")), AllOf(RefRange(Test2Code.range("Foo")), FileURI("unittest:///test2.cc")))))); + + Request.Limit = 1; + RefSlab::Builder Results2; + Merge.refs(Request, [&](const Ref &O) { Results2.insert(Foo.ID, O); }); + EXPECT_THAT(std::move(Results2).build(), + ElementsAre(Pair( + _, ElementsAre(AnyOf(FileURI("unittest:///test.cc"), + FileURI("unittest:///test2.cc")))))); } MATCHER_P2(IncludeHeaderWithRef, IncludeHeader, References, "") { Index: unittests/clangd/XRefsTests.cpp =================================================================== --- unittests/clangd/XRefsTests.cpp +++ unittests/clangd/XRefsTests.cpp @@ -1219,7 +1219,7 @@ std::vector> ExpectedLocations; for (const auto &R : T.ranges()) ExpectedLocations.push_back(RangeIs(R)); - EXPECT_THAT(findReferences(AST, T.point()), + EXPECT_THAT(findReferences(AST, T.point(), 0), ElementsAreArray(ExpectedLocations)) << Test; } @@ -1266,7 +1266,7 @@ std::vector> ExpectedLocations; for (const auto &R : T.ranges()) ExpectedLocations.push_back(RangeIs(R)); - EXPECT_THAT(findReferences(AST, T.point()), + EXPECT_THAT(findReferences(AST, T.point(), 0), ElementsAreArray(ExpectedLocations)) << Test; } @@ -1281,7 +1281,7 @@ auto AST = TU.build(); // References in main file are returned without index. - EXPECT_THAT(findReferences(AST, Main.point(), /*Index=*/nullptr), + EXPECT_THAT(findReferences(AST, Main.point(), 0, /*Index=*/nullptr), ElementsAre(RangeIs(Main.range()))); Annotations IndexedMain(R"cpp( int main() { [[f^oo]](); } @@ -1292,13 +1292,17 @@ IndexedTU.Code = IndexedMain.code(); IndexedTU.Filename = "Indexed.cpp"; IndexedTU.HeaderCode = Header; - EXPECT_THAT(findReferences(AST, Main.point(), IndexedTU.index().get()), + 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()); + // 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(), TU.index().get()), + EXPECT_THAT(findReferences(AST, Main.point(), 0, TU.index().get()), ElementsAre(RangeIs(Main.range()))); } @@ -1328,7 +1332,7 @@ Annotations File(T.AnnotatedCode); RecordingIndex Rec; auto AST = TestTU::withCode(File.code()).build(); - findReferences(AST, File.point(), &Rec); + findReferences(AST, File.point(), 0, &Rec); if (T.WantQuery) EXPECT_NE(Rec.RefIDs, None) << T.AnnotatedCode; else