[clangd] Add Limit parameter for xref.

Reviewers: sammccall

Subscribers: ilya-biryukov, ioeric, MaskRay, jkorous, arphaman, kadircet, cfe-commits

Differential Revision: https://reviews.llvm.org/D56597

12 files changed, 92 insertions, 40 deletions

diff --git a/clang-tools-extra/clangd/ClangdLSPServer.cpp b/clang-tools-extra/clangd/ClangdLSPServer.cpp
index e84adebac44..80792eddc69 100644
--- a/clang-tools-extra/clangd/ClangdLSPServer.cpp
+++ b/clang-tools-extra/clangd/ClangdLSPServer.cpp
@@ -710,7 +710,7 @@ void ClangdLSPServer::onChangeConfiguration(
 void ClangdLSPServer::onReference(const ReferenceParams &Params,
                                   Callback<std::vector<Location>> Reply) {
   Server->findReferences(Params.textDocument.uri.file(), Params.position,
-                         std::move(Reply));
+                         CCOpts.Limit, std::move(Reply));
 }
 
 void ClangdLSPServer::onSymbolInfo(const TextDocumentPositionParams &Params,
diff --git a/clang-tools-extra/clangd/ClangdServer.cpp b/clang-tools-extra/clangd/ClangdServer.cpp
index f799d6bc670..53a27bdb91d 100644
--- a/clang-tools-extra/clangd/ClangdServer.cpp
+++ b/clang-tools-extra/clangd/ClangdServer.cpp
@@ -500,13 +500,13 @@ void ClangdServer::documentSymbols(llvm::StringRef File,
                            Bind(Action, std::move(CB)));
 }
 
-void ClangdServer::findReferences(PathRef File, Position Pos,
+void ClangdServer::findReferences(PathRef File, Position Pos, uint32_t Limit,
                                   Callback<std::vector<Location>> CB) {
-  auto Action = [Pos, this](Callback<std::vector<Location>> CB,
-                            llvm::Expected<InputsAndAST> InpAST) {
+  auto Action = [Pos, Limit, this](Callback<std::vector<Location>> CB,
+                                   llvm::Expected<InputsAndAST> 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)));
diff --git a/clang-tools-extra/clangd/ClangdServer.h b/clang-tools-extra/clangd/ClangdServer.h
index 4ce1506dfce..213f042c83e 100644
--- a/clang-tools-extra/clangd/ClangdServer.h
+++ b/clang-tools-extra/clangd/ClangdServer.h
@@ -181,7 +181,7 @@ public:
                        Callback<std::vector<DocumentSymbol>> CB);
 
   /// Retrieve locations for symbol references.
-  void findReferences(PathRef File, Position Pos,
+  void findReferences(PathRef File, Position Pos, uint32_t Limit,
                       Callback<std::vector<Location>> CB);
 
   /// Run formatting for \p Rng inside \p File with content \p Code.
diff --git a/clang-tools-extra/clangd/XRefs.cpp b/clang-tools-extra/clangd/XRefs.cpp
index b07efe9d69d..29561dfe5a5 100644
--- a/clang-tools-extra/clangd/XRefs.cpp
+++ b/clang-tools-extra/clangd/XRefs.cpp
@@ -704,7 +704,9 @@ llvm::Optional<Hover> getHover(ParsedAST &AST, Position Pos) {
 }
 
 std::vector<Location> findReferences(ParsedAST &AST, Position Pos,
-                                     const SymbolIndex *Index) {
+                                     uint32_t Limit, const SymbolIndex *Index) {
+  if (!Limit)
+    Limit = std::numeric_limits<uint32_t>::max();
   std::vector<Location> Results;
   const SourceManager &SM = AST.getASTContext().getSourceManager();
   auto MainFilePath =
@@ -733,26 +735,30 @@ std::vector<Location> findReferences(ParsedAST &AST, Position Pos,
   }
 
   // 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;
 }
 
diff --git a/clang-tools-extra/clangd/XRefs.h b/clang-tools-extra/clangd/XRefs.h
index 77d602489f3..631879d63bf 100644
--- a/clang-tools-extra/clangd/XRefs.h
+++ b/clang-tools-extra/clangd/XRefs.h
@@ -35,7 +35,9 @@ std::vector<DocumentHighlight> findDocumentHighlights(ParsedAST &AST,
 llvm::Optional<Hover> 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<Location> findReferences(ParsedAST &AST, Position Pos,
+                                     uint32_t Limit,
                                      const SymbolIndex *Index = nullptr);
 
 /// Get info about symbols at \p Pos.
diff --git a/clang-tools-extra/clangd/index/Index.h b/clang-tools-extra/clangd/index/Index.h
index 3957a7f6830..cc4cf2a29c9 100644
--- a/clang-tools-extra/clangd/index/Index.h
+++ b/clang-tools-extra/clangd/index/Index.h
@@ -476,6 +476,10 @@ struct LookupRequest {
 struct RefsRequest {
   llvm::DenseSet<SymbolID> 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<uint32_t> Limit;
 };
 
 /// Interface for symbol indexes that can be used for searching or
diff --git a/clang-tools-extra/clangd/index/MemIndex.cpp b/clang-tools-extra/clangd/index/MemIndex.cpp
index 35352e6e686..e4dc56c5afa 100644
--- a/clang-tools-extra/clangd/index/MemIndex.cpp
+++ b/clang-tools-extra/clangd/index/MemIndex.cpp
@@ -69,13 +69,18 @@ void MemIndex::lookup(const LookupRequest &Req,
 void MemIndex::refs(const RefsRequest &Req,
                     llvm::function_ref<void(const Ref &)> Callback) const {
   trace::Span Tracer("MemIndex refs");
+  uint32_t Remaining =
+      Req.Limit.getValueOr(std::numeric_limits<uint32_t>::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<int>(Req.Filter & O.Kind))
+    for (const auto &O : SymRefs->second) {
+      if (Remaining > 0 && static_cast<int>(Req.Filter & O.Kind)) {
+        --Remaining;
         Callback(O);
+      }
+    }
   }
 }
 
diff --git a/clang-tools-extra/clangd/index/Merge.cpp b/clang-tools-extra/clangd/index/Merge.cpp
index f599d757aa8..52f54e39dde 100644
--- a/clang-tools-extra/clangd/index/Merge.cpp
+++ b/clang-tools-extra/clangd/index/Merge.cpp
@@ -87,6 +87,8 @@ void MergedIndex::lookup(
 void MergedIndex::refs(const RefsRequest &Req,
                        llvm::function_ref<void(const Ref &)> Callback) const {
   trace::Span Tracer("MergedIndex refs");
+  uint32_t Remaining =
+      Req.Limit.getValueOr(std::numeric_limits<uint32_t>::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 @@ void MergedIndex::refs(const RefsRequest &Req,
   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);
+    }
   });
 }
 
diff --git a/clang-tools-extra/clangd/index/dex/Dex.cpp b/clang-tools-extra/clangd/index/dex/Dex.cpp
index 157d4668856..c38706c65ce 100644
--- a/clang-tools-extra/clangd/index/dex/Dex.cpp
+++ b/clang-tools-extra/clangd/index/dex/Dex.cpp
@@ -236,10 +236,15 @@ void Dex::lookup(const LookupRequest &Req,
 void Dex::refs(const RefsRequest &Req,
                llvm::function_ref<void(const Ref &)> Callback) const {
   trace::Span Tracer("Dex refs");
+  uint32_t Remaining =
+      Req.Limit.getValueOr(std::numeric_limits<uint32_t>::max());
   for (const auto &ID : Req.IDs)
-    for (const auto &Ref : Refs.lookup(ID))
-      if (static_cast<int>(Req.Filter & Ref.Kind))
+    for (const auto &Ref : Refs.lookup(ID)) {
+      if (Remaining > 0 && static_cast<int>(Req.Filter & Ref.Kind)) {
+        --Remaining;
         Callback(Ref);
+      }
+    }
 }
 
 size_t Dex::estimateMemoryUsage() const {
diff --git a/clang-tools-extra/unittests/clangd/DexTests.cpp b/clang-tools-extra/unittests/clangd/DexTests.cpp
index 6c347b65052..078bae27bc7 100644
--- a/clang-tools-extra/unittests/clangd/DexTests.cpp
+++ b/clang-tools-extra/unittests/clangd/DexTests.cpp
@@ -667,18 +667,26 @@ TEST(DexTests, Refs) {
   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<std::string> Files;
   RefsRequest Req;
   Req.IDs.insert(Foo.ID);
   Req.Filter = RefKind::Declaration | RefKind::Definition;
+
+  std::vector<std::string> Files;
   Dex(std::vector<Symbol>{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<Symbol>{Foo, Bar}, Refs).refs(Req, [&](const Ref &R) {
+    Files.push_back(R.Location.FileURI);
+  });
+  EXPECT_THAT(Files, ElementsAre(AnyOf("foo.h", "foo.cc")));
 }
 
 } // namespace
diff --git a/clang-tools-extra/unittests/clangd/IndexTests.cpp b/clang-tools-extra/unittests/clangd/IndexTests.cpp
index 596366d5ba1..3b7d6ff2b96 100644
--- a/clang-tools-extra/unittests/clangd/IndexTests.cpp
+++ b/clang-tools-extra/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 @@ TEST(MergeIndexTest, Refs) {
   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 @@ TEST(MergeIndexTest, Refs) {
                                         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, "") {
diff --git a/clang-tools-extra/unittests/clangd/XRefsTests.cpp b/clang-tools-extra/unittests/clangd/XRefsTests.cpp
index 123049cf116..88394b6cc51 100644
--- a/clang-tools-extra/unittests/clangd/XRefsTests.cpp
+++ b/clang-tools-extra/unittests/clangd/XRefsTests.cpp
@@ -1219,7 +1219,7 @@ TEST(FindReferences, WithinAST) {
     std::vector<Matcher<Location>> 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 @@ TEST(FindReferences, ExplicitSymbols) {
     std::vector<Matcher<Location>> 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 @@ TEST(FindReferences, NeedsIndex) {
   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 @@ TEST(FindReferences, NeedsIndex) {
   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 @@ TEST(FindReferences, NoQueryForLocalSymbols) {
     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