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 @@ -165,6 +165,11 @@ std::unique_ptr createBoost(std::unique_ptr Child, float Factor); +/// Returns LIMIT iterator, which is exhausted as soon requested number of items +/// is consumed from the root of query tree. +std::unique_ptr createLimit(std::unique_ptr Child, + size_t Limit); + /// This allows createAnd(create(...), create(...)) syntax. template std::unique_ptr createAnd(Args... args) { std::vector> Children; 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 @@ -318,6 +318,43 @@ float Factor; }; +/// This iterator limits the number of items retrieved from the child iterator +/// on top of the query tree. To ensure that query tree with LIMIT iterators +/// inside works correctly, users have to call Root->consume(Root->peek()) each +/// time item is retrieved at the root of query tree. +class LimitIterator : public Iterator { +public: + LimitIterator(std::unique_ptr Child, size_t Limit) + : Child(move(Child)), ItemsLeft(Limit) {} + + bool reachedEnd() const override { + return ItemsLeft == 0 || Child->reachedEnd(); + } + + void advance() override { Child->advance(); } + + void advanceTo(DocID ID) override { Child->advanceTo(ID); } + + DocID peek() const override { return Child->peek(); } + + /// Decreases the limit in case the element consumed at top of the query tree + /// comes from the underlying iterator. + float consume(DocID ID) override { + if (!reachedEnd() && peek() == ID) + --ItemsLeft; + return Child->consume(ID); + } + +private: + llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override { + OS << "(LIMIT items left " << ItemsLeft << *Child << ")"; + return OS; + } + + std::unique_ptr Child; + size_t ItemsLeft; +}; + } // end namespace std::vector> consume(Iterator &It, size_t Limit) { @@ -353,6 +390,11 @@ return llvm::make_unique(move(Child), Factor); } +std::unique_ptr createLimit(std::unique_ptr Child, + size_t Size) { + return llvm::make_unique(move(Child), Size); +} + } // namespace dex } // namespace clangd } // namespace clang Index: clang-tools-extra/unittests/clangd/DexIndexTests.cpp =================================================================== --- clang-tools-extra/unittests/clangd/DexIndexTests.cpp +++ clang-tools-extra/unittests/clangd/DexIndexTests.cpp @@ -265,24 +265,27 @@ } 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; + const PostingList L0 = {3, 6, 7, 20, 42, 100}; + const PostingList L1 = {1, 3, 5, 6, 7, 30, 100}; + const PostingList L2 = {0, 3, 5, 7, 8, 100}; auto DocIterator = create(L0); - EXPECT_THAT(consumeIDs(*DocIterator, 42), ElementsAre(4, 7, 8, 20, 42, 100)); + EXPECT_THAT(consumeIDs(*DocIterator, 42), + ElementsAre(3, 6, 7, 20, 42, 100)); DocIterator = create(L0); - EXPECT_THAT(consumeIDs(*DocIterator), ElementsAre(4, 7, 8, 20, 42, 100)); + EXPECT_THAT(consumeIDs(*DocIterator), ElementsAre(3, 6, 7, 20, 42, 100)); DocIterator = create(L0); - EXPECT_THAT(consumeIDs(*DocIterator, 3), ElementsAre(4, 7, 8)); + EXPECT_THAT(consumeIDs(*DocIterator, 3), ElementsAre(3, 6, 7)); DocIterator = create(L0); EXPECT_THAT(consumeIDs(*DocIterator, 0), ElementsAre()); + + auto AndIterator = + createAnd(createLimit(createTrue(9000), 343), createLimit(create(L0), 2), + createLimit(create(L1), 3), createLimit(create(L2), 42)); + EXPECT_THAT(consumeIDs(*AndIterator), ElementsAre(3, 7)); } TEST(DexIndexIterators, True) {