[clangd] Implement LIMIT iterator

[kbobyrev]

Summary

This patch introduces LIMIT iterator, which is very important for improving the quality of search query. LIMIT iterators can be applied on top of BOOST iterators to prevent populating query request with a huge number of low-quality symbols.

[kbobyrev]

Patch is in the preview mode, I will add documentation and more tests before opening a review.

[kbobyrev]

Add comprehensive tests, improve documentation.

[kbobyrev]

Since it's bundled with the BOOST iterator and doesn't make too much sense without it, I should probable rebase on top of D50970 and add it as the parent revision.

[kbobyrev]

Rebase on top of BOOST iterator patch, update tests.

[kbobyrev]

Use better wording for internal LimitIterator documentation, perform minor code cleanup.

[sammccall]

first few comments, sorry I'll need to come back to this.

clang-tools-extra/clangd/index/dex/Iterator.h

89–90

Update documentation.

Informs the iterator that the current document was consumed, and returns its boost.

The rest of the doc seems a bit to far in the details, it's hard to follow.
Maybe just

If this iterator has any child iterators that contain the document, consume()
should be called on those and their boosts incorporated.
consume() must *not* be called on children that don't contain the current doc.

90–97

There's two possible formulations here:

• the DocID parameter is used by the Iterator to determine whether to actually decrement its limit (do we actually have that doc?)
• the DocID parameter is redundant, as the parent should only call consume() on children that have the doc.

We need to be explicit about which. It seems you want the second (based on offline discussion). In this case we should either assert that the DocID matches peek(), or just drop the DocID parameter and use peek() instead. (If this later means we need to optimize peek, so be it).

169

This comment is accurate, but maybe a bit too concise as what's going on here is a little subtle.
Maybe expand a little like

Returns LIMIT iterator, which yields up to N elements of its child iterator.
Elements only count towards the limit if they are part of the final result set.
Therefore the following iterator (AND (2) (LIMIT (1 2) 1)) yields (2), not ().

[kbobyrev]

Address a round comments from Sam.

[sammccall]

clang-tools-extra/clangd/index/dex/Iterator.cpp

109

this isn't possible, as the item being consumed is the value of peek().
You could assert !reachedEnd() or just delete this.

227

and again here

341

if condition is always true
you could replace the if with an assertion, but the child will assert, so better just to remove it.
replace if with assertion

348

ISTM this should show the immutable state at the start of the query, rather than the mutated state?
Or if you really think it's useful, both.

e.g. (LIMIT 5(3) <child>)

358

I think you can remove the Limit parameter now, replacing it with an outer Limit iterator

362

you can inline this now

clang-tools-extra/unittests/clangd/DexIndexTests.cpp

267

Is this testing the "limit" parameter to consume, or the limit iterator?
It shouldn't test both just because they happen to share a name :-)

(As noted above, I think you can remove the limit parameter)

[kbobyrev]

Address a round of comments & simplify code.

[sammccall]

clang-tools-extra/unittests/clangd/DexIndexTests.cpp

35

remove limit param here too?
Where you're *testing* limit, the tests will be clearer if you do it explicitly.

276

this doesn't seem to test the limit iterator (at least not more than it's tested elsewhere)

clang-tools-extra/clangd/index/dex/Iterator.h

public:
   /// Moves to the first valid DocID which is equal or higher than given ID. If
   /// it doesn't exist, the iterator is exhausted and proceeds to the END.
   ///
   /// Note: reachedEnd() must be false.
   virtual void advanceTo(DocID ID) = 0;
   /// Returns the current element this iterator points to.
   ///
   /// Note: reachedEnd() must be false.
   virtual DocID peek() const = 0;
-  /// Retrieves boosting score. Query tree root should pass Root->peek() to this
-  /// function, the parameter is needed to propagate through the tree. Given ID
-  /// should be compared against BOOST iterator peek()s: some of the iterators
-  /// would not point to the item which was propagated to the top of the query
-  /// tree (e.g. if these iterators are branches of OR iterator) and hence
-  /// shouldn't apply any boosting to the consumed item.
-  virtual float consume(DocID ID) = 0;
+  /// Informs the iterator that the current document was consumed, and returns
+  /// its boost.
+  ///
+  /// Note: If this iterator has any child iterators that contain the document,
+  /// consume() should be called on those and their boosts incorporated.
+  /// consume() must *not* be called on children that don't contain the current
+  /// doc.
+  virtual float consume() = 0;
 
   virtual ~Iterator() {}
 
   /// Prints a convenient human-readable iterator representation by recursively
   /// dumping iterators in the following format:
   ///
   /// (Type Child1 Child2 ...)
   ///

 
 /// Returns BOOST iterator which multiplies the score of each item by given
 /// factor. Boosting can be used as a computationally inexpensive filtering.
 /// Users can return significantly more items using consumeAndBoost() and then
 /// trim Top K using retrieval score.
 std::unique_ptr<Iterator> createBoost(std::unique_ptr<Iterator> Child,
                                       float Factor);
 
+/// Returns LIMIT iterator, which yields up to N elements of its child iterator.
+/// Elements only count towards the limit if they are part of the final result
+/// set. Therefore the following iterator (AND (2) (LIMIT (1 2) 1)) yields (2),
+/// not ().
+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;
   populateChildren(Children, args...);
   return createAnd(move(Children));
 }
 
 /// This allows createOr(create(...), create(...)) syntax.

clang-tools-extra/clangd/index/dex/Iterator.cpp

void advanceTo(DocID ID) override {
     Index = std::lower_bound(Index, std::end(Documents), ID);
   }
 
   DocID peek() const override {
     assert(!reachedEnd() && "DocumentIterator can't call peek() at the end.");
     return *Index;
   }
 
-  float consume(DocID ID) override { return DEFAULT_BOOST_SCORE; }
+  float consume() override { return DEFAULT_BOOST_SCORE; }
 
 private:
   llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override {
     OS << '[';
     auto Separator = "";
     for (auto It = std::begin(Documents); It != std::end(Documents); ++It) {
       OS << Separator;
       if (It == Index)
public:
   void advanceTo(DocID ID) override {
     assert(!reachedEnd() && "AndIterator can't call advanceTo() at the end.");
     Children.front()->advanceTo(ID);
     sync();
   }
 
   DocID peek() const override { return Children.front()->peek(); }
 
-  // If not exhausted and points to the given item, consume() returns the
-  // product of Children->consume(ID). Otherwise, DEFAULT_BOOST_SCORE is
-  // returned.
-  float consume(DocID ID) override {
-    if (reachedEnd() || peek() != ID)
+  float consume() override {
+    if (reachedEnd())
       return DEFAULT_BOOST_SCORE;
     return std::accumulate(
         begin(Children), end(Children), DEFAULT_BOOST_SCORE,
         [&](float Current, const std::unique_ptr<Iterator> &Child) {
-          return Current * Child->consume(ID);
+          return Current * Child->consume();
         });
   }
 
 private:
   llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override {
     OS << "(& ";
     auto Separator = "";
     for (const auto &Child : Children) {
for (const auto &Child : Children)
       if (!Child->reachedEnd())
         Result = std::min(Result, Child->peek());
 
     return Result;
   }
 
   // Returns the maximum boosting score among all Children when iterator is not
   // exhausted and points to the given ID, DEFAULT_BOOST_SCORE otherwise.
-  float consume(DocID ID) override {
-    if (reachedEnd() || peek() != ID)
+  float consume() override {
+    if (reachedEnd())
       return DEFAULT_BOOST_SCORE;
+    const DocID ID = peek();
     return std::accumulate(
         begin(Children), end(Children), DEFAULT_BOOST_SCORE,
         [&](float Current, const std::unique_ptr<Iterator> &Child) {
           return (!Child->reachedEnd() && Child->peek() == ID)
-                     ? std::max(Current, Child->consume(ID))
+                     ? std::max(Current, Child->consume())
                      : Current;
         });
   }
 
 private:
   llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override {
     OS << "(| ";
     auto Separator = "";
void advanceTo(DocID ID) override {
     Index = std::min(ID, Size);
   }
 
   DocID peek() const override {
     assert(!reachedEnd() && "TrueIterator can't call peek() at the end.");
     return Index;
   }
 
-  float consume(DocID) override { return DEFAULT_BOOST_SCORE; }
+  float consume() override { return DEFAULT_BOOST_SCORE; }
 
 private:
   llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override {
     OS << "(TRUE {" << Index << "} out of " << Size << ")";
     return OS;
   }
 
   DocID Index = 0;
public:
   bool reachedEnd() const override { return Child->reachedEnd(); }
 
   void advance() override { Child->advance(); }
 
   void advanceTo(DocID ID) override { Child->advanceTo(ID); }
 
   DocID peek() const override { return Child->peek(); }
 
-  float consume(DocID ID) override { return Child->consume(ID) * Factor; }
+  float consume() override { return Child->consume() * Factor; }
 
 private:
   llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override {
     OS << "(BOOST " << Factor << ' ' << *Child << ')';
     return OS;
   }
 
   std::unique_ptr<Iterator> Child;
   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<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.
+  float consume() override {
+    if (!reachedEnd())
+      --ItemsLeft;
+    return Child->consume();
+  }
+
+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<std::pair<DocID, float>> consume(Iterator &It, size_t Limit) {
   std::vector<std::pair<DocID, float>> Result;
   for (size_t Retrieved = 0; !It.reachedEnd() && Retrieved < Limit;
        It.advance(), ++Retrieved) {
     DocID Document = It.peek();
-    Result.push_back(std::make_pair(Document, It.consume(Document)));
+    Result.push_back(std::make_pair(Document, It.consume()));
   }
   return Result;
 }
 
 std::unique_ptr<Iterator> create(PostingListRef Documents) {
   return llvm::make_unique<DocumentIterator>(Documents);
 }
 
std::unique_ptr<Iterator> createTrue(DocID Size) {
   return llvm::make_unique<TrueIterator>(Size);
 }
 
 std::unique_ptr<Iterator> createBoost(std::unique_ptr<Iterator> Child,
                                       float Factor) {
   return llvm::make_unique<BoostIterator>(move(Child), Factor);
 }
 
+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

clang-tools-extra/unittests/clangd/DexIndexTests.cpp

 namespace clang {
 namespace clangd {
 namespace dex {
 namespace {
 
 std::vector<DocID>
 consumeIDs(Iterator &It, size_t Limit = std::numeric_limits<size_t>::max()) {
   auto IDAndScore = consume(It, Limit);
   std::vector<DocID> IDs(IDAndScore.size());
   for (size_t I = 0; I < IDAndScore.size(); ++I)
     IDs[I] = IDAndScore[I].first;
   return IDs;
 }
 
 TEST(DexIndexIterators, DocumentIterator) {
   const PostingList L = {4, 7, 8, 20, 42, 100};
   auto DocIterator = create(L);
TEST(DexIndexIterators, QueryTree) {
 
   EXPECT_FALSE(Root->reachedEnd());
   EXPECT_EQ(Root->peek(), 1U);
   Root->advanceTo(0);
   // Advance multiple times. Shouldn't do anything.
   Root->advanceTo(1);
   Root->advanceTo(0);
   EXPECT_EQ(Root->peek(), 1U);
-  auto ElementBoost = Root->consume(Root->peek());
+  auto ElementBoost = Root->consume();
   EXPECT_THAT(ElementBoost, 6);
   Root->advance();
   EXPECT_EQ(Root->peek(), 5U);
   Root->advanceTo(5);
   EXPECT_EQ(Root->peek(), 5U);
-  ElementBoost = Root->consume(Root->peek());
+  ElementBoost = Root->consume();
   EXPECT_THAT(ElementBoost, 8);
   Root->advanceTo(9000);
   EXPECT_TRUE(Root->reachedEnd());
 }
 
 TEST(DexIndexIterators, StringRepresentation) {
   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;
 
   EXPECT_EQ(llvm::to_string(*(create(L0))), "[{4}, 7, 8, 20, 42, 100, END]");
 
   auto Nested = createAnd(createAnd(create(L1), create(L2)),
                           createOr(create(L3), create(L4), create(L5)));
 
   EXPECT_EQ(llvm::to_string(*Nested),
             "(& (& [{1}, 3, 5, 8, 9, END] [{1}, 5, 7, 9, END]) (| [0, {5}, "
             "END] [0, {1}, 5, END] [{END}]))");
 }
 
 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) {
   auto TrueIterator = createTrue(0U);
   EXPECT_TRUE(TrueIterator->reachedEnd());
   EXPECT_THAT(consumeIDs(*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(consumeIDs(*AndIterator), ElementsAre(1, 2, 5));
 }
 
 TEST(DexIndexIterators, Boost) {
   auto BoostIterator = createBoost(createTrue(5U), 42U);
   EXPECT_FALSE(BoostIterator->reachedEnd());
-  auto ElementBoost = BoostIterator->consume(BoostIterator->peek());
+  auto ElementBoost = BoostIterator->consume();
   EXPECT_THAT(ElementBoost, 42U);
 
   const PostingList L0 = {2, 4};
   const PostingList L1 = {1, 4};
   auto Root = createOr(createTrue(5U), createBoost(create(L0), 2U),
                        createBoost(create(L1), 3U));
 
-  ElementBoost = Root->consume(Root->peek());
+  ElementBoost = Root->consume();
   EXPECT_THAT(ElementBoost, Iterator::DEFAULT_BOOST_SCORE);
   Root->advance();
   EXPECT_THAT(Root->peek(), 1U);
-  ElementBoost = Root->consume(Root->peek());
+  ElementBoost = Root->consume();
   EXPECT_THAT(ElementBoost, 3);
 
   Root->advance();
   EXPECT_THAT(Root->peek(), 2U);
-  ElementBoost = Root->consume(Root->peek());
+  ElementBoost = Root->consume();
   EXPECT_THAT(ElementBoost, 2);
 
   Root->advanceTo(4);
-  ElementBoost = Root->consume(Root->peek());
+  ElementBoost = Root->consume();
   EXPECT_THAT(ElementBoost, 3);
 }
 
 testing::Matcher<std::vector<Token>>
 trigramsAre(std::initializer_list<std::string> Trigrams) {
   std::vector<Token> Tokens;
   for (const auto &Symbols : Trigrams) {
     Tokens.push_back(Token(Token::Kind::Trigram, Symbols));