rust-rocksdb/db/compaction/clipping_iterator.h

//  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).

#pragma once

#include <cassert>

#include "rocksdb/comparator.h"
#include "table/internal_iterator.h"

namespace ROCKSDB_NAMESPACE {

// An internal iterator that wraps another one and ensures that any keys
// returned are strictly within a range [start, end). If the underlying
// iterator has already performed the bounds checking, it relies on that result;
// otherwise, it performs the necessary key comparisons itself. Both bounds
// are optional.
class ClippingIterator : public InternalIterator {
 public:
  ClippingIterator(InternalIterator* iter, const Slice* start, const Slice* end,
                   const CompareInterface* cmp)
      : iter_(iter), start_(start), end_(end), cmp_(cmp), valid_(false) {
    assert(iter_);
    assert(cmp_);
    assert(!start_ || !end_ || cmp_->Compare(*start_, *end_) <= 0);

    UpdateAndEnforceBounds();
  }

  bool Valid() const override { return valid_; }

  void SeekToFirst() override {
    if (start_) {
      iter_->Seek(*start_);
    } else {
      iter_->SeekToFirst();
    }

    UpdateAndEnforceUpperBound();
  }

  void SeekToLast() override {
    if (end_) {
      iter_->SeekForPrev(*end_);

      // Upper bound is exclusive, so we need a key which is strictly smaller
      if (iter_->Valid() && cmp_->Compare(iter_->key(), *end_) == 0) {
        iter_->Prev();
      }
    } else {
      iter_->SeekToLast();
    }

    UpdateAndEnforceLowerBound();
  }

  void Seek(const Slice& target) override {
    if (start_ && cmp_->Compare(target, *start_) < 0) {
      iter_->Seek(*start_);
      UpdateAndEnforceUpperBound();
      return;
    }

    if (end_ && cmp_->Compare(target, *end_) >= 0) {
      valid_ = false;
      return;
    }

    iter_->Seek(target);
    UpdateAndEnforceUpperBound();
  }

  void SeekForPrev(const Slice& target) override {
    if (start_ && cmp_->Compare(target, *start_) < 0) {
      valid_ = false;
      return;
    }

    if (end_ && cmp_->Compare(target, *end_) >= 0) {
      iter_->SeekForPrev(*end_);

      // Upper bound is exclusive, so we need a key which is strictly smaller
      if (iter_->Valid() && cmp_->Compare(iter_->key(), *end_) == 0) {
        iter_->Prev();
      }

      UpdateAndEnforceLowerBound();
      return;
    }

    iter_->SeekForPrev(target);
    UpdateAndEnforceLowerBound();
  }

  void Next() override {
    assert(valid_);
    iter_->Next();
    UpdateAndEnforceUpperBound();
  }

  bool NextAndGetResult(IterateResult* result) override {
    assert(valid_);
    assert(result);

    IterateResult res;
    valid_ = iter_->NextAndGetResult(&res);

    if (!valid_) {
      return false;
    }

    if (end_) {
      EnforceUpperBoundImpl(res.bound_check_result);

      if (!valid_) {
        return false;
      }
    }

    res.bound_check_result = IterBoundCheck::kInbound;
    *result = res;

    return true;
  }

  void Prev() override {
    assert(valid_);
    iter_->Prev();
    UpdateAndEnforceLowerBound();
  }

  Slice key() const override {
    assert(valid_);
    return iter_->key();
  }

  Slice user_key() const override {
    assert(valid_);
    return iter_->user_key();
  }

  Slice value() const override {
    assert(valid_);
    return iter_->value();
  }

  Status status() const override { return iter_->status(); }

  bool PrepareValue() override {
    assert(valid_);

    if (iter_->PrepareValue()) {
      return true;
    }

    assert(!iter_->Valid());
    valid_ = false;
    return false;
  }

  bool MayBeOutOfLowerBound() override {
    assert(valid_);
    return false;
  }

  IterBoundCheck UpperBoundCheckResult() override {
    assert(valid_);
    return IterBoundCheck::kInbound;
  }

  void SetPinnedItersMgr(PinnedIteratorsManager* pinned_iters_mgr) override {
    iter_->SetPinnedItersMgr(pinned_iters_mgr);
  }

  bool IsKeyPinned() const override {
    assert(valid_);
    return iter_->IsKeyPinned();
  }

  bool IsValuePinned() const override {
    assert(valid_);
    return iter_->IsValuePinned();
  }

  Status GetProperty(std::string prop_name, std::string* prop) override {
    return iter_->GetProperty(prop_name, prop);
  }

  bool IsDeleteRangeSentinelKey() const override {
    assert(valid_);
    return iter_->IsDeleteRangeSentinelKey();
  }

 private:
  void UpdateValid() {
    assert(!iter_->Valid() || iter_->status().ok());

    valid_ = iter_->Valid();
  }

  void EnforceUpperBoundImpl(IterBoundCheck bound_check_result) {
    if (bound_check_result == IterBoundCheck::kInbound) {
      return;
    }

    if (bound_check_result == IterBoundCheck::kOutOfBound) {
      valid_ = false;
      return;
    }

    assert(bound_check_result == IterBoundCheck::kUnknown);

    if (cmp_->Compare(key(), *end_) >= 0) {
      valid_ = false;
    }
  }

  void EnforceUpperBound() {
    if (!valid_) {
      return;
    }

    if (!end_) {
      return;
    }

    EnforceUpperBoundImpl(iter_->UpperBoundCheckResult());
  }

  void EnforceLowerBound() {
    if (!valid_) {
      return;
    }

    if (!start_) {
      return;
    }

    if (!iter_->MayBeOutOfLowerBound()) {
      return;
    }

    if (cmp_->Compare(key(), *start_) < 0) {
      valid_ = false;
    }
  }

  void AssertBounds() {
    assert(!valid_ || !start_ || cmp_->Compare(key(), *start_) >= 0);
    assert(!valid_ || !end_ || cmp_->Compare(key(), *end_) < 0);
  }

  void UpdateAndEnforceBounds() {
    UpdateValid();
    EnforceUpperBound();
    EnforceLowerBound();
    AssertBounds();
  }

  void UpdateAndEnforceUpperBound() {
    UpdateValid();
    EnforceUpperBound();
    AssertBounds();
  }

  void UpdateAndEnforceLowerBound() {
    UpdateValid();
    EnforceLowerBound();
    AssertBounds();
  }

  InternalIterator* iter_;
  const Slice* start_;
  const Slice* end_;
  const CompareInterface* cmp_;
  bool valid_;
};

}  // namespace ROCKSDB_NAMESPACE