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,7 @@ } std::vector findReferences(ParsedAST &AST, Position Pos, - const SymbolIndex *Index) { + uint32_t Limit, const SymbolIndex *Index) { std::vector Results; const SourceManager &SM = AST.getASTContext().getSourceManager(); auto MainFilePath = @@ -721,11 +721,14 @@ if (DI.IsReferencedExplicitly) TargetDecls.push_back(DI.D); } - + if (!Limit) + Limit = std::numeric_limits::max(); // We traverse the AST to find references in the main file. // TODO: should we handle macros, too? auto MainFileRefs = findRefs(TargetDecls, AST); for (const auto &Ref : MainFileRefs) { + if (Results.size() == Limit) + break; Location Result; Result.range = getTokenRange(AST, Ref.Loc); Result.uri = URIForFile::canonicalize(*MainFilePath, *MainFilePath); @@ -733,9 +736,11 @@ } // Now query the index for references from other files. - if (!Index) + if (!Index || Results.size() == Limit) return Results; 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 @@ -748,6 +753,8 @@ if (Req.IDs.empty()) return Results; Index->refs(Req, [&](const Ref &R) { + if (Results.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) 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,20 @@ void MemIndex::refs(const RefsRequest &Req, llvm::function_ref Callback) const { trace::Span Tracer("MemIndex refs"); + uint32_t Limit = + Req.Limit ? *Req.Limit : 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 (Limit == 0) + return; + if (static_cast(Req.Filter & O.Kind)) { + --Limit; Callback(O); + } + } } } Index: clangd/index/Merge.cpp =================================================================== --- clangd/index/Merge.cpp +++ clangd/index/Merge.cpp @@ -95,14 +95,27 @@ // FIXME: The heuristic fails if the dynamic index contains a file, but all // refs were removed (we will report stale ones from the static index). // Ultimately we should explicit check which index has the file instead. + uint32_t DynamicCount = 0; + uint32_t Limit = + Req.Limit ? *Req.Limit : std::numeric_limits::max(); llvm::StringSet<> DynamicIndexFileURIs; Dynamic->refs(Req, [&](const Ref &O) { DynamicIndexFileURIs.insert(O.Location.FileURI); Callback(O); + ++DynamicCount; }); + if (DynamicCount == Limit) + return; + uint32_t StaticCount = 0; + // 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 (StaticCount + DynamicCount == Limit) + return; + if (!DynamicIndexFileURIs.count(O.Location.FileURI)) { + ++StaticCount; Callback(O); + } }); } Index: clangd/index/dex/Dex.cpp =================================================================== --- clangd/index/dex/Dex.cpp +++ clangd/index/dex/Dex.cpp @@ -236,10 +236,17 @@ void Dex::refs(const RefsRequest &Req, llvm::function_ref Callback) const { trace::Span Tracer("Dex refs"); + uint32_t Limit = + Req.Limit ? *Req.Limit : 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 (Limit == 0) + return; + if (static_cast(Req.Filter & Ref.Kind)) { + --Limit; Callback(Ref); + } + } } size_t Dex::estimateMemoryUsage() const { Index: unittests/clangd/DexTests.cpp =================================================================== --- unittests/clangd/DexTests.cpp +++ unittests/clangd/DexTests.cpp @@ -670,15 +670,26 @@ 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; - Dex(std::vector{Foo, Bar}, Refs).refs(Req, [&](const Ref &R) { - Files.push_back(R.Location.FileURI); - }); - - EXPECT_THAT(Files, ElementsAre("foo.h")); + { + std::vector Files; + RefsRequest Req; + Req.IDs.insert(Foo.ID); + Req.Filter = RefKind::Declaration | RefKind::Definition; + Dex(std::vector{Foo, Bar}, Refs).refs(Req, [&](const Ref &R) { + Files.push_back(R.Location.FileURI); + }); + EXPECT_THAT(Files, ElementsAre("foo.h")); + } + { + RefsRequest Req; + Req.IDs.insert(Foo.ID); + Req.Limit = 1; + size_t RefCount = 0; + Dex(std::vector{Foo, Bar}, Refs).refs(Req, [&](const Ref &R) { + ++RefCount; + }); + EXPECT_THAT(*Req.Limit, RefCount); + } } } // namespace Index: unittests/clangd/IndexTests.cpp =================================================================== --- unittests/clangd/IndexTests.cpp +++ unittests/clangd/IndexTests.cpp @@ -290,16 +290,24 @@ RefsRequest Request; 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( - _, UnorderedElementsAre(AllOf(RefRange(Test1Code.range("Foo")), - FileURI("unittest:///test.cc")), - AllOf(RefRange(Test2Code.range("Foo")), - FileURI("unittest:///test2.cc")))))); + { + RefSlab::Builder Results; + Merge.refs(Request, [&](const Ref &O) { Results.insert(Foo.ID, O); }); + + EXPECT_THAT( + std::move(Results).build(), + ElementsAre(Pair( + _, UnorderedElementsAre(AllOf(RefRange(Test1Code.range("Foo")), + FileURI("unittest:///test.cc")), + AllOf(RefRange(Test2Code.range("Foo")), + FileURI("unittest:///test2.cc")))))); + } + { + Request.Limit = 1; + size_t RefsCount = 0; + Merge.refs(Request, [&](const Ref &O) { ++RefsCount; }); + EXPECT_EQ(*Request.Limit, RefsCount); + } } 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