use crate::database::{
Error as DatabaseError,
Result as DatabaseResult,
};
use anyhow::anyhow;
use fuel_core_storage::iter::IterDirection;
use rocksdb::{
DBAccess,
DBRawIteratorWithThreadMode,
IteratorMode,
};
pub struct RocksDBKeyIterator<'a, D: DBAccess, R> {
raw: DBRawIteratorWithThreadMode<'a, D>,
direction: IterDirection,
done: bool,
_marker: core::marker::PhantomData<R>,
}
pub trait ExtractItem: Send + Sync + 'static {
type Item: Send + Sync;
fn extract_item<D>(
raw_iterator: &DBRawIteratorWithThreadMode<D>,
) -> Option<Self::Item>
where
D: DBAccess;
fn size(item: &Self::Item) -> u64;
fn starts_with(item: &Self::Item, prefix: &[u8]) -> bool;
}
impl<D: DBAccess, R> Iterator for RocksDBKeyIterator<'_, D, R>
where
R: ExtractItem,
{
type Item = DatabaseResult<R::Item>;
fn next(&mut self) -> Option<Self::Item> {
if self.done {
None
} else {
match R::extract_item(&self.raw) {
Some(item) => {
match self.direction {
IterDirection::Forward => self.raw.next(),
IterDirection::Reverse => self.raw.prev(),
}
Some(DatabaseResult::Ok(item))
}
_ => {
self.done = true;
self.raw.status().err().map(|e| {
Some(DatabaseResult::Err(DatabaseError::Other(anyhow!(e))))
})?
}
}
}
}
}
impl<'a, D: DBAccess, R> RocksDBKeyIterator<'a, D, R> {
fn set_mode(&mut self, mode: IteratorMode) {
self.done = false;
self.direction = match mode {
IteratorMode::Start => {
self.raw.seek_to_first();
IterDirection::Forward
}
IteratorMode::End => {
self.raw.seek_to_last();
IterDirection::Reverse
}
IteratorMode::From(key, rocksdb::Direction::Forward) => {
self.raw.seek(key);
IterDirection::Forward
}
IteratorMode::From(key, rocksdb::Direction::Reverse) => {
self.raw.seek_for_prev(key);
IterDirection::Reverse
}
};
}
pub fn new(
raw: DBRawIteratorWithThreadMode<'a, D>,
mode: IteratorMode,
) -> RocksDBKeyIterator<'a, D, R> {
let mut raw_iterator = Self {
raw,
direction: IterDirection::Forward, done: false,
_marker: Default::default(),
};
raw_iterator.set_mode(mode);
raw_iterator
}
}