[clangd] Implement iterator cost

[kbobyrev]

Summary

This patch introduces iterator cost concept to improve the performance of Dex query iterators (mainly, AND iterator). Benchmarks show that the queries become ~10% faster.

Before

-------------------------------------------------------
Benchmark                Time           CPU Iteration
-------------------------------------------------------
DexAdHocQueries    5883074 ns    5883018 ns        117
DexRealQ         959904457 ns  959898507 ns          1

After

-------------------------------------------------------
Benchmark                Time           CPU Iteration
-------------------------------------------------------
DexAdHocQueries    5238403 ns    5238361 ns        130
DexRealQ         873275207 ns  873269453 ns          1

[kbobyrev]

It's probably better to roll out proximity path boosting & actual two-stage filtering before rolling this out.

[sammccall]

In D51310#1214548, @kbobyrev wrote:

It's probably better to roll out proximity path boosting & actual two-stage filtering before rolling this out.

Up to you (I don't understand the interaction), but this looks good to go as-is. I'd expect a 10% speedup to be pretty robust to the sorts of changes you're talking about.

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

90 ↗

(On Diff #162686)

this is a performance heuristic and deserves a brief comment (about *why* we're best to have the shortest first). Not too many details, handwavy is fine :-)

93 ↗

(On Diff #162686)

(I'd actually suggest declaring operator< in iterator.h and making it compare by cost, it seems natural/important enough, and saves duplicating this lambda)

125 ↗

(On Diff #162686)

The use of std::accumulate here seems less clear than the direct:

size_t Smallest = ...;
for (const auto& Child : Children)
  Smallest = std::min(Smallest, Child->cost());
return Smallest;

For better or worse, functional styles don't have the most natural syntax in C++, and (IMO) shouldn't be used just because they fit.

(This is consistent with other places, but I also think that the same applies there)

125 ↗

(On Diff #162686)

actually, we've sorted by cost, so don't we just want Children.front().cost() here?

258 ↗

(On Diff #162686)

as above, we've sorted, so just Children.back()->cost()?

377 ↗

(On Diff #162686)

min() this with Child->cost()?

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

99 ↗

(On Diff #162686)

const?

99 ↗

(On Diff #162686)

I know this terminology comes from elsewhere, but I struggle with "cost" as a metaphor because I expect it to aggregate its children as a sum (at least in cases like "or" when all children will be advanced until end).

Maybe estimateSize()?

[kbobyrev]

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

93 ↗

(On Diff #162686)

Actually, I think sorting the children of OR iterator doesn't improve performance so I should probably just have a lambda in one place (AND) as it won't be used anywhere else.

125 ↗

(On Diff #162686)

It's a sad day for the noble λ knights, but I agree.

I think I should push it like this in this patch (to comply with the common style) and then send out a separate patch with refactorings (I had few concerns about the inconsistency in the Iterator.cpp).

clangd/index/dex/Iterator.h

public:
   /// 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;
+  /// Returns an estimate of advance() calls before the iterator is exhausted.
+  virtual size_t estimateSize() const = 0;
 
   virtual ~Iterator() {}
 
   /// Prints a convenient human-readable iterator representation by recursively
   /// dumping iterators in the following format:
   ///
   /// (Type Child1 Child2 ...)
   ///

clangd/index/dex/Iterator.cpp

public:
 
   DocID peek() const override {
     assert(!reachedEnd() && "DocumentIterator can't call peek() at the end.");
     return *Index;
   }
 
   float consume() override { return DEFAULT_BOOST_SCORE; }
 
+  size_t estimateSize() const override { return Documents.size(); }
+
 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 << '}';

 /// iterator restores the invariant: all children must point to the same item.
 class AndIterator : public Iterator {
 public:
   AndIterator(std::vector<std::unique_ptr<Iterator>> AllChildren)
       : Children(std::move(AllChildren)) {
     assert(!Children.empty() && "AndIterator should have at least one child.");
     // Establish invariants.
     sync();
+    // When children are sorted by the estimateSize(), sync() calls are more
+    // effective. Each sync() starts with the first child and makes sure all
+    // children point to the same element. If any child is "above" the previous
+    // ones, the algorithm resets and and advances the children to the next
+    // highest element starting from the front. When child iterators in the
+    // beginning have smaller estimated size, the sync() will have less restarts
+    // and become more effective.
+    std::sort(begin(Children), end(Children),
+              [](const std::unique_ptr<Iterator> &LHS,
+                 const std::unique_ptr<Iterator> &RHS) {
+                return LHS->estimateSize() < RHS->estimateSize();
+              });
   }
 
   bool reachedEnd() const override { return ReachedEnd; }
 
   /// Advances all children to the next common item.
   void advance() override {
     assert(!reachedEnd() && "AndIterator can't call advance() at the end.");
     Children.front()->advance();
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();
         });
   }
 
+  size_t estimateSize() const override {
+    return Children.front()->estimateSize();
+  }
+
 private:
   llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override {
     OS << "(& ";
     auto Separator = "";
     for (const auto &Child : Children) {
       OS << Separator << *Child;
       Separator = " ";
     }
do {
         Child->advanceTo(SyncID);
         ReachedEnd |= Child->reachedEnd();
         // If any child reaches end And iterator can not match any other items.
         // In this case, just terminate the process.
         if (ReachedEnd)
           return;
         // If any child goes beyond given ID (i.e. ID is not the common item),
         // all children should be advanced to the next common item.
-        // FIXME(kbobyrev): This is not a very optimized version; after costs
-        // are introduced, cycle should break whenever ID exceeds current one
-        // and cheapest children should be advanced over again.
         if (Child->peek() > SyncID) {
           SyncID = Child->peek();
           NeedsAdvance = true;
         }
       }
     } while (NeedsAdvance);
   }
 

 /// one child. To preserve the sorted order, this iterator always advances the
 /// child with smallest Child->peek() value. OrIterator becomes exhausted as
 /// soon as all of its children are exhausted.
 class OrIterator : public Iterator {
 public:
   OrIterator(std::vector<std::unique_ptr<Iterator>> AllChildren)
       : Children(std::move(AllChildren)) {
     assert(Children.size() > 0 && "Or Iterator must have at least one child.");
+    std::sort(begin(Children), end(Children));
   }
 
   /// Returns true if all children are exhausted.
   bool reachedEnd() const override {
     return std::all_of(begin(Children), end(Children),
                        [](const std::unique_ptr<Iterator> &Child) {
                          return Child->reachedEnd();
                        });
return std::accumulate(
         begin(Children), end(Children), DEFAULT_BOOST_SCORE,
         [&](float Boost, const std::unique_ptr<Iterator> &Child) {
           return (!Child->reachedEnd() && Child->peek() == ID)
                      ? std::max(Boost, Child->consume())
                      : Boost;
         });
   }
 
+  size_t estimateSize() const override {
+    return std::accumulate(
+        begin(Children), end(Children), Children.front()->estimateSize(),
+        [&](size_t Current, const std::unique_ptr<Iterator> &Child) {
+          return std::max(Current, Child->estimateSize());
+        });
+  }
+
 private:
   llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override {
     OS << "(| ";
     auto Separator = "";
     for (const auto &Child : Children) {
       OS << Separator << *Child;
       Separator = " ";
     }
public:
 
   DocID peek() const override {
     assert(!reachedEnd() && "TrueIterator can't call peek() at the end.");
     return Index;
   }
 
   float consume() override { return DEFAULT_BOOST_SCORE; }
 
+  size_t estimateSize() const override { return Size; }
+
 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.
public:
   void advance() override { Child->advance(); }
 
   void advanceTo(DocID ID) override { Child->advanceTo(ID); }
 
   DocID peek() const override { return Child->peek(); }
 
   float consume() override { return Child->consume() * Factor; }
 
+  size_t estimateSize() const override { return Child->estimateSize(); }
+
 private:
   llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override {
     OS << "(BOOST " << Factor << ' ' << *Child << ')';
     return OS;
   }
 
   std::unique_ptr<Iterator> Child;
   float Factor;
public:
   /// 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();
   }
 
+  size_t estimateSize() const override {
+    return std::min(Child->estimateSize(), Limit);
+  }
+
 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;

unittests/clangd/DexIndexTests.cpp

TEST(DexIndexIterators, StringRepresentation) {
   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}]))");
+            "(& (| [{END}] [0, {5}, END] [0, {1}, 5, END]) (& [{1}, 5, 7, 9, "
+            "END] [{1}, 3, 5, 8, 9, END]))");
 }
 
 TEST(DexIndexIterators, Limit) {
   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);