[clangd] Add Limit parameter for xref.

[hokein]

Fix an issue and add comments.

[sammccall]

clangd/XRefs.cpp

726

nit: move this to the top, since it's just adjusting input params?

731

I don't think we need to break up the flow here to early exit.

I'd suggest reducing the number of places we deal with limits, we can get away with 3 instead of 5:

... collect results as before ...
if (Index && Limit && Results.size() < Limit) {
  ... query index ...
  Req.Limit = Limit;
  ... append results as usual ...
}
if (Limit && Results.size() > Limit)
  Results.resize(limit)

clangd/index/MemIndex.cpp

72

Req.Limit.getValueOr(...)

72

if you're going to mutate this, maybe call it "remaining" or so?

clangd/index/Merge.cpp

100

why two separate count variables?

100

you've inserted the new code between existing code and the comment that describes it.

unittests/clangd/DexTests.cpp

684–689

If you want these to be the same test, there should be some connection between them.
A natural one is that it's the same query with a different limit - you could extend the original test to have multiple results.

697

As far as I can tell, this limit makes no difference, there's only one result anyway.

702

In tests, we know what the number is, and it's clearer to just spell it out: EXPECT_EQ(1, RefCount)

702

generally prefer to use EXPECT_THAT only with matchers. EXPECT_EQ gives clearer error messages, reads more naturally, and is more common.

Better still would be to assert the actual outcomes, e.g. ElementsAre(anyOf("foo1.h", "foo2.h"))

unittests/clangd/IndexTests.cpp

315

why new scope here?

[hokein]

Address comments.

[hokein]

unittests/clangd/DexTests.cpp

684–689

Done, made the query request the same in this test except the limit.

697

Actually, the filter of the request was different than the case above, we would get 2 refs. Used the same request to avoid confusion.

unittests/clangd/IndexTests.cpp

315

To align with the case above.

[sammccall]

unittests/clangd/DexTests.cpp

691

instead of splitting scopes here, might be more obvious just to Files.clear() between?

unittests/clangd/IndexTests.cpp

315

FWIW I'd find it easier to read if neither this nor the one above used new scopes, which tend to suggest some spooky scoped object side effects to me.
Up to you, though.

[hokein]

Remove scopes.

clangd/ClangdLSPServer.cpp

 void ClangdLSPServer::onChangeConfiguration(
     const DidChangeConfigurationParams &Params) {
   applyConfiguration(Params.settings);
 }
 
 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,
                                    Callback<std::vector<SymbolDetails>> Reply) {
   Server->symbolInfo(Params.textDocument.uri.file(), Params.position,
                      std::move(Reply));
 }
 

clangd/ClangdServer.h

public:
   void workspaceSymbols(StringRef Query, int Limit,
                         Callback<std::vector<SymbolInformation>> CB);
 
   /// Retrieve the symbols within the specified file.
   void documentSymbols(StringRef File,
                        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.
   llvm::Expected<tooling::Replacements> formatRange(StringRef Code,
                                                     PathRef File, Range Rng);
 
   /// Run formatting for the whole \p File with content \p Code.
   llvm::Expected<tooling::Replacements> formatFile(StringRef Code,

clangd/ClangdServer.cpp

auto Action = [](Callback<std::vector<DocumentSymbol>> CB,
     if (!InpAST)
       return CB(InpAST.takeError());
     CB(clangd::getDocumentSymbols(InpAST->AST));
   };
   WorkScheduler.runWithAST("documentSymbols", 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,
+  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)));
 }
 
 void ClangdServer::symbolInfo(PathRef File, Position Pos,
                               Callback<std::vector<SymbolDetails>> CB) {
   auto Action = [Pos](Callback<std::vector<SymbolDetails>> CB,

clangd/XRefs.h

 /// Returns highlights for all usages of a symbol at \p Pos.
 std::vector<DocumentHighlight> findDocumentHighlights(ParsedAST &AST,
                                                       Position Pos);
 
 /// Get the hover information when hovering at \p Pos.
 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.
 std::vector<SymbolDetails> getSymbolInfo(ParsedAST &AST, Position Pos);
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_XREFS_H

clangd/XRefs.cpp

llvm::Optional<Hover> getHover(ParsedAST &AST, Position Pos) {
   auto DeducedType = getDeducedType(AST, SourceLocationBeg);
   if (DeducedType && !DeducedType->isNull())
     return getHoverContents(*DeducedType, AST.getASTContext());
 
   return None;
 }
 
 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 =
       getCanonicalPath(SM.getFileEntryForID(SM.getMainFileID()), SM);
   if (!MainFilePath) {
     elog("Failed to get a path for the main file, so no references");
     return Results;
   }
   auto Loc = getBeginningOfIdentifier(AST, Pos, SM.getMainFileID());
   auto Symbols = getSymbolAtPosition(AST, Loc);
 
   std::vector<const Decl *> TargetDecls;
   for (const DeclInfo &DI : Symbols.Decls) {
     if (DI.IsReferencedExplicitly)
       TargetDecls.push_back(DI.D);
   }
 
   // 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) {
     Location Result;
     Result.range = getTokenRange(AST, Ref.Loc);
     Result.uri = URIForFile::canonicalize(*MainFilePath, *MainFilePath);
     Results.push_back(std::move(Result));
   }
 
   // Now query the index for references from other files.
-  if (!Index)
-    return Results;
+  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 (Results.size() > Limit)
+    Results.resize(Limit);
   return Results;
 }
 
 std::vector<SymbolDetails> getSymbolInfo(ParsedAST &AST, Position Pos) {
   const SourceManager &SM = AST.getASTContext().getSourceManager();
 
   auto Loc = getBeginningOfIdentifier(AST, Pos, SM.getMainFileID());
   auto Symbols = getSymbolAtPosition(AST, Loc);

clangd/index/Index.h

 
 struct LookupRequest {
   llvm::DenseSet<SymbolID> IDs;
 };
 
 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
 /// matching symbols among a set of symbols based on names or unique IDs.
 class SymbolIndex {
 public:
   virtual ~SymbolIndex() = default;
 

clangd/index/MemIndex.cpp

for (const auto &ID : Req.IDs) {
     if (I != Index.end())
       Callback(*I->second);
   }
 }
 
 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);
       }
     }
+  }
+}
 
 size_t MemIndex::estimateMemoryUsage() const {
   return Index.getMemorySize() + Refs.getMemorySize() + BackingDataSize;
 }
 
 } // namespace clangd
 } // namespace clang

clangd/index/Merge.cpp

void MergedIndex::lookup(
   for (const auto &ID : RemainingIDs)
     if (const Symbol *Sym = B.find(ID))
       Callback(*Sym);
 }
 
 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,
   // and only report static index refs from other files.
   //
   // 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.
   llvm::StringSet<> DynamicIndexFileURIs;
   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);
+    }
   });
 }
 
 Symbol mergeSymbol(const Symbol &L, const Symbol &R) {
   assert(L.ID == R.ID);
   // We prefer information from TUs that saw the definition.
   // Classes: this is the def itself. Functions: hopefully the header decl.
   // If both did (or both didn't), continue to prefer L over R.

clangd/index/dex/Dex.cpp

for (const auto &ID : Req.IDs) {
     if (I != LookupTable.end())
       Callback(*I->second);
   }
 }
 
 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 {
   size_t Bytes = Symbols.size() * sizeof(const Symbol *);
   Bytes += SymbolQuality.size() * sizeof(float);
   Bytes += LookupTable.getMemorySize();
   Bytes += InvertedIndex.getMemorySize();
   for (const auto &TokenToPostingList : InvertedIndex)
     Bytes += TokenToPostingList.second.bytes();

unittests/clangd/DexTests.cpp

auto AddRef = [&](const Symbol &Sym, const char *Filename, RefKind Kind) {
     auto &SymbolRefs = Refs[Sym.ID];
     SymbolRefs.emplace_back();
     SymbolRefs.back().Kind = Kind;
     SymbolRefs.back().Location.FileURI = Filename;
   };
   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
 } // namespace dex
 } // namespace clangd
 } // namespace clang

unittests/clangd/IndexTests.cpp

 #include "index/Index.h"
 #include "index/MemIndex.h"
 #include "index/Merge.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 using testing::_;
 using testing::AllOf;
+using testing::AnyOf;
 using testing::ElementsAre;
 using testing::Pair;
 using testing::Pointee;
 using testing::UnorderedElementsAre;
 
 namespace clang {
 namespace clangd {
 namespace {
TEST(MergeIndexTest, Refs) {
   Test2.Filename = "test2.cc";
   StaticAST = Test2.build();
   StaticIndex.updateMain(Test2.Filename, StaticAST);
 
   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"))))));
+
+  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, "") {
   return (arg.IncludeHeader == IncludeHeader) && (arg.References == References);
 }
 
 TEST(MergeTest, MergeIncludesOnDifferentDefinitions) {
   Symbol L, R;
   L.Name = "left";
   R.Name = "right";
   L.ID = R.ID = SymbolID("hello");
   L.IncludeHeaders.emplace_back("common", 1);
   R.IncludeHeaders.emplace_back("common", 1);

unittests/clangd/XRefsTests.cpp

const char *Tests[] = {
       )cpp",
   };
   for (const char *Test : Tests) {
     Annotations T(Test);
     auto AST = TestTU::withCode(T.code()).build();
     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;
   }
 }
 
 TEST(FindReferences, ExplicitSymbols) {
   const char *Tests[] = {
       R"cpp(
const char *Tests[] = {
       )cpp",
   };
   for (const char *Test : Tests) {
     Annotations T(Test);
     auto AST = TestTU::withCode(T.code()).build();
     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;
   }
 }
 
 TEST(FindReferences, NeedsIndex) {
   const char *Header = "int foo();";
   Annotations Main("int main() { [[f^oo]](); }");
   TestTU TU;
   TU.Code = Main.code();
   TU.HeaderCode = Header;
   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]](); }
   )cpp");
 
   // References from indexed files are included.
   TestTU IndexedTU;
   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())));
 }
 
 TEST(FindReferences, NoQueryForLocalSymbols) {
   struct RecordingIndex : public MemIndex {
     mutable Optional<llvm::DenseSet<SymbolID>> RefIDs;
     void refs(const RefsRequest &Req,
               llvm::function_ref<void(const Ref &)>) const override {
bool WantQuery;
       {"void foo() { int ^x; }", false},
       {"void foo() { struct ^x{}; }", false},
       {"auto lambda = []{ int ^x; };", false},
   };
   for (Test T : Tests) {
     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
       EXPECT_EQ(Rec.RefIDs, None) << T.AnnotatedCode;
   }
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang