[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

90 ↗

(On Diff #161952)

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.

96 ↗

(On Diff #161952)

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).

168 ↗

(On Diff #161952)

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 ↗

(On Diff #162183)

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

227 ↗

(On Diff #162183)

and again here

341 ↗

(On Diff #162183)

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 ↗

(On Diff #162183)

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 ↗

(On Diff #162183)

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

362 ↗

(On Diff #162183)

you can inline this now

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

267 ↗

(On Diff #162183)

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

34 ↗

(On Diff #162341)

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

279 ↗

(On Diff #162341)

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

clang-tools-extra/trunk/clangd/index/dex/DexIndex.cpp

auto QueryIterator = TopLevelChildren.empty()
                              ? createTrue(Symbols->size())
                              : createAnd(move(TopLevelChildren));
     // Retrieve more items than it was requested: some of  the items with high
     // final score might not be retrieved otherwise.
     // FIXME(kbobyrev): Pre-scoring retrieval threshold should be adjusted as
     // 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;
+    auto Root = createLimit(move(QueryIterator), ItemsToRetrieve);
     // FIXME(kbobyrev): Add boosting to the query and utilize retrieved
     // boosting scores.
-    std::vector<std::pair<DocID, float>> SymbolDocIDs =
-        consume(*QueryIterator, ItemsToRetrieve);
+    std::vector<std::pair<DocID, float>> SymbolDocIDs = consume(*Root);
 
     // Retrieve top Req.MaxCandidateCount items.
     std::priority_queue<std::pair<float, const Symbol *>> Top;
     for (const auto &P : SymbolDocIDs) {
       const DocID SymbolDocID = P.first;
       const auto *Sym = (*Symbols)[SymbolDocID];
       const llvm::Optional<float> Score = Filter.match(Sym->Name);
       if (!Score)

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

public:
   /// 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 ...)
   ///
   /// Where Type is the iterator type representation: "&" for And, "|" for Or,
   /// ChildN is N-th iterator child. Raw iterators over PostingList are
   /// represented as "[ID1, ID2, ..., {IDN}, ... END]" where IDN is N-th
   /// PostingList entry and the element which is pointed to by the PostingList
   /// iterator is enclosed in {} braces.
   friend llvm::raw_ostream &operator<<(llvm::raw_ostream &OS,
                                        const Iterator &Iterator) {
     return Iterator.dump(OS);
   }
 
   constexpr static float DEFAULT_BOOST_SCORE = 1;
 
 private:
   virtual llvm::raw_ostream &dump(llvm::raw_ostream &OS) const = 0;
 };
 
-/// Advances the iterator until it is either exhausted or the number of
-/// requested items is reached. Returns pairs of document IDs with the
-/// corresponding boosting score.
+/// Advances the iterator until it is exhausted. Returns pairs of document IDs
+/// with the corresponding boosting score.
 ///
 /// Boosting can be seen as a compromise between retrieving too many items and
 /// calculating finals score for each of them (which might be very expensive)
 /// and not retrieving enough items so that items with very high final score
 /// would not be processed. Boosting score is a computationally efficient way
 /// to acquire preliminary scores of requested items.
-std::vector<std::pair<DocID, float>>
-consume(Iterator &It, size_t Limit = std::numeric_limits<size_t>::max());
+std::vector<std::pair<DocID, float>> consume(Iterator &It);
 
 /// Returns a document iterator over given PostingList.
 ///
 /// DocumentIterator returns DEFAULT_BOOST_SCORE for each processed item.
 std::unique_ptr<Iterator> create(PostingListRef Documents);
 
 /// Returns AND Iterator which performs the intersection of the PostingLists of
 /// its children.

 
 /// 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/trunk/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)
-      return DEFAULT_BOOST_SCORE;
+  float consume() override {
+    assert(!reachedEnd() && "AndIterator can't consume() at the end.");
     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)
-      return DEFAULT_BOOST_SCORE;
+  float consume() override {
+    assert(!reachedEnd() &&
+           "OrIterator can't consume() after it reached the end.");
+    const DocID ID = peek();
     return std::accumulate(
         begin(Children), end(Children), DEFAULT_BOOST_SCORE,
-        [&](float Current, const std::unique_ptr<Iterator> &Child) {
+        [&](float Boost, const std::unique_ptr<Iterator> &Child) {
           return (!Child->reachedEnd() && Child->peek() == ID)
-                     ? std::max(Current, Child->consume(ID))
-                     : Current;
+                     ? std::max(Boost, Child->consume())
+                     : Boost;
         });
   }
 
 private:
   llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override {
     OS << "(| ";
     auto Separator = "";
     for (const auto &Child : Children) {
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)), Limit(Limit), 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 {
+    assert(!reachedEnd() && "LimitIterator can't consume at the end.");
+    --ItemsLeft;
+    return Child->consume();
+  }
+
+private:
+  llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override {
+    OS << "(LIMIT " << Limit << '(' << ItemsLeft << ") " << *Child << ')';
+    return OS;
+  }
+
+  std::unique_ptr<Iterator> Child;
+  size_t Limit;
+  size_t ItemsLeft;
+};
+
 } // end namespace
 
-std::vector<std::pair<DocID, float>> consume(Iterator &It, size_t Limit) {
+std::vector<std::pair<DocID, float>> consume(Iterator &It) {
   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)));
-  }
+  for (; !It.reachedEnd(); It.advance())
+    Result.push_back(std::make_pair(It.peek(), It.consume()));
   return Result;
 }
 
 std::unique_ptr<Iterator> create(PostingListRef Documents) {
   return llvm::make_unique<DocumentIterator>(Documents);
 }
 
 std::unique_ptr<Iterator>
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/trunk/unittests/clangd/DexIndexTests.cpp

 using ::testing::ElementsAre;
 using ::testing::UnorderedElementsAre;
 
 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> consumeIDs(Iterator &It) {
+  auto IDAndScore = consume(It);
   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};
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;
-
-  auto DocIterator = create(L0);
-  EXPECT_THAT(consumeIDs(*DocIterator, 42), ElementsAre(4, 7, 8, 20, 42, 100));
-
-  DocIterator = create(L0);
-  EXPECT_THAT(consumeIDs(*DocIterator), ElementsAre(4, 7, 8, 20, 42, 100));
-
-  DocIterator = create(L0);
-  EXPECT_THAT(consumeIDs(*DocIterator, 3), ElementsAre(4, 7, 8));
-
-  DocIterator = create(L0);
-  EXPECT_THAT(consumeIDs(*DocIterator, 0), ElementsAre());
+  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 = createLimit(create(L0), 42);
+  EXPECT_THAT(consumeIDs(*DocIterator), ElementsAre(3, 6, 7, 20, 42, 100));
+
+  DocIterator = createLimit(create(L0), 3);
+  EXPECT_THAT(consumeIDs(*DocIterator), ElementsAre(3, 6, 7));
+
+  DocIterator = createLimit(create(L0), 0);
+  EXPECT_THAT(consumeIDs(*DocIterator), 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));