use crate::{
db::{DBAccess, DB},
ffi, Error, ReadOptions, WriteBatch,
};
use libc::{c_char, c_uchar, size_t};
use std::{marker::PhantomData, slice};
pub type DBRawIterator<'a> = DBRawIteratorWithThreadMode<'a, DB>;
pub struct DBRawIteratorWithThreadMode<'a, D: DBAccess> {
inner: std::ptr::NonNull<ffi::rocksdb_iterator_t>,
_readopts: ReadOptions,
db: PhantomData<&'a D>,
}
impl<'a, D: DBAccess> DBRawIteratorWithThreadMode<'a, D> {
pub(crate) fn new(db: &D, readopts: ReadOptions) -> Self {
let inner = unsafe { db.create_iterator(&readopts) };
Self::from_inner(inner, readopts)
}
pub(crate) fn new_cf(
db: &'a D,
cf_handle: *mut ffi::rocksdb_column_family_handle_t,
readopts: ReadOptions,
) -> Self {
let inner = unsafe { db.create_iterator_cf(cf_handle, &readopts) };
Self::from_inner(inner, readopts)
}
fn from_inner(inner: *mut ffi::rocksdb_iterator_t, readopts: ReadOptions) -> Self {
let inner = std::ptr::NonNull::new(inner).unwrap();
Self {
inner,
_readopts: readopts,
db: PhantomData,
}
}
pub fn valid(&self) -> bool {
unsafe { ffi::rocksdb_iter_valid(self.inner.as_ptr()) != 0 }
}
pub fn status(&self) -> Result<(), Error> {
unsafe {
ffi_try!(ffi::rocksdb_iter_get_error(self.inner.as_ptr()));
}
Ok(())
}
pub fn seek_to_first(&mut self) {
unsafe {
ffi::rocksdb_iter_seek_to_first(self.inner.as_ptr());
}
}
pub fn seek_to_last(&mut self) {
unsafe {
ffi::rocksdb_iter_seek_to_last(self.inner.as_ptr());
}
}
pub fn seek<K: AsRef<[u8]>>(&mut self, key: K) {
let key = key.as_ref();
unsafe {
ffi::rocksdb_iter_seek(
self.inner.as_ptr(),
key.as_ptr() as *const c_char,
key.len() as size_t,
);
}
}
pub fn seek_for_prev<K: AsRef<[u8]>>(&mut self, key: K) {
let key = key.as_ref();
unsafe {
ffi::rocksdb_iter_seek_for_prev(
self.inner.as_ptr(),
key.as_ptr() as *const c_char,
key.len() as size_t,
);
}
}
pub fn next(&mut self) {
unsafe {
ffi::rocksdb_iter_next(self.inner.as_ptr());
}
}
pub fn prev(&mut self) {
unsafe {
ffi::rocksdb_iter_prev(self.inner.as_ptr());
}
}
pub fn key(&self) -> Option<&[u8]> {
if self.valid() {
Some(self.key_impl())
} else {
None
}
}
pub fn value(&self) -> Option<&[u8]> {
if self.valid() {
Some(self.value_impl())
} else {
None
}
}
pub fn item(&self) -> Option<(&[u8], &[u8])> {
if self.valid() {
Some((self.key_impl(), self.value_impl()))
} else {
None
}
}
fn key_impl(&self) -> &[u8] {
unsafe {
let mut key_len: size_t = 0;
let key_len_ptr: *mut size_t = &mut key_len;
let key_ptr = ffi::rocksdb_iter_key(self.inner.as_ptr(), key_len_ptr);
slice::from_raw_parts(key_ptr as *const c_uchar, key_len as usize)
}
}
fn value_impl(&self) -> &[u8] {
unsafe {
let mut val_len: size_t = 0;
let val_len_ptr: *mut size_t = &mut val_len;
let val_ptr = ffi::rocksdb_iter_value(self.inner.as_ptr(), val_len_ptr);
slice::from_raw_parts(val_ptr as *const c_uchar, val_len as usize)
}
}
}
impl<'a, D: DBAccess> Drop for DBRawIteratorWithThreadMode<'a, D> {
fn drop(&mut self) {
unsafe {
ffi::rocksdb_iter_destroy(self.inner.as_ptr());
}
}
}
unsafe impl<'a, D: DBAccess> Send for DBRawIteratorWithThreadMode<'a, D> {}
unsafe impl<'a, D: DBAccess> Sync for DBRawIteratorWithThreadMode<'a, D> {}
pub type DBIterator<'a> = DBIteratorWithThreadMode<'a, DB>;
pub struct DBIteratorWithThreadMode<'a, D: DBAccess> {
raw: DBRawIteratorWithThreadMode<'a, D>,
direction: Direction,
done: bool,
}
#[derive(Copy, Clone)]
pub enum Direction {
Forward,
Reverse,
}
pub type KVBytes = (Box<[u8]>, Box<[u8]>);
#[derive(Copy, Clone)]
pub enum IteratorMode<'a> {
Start,
End,
From(&'a [u8], Direction),
}
impl<'a, D: DBAccess> DBIteratorWithThreadMode<'a, D> {
pub(crate) fn new(db: &D, readopts: ReadOptions, mode: IteratorMode) -> Self {
Self::from_raw(DBRawIteratorWithThreadMode::new(db, readopts), mode)
}
pub(crate) fn new_cf(
db: &'a D,
cf_handle: *mut ffi::rocksdb_column_family_handle_t,
readopts: ReadOptions,
mode: IteratorMode,
) -> Self {
Self::from_raw(
DBRawIteratorWithThreadMode::new_cf(db, cf_handle, readopts),
mode,
)
}
fn from_raw(raw: DBRawIteratorWithThreadMode<'a, D>, mode: IteratorMode) -> Self {
let mut rv = DBIteratorWithThreadMode {
raw,
direction: Direction::Forward, done: false,
};
rv.set_mode(mode);
rv
}
pub fn set_mode(&mut self, mode: IteratorMode) {
self.done = false;
self.direction = match mode {
IteratorMode::Start => {
self.raw.seek_to_first();
Direction::Forward
}
IteratorMode::End => {
self.raw.seek_to_last();
Direction::Reverse
}
IteratorMode::From(key, Direction::Forward) => {
self.raw.seek(key);
Direction::Forward
}
IteratorMode::From(key, Direction::Reverse) => {
self.raw.seek_for_prev(key);
Direction::Reverse
}
};
}
}
impl<'a, D: DBAccess> Iterator for DBIteratorWithThreadMode<'a, D> {
type Item = Result<KVBytes, Error>;
fn next(&mut self) -> Option<Result<KVBytes, Error>> {
if self.done {
None
} else if let Some((key, value)) = self.raw.item() {
let item = (Box::from(key), Box::from(value));
match self.direction {
Direction::Forward => self.raw.next(),
Direction::Reverse => self.raw.prev(),
}
Some(Ok(item))
} else {
self.done = true;
self.raw.status().err().map(Result::Err)
}
}
}
impl<'a, D: DBAccess> std::iter::FusedIterator for DBIteratorWithThreadMode<'a, D> {}
impl<'a, D: DBAccess> Into<DBRawIteratorWithThreadMode<'a, D>> for DBIteratorWithThreadMode<'a, D> {
fn into(self) -> DBRawIteratorWithThreadMode<'a, D> {
self.raw
}
}
pub struct DBWALIterator {
pub(crate) inner: *mut ffi::rocksdb_wal_iterator_t,
}
impl DBWALIterator {
pub fn valid(&self) -> bool {
unsafe { ffi::rocksdb_wal_iter_valid(self.inner) != 0 }
}
pub fn status(&self) -> Result<(), Error> {
unsafe {
ffi_try!(ffi::rocksdb_wal_iter_status(self.inner));
}
Ok(())
}
}
impl Iterator for DBWALIterator {
type Item = Result<(u64, WriteBatch), Error>;
fn next(&mut self) -> Option<Self::Item> {
if !self.valid() {
return None;
}
unsafe {
ffi::rocksdb_wal_iter_next(self.inner);
}
if self.valid() {
let mut seq: u64 = 0;
let inner = unsafe { ffi::rocksdb_wal_iter_get_batch(self.inner, &mut seq) };
Some(Ok((seq, WriteBatch { inner })))
} else {
self.status().err().map(Result::Err)
}
}
}
impl Drop for DBWALIterator {
fn drop(&mut self) {
unsafe {
ffi::rocksdb_wal_iter_destroy(self.inner);
}
}
}