hokein created this revision.
hokein added a reviewer: sammccall.
Herald added subscribers: kadircet, arphaman, jkorous, MaskRay, ioeric, 
ilya-biryukov.

Repository:
  rCTE Clang Tools Extra

https://reviews.llvm.org/D56597

Files:
  clangd/ClangdLSPServer.cpp
  clangd/ClangdServer.cpp
  clangd/ClangdServer.h
  clangd/XRefs.cpp
  clangd/XRefs.h
  clangd/index/Index.h
  clangd/index/MemIndex.cpp
  clangd/index/Merge.cpp
  clangd/index/dex/Dex.cpp
  unittests/clangd/DexTests.cpp
  unittests/clangd/IndexTests.cpp
  unittests/clangd/XRefsTests.cpp

Index: unittests/clangd/XRefsTests.cpp
===================================================================
--- unittests/clangd/XRefsTests.cpp
+++ unittests/clangd/XRefsTests.cpp
@@ -1219,7 +1219,7 @@
     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 @@
     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 @@
   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(
+      1, 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
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/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<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"));
+  {
+    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
Index: clangd/index/dex/Dex.cpp
===================================================================
--- clangd/index/dex/Dex.cpp
+++ clangd/index/dex/Dex.cpp
@@ -236,10 +236,16 @@
 void Dex::refs(const RefsRequest &Req,
                llvm::function_ref<void(const Ref &)> Callback) const {
   trace::Span Tracer("Dex refs");
+  size_t Limit = Req.Limit ? *Req.Limit : std::numeric_limits<size_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 {
Index: clangd/index/Merge.cpp
===================================================================
--- clangd/index/Merge.cpp
+++ clangd/index/Merge.cpp
@@ -95,14 +95,20 @@
   // 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.
+  size_t Limit = Req.Limit ? *Req.Limit : std::numeric_limits<size_t>::max();
   llvm::StringSet<> DynamicIndexFileURIs;
   Dynamic->refs(Req, [&](const Ref &O) {
     DynamicIndexFileURIs.insert(O.Location.FileURI);
     Callback(O);
+    --Limit;
   });
+  if (Limit == 0)
+    return;
   Static->refs(Req, [&](const Ref &O) {
-    if (!DynamicIndexFileURIs.count(O.Location.FileURI))
+    if (Limit > 0 && !DynamicIndexFileURIs.count(O.Location.FileURI)) {
+      --Limit;
       Callback(O);
+    }
   });
 }
 
Index: clangd/index/MemIndex.cpp
===================================================================
--- clangd/index/MemIndex.cpp
+++ clangd/index/MemIndex.cpp
@@ -69,13 +69,19 @@
 void MemIndex::refs(const RefsRequest &Req,
                     llvm::function_ref<void(const Ref &)> Callback) const {
   trace::Span Tracer("MemIndex refs");
+  size_t Limit = Req.Limit ? *Req.Limit : std::numeric_limits<size_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);
+      }
+    }
   }
 }
 
Index: clangd/index/Index.h
===================================================================
--- clangd/index/Index.h
+++ clangd/index/Index.h
@@ -474,6 +474,8 @@
 struct RefsRequest {
   llvm::DenseSet<SymbolID> IDs;
   RefKind Filter = RefKind::All;
+  /// If set, limit the number of refers returned from the index.
+  llvm::Optional<uint32_t> Limit;
 };
 
 /// Interface for symbol indexes that can be used for searching or
Index: clangd/XRefs.h
===================================================================
--- clangd/XRefs.h
+++ clangd/XRefs.h
@@ -35,7 +35,8 @@
 llvm::Optional<Hover> getHover(ParsedAST &AST, Position Pos);
 
 /// Returns reference locations of the symbol at a specified \p Pos.
-std::vector<Location> findReferences(ParsedAST &AST, Position Pos,
+/// \p Limit limits the number of results returned (0 means no limit).
+std::vector<Location> findReferences(ParsedAST &AST, Position Pos, int Limit,
                                      const SymbolIndex *Index = nullptr);
 
 /// Get info about symbols at \p Pos.
Index: clangd/XRefs.cpp
===================================================================
--- clangd/XRefs.cpp
+++ clangd/XRefs.cpp
@@ -703,7 +703,7 @@
   return None;
 }
 
-std::vector<Location> findReferences(ParsedAST &AST, Position Pos,
+std::vector<Location> findReferences(ParsedAST &AST, Position Pos, int Limit,
                                      const SymbolIndex *Index) {
   std::vector<Location> Results;
   const SourceManager &SM = AST.getASTContext().getSourceManager();
@@ -721,11 +721,15 @@
     if (DI.IsReferencedExplicitly)
       TargetDecls.push_back(DI.D);
   }
-
+  if (!Limit)
+    Limit = std::numeric_limits<int>::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 (Limit == 0)
+      break;
+    --Limit;
     Location Result;
     Result.range = getTokenRange(AST, Ref.Loc);
     Result.uri = URIForFile::canonicalize(*MainFilePath, *MainFilePath);
@@ -733,9 +737,11 @@
   }
 
   // Now query the index for references from other files.
-  if (!Index)
+  if (!Index || Limit == 0)
     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
@@ -750,8 +756,10 @@
   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)
+    if (Limit > 0 && LSPLoc && LSPLoc->uri.file() != *MainFilePath) {
+      --Limit;
       Results.push_back(std::move(*LSPLoc));
+    }
   });
   return Results;
 }
Index: clangd/ClangdServer.h
===================================================================
--- clangd/ClangdServer.h
+++ clangd/ClangdServer.h
@@ -181,7 +181,7 @@
                        Callback<std::vector<DocumentSymbol>> CB);
 
   /// Retrieve locations for symbol references.
-  void findReferences(PathRef File, Position Pos,
+  void findReferences(PathRef File, Position Pos, int Limit,
                       Callback<std::vector<Location>> 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, int 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)));
Index: clangd/ClangdLSPServer.cpp
===================================================================
--- clangd/ClangdLSPServer.cpp
+++ clangd/ClangdLSPServer.cpp
@@ -710,7 +710,7 @@
 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,
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
http://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to