use std::{iter::FusedIterator, sync::Arc}; use conduit::Result; use rocksdb::{ColumnFamily, DBRawIteratorWithThreadMode, Direction, IteratorMode, ReadOptions}; use crate::{ engine::Db, result, slice::{OwnedKeyVal, OwnedKeyValPair}, Engine, }; type Cursor<'cursor> = DBRawIteratorWithThreadMode<'cursor, Db>; struct State<'cursor> { cursor: Cursor<'cursor>, direction: Direction, valid: bool, init: bool, } impl<'cursor> State<'cursor> { pub(crate) fn new( db: &'cursor Arc, cf: &'cursor Arc, opts: ReadOptions, mode: &IteratorMode<'_>, ) -> Self { let mut cursor = db.db.raw_iterator_cf_opt(&**cf, opts); let direction = into_direction(mode); let valid = seek_init(&mut cursor, mode); Self { cursor, direction, valid, init: true, } } } pub struct Iter<'cursor> { state: State<'cursor>, } impl<'cursor> Iter<'cursor> { pub(crate) fn new( db: &'cursor Arc, cf: &'cursor Arc, opts: ReadOptions, mode: &IteratorMode<'_>, ) -> Self { Self { state: State::new(db, cf, opts, mode), } } } impl Iterator for Iter<'_> { type Item = OwnedKeyValPair; fn next(&mut self) -> Option { if !self.state.init && self.state.valid { seek_next(&mut self.state.cursor, self.state.direction); } else if self.state.init { self.state.init = false; } self.state .cursor .item() .map(OwnedKeyVal::from) .map(OwnedKeyVal::to_tuple) .or_else(|| { when_invalid(&mut self.state).expect("iterator invalidated due to error"); None }) } } impl FusedIterator for Iter<'_> {} fn when_invalid(state: &mut State<'_>) -> Result<()> { state.valid = false; result(state.cursor.status()) } fn seek_next(cursor: &mut Cursor<'_>, direction: Direction) { match direction { Direction::Forward => cursor.next(), Direction::Reverse => cursor.prev(), } } fn seek_init(cursor: &mut Cursor<'_>, mode: &IteratorMode<'_>) -> bool { use Direction::{Forward, Reverse}; use IteratorMode::{End, From, Start}; match mode { Start => cursor.seek_to_first(), End => cursor.seek_to_last(), From(key, Forward) => cursor.seek(key), From(key, Reverse) => cursor.seek_for_prev(key), }; cursor.valid() } fn into_direction(mode: &IteratorMode<'_>) -> Direction { use Direction::{Forward, Reverse}; use IteratorMode::{End, From, Start}; match mode { Start | From(_, Forward) => Forward, End | From(_, Reverse) => Reverse, } }