kbobyrev updated this revision to Diff 161677.
kbobyrev added a comment.

Add comprehensive tests, improve documentation.


https://reviews.llvm.org/D51029

Files:
  clang-tools-extra/clangd/index/dex/DexIndex.cpp
  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
@@ -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<std::vector<Token>>
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<DocID> consume(Iterator &It,
+std::vector<DocID> consume(std::unique_ptr<Iterator> It,
                            size_t Limit = std::numeric_limits<size_t>::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<Iterator> 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<Iterator> createLimit(std::unique_ptr<Iterator> Child,
+                                      size_t Limit);
+
 /// This allows createAnd(create(...), create(...)) syntax.
 template <typename... Args> std::unique_ptr<Iterator> createAnd(Args... args) {
   std::vector<std::unique_ptr<Iterator>> 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<Iterator> 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<Iterator> Child;
+  size_t ItemsLeft;
+};
+
 } // end namespace
 
-std::vector<DocID> consume(Iterator &It, size_t Limit) {
+std::vector<DocID> consume(std::unique_ptr<Iterator> It, size_t Limit) {
   std::vector<DocID> 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<TrueIterator>(Size);
 }
 
+std::unique_ptr<Iterator> createLimit(std::unique_ptr<Iterator> Child,
+                                      size_t Size) {
+  return llvm::make_unique<LimitIterator>(move(Child), Size);
+}
+
 } // namespace dex
 } // namespace clangd
 } // namespace clang
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<DocID> SymbolDocIDs = consume(*QueryIterator, ItemsToRetrieve);
+    std::vector<DocID> SymbolDocIDs =
+        consume(move(QueryIterator), ItemsToRetrieve);
 
     // Retrieve top Req.MaxCandidateCount items.
     std::priority_queue<std::pair<float, const Symbol *>> Top;
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
http://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to