Use Member Functions and remove directions

jaykorean · Mar 21, 2024 · c2c97cf · c2c97cf
1 parent 6ddfa5f
commit c2c97cf
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 63 deletions.
diff --git a/db/multi_cf_iterator.cc b/db/multi_cf_iterator.cc
@@ -9,35 +9,28 @@
 
 namespace ROCKSDB_NAMESPACE {
 
-void MultiCfIterator::SeekCommon(
-    const std::function<void(Iterator*)>& child_seek_func,
-    Direction direction) {
-  direction_ = direction;
-  Reset();
+template <typename BinaryHeap, typename ChildSeekFuncType>
+void MultiCfIterator::SeekCommon(BinaryHeap& heap,
+                                 const ChildSeekFuncType& child_seek_func) {
+  heap.clear();
   int i = 0;
   for (auto& cfh_iter_pair : cfh_iter_pairs_) {
     auto& cfh = cfh_iter_pair.first;
     auto& iter = cfh_iter_pair.second;
     child_seek_func(iter.get());
     if (iter->Valid()) {
       assert(iter->status().ok());
-      if (direction_ == kReverse) {
-        auto& max_heap = std::get<MultiCfMaxHeap>(heap_);
-        max_heap.push(MultiCfIteratorInfo{iter.get(), cfh, i});
-      } else {
-        auto& min_heap = std::get<MultiCfMinHeap>(heap_);
-        min_heap.push(MultiCfIteratorInfo{iter.get(), cfh, i});
-      }
+      heap.push(MultiCfIteratorInfo{iter.get(), cfh, i});
     } else {
       considerStatus(iter->status());
     }
     ++i;
   }
 }
 
-template <typename BinaryHeap>
-void MultiCfIterator::AdvanceIterator(
-    BinaryHeap& heap, const std::function<void(Iterator*)>& advance_func) {
+template <typename BinaryHeap, typename AdvancedFuncType>
+void MultiCfIterator::AdvanceIterator(BinaryHeap& heap,
+                                      const AdvancedFuncType& advance_func) {
   // 1. Keep the top iterator (by popping it from the heap)
   // 2. Make sure all others have iterated past the top iterator key slice
   // 3. Advance the top iterator, and add it back to the heap if valid
@@ -70,33 +63,39 @@ void MultiCfIterator::AdvanceIterator(
 }
 
 void MultiCfIterator::SeekToFirst() {
-  SeekCommon([](Iterator* iter) { iter->SeekToFirst(); }, kForward);
+  auto& min_heap = GetHeap<MultiCfMinHeap>([this]() { InitMinHeap(); });
+  SeekCommon(min_heap, [](Iterator* iter) { iter->SeekToFirst(); });
 }
 void MultiCfIterator::Seek(const Slice& target) {
-  SeekCommon([&target](Iterator* iter) { iter->Seek(target); }, kForward);
+  auto& min_heap = GetHeap<MultiCfMinHeap>([this]() { InitMinHeap(); });
+  SeekCommon(min_heap, [&target](Iterator* iter) { iter->Seek(target); });
 }
 void MultiCfIterator::SeekToLast() {
-  SeekCommon([](Iterator* iter) { iter->SeekToLast(); }, kReverse);
+  auto& max_heap = GetHeap<MultiCfMaxHeap>([this]() { InitMaxHeap(); });
+  SeekCommon(max_heap, [](Iterator* iter) { iter->SeekToLast(); });
 }
 void MultiCfIterator::SeekForPrev(const Slice& target) {
-  SeekCommon([&target](Iterator* iter) { iter->SeekForPrev(target); },
-             kReverse);
+  auto& max_heap = GetHeap<MultiCfMaxHeap>([this]() { InitMaxHeap(); });
+  SeekCommon(max_heap,
+             [&target](Iterator* iter) { iter->SeekForPrev(target); });
 }
 
 void MultiCfIterator::Next() {
   assert(Valid());
-  if (direction_ != kForward) {
-    SwitchToDirection(kForward);
-  }
-  auto& min_heap = std::get<MultiCfMinHeap>(heap_);
+  auto& min_heap = GetHeap<MultiCfMinHeap>([this]() {
+    Slice target = key();
+    InitMinHeap();
+    Seek(target);
+  });
   AdvanceIterator(min_heap, [](Iterator* iter) { iter->Next(); });
 }
 void MultiCfIterator::Prev() {
   assert(Valid());
-  if (direction_ != kReverse) {
-    SwitchToDirection(kReverse);
-  }
-  auto& max_heap = std::get<MultiCfMaxHeap>(heap_);
+  auto& max_heap = GetHeap<MultiCfMaxHeap>([this]() {
+    Slice target = key();
+    InitMaxHeap();
+    SeekForPrev(target);
+  });
   AdvanceIterator(max_heap, [](Iterator* iter) { iter->Prev(); });
 }
 

diff --git a/db/multi_cf_iterator.h b/db/multi_cf_iterator.h
@@ -86,13 +86,10 @@ class MultiCfIterator : public Iterator {
 
   MultiCfIterHeap heap_;
 
-  enum Direction : uint8_t { kForward, kReverse };
-  Direction direction_ = kForward;
-
   // TODO: Lower and Upper bounds
 
   Iterator* current() const {
-    if (direction_ == kReverse) {
+    if (std::holds_alternative<MultiCfMaxHeap>(heap_)) {
       auto& max_heap = std::get<MultiCfMaxHeap>(heap_);
       return max_heap.top().iterator;
     }
@@ -114,7 +111,7 @@ class MultiCfIterator : public Iterator {
   }
 
   bool Valid() const override {
-    if (direction_ == kReverse) {
+    if (std::holds_alternative<MultiCfMaxHeap>(heap_)) {
       auto& max_heap = std::get<MultiCfMaxHeap>(heap_);
       return !max_heap.empty() && status_.ok();
     }
@@ -128,21 +125,13 @@ class MultiCfIterator : public Iterator {
       status_ = std::move(s);
     }
   }
-  void Reset() {
-    std::visit(overload{[&](MultiCfMinHeap& min_heap) -> void {
-                          min_heap.clear();
-                          if (direction_ == kReverse) {
-                            InitMaxHeap();
-                          }
-                        },
-                        [&](MultiCfMaxHeap& max_heap) -> void {
-                          max_heap.clear();
-                          if (direction_ == kForward) {
-                            InitMinHeap();
-                          }
-                        }},
-               heap_);
-    status_ = Status::OK();
+
+  template <typename HeapType, typename InitFunc>
+  HeapType& GetHeap(InitFunc initFunc) {
+    if (!std::holds_alternative<HeapType>(heap_)) {
+      initFunc();
+    }
+    return std::get<HeapType>(heap_);
   }
 
   void InitMinHeap() {
@@ -154,21 +143,10 @@ class MultiCfIterator : public Iterator {
         MultiCfHeapItemComparator<std::less<int>>(comparator_));
   }
 
-  void SwitchToDirection(Direction new_direction) {
-    assert(direction_ != new_direction);
-    Slice target = key();
-    if (new_direction == kForward) {
-      Seek(target);
-    } else {
-      SeekForPrev(target);
-    }
-  }
-
-  void SeekCommon(const std::function<void(Iterator*)>& child_seek_func,
-                  Direction direction);
-  template <typename BinaryHeap>
-  void AdvanceIterator(BinaryHeap& heap,
-                       const std::function<void(Iterator*)>& advance_func);
+  template <typename BinaryHeap, typename ChildSeekFuncType>
+  void SeekCommon(BinaryHeap& heap, const ChildSeekFuncType& child_seek_func);
+  template <typename BinaryHeap, typename AdvancedFuncType>
+  void AdvanceIterator(BinaryHeap& heap, const AdvancedFuncType& advance_func);
 
   void SeekToFirst() override;
   void SeekToLast() override;