[clangd] Dex: FALSE iterator, peephole optimizations, fix AND bug

[sammccall]

Summary

The FALSE iterator will be used in a followup patch to fix a logic bug in Dex
(currently, tokens that don't have posting lists in the index are simply dropped
from the query, changing semantics).

It can usually be optimized away, so added the following opmitizations:

• simplify booleans inside AND/OR
• replace effectively-empty AND/OR with booleans
• flatten nested AND/ORs

While working on this, found a bug in the AND iterator: its constructor sync()
assumes that ReachedEnd is set if applicable, but the constructor never sets it.
This crashes if a non-first iterator is nonempty.

[ilya-biryukov]

LGTM and a few NITs.

clangd/index/dex/Iterator.cpp

376

Maybe replace with case Kind::Other to make sure compiler gives a warning to update the switch when new kinds are added?
Same for other switches on kind

378

Add braces for for statement to make sure there's a closing brace at the right indentation?

385

Maybe use the qualified std::move for consistency with the previous line?

408

Again, for () switch pattern ends up putting the closing brace at the wrong indent level. Consider adding the braces around for the loop body.

415

Maybe use the qualified std::move for consistency with the previous line?

clangd/index/dex/Iterator.h

109

Maybe put data members after functions to follow the pattern we have in most (all?) other places in clangd?

[sammccall]

clangd/index/dex/Iterator.cpp

376

Actually this is deliberate, there's no particular need to specialize on every kind. (and in fact we have Or specialized on one fewer than And, but I wrote the case out in order to leave a comment)

385

Updated throughout.
(This file was inconsistent and I was trying to preserve... something, but I'm not sure what)

clangd/index/dex/Iterator.h

public:
   /// ChildN is N-th iterator child. Raw iterators over PostingList are
   /// represented as "[... CurID ...]" where CurID is the current PostingList
   /// entry being inspected.
   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; }
+
+protected:
+  Iterator(Kind MyKind = Kind::Other) : MyKind(MyKind) {}
+
 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.
 ///
 /// 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
public:
   /// containing the requested item.
   std::unique_ptr<Iterator>
   unionOf(std::vector<std::unique_ptr<Iterator>> Children) const;
 
   /// Returns TRUE Iterator which iterates over "virtual" PostingList
   /// containing all items in range [0, Size) in an efficient manner.
   std::unique_ptr<Iterator> all() const;
 
+  /// Returns FALSE Iterator which iterates over no documents.
+  std::unique_ptr<Iterator> none() const;
+
   /// 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> boost(std::unique_ptr<Iterator> Child,
                                   float Factor) const;
 
   /// Returns LIMIT iterator, which yields up to N elements of its child

clangd/index/dex/Iterator.cpp

 //===--- Iterator.cpp - Query Symbol Retrieval ------------------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "Iterator.h"
+#include "llvm/Support/Casting.h"
 #include <algorithm>
 #include <cassert>
 #include <numeric>
 
 namespace clang {
 namespace clangd {
 namespace dex {
 
 namespace {
 
 /// 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)
-      : Children(std::move(AllChildren)) {
+      : Iterator(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.
private:
   /// 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)
-      : Children(std::move(AllChildren)) {
+      : Iterator(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;
for (const auto &Child : Children) {
       Separator = " ";
     }
     OS << ')';
     return OS;
   }
 
   // FIXME(kbobyrev): Would storing Children in min-heap be faster?
   std::vector<std::unique_ptr<Iterator>> Children;
+  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) : Size(Size) {}
+  explicit TrueIterator(DocID Size) : Iterator(Kind::True), Size(Size) {}
 
   bool reachedEnd() const override { return Index >= Size; }
 
   void advance() override {
     assert(!reachedEnd() && "TRUE iterator can't advance() at the end.");
     ++Index;
   }
 
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; }
+  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(move(Child)), Factor(Factor) {}
+      : Child(std::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(); }

 
 /// 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) {}
+      : Child(std::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); }
std::vector<std::pair<DocID, float>> consume(Iterator &It) {
   std::vector<std::pair<DocID, float>> Result;
   for (; !It.reachedEnd(); It.advance())
     Result.emplace_back(It.peek(), It.consume());
   return Result;
 }
 
 std::unique_ptr<Iterator>
 Corpus::intersect(std::vector<std::unique_ptr<Iterator>> Children) const {
-  // If there is exactly one child, pull it one level up: AND(Child) -> Child.
-  return Children.size() == 1 ? std::move(Children.front())
-                              : llvm::make_unique<AndIterator>(move(Children));
+  std::vector<std::unique_ptr<Iterator>> RealChildren;
+  for (auto &Child : Children) {
+    switch (Child->kind()) {
+    case Iterator::Kind::True:
+      break; // No effect, drop the iterator.
+    case Iterator::Kind::False:
+      return std::move(Child); // Intersection is empty.
+    case Iterator::Kind::And: {
+      // Inline nested AND into parent AND.
+      auto &NewChildren = static_cast<AndIterator *>(Child.get())->Children;
+      std::move(NewChildren.begin(), NewChildren.end(),
+                std::back_inserter(RealChildren));
+      break;
+    }
+    default:
+      RealChildren.push_back(std::move(Child));
+    }
+  }
+  switch (RealChildren.size()) {
+  case 0:
+    return all();
+  case 1:
+    return std::move(RealChildren.front());
+  default:
+    return llvm::make_unique<AndIterator>(std::move(RealChildren));
+  }
 }
 
 std::unique_ptr<Iterator>
 Corpus::unionOf(std::vector<std::unique_ptr<Iterator>> Children) const {
-  // If there is exactly one child, pull it one level up: OR(Child) -> Child.
-  return Children.size() == 1 ? std::move(Children.front())
-                              : llvm::make_unique<OrIterator>(move(Children));
+  std::vector<std::unique_ptr<Iterator>> RealChildren;
+  for (auto &Child : Children) {
+    switch (Child->kind()) {
+    case Iterator::Kind::False:
+      break; // No effect, drop the iterator.
+    case Iterator::Kind::Or: {
+      // Inline nested OR into parent OR.
+      auto &NewChildren = static_cast<OrIterator *>(Child.get())->Children;
+      std::move(NewChildren.begin(), NewChildren.end(),
+                std::back_inserter(RealChildren));
+      break;
+    }
+    case Iterator::Kind::True:
+      // Don't return all(), which would discard sibling boosts.
+    default:
+      RealChildren.push_back(std::move(Child));
+    }
+  }
+  switch (RealChildren.size()) {
+  case 0:
+    return none();
+  case 1:
+    return std::move(RealChildren.front());
+  default:
+    return llvm::make_unique<OrIterator>(std::move(RealChildren));
+  }
 }
 
 std::unique_ptr<Iterator> Corpus::all() const {
   return llvm::make_unique<TrueIterator>(Size);
 }
 
+std::unique_ptr<Iterator> Corpus::none() const {
+  return llvm::make_unique<FalseIterator>();
+}
+
 std::unique_ptr<Iterator> Corpus::boost(std::unique_ptr<Iterator> Child,
                                         float Factor) const {
-  return llvm::make_unique<BoostIterator>(move(Child), Factor);
+  if (Factor == 1)
+    return Child;
+  if (Child->kind() == Iterator::Kind::False)
+    return Child;
+  return llvm::make_unique<BoostIterator>(std::move(Child), Factor);
 }
 
 std::unique_ptr<Iterator> Corpus::limit(std::unique_ptr<Iterator> Child,
                                         size_t Limit) const {
-  return llvm::make_unique<LimitIterator>(move(Child), Limit);
+  if (Child->kind() == Iterator::Kind::False)
+    return Child;
+  return llvm::make_unique<LimitIterator>(std::move(Child), Limit);
 }
 
 } // namespace dex
 } // namespace clangd
 } // namespace clang

unittests/clangd/DexTests.cpp

TEST(DexIterators, AndThreeLists) {
   EXPECT_EQ(And->peek(), 7U);
   And->advanceTo(300);
   EXPECT_EQ(And->peek(), 320U);
   And->advanceTo(100000);
 
   EXPECT_TRUE(And->reachedEnd());
 }
 
+TEST(DexIterators, AndEmpty) {
+  Corpus C{10000};
+  const PostingList L1({1});
+  const PostingList L2({2});
+  // These iterators are empty, but the optimizer can't tell.
+  auto Empty1 = C.intersect(L1.iterator(), L2.iterator());
+  auto Empty2 = C.intersect(L1.iterator(), L2.iterator());
+  // And syncs iterators on construction, and used to fail on empty children.
+  auto And = C.intersect(std::move(Empty1), std::move(Empty2));
+  EXPECT_TRUE(And->reachedEnd());
+}
+
 TEST(DexIterators, OrTwoLists) {
   Corpus C{10000};
   const PostingList L0({0, 5, 7, 10, 42, 320, 9000});
   const PostingList L1({0, 4, 7, 10, 30, 60, 320, 9000});
 
   auto Or = C.unionOf(L0.iterator(), L1.iterator());
 
   EXPECT_FALSE(Or->reachedEnd());
TEST(DexIterators, Limit) {
 
   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));
 }
 
 TEST(DexIterators, True) {
-  Corpus C{0};
-  auto TrueIterator = C.all();
-  EXPECT_TRUE(TrueIterator->reachedEnd());
-  EXPECT_THAT(consumeIDs(*TrueIterator), ElementsAre());
-
-  C = Corpus{7};
-  const PostingList L0({1, 2, 5, 7});
-  TrueIterator = C.all();
-  EXPECT_THAT(TrueIterator->peek(), 0);
-  auto AndIterator = C.intersect(L0.iterator(), move(TrueIterator));
-  EXPECT_FALSE(AndIterator->reachedEnd());
-  EXPECT_THAT(consumeIDs(*AndIterator), ElementsAre(1, 2, 5));
+  EXPECT_TRUE(Corpus{0}.all()->reachedEnd());
+  EXPECT_THAT(consumeIDs(*Corpus{4}.all()), ElementsAre(0, 1, 2, 3));
 }
 
 TEST(DexIterators, Boost) {
   Corpus C{5};
   auto BoostIterator = C.boost(C.all(), 42U);
   EXPECT_FALSE(BoostIterator->reachedEnd());
   auto ElementBoost = BoostIterator->consume();
   EXPECT_THAT(ElementBoost, 42U);
TEST(DexIterators, Boost) {
   ElementBoost = Root->consume();
   EXPECT_THAT(ElementBoost, 2);
 
   Root->advanceTo(4);
   ElementBoost = Root->consume();
   EXPECT_THAT(ElementBoost, 3);
 }
 
+TEST(DexIterators, Optimizations) {
+  Corpus C{5};
+  const PostingList L1({1});
+  const PostingList L2({2});
+  const PostingList L3({3});
+
+  // empty and/or yield true/false
+  EXPECT_EQ(llvm::to_string(*C.intersect()), "true");
+  EXPECT_EQ(llvm::to_string(*C.unionOf()), "false");
+
+  // true/false inside and/or short-circuit
+  EXPECT_EQ(llvm::to_string(*C.intersect(L1.iterator(), C.all())), "[1]");
+  EXPECT_EQ(llvm::to_string(*C.intersect(L1.iterator(), C.none())), "false");
+  // Not optimized to avoid breaking boosts.
+  EXPECT_EQ(llvm::to_string(*C.unionOf(L1.iterator(), C.all())),
+            "(| [1] true)");
+  EXPECT_EQ(llvm::to_string(*C.unionOf(L1.iterator(), C.none())), "[1]");
+
+  // and/or nested inside and/or are flattened
+  EXPECT_EQ(llvm::to_string(*C.intersect(
+                L1.iterator(), C.intersect(L1.iterator(), L1.iterator()))),
+            "(& [1] [1] [1])");
+  EXPECT_EQ(llvm::to_string(*C.unionOf(
+                L1.iterator(), C.unionOf(L2.iterator(), L3.iterator()))),
+            "(| [1] [2] [3])");
+
+  // optimizations combine over multiple levels
+  EXPECT_EQ(llvm::to_string(*C.intersect(
+                C.intersect(L1.iterator(), C.intersect()), C.unionOf(C.all()))),
+            "[1]");
+}
+
 //===----------------------------------------------------------------------===//
 // Search token tests.
 //===----------------------------------------------------------------------===//
 
 testing::Matcher<std::vector<Token>>
 tokensAre(std::initializer_list<std::string> Strings, Token::Kind Kind) {
   std::vector<Token> Tokens;
   for (const auto &TokenData : Strings) {