[clangd] Add Limit parameter for xref.

[hokein]

Fix an issue and add comments.

[sammccall]

clangd/XRefs.cpp

724

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

730

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

98

why two separate count variables?

98

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

unittests/clangd/DexTests.cpp

673

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.

686

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

691

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

691

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

306

why new scope here?

[hokein]

Address comments.

[hokein]

unittests/clangd/DexTests.cpp

673

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

686

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

306

To align with the case above.

[sammccall]

unittests/clangd/DexTests.cpp

680

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

unittests/clangd/IndexTests.cpp

306

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) {
   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);
   }
-
+  if (!Limit)
+    Limit = std::numeric_limits<uint32_t>::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);
     Results.push_back(std::move(Result));
   }
 
   // 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
     // 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) {
+    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)
       Results.push_back(std::move(*LSPLoc));
   });
   return Results;
 }
 

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 Limit =
+      Req.Limit ? *Req.Limit : 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 (Limit == 0)
+        return;
+      if (static_cast<int>(Req.Filter & O.Kind)) {
+        --Limit;
         Callback(O);
       }
     }
+  }
+}
 
 size_t MemIndex::estimateMemoryUsage() const {
   return Index.getMemorySize() + Refs.getMemorySize() + BackingDataSize;
 }
 
 } // namespace clangd
 } // namespace clang

clangd/index/Merge.cpp

void MergedIndex::refs(const RefsRequest &Req,
   // 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.
+  uint32_t DynamicCount = 0;
+  uint32_t Limit =
+      Req.Limit ? *Req.Limit : std::numeric_limits<uint32_t>::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);
+    }
   });
 }
 
 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 Limit =
+      Req.Limit ? *Req.Limit : 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 (Limit == 0)
+        return;
+      if (static_cast<int>(Req.Filter & Ref.Kind)) {
+        --Limit;
         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) {
     SymbolRefs.back().Location.FileURI = Filename;
   };
   auto Foo = symbol("foo");
   auto Bar = symbol("bar");
   AddRef(Foo, "foo.h", RefKind::Declaration);
   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;
     Dex(std::vector<Symbol>{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<Symbol>{Foo, Bar}, Refs).refs(Req, [&](const Ref &R) {
+      ++RefCount;
+    });
+    EXPECT_THAT(*Req.Limit, RefCount);
+  }
+}
 
 } // namespace
 } // namespace dex
 } // namespace clangd
 } // namespace clang

unittests/clangd/IndexTests.cpp

TEST(MergeIndexTest, Refs) {
   Test2.HeaderCode = HeaderCode;
   Test2.Code = Test2Code.code();
   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;
+    size_t RefsCount = 0;
+    Merge.refs(Request, [&](const Ref &O) { ++RefsCount; });
+    EXPECT_EQ(*Request.Limit, RefsCount);
+  }
+}
 
 MATCHER_P2(IncludeHeaderWithRef, IncludeHeader, References, "") {
   return (arg.IncludeHeader == IncludeHeader) && (arg.References == References);
 }
 
 TEST(MergeTest, MergeIncludesOnDifferentDefinitions) {
   Symbol L, R;
   L.Name = "left";

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