Index: clang-tools-extra/trunk/clangd/index/dex/Iterator.h =================================================================== --- clang-tools-extra/trunk/clangd/index/dex/Iterator.h +++ clang-tools-extra/trunk/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 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 consume(Iterator &It, + size_t Limit = std::numeric_limits::max()); /// Returns a document iterator over given PostingList. std::unique_ptr create(PostingListRef Documents); Index: clang-tools-extra/trunk/clangd/index/dex/Iterator.cpp =================================================================== --- clang-tools-extra/trunk/clangd/index/dex/Iterator.cpp +++ clang-tools-extra/trunk/clangd/index/dex/Iterator.cpp @@ -218,9 +218,10 @@ } // end namespace -std::vector consume(Iterator &It) { +std::vector consume(Iterator &It, size_t Limit) { std::vector 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/trunk/unittests/clangd/DexIndexTests.cpp =================================================================== --- clang-tools-extra/trunk/unittests/clangd/DexIndexTests.cpp +++ clang-tools-extra/trunk/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> trigramsAre(std::initializer_list Trigrams) { std::vector Tokens;