kbobyrev updated this revision to Diff 160088.
kbobyrev marked 2 inline comments as done.

https://reviews.llvm.org/D50500

Files:
  clang-tools-extra/clangd/index/dex/Iterator.cpp
  clang-tools-extra/clangd/index/dex/Iterator.h
  clang-tools-extra/unittests/clangd/DexIndexTests.cpp


Index: clang-tools-extra/unittests/clangd/DexIndexTests.cpp
===================================================================
--- clang-tools-extra/unittests/clangd/DexIndexTests.cpp
+++ clang-tools-extra/unittests/clangd/DexIndexTests.cpp
@@ -240,6 +240,27 @@
             "(& (& [1, 3, 5, 8, 9] [1, 5, 7, 9]) (| [0, 5] [0, 1, 5] []))");
 }
 
+TEST(DexIndexIterators, Limit) {
+  const PostingList L0 = {4, 7, 8, 20, 42, 100};
+  const PostingList L1 = {1, 3, 5, 8, 9};
+  const PostingList L2 = {1, 5, 7, 9};
+  const PostingList L3 = {0, 5};
+  const PostingList L4 = {0, 1, 5};
+  const PostingList L5;
+
+  auto DocIterator = create(L0);
+  EXPECT_THAT(consume(*DocIterator, 42), ElementsAre(4, 7, 8, 20, 42, 100));
+
+  DocIterator = create(L0);
+  EXPECT_THAT(consume(*DocIterator), ElementsAre(4, 7, 8, 20, 42, 100));
+
+  DocIterator = create(L0);
+  EXPECT_THAT(consume(*DocIterator, 3), ElementsAre(4, 7, 8));
+
+  DocIterator = create(L0);
+  EXPECT_THAT(consume(*DocIterator, 0), ElementsAre());
+}
+
 testing::Matcher<std::vector<Token>>
 trigramsAre(std::initializer_list<std::string> Trigrams) {
   std::vector<Token> Tokens;
Index: clang-tools-extra/clangd/index/dex/Iterator.h
===================================================================
--- clang-tools-extra/clangd/index/dex/Iterator.h
+++ clang-tools-extra/clangd/index/dex/Iterator.h
@@ -101,9 +101,10 @@
   virtual llvm::raw_ostream &dump(llvm::raw_ostream &OS) const = 0;
 };
 
-/// Exhausts given iterator and returns all processed DocIDs. The result
-/// contains sorted DocumentIDs.
-std::vector<DocID> consume(Iterator &It);
+/// Advances the iterator until it is either exhausted or the number of
+/// requested items is reached. The result contains sorted DocumentIDs.
+std::vector<DocID> consume(Iterator &It,
+                           size_t Limit = std::numeric_limits<size_t>::max());
 
 /// Returns a document iterator over given PostingList.
 std::unique_ptr<Iterator> create(PostingListRef Documents);
Index: clang-tools-extra/clangd/index/dex/Iterator.cpp
===================================================================
--- clang-tools-extra/clangd/index/dex/Iterator.cpp
+++ clang-tools-extra/clangd/index/dex/Iterator.cpp
@@ -218,9 +218,10 @@
 
 } // end namespace
 
-std::vector<DocID> consume(Iterator &It) {
+std::vector<DocID> consume(Iterator &It, size_t Limit) {
   std::vector<DocID> Result;
-  for (; !It.reachedEnd(); It.advance())
+  for (size_t Retrieved = 0; !It.reachedEnd() && Retrieved < Limit;
+       It.advance(), ++Retrieved)
     Result.push_back(It.peek());
   return Result;
 }


Index: clang-tools-extra/unittests/clangd/DexIndexTests.cpp
===================================================================
--- clang-tools-extra/unittests/clangd/DexIndexTests.cpp
+++ clang-tools-extra/unittests/clangd/DexIndexTests.cpp
@@ -240,6 +240,27 @@
             "(& (& [1, 3, 5, 8, 9] [1, 5, 7, 9]) (| [0, 5] [0, 1, 5] []))");
 }
 
+TEST(DexIndexIterators, Limit) {
+  const PostingList L0 = {4, 7, 8, 20, 42, 100};
+  const PostingList L1 = {1, 3, 5, 8, 9};
+  const PostingList L2 = {1, 5, 7, 9};
+  const PostingList L3 = {0, 5};
+  const PostingList L4 = {0, 1, 5};
+  const PostingList L5;
+
+  auto DocIterator = create(L0);
+  EXPECT_THAT(consume(*DocIterator, 42), ElementsAre(4, 7, 8, 20, 42, 100));
+
+  DocIterator = create(L0);
+  EXPECT_THAT(consume(*DocIterator), ElementsAre(4, 7, 8, 20, 42, 100));
+
+  DocIterator = create(L0);
+  EXPECT_THAT(consume(*DocIterator, 3), ElementsAre(4, 7, 8));
+
+  DocIterator = create(L0);
+  EXPECT_THAT(consume(*DocIterator, 0), ElementsAre());
+}
+
 testing::Matcher<std::vector<Token>>
 trigramsAre(std::initializer_list<std::string> Trigrams) {
   std::vector<Token> Tokens;
Index: clang-tools-extra/clangd/index/dex/Iterator.h
===================================================================
--- clang-tools-extra/clangd/index/dex/Iterator.h
+++ clang-tools-extra/clangd/index/dex/Iterator.h
@@ -101,9 +101,10 @@
   virtual llvm::raw_ostream &dump(llvm::raw_ostream &OS) const = 0;
 };
 
-/// Exhausts given iterator and returns all processed DocIDs. The result
-/// contains sorted DocumentIDs.
-std::vector<DocID> consume(Iterator &It);
+/// Advances the iterator until it is either exhausted or the number of
+/// requested items is reached. The result contains sorted DocumentIDs.
+std::vector<DocID> consume(Iterator &It,
+                           size_t Limit = std::numeric_limits<size_t>::max());
 
 /// Returns a document iterator over given PostingList.
 std::unique_ptr<Iterator> create(PostingListRef Documents);
Index: clang-tools-extra/clangd/index/dex/Iterator.cpp
===================================================================
--- clang-tools-extra/clangd/index/dex/Iterator.cpp
+++ clang-tools-extra/clangd/index/dex/Iterator.cpp
@@ -218,9 +218,10 @@
 
 } // end namespace
 
-std::vector<DocID> consume(Iterator &It) {
+std::vector<DocID> consume(Iterator &It, size_t Limit) {
   std::vector<DocID> Result;
-  for (; !It.reachedEnd(); It.advance())
+  for (size_t Retrieved = 0; !It.reachedEnd() && Retrieved < Limit;
+       It.advance(), ++Retrieved)
     Result.push_back(It.peek());
   return Result;
 }
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
http://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to