[clangd] Implement BOOST iterator

[kbobyrev]

Summary

This patch introduces BOOST iterator - a substantial block for efficient and high-quality symbol retrieval. The concept of boosting allows performing computationally inexpensive scoring on the query side so that the final (expensive) scoring can only be applied on the items with the highest preliminary score while eliminating the need to score too many items.

[kbobyrev]

This patch is in preview mode, documentation overhaul and minor cleanup incoming.

[kbobyrev]

Add documentation, cleanup tests.

[ilya-biryukov]

Looking at this one. (to make sure we don't clash with @ioeric)

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

111 ↗

(On Diff #161477)

Maybe make it constexpr and put the value into the header?

[ilya-biryukov]

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

92 ↗

(On Diff #161477)

Maybe add a note to the comment on why an ID parameter actually is here?
IIUC, we need to because various iterators in the tree may point to different elements and we need to know which one we've actually matched.

129 ↗

(On Diff #161477)

Could we describe the rationale for keeping both consume and consumeAndBoost somewhere in the comments?

From the offline conversation, it seems consumeAndBoost is more expensive, but our clients will all use it at some point in the future.
The idea of paying for boosting without actually using it seems bad, so keeping this function separate makes sense.

164 ↗

(On Diff #161477)

Maybe use float scores to align with the scoring code we have for completion?

[kbobyrev]

• Add more comments explaining the difference between consume() and consumeAndBoost() and their potential usecases for the clients
• Move DEFAULT_BOOSTING_SCORE to the header by making it constexpr and change its type to float

[ilya-biryukov]

Summarizing the offline discussion with @kbobyrev, @ioeric and @sammccall in two comments.

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

96 ↗

(On Diff #161909)

With limit iterator in mind, let's rename this to consume() and make it non-const.

137 ↗

(On Diff #161909)

Let's remove this function and change the interface of consume to return a vector of pairs instead.

[kbobyrev]

Address remaining comments.

[ilya-biryukov]

LGTM

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;
-    std::vector<DocID> SymbolDocIDs = consume(*QueryIterator, ItemsToRetrieve);
+    // FIXME(kbobyrev): Add boosting to the query and utilize retrieved
+    // boosting scores.
+    std::vector<std::pair<DocID, float>> SymbolDocIDs =
+        consume(*QueryIterator, ItemsToRetrieve);
 
     // Retrieve top Req.MaxCandidateCount items.
     std::priority_queue<std::pair<float, const Symbol *>> Top;
-    for (const auto &SymbolDocID : SymbolDocIDs) {
+    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)
         continue;
       // Multiply score by a negative factor so that Top stores items with the
       // highest actual score.
       Top.emplace(-(*Score) * SymbolQuality.find(Sym)->second, Sym);
       if (Top.size() > Req.MaxCandidateCount) {
void DexIndex::lookup(const LookupRequest &Req,
   std::lock_guard<std::mutex> Lock(Mutex);
   for (const auto &ID : Req.IDs) {
     auto I = LookupTable.find(ID);
     if (I != LookupTable.end())
       Callback(*I->second);
   }
 }
 
-
 void DexIndex::findOccurrences(
     const OccurrencesRequest &Req,
     llvm::function_ref<void(const SymbolOccurrence &)> Callback) const {
   log("findOccurrences is not implemented.");
 }
 
 } // namespace dex
 } // namespace clangd
 } // namespace clang

clangd/index/dex/Iterator.h

 
 /// Iterator is the interface for Query Tree node. The simplest type of Iterator
 /// is DocumentIterator which is simply a wrapper around PostingList iterator
 /// and serves as the Query Tree leaf. More sophisticated examples of iterators
 /// can manage intersection, union of the elements produced by other iterators
 /// (their children) to form a multi-level Query Tree. The interface is designed
 /// to be extensible in order to support multiple types of iterators.
 class Iterator {
-  // FIXME(kbobyrev): Provide callback for matched documents.
-  // FIXME(kbobyrev): Implement new types of iterators: Label, Boost (with
-  // scoring), Limit.
   // FIXME(kbobyrev): Implement iterator cost, an estimate of advance() calls
   // before iterator exhaustion.
+  // FIXME(kbobyrev): Implement Limit iterator.
 public:
   /// Returns true if all valid DocIDs were processed and hence the iterator is
   /// exhausted.
   virtual bool reachedEnd() const = 0;
   /// Moves to next valid DocID. If it doesn't exist, the iterator is exhausted
   /// and proceeds to the END.
   ///
   /// Note: reachedEnd() must be false.
   virtual void advance() = 0;
   /// 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;
 
   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. The result contains sorted DocumentIDs.
-std::vector<DocID> consume(Iterator &It,
-                           size_t Limit = std::numeric_limits<size_t>::max());
+/// requested items is reached. 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());
 
 /// 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.
+///
+/// consume(): AND Iterator returns the product of Childrens' boosting scores
+/// when not exhausted and DEFAULT_BOOST_SCORE otherwise.
 std::unique_ptr<Iterator>
 createAnd(std::vector<std::unique_ptr<Iterator>> Children);
 
 /// Returns OR Iterator which performs the union of the PostingLists of its
 /// children.
+///
+/// consume(): OR Iterator returns the highest boost value among children
+/// pointing to requested item when not exhausted and DEFAULT_BOOST_SCORE
+/// otherwise.
 std::unique_ptr<Iterator>
 createOr(std::vector<std::unique_ptr<Iterator>> Children);
 
 /// Returns TRUE Iterator which iterates over "virtual" PostingList containing
 /// all items in range [0, Size) in an efficient manner.
+///
+/// TRUE returns DEFAULT_BOOST_SCORE for each processed item.
 std::unique_ptr<Iterator> createTrue(DocID Size);
 
+/// 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);
+
 /// 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.

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; }
+
 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)
         OS << '{' << *It << '}';
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;
+    return std::accumulate(
+        begin(Children), end(Children), DEFAULT_BOOST_SCORE,
+        [&](float Current, const std::unique_ptr<Iterator> &Child) {
+          return Current * Child->consume(ID);
+        });
+  }
+
 private:
   llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override {
     OS << "(& ";
     auto Separator = "";
     for (const auto &Child : Children) {
       OS << Separator << *Child;
       Separator = " ";
     }
DocID peek() const override {
 
     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;
+    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))
+                     : Current;
+        });
+  }
+
 private:
   llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override {
     OS << "(| ";
     auto Separator = "";
     for (const auto &Child : Children) {
       OS << Separator << *Child;
       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; }
+
 private:
   llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override {
     OS << "(TRUE {" << Index << "} out of " << Size << ")";
     return OS;
   }
 
   DocID Index = 0;
   /// Size of the underlying virtual PostingList.
   DocID Size;
 };
 
+/// Boost iterator is a wrapper around its child which multiplies scores of
+/// each retrieved item by a given factor.
+class BoostIterator : public Iterator {
+public:
+  BoostIterator(std::unique_ptr<Iterator> Child, float Factor)
+      : Child(move(Child)), Factor(Factor) {}
+
+  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; }
+
+private:
+  llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override {
+    OS << "(BOOST " << Factor << ' ' << *Child << ')';
+    return OS;
+  }
+
+  std::unique_ptr<Iterator> Child;
+  float Factor;
+};
+
 } // end namespace
 
-std::vector<DocID> consume(Iterator &It, size_t Limit) {
-  std::vector<DocID> Result;
+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)
-    Result.push_back(It.peek());
+       It.advance(), ++Retrieved) {
+    DocID Document = It.peek();
+    Result.push_back(std::make_pair(Document, It.consume(Document)));
+  }
   return Result;
 }
 
 std::unique_ptr<Iterator> create(PostingListRef Documents) {
   return llvm::make_unique<DocumentIterator>(Documents);
 }
 
 std::unique_ptr<Iterator>
 createAnd(std::vector<std::unique_ptr<Iterator>> Children) {
   return llvm::make_unique<AndIterator>(move(Children));
 }
 
 std::unique_ptr<Iterator>
 createOr(std::vector<std::unique_ptr<Iterator>> Children) {
   return llvm::make_unique<OrIterator>(move(Children));
 }
 
 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);
+}
+
 } // namespace dex
 } // namespace clangd
 } // namespace clang

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> 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);
 
   EXPECT_EQ(DocIterator->peek(), 4U);
   EXPECT_FALSE(DocIterator->reachedEnd());
 
   DocIterator->advance();
TEST(DexIndexIterators, AndWithEmpty) {
   const PostingList L1 = {0, 5, 7, 10, 42, 320, 9000};
 
   auto AndEmpty = createAnd(create(L0));
   EXPECT_TRUE(AndEmpty->reachedEnd());
 
   auto AndWithEmpty = createAnd(create(L0), create(L1));
   EXPECT_TRUE(AndWithEmpty->reachedEnd());
 
-  EXPECT_THAT(consume(*AndWithEmpty), ElementsAre());
+  EXPECT_THAT(consumeIDs(*AndWithEmpty), ElementsAre());
 }
 
 TEST(DexIndexIterators, AndTwoLists) {
   const PostingList L0 = {0, 5, 7, 10, 42, 320, 9000};
   const PostingList L1 = {0, 4, 7, 10, 30, 60, 320, 9000};
 
   auto And = createAnd(create(L1), create(L0));
 
   EXPECT_FALSE(And->reachedEnd());
-  EXPECT_THAT(consume(*And), ElementsAre(0U, 7U, 10U, 320U, 9000U));
+  EXPECT_THAT(consumeIDs(*And), ElementsAre(0U, 7U, 10U, 320U, 9000U));
 
   And = createAnd(create(L0), create(L1));
 
   And->advanceTo(0);
   EXPECT_EQ(And->peek(), 0U);
   And->advanceTo(5);
   EXPECT_EQ(And->peek(), 7U);
   And->advanceTo(10);
TEST(DexIndexIterators, OrWithEmpty) {
   const PostingList L1 = {0, 5, 7, 10, 42, 320, 9000};
 
   auto OrEmpty = createOr(create(L0));
   EXPECT_TRUE(OrEmpty->reachedEnd());
 
   auto OrWithEmpty = createOr(create(L0), create(L1));
   EXPECT_FALSE(OrWithEmpty->reachedEnd());
 
-  EXPECT_THAT(consume(*OrWithEmpty),
+  EXPECT_THAT(consumeIDs(*OrWithEmpty),
               ElementsAre(0U, 5U, 7U, 10U, 42U, 320U, 9000U));
 }
 
 TEST(DexIndexIterators, OrTwoLists) {
   const PostingList L0 = {0, 5, 7, 10, 42, 320, 9000};
   const PostingList L1 = {0, 4, 7, 10, 30, 60, 320, 9000};
 
   auto Or = createOr(create(L0), create(L1));
TEST(DexIndexIterators, OrTwoLists) {
   EXPECT_EQ(Or->peek(), 320U);
   Or->advanceTo(9000);
   EXPECT_EQ(Or->peek(), 9000U);
   Or->advanceTo(9001);
   EXPECT_TRUE(Or->reachedEnd());
 
   Or = createOr(create(L0), create(L1));
 
-  EXPECT_THAT(consume(*Or),
+  EXPECT_THAT(consumeIDs(*Or),
               ElementsAre(0U, 4U, 5U, 7U, 10U, 30U, 42U, 60U, 320U, 9000U));
 }
 
 TEST(DexIndexIterators, OrThreeLists) {
   const PostingList L0 = {0, 5, 7, 10, 42, 320, 9000};
   const PostingList L1 = {0, 4, 7, 10, 30, 60, 320, 9000};
   const PostingList L2 = {1, 4, 7, 11, 30, 60, 320, 9000};
 
TEST(DexIndexIterators, OrThreeLists) {
 
   Or->advanceTo(9001);
   EXPECT_TRUE(Or->reachedEnd());
 }
 
 // FIXME(kbobyrev): The testcase below is similar to what is expected in real
 // queries. It should be updated once new iterators (such as boosting, limiting,
 // etc iterators) appear. However, it is not exhaustive and it would be
-// beneficial to implement automatic generation of query trees for more
-// comprehensive testing.
+// beneficial to implement automatic generation (e.g. fuzzing) of query trees
+// for more comprehensive testing.
 TEST(DexIndexIterators, QueryTree) {
-  // An example of more complicated query
   //
   //                      +-----------------+
   //                      |And Iterator:1, 5|
   //                      +--------+--------+
   //                               |
   //                               |
-  //                 +------------------------------------+
+  //                 +-------------+----------------------+
   //                 |                                    |
   //                 |                                    |
-  //      +----------v----------+              +----------v---------+
-  //      |And Iterator: 1, 5, 9|              |Or Iterator: 0, 1, 5|
-  //      +----------+----------+              +----------+---------+
+  //      +----------v----------+              +----------v------------+
+  //      |And Iterator: 1, 5, 9|              |Or Iterator: 0, 1, 3, 5|
+  //      +----------+----------+              +----------+------------+
   //                 |                                    |
-  //          +------+-----+                    +---------+-----------+
+  //          +------+-----+                    +---------------------+
   //          |            |                    |         |           |
-  //  +-------v-----+ +----v-----+           +--v--+    +-V--+    +---v---+
-  //  |1, 3, 5, 8, 9| |1, 5, 7, 9|           |Empty|    |0, 5|    |0, 1, 5|
-  //  +-------------+ +----------+           +-----+    +----+    +-------+
-
+  //  +-------v-----+ +----+---+             +--v--+  +---v----+ +----v---+
+  //  |1, 3, 5, 8, 9| |Boost: 2|             |Empty|  |Boost: 3| |Boost: 4|
+  //  +-------------+ +----+---+             +-----+  +---+----+ +----+---+
+  //                       |                              |           |
+  //                  +----v-----+                      +-v--+    +---v---+
+  //                  |1, 5, 7, 9|                      |1, 5|    |0, 3, 5|
+  //                  +----------+                      +----+    +-------+
+  //
   const PostingList L0 = {1, 3, 5, 8, 9};
   const PostingList L1 = {1, 5, 7, 9};
-  const PostingList L2 = {0, 5};
-  const PostingList L3 = {0, 1, 5};
-  const PostingList L4;
+  const PostingList L3;
+  const PostingList L4 = {1, 5};
+  const PostingList L5 = {0, 3, 5};
 
   // Root of the query tree: [1, 5]
   auto Root = createAnd(
       // Lower And Iterator: [1, 5, 9]
-      createAnd(create(L0), create(L1)),
+      createAnd(create(L0), createBoost(create(L1), 2U)),
       // Lower Or Iterator: [0, 1, 5]
-      createOr(create(L2), create(L3), create(L4)));
+      createOr(create(L3), createBoost(create(L4), 3U),
+               createBoost(create(L5), 4U)));
 
   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());
+  EXPECT_THAT(ElementBoost, 6);
   Root->advance();
   EXPECT_EQ(Root->peek(), 5U);
   Root->advanceTo(5);
   EXPECT_EQ(Root->peek(), 5U);
+  ElementBoost = Root->consume(Root->peek());
+  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};
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(consume(*DocIterator, 42), ElementsAre(4, 7, 8, 20, 42, 100));
+  EXPECT_THAT(consumeIDs(*DocIterator, 42), ElementsAre(4, 7, 8, 20, 42, 100));
 
   DocIterator = create(L0);
-  EXPECT_THAT(consume(*DocIterator), ElementsAre(4, 7, 8, 20, 42, 100));
+  EXPECT_THAT(consumeIDs(*DocIterator), ElementsAre(4, 7, 8, 20, 42, 100));
 
   DocIterator = create(L0);
-  EXPECT_THAT(consume(*DocIterator, 3), ElementsAre(4, 7, 8));
+  EXPECT_THAT(consumeIDs(*DocIterator, 3), ElementsAre(4, 7, 8));
 
   DocIterator = create(L0);
-  EXPECT_THAT(consume(*DocIterator, 0), ElementsAre());
+  EXPECT_THAT(consumeIDs(*DocIterator, 0), ElementsAre());
 }
 
 TEST(DexIndexIterators, True) {
   auto TrueIterator = createTrue(0U);
   EXPECT_TRUE(TrueIterator->reachedEnd());
-  EXPECT_THAT(consume(*TrueIterator), ElementsAre());
+  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(consume(*AndIterator), ElementsAre(1, 2, 5));
+  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());
+  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());
+  EXPECT_THAT(ElementBoost, Iterator::DEFAULT_BOOST_SCORE);
+  Root->advance();
+  EXPECT_THAT(Root->peek(), 1U);
+  ElementBoost = Root->consume(Root->peek());
+  EXPECT_THAT(ElementBoost, 3);
+
+  Root->advance();
+  EXPECT_THAT(Root->peek(), 2U);
+  ElementBoost = Root->consume(Root->peek());
+  EXPECT_THAT(ElementBoost, 2);
+
+  Root->advanceTo(4);
+  ElementBoost = Root->consume(Root->peek());
+  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));
   }