1 files changed, 220 insertions, 0 deletions
diff --git a/src/rocksdb/table/two_level_iterator.cc b/src/rocksdb/table/two_level_iterator.cc
new file mode 100644
index 000000000..4b6634e5c
--- /dev/null
+++ b/src/rocksdb/table/two_level_iterator.cc
@@ -0,0 +1,220 @@
+//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+//
+// Copyright (c) 2011 The LevelDB Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file. See the AUTHORS file for names of contributors.
+
+#include "table/two_level_iterator.h"
+
+#include "db/pinned_iterators_manager.h"
+#include "memory/arena.h"
+#include "rocksdb/options.h"
+#include "rocksdb/table.h"
+#include "table/block_based/block.h"
+#include "table/format.h"
+
+namespace ROCKSDB_NAMESPACE {
+
+namespace {
+
+class TwoLevelIndexIterator : public InternalIteratorBase<IndexValue> {
+ public:
+  explicit TwoLevelIndexIterator(
+      TwoLevelIteratorState* state,
+      InternalIteratorBase<IndexValue>* first_level_iter);
+
+  ~TwoLevelIndexIterator() override {
+    first_level_iter_.DeleteIter(false /* is_arena_mode */);
+    second_level_iter_.DeleteIter(false /* is_arena_mode */);
+    delete state_;
+  }
+
+  void Seek(const Slice& target) override;
+  void SeekForPrev(const Slice& target) override;
+  void SeekToFirst() override;
+  void SeekToLast() override;
+  void Next() override;
+  void Prev() override;
+
+  bool Valid() const override { return second_level_iter_.Valid(); }
+  Slice key() const override {
+    assert(Valid());
+    return second_level_iter_.key();
+  }
+  Slice user_key() const override {
+    assert(Valid());
+    return second_level_iter_.user_key();
+  }
+  IndexValue value() const override {
+    assert(Valid());
+    return second_level_iter_.value();
+  }
+  Status status() const override {
+    if (!first_level_iter_.status().ok()) {
+      assert(second_level_iter_.iter() == nullptr);
+      return first_level_iter_.status();
+    } else if (second_level_iter_.iter() != nullptr &&
+               !second_level_iter_.status().ok()) {
+      return second_level_iter_.status();
+    } else {
+      return status_;
+    }
+  }
+  void SetPinnedItersMgr(
+      PinnedIteratorsManager* /*pinned_iters_mgr*/) override {}
+  bool IsKeyPinned() const override { return false; }
+  bool IsValuePinned() const override { return false; }
+
+ private:
+  void SaveError(const Status& s) {
+    if (status_.ok() && !s.ok()) status_ = s;
+  }
+  void SkipEmptyDataBlocksForward();
+  void SkipEmptyDataBlocksBackward();
+  void SetSecondLevelIterator(InternalIteratorBase<IndexValue>* iter);
+  void InitDataBlock();
+
+  TwoLevelIteratorState* state_;
+  IteratorWrapperBase<IndexValue> first_level_iter_;
+  IteratorWrapperBase<IndexValue> second_level_iter_;  // May be nullptr
+  Status status_;
+  // If second_level_iter is non-nullptr, then "data_block_handle_" holds the
+  // "index_value" passed to block_function_ to create the second_level_iter.
+  BlockHandle data_block_handle_;
+};
+
+TwoLevelIndexIterator::TwoLevelIndexIterator(
+    TwoLevelIteratorState* state,
+    InternalIteratorBase<IndexValue>* first_level_iter)
+    : state_(state), first_level_iter_(first_level_iter) {}
+
+void TwoLevelIndexIterator::Seek(const Slice& target) {
+  first_level_iter_.Seek(target);
+
+  InitDataBlock();
+  if (second_level_iter_.iter() != nullptr) {
+    second_level_iter_.Seek(target);
+  }
+  SkipEmptyDataBlocksForward();
+}
+
+void TwoLevelIndexIterator::SeekForPrev(const Slice& target) {
+  first_level_iter_.Seek(target);
+  InitDataBlock();
+  if (second_level_iter_.iter() != nullptr) {
+    second_level_iter_.SeekForPrev(target);
+  }
+  if (!Valid()) {
+    if (!first_level_iter_.Valid() && first_level_iter_.status().ok()) {
+      first_level_iter_.SeekToLast();
+      InitDataBlock();
+      if (second_level_iter_.iter() != nullptr) {
+        second_level_iter_.SeekForPrev(target);
+      }
+    }
+    SkipEmptyDataBlocksBackward();
+  }
+}
+
+void TwoLevelIndexIterator::SeekToFirst() {
+  first_level_iter_.SeekToFirst();
+  InitDataBlock();
+  if (second_level_iter_.iter() != nullptr) {
+    second_level_iter_.SeekToFirst();
+  }
+  SkipEmptyDataBlocksForward();
+}
+
+void TwoLevelIndexIterator::SeekToLast() {
+  first_level_iter_.SeekToLast();
+  InitDataBlock();
+  if (second_level_iter_.iter() != nullptr) {
+    second_level_iter_.SeekToLast();
+  }
+  SkipEmptyDataBlocksBackward();
+}
+
+void TwoLevelIndexIterator::Next() {
+  assert(Valid());
+  second_level_iter_.Next();
+  SkipEmptyDataBlocksForward();
+}
+
+void TwoLevelIndexIterator::Prev() {
+  assert(Valid());
+  second_level_iter_.Prev();
+  SkipEmptyDataBlocksBackward();
+}
+
+void TwoLevelIndexIterator::SkipEmptyDataBlocksForward() {
+  while (second_level_iter_.iter() == nullptr ||
+         (!second_level_iter_.Valid() && second_level_iter_.status().ok())) {
+    // Move to next block
+    if (!first_level_iter_.Valid()) {
+      SetSecondLevelIterator(nullptr);
+      return;
+    }
+    first_level_iter_.Next();
+    InitDataBlock();
+    if (second_level_iter_.iter() != nullptr) {
+      second_level_iter_.SeekToFirst();
+    }
+  }
+}
+
+void TwoLevelIndexIterator::SkipEmptyDataBlocksBackward() {
+  while (second_level_iter_.iter() == nullptr ||
+         (!second_level_iter_.Valid() && second_level_iter_.status().ok())) {
+    // Move to next block
+    if (!first_level_iter_.Valid()) {
+      SetSecondLevelIterator(nullptr);
+      return;
+    }
+    first_level_iter_.Prev();
+    InitDataBlock();
+    if (second_level_iter_.iter() != nullptr) {
+      second_level_iter_.SeekToLast();
+    }
+  }
+}
+
+void TwoLevelIndexIterator::SetSecondLevelIterator(
+    InternalIteratorBase<IndexValue>* iter) {
+  InternalIteratorBase<IndexValue>* old_iter = second_level_iter_.Set(iter);
+  delete old_iter;
+}
+
+void TwoLevelIndexIterator::InitDataBlock() {
+  if (!first_level_iter_.Valid()) {
+    SetSecondLevelIterator(nullptr);
+  } else {
+    BlockHandle handle = first_level_iter_.value().handle;
+    if (second_level_iter_.iter() != nullptr &&
+        !second_level_iter_.status().IsIncomplete() &&
+        handle.offset() == data_block_handle_.offset()) {
+      // second_level_iter is already constructed with this iterator, so
+      // no need to change anything
+    } else {
+      InternalIteratorBase<IndexValue>* iter =
+          state_->NewSecondaryIterator(handle);
+      data_block_handle_ = handle;
+      SetSecondLevelIterator(iter);
+      if (iter == nullptr) {
+        status_ = Status::Corruption("Missing block for partition " +
+                                     handle.ToString());
+      }
+    }
+  }
+}
+
+}  // namespace
+
+InternalIteratorBase<IndexValue>* NewTwoLevelIterator(
+    TwoLevelIteratorState* state,
+    InternalIteratorBase<IndexValue>* first_level_iter) {
+  return new TwoLevelIndexIterator(state, first_level_iter);
+}
+}  // namespace ROCKSDB_NAMESPACE