Index: clang-tools-extra/clangd/index/dex/DexIndex.cpp =================================================================== --- clang-tools-extra/clangd/index/dex/DexIndex.cpp +++ clang-tools-extra/clangd/index/dex/DexIndex.cpp @@ -119,7 +119,8 @@ // using 100x of the requested number might not be good in practice, e.g. // when the requested number of items is small. const unsigned ItemsToRetrieve = 100 * Req.MaxCandidateCount; - std::vector SymbolDocIDs = consume(*QueryIterator, ItemsToRetrieve); + std::vector SymbolDocIDs = + consume(move(QueryIterator), ItemsToRetrieve); // Retrieve top Req.MaxCandidateCount items. std::priority_queue> Top; 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 @@ -89,6 +89,13 @@ /// /// Note: reachedEnd() must be false. virtual DocID peek() const = 0; + /// Marks an element as "consumed" and triggers limit decrement for every + /// LIMIT iterator which yields given item. Root iterator should call + /// consume(peek) as the argument is used to propagate the tree, otherwise + /// the behavior will be incorrect. + virtual void consume(DocID ID) = 0; + /// Returns true if the iterator can retrieve more items upon request. + virtual bool canRetrieveMore() const { return !reachedEnd(); } virtual ~Iterator() {} @@ -113,7 +120,7 @@ /// 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, +std::vector consume(std::unique_ptr It, size_t Limit = std::numeric_limits::max()); /// Returns a document iterator over given PostingList. @@ -133,6 +140,11 @@ /// all items in range [0, Size) in an efficient manner. std::unique_ptr createTrue(DocID Size); +/// 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 @@ -46,6 +46,8 @@ return *Index; } + void consume(DocID ID) override {} + private: llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override { OS << '['; @@ -103,6 +105,13 @@ DocID peek() const override { return Children.front()->peek(); } + void consume(DocID ID) override { + if (peek() != ID) + return; + for (const auto &Child : Children) + Child->consume(ID); + } + private: llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override { OS << "(& "; @@ -209,6 +218,13 @@ return Result; } + void consume(DocID ID) override { + if (peek() != ID) + return; + for (const auto &Child : Children) + Child->consume(ID); + } + private: llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override { OS << "(| "; @@ -249,6 +265,8 @@ return Index; } + void consume(DocID ID) override {} + private: llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override { OS << "(TRUE {" << Index << "} out of " << Size << ")"; @@ -260,13 +278,54 @@ DocID Size; }; +/// This iterator is a wrapper around the underlying child which limits the +/// number of elements which are retrieved from the whole iterator 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. + void consume(DocID ID) override { + if (!reachedEnd() && peek() == ID) + --ItemsLeft; + } + + /// Unlike most iterators, LIMIT iterator can retrieve more items after being + /// exhausted. Hence, it only matters whether the child is exhausted or not. + bool canRetrieveMore() const override { return !Child->reachedEnd(); } + +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) { +std::vector consume(std::unique_ptr It, size_t Limit) { std::vector Result; - for (size_t Retrieved = 0; !It.reachedEnd() && Retrieved < Limit; - It.advance(), ++Retrieved) - Result.push_back(It.peek()); + auto Root = createLimit(std::move(It), Limit); + for (; !Root->reachedEnd(); Root->advance()) { + const DocID Item = Root->peek(); + Root->consume(Item); + Result.push_back(Item); + } return Result; } @@ -288,6 +347,11 @@ return llvm::make_unique(Size); } +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 @@ -62,7 +62,7 @@ auto AndWithEmpty = createAnd(create(L0), create(L1)); EXPECT_TRUE(AndWithEmpty->reachedEnd()); - EXPECT_THAT(consume(*AndWithEmpty), ElementsAre()); + EXPECT_THAT(consume(move(AndWithEmpty)), ElementsAre()); } TEST(DexIndexIterators, AndTwoLists) { @@ -72,7 +72,7 @@ auto And = createAnd(create(L1), create(L0)); EXPECT_FALSE(And->reachedEnd()); - EXPECT_THAT(consume(*And), ElementsAre(0U, 7U, 10U, 320U, 9000U)); + EXPECT_THAT(consume(move(And)), ElementsAre(0U, 7U, 10U, 320U, 9000U)); And = createAnd(create(L0), create(L1)); @@ -113,7 +113,7 @@ auto OrWithEmpty = createOr(create(L0), create(L1)); EXPECT_FALSE(OrWithEmpty->reachedEnd()); - EXPECT_THAT(consume(*OrWithEmpty), + EXPECT_THAT(consume(move(OrWithEmpty)), ElementsAre(0U, 5U, 7U, 10U, 42U, 320U, 9000U)); } @@ -146,7 +146,7 @@ Or = createOr(create(L0), create(L1)); - EXPECT_THAT(consume(*Or), + EXPECT_THAT(consume(move(Or)), ElementsAre(0U, 4U, 5U, 7U, 10U, 30U, 42U, 60U, 320U, 9000U)); } @@ -248,37 +248,40 @@ } 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(consume(*DocIterator, 42), ElementsAre(4, 7, 8, 20, 42, 100)); + EXPECT_THAT(consume(move(DocIterator), 42), + ElementsAre(3, 6, 7, 20, 42, 100)); DocIterator = create(L0); - EXPECT_THAT(consume(*DocIterator), ElementsAre(4, 7, 8, 20, 42, 100)); + EXPECT_THAT(consume(move(DocIterator)), ElementsAre(3, 6, 7, 20, 42, 100)); DocIterator = create(L0); - EXPECT_THAT(consume(*DocIterator, 3), ElementsAre(4, 7, 8)); + EXPECT_THAT(consume(move(DocIterator), 3), ElementsAre(3, 6, 7)); DocIterator = create(L0); - EXPECT_THAT(consume(*DocIterator, 0), ElementsAre()); + EXPECT_THAT(consume(move(DocIterator), 0), ElementsAre()); + + auto AndIterator = + createAnd(createLimit(createTrue(9000), 343), createLimit(create(L0), 2), + createLimit(create(L1), 3), createLimit(create(L2), 42)); + EXPECT_THAT(consume(move(AndIterator)), ElementsAre(3, 7)); } TEST(DexIndexIterators, True) { auto TrueIterator = createTrue(0U); EXPECT_TRUE(TrueIterator->reachedEnd()); - EXPECT_THAT(consume(*TrueIterator), ElementsAre()); + EXPECT_THAT(consume(move(TrueIterator)), ElementsAre()); PostingList L0 = {1, 2, 5, 7}; TrueIterator = createTrue(7U); EXPECT_THAT(TrueIterator->peek(), 0); auto AndIterator = createAnd(create(L0), move(TrueIterator)); EXPECT_FALSE(AndIterator->reachedEnd()); - EXPECT_THAT(consume(*AndIterator), ElementsAre(1, 2, 5)); + EXPECT_THAT(consume(move(AndIterator)), ElementsAre(1, 2, 5)); } testing::Matcher>