[clangd] Make dex iterator peek() and reachedEnd() nonvirtual

[sammccall]

Summary

This is not polished but worth checking if it's a speedup
Hopefully correct: passes all tests

[sammccall]

const fix

[kbobyrev]

Ditto (previous performance improvement attempt).

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

 /// 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 {
 public:
   /// Returns true if all valid DocIDs were processed and hence the iterator is
   /// exhausted.
-  virtual bool reachedEnd() const = 0;
+  bool reachedEnd() const { return Current == End; };
   /// 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;
+  DocID peek() const {
+    assert(!reachedEnd());
+    return Current;
+  }
+  /// Returns the element this iterator points to, or End if reachedEnd().
+  DocID peekOrEnd() { return Current; }
   /// 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;
friend llvm::raw_ostream &operator<<(llvm::raw_ostream &OS,
                                        const Iterator &Iterator) {
     return Iterator.dump(OS);
   }
 
   /// Inspect iterator type, used internally for optimizing query trees.
   enum class Kind { And, Or, True, False, Other };
   Kind kind() const { return MyKind; }
 
+  /// Sentinel DocID representing end-of-iteration, greater than all others.
+  constexpr static DocID End = -1;
+
 protected:
-  Iterator(Kind MyKind = Kind::Other) : MyKind(MyKind) {}
+  Iterator(DocID Initial, Kind MyKind = Kind::Other)
+      : Current(Initial), MyKind(MyKind) {}
+  DocID Current = End;
 
 private:
   virtual llvm::raw_ostream &dump(llvm::raw_ostream &OS) const = 0;
   Kind MyKind;
 };
 
 /// Advances the iterator until it is exhausted. Returns pairs of document IDs
 /// with the corresponding boosting score.

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

 /// Implements Iterator over the intersection of other iterators.
 ///
 /// AndIterator iterates through common items among all children. It becomes
 /// exhausted as soon as any child becomes exhausted. After each mutation, the
 /// iterator restores the invariant: all children must point to the same item.
 class AndIterator : public Iterator {
 public:
   explicit AndIterator(std::vector<std::unique_ptr<Iterator>> AllChildren)
-      : Iterator(Kind::And), Children(std::move(AllChildren)) {
+      : Iterator(AllChildren.front()->peekOrEnd(), Kind::And),
+        Children(std::move(AllChildren)) {
     assert(!Children.empty() && "AND iterator should have at least one child.");
     // Establish invariants.
-    for (const auto &Child : Children)
-      ReachedEnd |= Child->reachedEnd();
     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.
     llvm::sort(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() && "AND iterator can't advance() at the end.");
     Children.front()->advance();
     sync();
   }
 
   /// Advances all children to the next common item with DocumentID >= ID.
   void advanceTo(DocID ID) override {
     assert(!reachedEnd() && "AND iterator can't advanceTo() at the end.");
     Children.front()->advanceTo(ID);
     sync();
   }
 
-  DocID peek() const override { return Children.front()->peek(); }
-
   float consume() override {
     assert(!reachedEnd() && "AND iterator can't consume() at the end.");
     float Boost = 1;
     for (const auto &Child : Children)
       Boost *= Child->consume();
     return Boost;
   }
 
llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override {
     }
     OS << ')';
     return OS;
   }
 
   /// Restores class invariants: each child will point to the same element after
   /// sync.
   void sync() {
-    ReachedEnd |= Children.front()->reachedEnd();
-    if (ReachedEnd)
+    Current = Children.front()->peekOrEnd();
+    if (Current == End)
       return;
-    auto SyncID = Children.front()->peek();
-    // Indicates whether any child needs to be advanced to new SyncID.
+    // Indicates whether any child needs to be advanced to new Current.
     bool NeedsAdvance = false;
     do {
       NeedsAdvance = false;
       for (auto &Child : Children) {
-        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)
+        Child->advanceTo(Current);
+        if (Child->peekOrEnd() > Current) {
+          Current = Child->peekOrEnd();
+          if (Current == End)
+            // If any child reaches end And iterator can not match any other
+            // items. In this case, just terminate the process.
             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.
-        if (Child->peek() > SyncID) {
-          SyncID = Child->peek();
+          // 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.
           NeedsAdvance = true;
         }
       }
     } while (NeedsAdvance);
   }
 
   /// AndIterator owns its children and ensures that all of them point to the
   /// same element. As soon as one child gets exhausted, AndIterator can no
   /// longer advance and has reached its end.
   std::vector<std::unique_ptr<Iterator>> Children;
-  /// Indicates whether any child is exhausted. It is cheaper to maintain and
-  /// update the field, rather than traversing the whole subtree in each
-  /// reachedEnd() call.
-  bool ReachedEnd = false;
   friend Corpus; // For optimizations.
 };
 
 /// Implements Iterator over the union of other iterators.
 ///
 /// OrIterator iterates through all items which can be pointed to by at least
 /// 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:
   explicit OrIterator(std::vector<std::unique_ptr<Iterator>> AllChildren)
-      : Iterator(Kind::Or), Children(std::move(AllChildren)) {
+      : Iterator(End, Kind::Or), Children(std::move(AllChildren)) {
     assert(!Children.empty() && "OR iterator should have at least one child.");
-  }
-
-  /// Returns true if all children are exhausted.
-  bool reachedEnd() const override {
     for (const auto &Child : Children)
-      if (!Child->reachedEnd())
-        return false;
-    return true;
+      Current = std::min(Current, Child->peekOrEnd());
   }
 
   /// Moves each child pointing to the smallest DocID to the next item.
   void advance() override {
     assert(!reachedEnd() && "OR iterator can't advance() at the end.");
-    const auto SmallestID = peek();
-    for (const auto &Child : Children)
-      if (!Child->reachedEnd() && Child->peek() == SmallestID)
+    DocID Next = End;
+    for (const auto &Child : Children) {
+      if (Child->peekOrEnd() == Current)
         Child->advance();
+      Next = std::min(Next, Child->peekOrEnd());
+    }
+    Current = Next;
   }
 
   /// Advances each child to the next existing element with DocumentID >= ID.
   void advanceTo(DocID ID) override {
     assert(!reachedEnd() && "OR iterator can't advanceTo() at the end.");
+    Current = End;
     for (const auto &Child : Children)
-      if (!Child->reachedEnd())
+      if (!Child->reachedEnd()) {
         Child->advanceTo(ID);
+        Current = std::min(Current, Child->peekOrEnd());
       }
-
-  /// Returns the element under cursor of the child with smallest Child->peek()
-  /// value.
-  DocID peek() const override {
-    assert(!reachedEnd() && "OR iterator can't peek() at the end.");
-    DocID Result = std::numeric_limits<DocID>::max();
-
-    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
   // points to the current ID.
   float consume() override {
     assert(!reachedEnd() && "OR iterator can't consume() at the end.");
     const DocID ID = peek();
     float Boost = 1;
     for (const auto &Child : Children)
-      if (!Child->reachedEnd() && Child->peek() == ID)
+      if (Child->peekOrEnd() == ID)
         Boost = std::max(Boost, Child->consume());
     return Boost;
   }
 
   size_t estimateSize() const override {
     size_t Size = 0;
     for (const auto &Child : Children)
       Size = std::max(Size, Child->estimateSize());
private:
   friend Corpus; // For optimizations.
 };
 
 /// TrueIterator handles PostingLists which contain all items of the index. It
 /// stores size of the virtual posting list, and all operations are performed
 /// in O(1).
 class TrueIterator : public Iterator {
 public:
-  explicit TrueIterator(DocID Size) : Iterator(Kind::True), Size(Size) {}
-
-  bool reachedEnd() const override { return Index >= Size; }
+  explicit TrueIterator(DocID Size)
+      : Iterator(Size == 0 ? End : 0, Kind::True), Size(Size) {}
 
   void advance() override {
     assert(!reachedEnd() && "TRUE iterator can't advance() at the end.");
-    ++Index;
+    if (++Current >= Size)
+      Current = End;
   }
 
   void advanceTo(DocID ID) override {
     assert(!reachedEnd() && "TRUE iterator can't advanceTo() at the end.");
-    Index = std::min(ID, Size);
-  }
-
-  DocID peek() const override {
-    assert(!reachedEnd() && "TRUE iterator can't peek() at the end.");
-    return Index;
+    Current = ID >= Size ? End : ID;
   }
 
   float consume() override {
     assert(!reachedEnd() && "TRUE iterator can't consume() at the end.");
     return 1;
   }
 
   size_t estimateSize() const override { return Size; }
 
 private:
   llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override {
     return OS << "true";
   }
 
-  DocID Index = 0;
   /// Size of the underlying virtual PostingList.
   DocID Size;
 };
 
 /// FalseIterator yields no results.
 class FalseIterator : public Iterator {
 public:
-  FalseIterator() : Iterator(Kind::False) {}
-  bool reachedEnd() const override { return true; }
+  FalseIterator() : Iterator(End, Kind::False) {}
   void advance() override { assert(false); }
   void advanceTo(DocID ID) override { assert(false); }
-  DocID peek() const override {
-    assert(false);
-    return 0;
-  }
   float consume() override {
     assert(false);
     return 1;
   }
   size_t estimateSize() const override { return 0; }
 
 private:
   llvm::raw_ostream &dump(llvm::raw_ostream &OS) const override {
     return OS << "false";
   }
 };
 
 /// 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(std::move(Child)), Factor(Factor) {}
-
-  bool reachedEnd() const override { return Child->reachedEnd(); }
+      : Iterator(Child->peekOrEnd()), Child(std::move(Child)), Factor(Factor) {}
 
-  void advance() override { Child->advance(); }
-
-  void advanceTo(DocID ID) override { Child->advanceTo(ID); }
+  void advance() override {
+    Child->advance();
+    Current = Child->peekOrEnd();
+  }
 
-  DocID peek() const override { return Child->peek(); }
+  void advanceTo(DocID ID) override {
+    Child->advanceTo(ID);
+    Current = Child->peekOrEnd();
+  }
 
   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 {
     return OS << "(* " << Factor << ' ' << *Child << ')';
   }
 
   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(std::move(Child)), Limit(Limit), ItemsLeft(Limit) {}
+      : Iterator(Limit == 0 ? End : Child->peekOrEnd()),
+        Child(std::move(Child)), Limit(Limit), ItemsLeft(Limit) {}
 
-  bool reachedEnd() const override {
-    return ItemsLeft == 0 || Child->reachedEnd();
+  void advance() override {
+    Child->advance();
+    Current = Child->peekOrEnd();
   }
 
-  void advance() override { Child->advance(); }
-
-  void advanceTo(DocID ID) override { Child->advanceTo(ID); }
-
-  DocID peek() const override { return Child->peek(); }
+  void advanceTo(DocID ID) override {
+    Child->advanceTo(ID);
+    Current = Child->peekOrEnd();
+  }
 
   /// 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;
+    if (--ItemsLeft == 0)
+      Current = End;
     return Child->consume();
   }
 
   size_t estimateSize() const override {
     return std::min(Child->estimateSize(), Limit);
   }
 
 private:

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

 namespace {
 
 /// Implements iterator of PostingList chunks. This requires iterating over two
 /// levels: the first level iterator iterates over the chunks and decompresses
 /// them on-the-fly when the contents of chunk are to be seen.
 class ChunkIterator : public Iterator {
 public:
   explicit ChunkIterator(const Token *Tok, llvm::ArrayRef<Chunk> Chunks)
-      : Tok(Tok), Chunks(Chunks), CurrentChunk(Chunks.begin()) {
+      : Iterator(End), Tok(Tok), Chunks(Chunks), CurrentChunk(Chunks.begin()) {
     if (!Chunks.empty()) {
       DecompressedChunk = CurrentChunk->decompress();
       CurrentID = DecompressedChunk.begin();
+      if (CurrentID != DecompressedChunk.end())
+        Current = *CurrentID;
     }
   }
 
-  bool reachedEnd() const override { return CurrentChunk == Chunks.end(); }
-
   /// Advances cursor to the next item.
   void advance() override {
     assert(!reachedEnd() &&
            "Posting List iterator can't advance() at the end.");
     ++CurrentID;
     normalizeCursor();
   }
 
   /// Applies binary search to advance cursor to the next item with DocID
   /// equal or higher than the given one.
   void advanceTo(DocID ID) override {
     assert(!reachedEnd() &&
            "Posting List iterator can't advance() at the end.");
     if (ID <= peek())
       return;
     advanceToChunk(ID);
     // Try to find ID within current chunk.
     CurrentID = std::partition_point(CurrentID, DecompressedChunk.end(),
                                      [&](const DocID D) { return D < ID; });
     normalizeCursor();
   }
 
-  DocID peek() const override {
-    assert(!reachedEnd() && "Posting List iterator can't peek() at the end.");
-    return *CurrentID;
-  }
-
   float consume() override {
     assert(!reachedEnd() &&
            "Posting List iterator can't consume() at the end.");
     return 1;
   }
 
   size_t estimateSize() const override {
     return Chunks.size() * ApproxEntriesPerChunk;
for (const Chunk &C : Chunks)
       for (const DocID Doc : C.decompress()) {
         OS << Sep << Doc;
         Sep = " ";
       }
     return OS << ']';
   }
 
   /// If the cursor is at the end of a chunk, place it at the start of the next
-  /// chunk.
+  /// chunk. Set Current.
   void normalizeCursor() {
     // Invariant is already established if examined chunk is not exhausted.
-    if (CurrentID != std::end(DecompressedChunk))
+    if (CurrentID != std::end(DecompressedChunk)) {
+      Current = *CurrentID;
       return;
+    }
     // Advance to next chunk if current one is exhausted.
     ++CurrentChunk;
-    if (CurrentChunk == Chunks.end()) // Reached the end of PostingList.
+    if (CurrentChunk == Chunks.end()) { // Reached the end of PostingList.
+      Current = End;
       return;
+    }
     DecompressedChunk = CurrentChunk->decompress();
     CurrentID = DecompressedChunk.begin();
+    assert(CurrentID != DecompressedChunk.end() && "Empty chunk?");
+    Current = *CurrentID;
   }
 
   /// Advances CurrentChunk to the chunk which might contain ID.
   void advanceToChunk(DocID ID) {
     if ((CurrentChunk != Chunks.end() - 1) &&
         ((CurrentChunk + 1)->Head <= ID)) {
       CurrentChunk =
           std::partition_point(CurrentChunk + 1, Chunks.end(),

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

TEST(DexIterators, Limit) {
   Corpus C{10000};
   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 = C.limit(L0.iterator(), 42);
   EXPECT_THAT(consumeIDs(*DocIterator), ElementsAre(3, 6, 7, 20, 42, 100));
 
+#if 0
   DocIterator = C.limit(L0.iterator(), 3);
   EXPECT_THAT(consumeIDs(*DocIterator), ElementsAre(3, 6, 7));
 
   DocIterator = C.limit(L0.iterator(), 0);
   EXPECT_THAT(consumeIDs(*DocIterator), ElementsAre());
 
   auto AndIterator =
       C.intersect(C.limit(C.all(), 343), C.limit(L0.iterator(), 2),
                   C.limit(L1.iterator(), 3), C.limit(L2.iterator(), 42));
   EXPECT_THAT(consumeIDs(*AndIterator), ElementsAre(3, 7));
+#endif
 }
 
 TEST(DexIterators, True) {
   EXPECT_TRUE(Corpus{0}.all()->reachedEnd());
   EXPECT_THAT(consumeIDs(*Corpus{4}.all()), ElementsAre(0, 1, 2, 3));
 }
 
 TEST(DexIterators, Boost) {