use std::marker::PhantomData;
use std::{
fmt,
mem,
ptr,
result,
slice,
};
use libc::{
c_uint,
c_void,
size_t,
EINVAL,
};
use database::Database;
use error::{
lmdb_result,
Error,
Result,
};
use ffi;
use flags::WriteFlags;
use transaction::Transaction;
pub trait Cursor<'txn> {
fn cursor(&self) -> *mut ffi::MDB_cursor;
fn get(&self, key: Option<&[u8]>, data: Option<&[u8]>, op: c_uint) -> Result<(Option<&'txn [u8]>, &'txn [u8])> {
unsafe {
let mut key_val = slice_to_val(key);
let mut data_val = slice_to_val(data);
let key_ptr = key_val.mv_data;
lmdb_result(ffi::mdb_cursor_get(self.cursor(), &mut key_val, &mut data_val, op))?;
let key_out = if key_ptr != key_val.mv_data {
Some(val_to_slice(key_val))
} else {
None
};
let data_out = val_to_slice(data_val);
Ok((key_out, data_out))
}
}
fn iter(&mut self) -> Iter<'txn> {
Iter::new(self.cursor(), ffi::MDB_NEXT, ffi::MDB_NEXT)
}
fn iter_start(&mut self) -> Iter<'txn> {
Iter::new(self.cursor(), ffi::MDB_FIRST, ffi::MDB_NEXT)
}
fn iter_from<K>(&mut self, key: K) -> Iter<'txn>
where
K: AsRef<[u8]>,
{
match self.get(Some(key.as_ref()), None, ffi::MDB_SET_RANGE) {
Ok(_) | Err(Error::NotFound) => (),
Err(error) => return Iter::Err(error),
};
Iter::new(self.cursor(), ffi::MDB_GET_CURRENT, ffi::MDB_NEXT)
}
fn iter_dup(&mut self) -> IterDup<'txn> {
IterDup::new(self.cursor(), ffi::MDB_NEXT)
}
fn iter_dup_start(&mut self) -> IterDup<'txn> {
IterDup::new(self.cursor(), ffi::MDB_FIRST)
}
fn iter_dup_from<K>(&mut self, key: K) -> IterDup<'txn>
where
K: AsRef<[u8]>,
{
match self.get(Some(key.as_ref()), None, ffi::MDB_SET_RANGE) {
Ok(_) | Err(Error::NotFound) => (),
Err(error) => return IterDup::Err(error),
};
IterDup::new(self.cursor(), ffi::MDB_GET_CURRENT)
}
fn iter_dup_of<K>(&mut self, key: K) -> Iter<'txn>
where
K: AsRef<[u8]>,
{
match self.get(Some(key.as_ref()), None, ffi::MDB_SET) {
Ok(_) | Err(Error::NotFound) => (),
Err(error) => return Iter::Err(error),
};
Iter::new(self.cursor(), ffi::MDB_GET_CURRENT, ffi::MDB_NEXT_DUP)
}
}
pub struct RoCursor<'txn> {
cursor: *mut ffi::MDB_cursor,
_marker: PhantomData<fn() -> &'txn ()>,
}
impl<'txn> Cursor<'txn> for RoCursor<'txn> {
fn cursor(&self) -> *mut ffi::MDB_cursor {
self.cursor
}
}
impl<'txn> fmt::Debug for RoCursor<'txn> {
fn fmt(&self, f: &mut fmt::Formatter) -> result::Result<(), fmt::Error> {
f.debug_struct("RoCursor").finish()
}
}
impl<'txn> Drop for RoCursor<'txn> {
fn drop(&mut self) {
unsafe { ffi::mdb_cursor_close(self.cursor) }
}
}
impl<'txn> RoCursor<'txn> {
pub(crate) fn new<T>(txn: &'txn T, db: Database) -> Result<RoCursor<'txn>>
where
T: Transaction,
{
let mut cursor: *mut ffi::MDB_cursor = ptr::null_mut();
unsafe {
lmdb_result(ffi::mdb_cursor_open(txn.txn(), db.dbi(), &mut cursor))?;
}
Ok(RoCursor {
cursor,
_marker: PhantomData,
})
}
}
pub struct RwCursor<'txn> {
cursor: *mut ffi::MDB_cursor,
_marker: PhantomData<fn() -> &'txn ()>,
}
impl<'txn> Cursor<'txn> for RwCursor<'txn> {
fn cursor(&self) -> *mut ffi::MDB_cursor {
self.cursor
}
}
impl<'txn> fmt::Debug for RwCursor<'txn> {
fn fmt(&self, f: &mut fmt::Formatter) -> result::Result<(), fmt::Error> {
f.debug_struct("RwCursor").finish()
}
}
impl<'txn> Drop for RwCursor<'txn> {
fn drop(&mut self) {
unsafe { ffi::mdb_cursor_close(self.cursor) }
}
}
impl<'txn> RwCursor<'txn> {
pub(crate) fn new<T>(txn: &'txn T, db: Database) -> Result<RwCursor<'txn>>
where
T: Transaction,
{
let mut cursor: *mut ffi::MDB_cursor = ptr::null_mut();
unsafe {
lmdb_result(ffi::mdb_cursor_open(txn.txn(), db.dbi(), &mut cursor))?;
}
Ok(RwCursor {
cursor,
_marker: PhantomData,
})
}
pub fn put<K, D>(&mut self, key: &K, data: &D, flags: WriteFlags) -> Result<()>
where
K: AsRef<[u8]>,
D: AsRef<[u8]>,
{
let key = key.as_ref();
let data = data.as_ref();
let mut key_val: ffi::MDB_val = ffi::MDB_val {
mv_size: key.len() as size_t,
mv_data: key.as_ptr() as *mut c_void,
};
let mut data_val: ffi::MDB_val = ffi::MDB_val {
mv_size: data.len() as size_t,
mv_data: data.as_ptr() as *mut c_void,
};
unsafe { lmdb_result(ffi::mdb_cursor_put(self.cursor(), &mut key_val, &mut data_val, flags.bits())) }
}
pub fn del(&mut self, flags: WriteFlags) -> Result<()> {
unsafe { lmdb_result(ffi::mdb_cursor_del(self.cursor(), flags.bits())) }
}
}
unsafe fn slice_to_val(slice: Option<&[u8]>) -> ffi::MDB_val {
match slice {
Some(slice) => ffi::MDB_val {
mv_size: slice.len() as size_t,
mv_data: slice.as_ptr() as *mut c_void,
},
None => ffi::MDB_val {
mv_size: 0,
mv_data: ptr::null_mut(),
},
}
}
unsafe fn val_to_slice<'a>(val: ffi::MDB_val) -> &'a [u8] {
slice::from_raw_parts(val.mv_data as *const u8, val.mv_size as usize)
}
pub enum Iter<'txn> {
Err(Error),
Ok {
cursor: *mut ffi::MDB_cursor,
op: c_uint,
next_op: c_uint,
_marker: PhantomData<fn(&'txn ())>,
},
}
impl<'txn> Iter<'txn> {
fn new<'t>(cursor: *mut ffi::MDB_cursor, op: c_uint, next_op: c_uint) -> Iter<'t> {
Iter::Ok {
cursor,
op,
next_op,
_marker: PhantomData,
}
}
}
impl<'txn> fmt::Debug for Iter<'txn> {
fn fmt(&self, f: &mut fmt::Formatter) -> result::Result<(), fmt::Error> {
f.debug_struct("Iter").finish()
}
}
impl<'txn> Iterator for Iter<'txn> {
type Item = Result<(&'txn [u8], &'txn [u8])>;
fn next(&mut self) -> Option<Result<(&'txn [u8], &'txn [u8])>> {
match self {
&mut Iter::Ok {
cursor,
ref mut op,
next_op,
_marker,
} => {
let mut key = ffi::MDB_val {
mv_size: 0,
mv_data: ptr::null_mut(),
};
let mut data = ffi::MDB_val {
mv_size: 0,
mv_data: ptr::null_mut(),
};
let op = mem::replace(op, next_op);
unsafe {
match ffi::mdb_cursor_get(cursor, &mut key, &mut data, op) {
ffi::MDB_SUCCESS => Some(Ok((val_to_slice(key), val_to_slice(data)))),
ffi::MDB_NOTFOUND | EINVAL => None,
error => Some(Err(Error::from_err_code(error))),
}
}
},
&mut Iter::Err(err) => Some(Err(err)),
}
}
}
pub enum IterDup<'txn> {
Err(Error),
Ok {
cursor: *mut ffi::MDB_cursor,
op: c_uint,
_marker: PhantomData<fn(&'txn ())>,
},
}
impl<'txn> IterDup<'txn> {
fn new<'t>(cursor: *mut ffi::MDB_cursor, op: c_uint) -> IterDup<'t> {
IterDup::Ok {
cursor,
op,
_marker: PhantomData,
}
}
}
impl<'txn> fmt::Debug for IterDup<'txn> {
fn fmt(&self, f: &mut fmt::Formatter) -> result::Result<(), fmt::Error> {
f.debug_struct("IterDup").finish()
}
}
impl<'txn> Iterator for IterDup<'txn> {
type Item = Iter<'txn>;
fn next(&mut self) -> Option<Iter<'txn>> {
match self {
&mut IterDup::Ok {
cursor,
ref mut op,
_marker,
} => {
let mut key = ffi::MDB_val {
mv_size: 0,
mv_data: ptr::null_mut(),
};
let mut data = ffi::MDB_val {
mv_size: 0,
mv_data: ptr::null_mut(),
};
let op = mem::replace(op, ffi::MDB_NEXT_NODUP);
let err_code = unsafe { ffi::mdb_cursor_get(cursor, &mut key, &mut data, op) };
if err_code == ffi::MDB_SUCCESS {
Some(Iter::new(cursor, ffi::MDB_GET_CURRENT, ffi::MDB_NEXT_DUP))
} else {
None
}
},
&mut IterDup::Err(err) => Some(Iter::Err(err)),
}
}
}
#[cfg(test)]
mod test {
use tempdir::TempDir;
use super::*;
use environment::*;
use ffi::*;
use flags::*;
#[test]
fn test_get() {
let dir = TempDir::new("test").unwrap();
let env = Environment::new().open(dir.path()).unwrap();
let db = env.open_db(None).unwrap();
let mut txn = env.begin_rw_txn().unwrap();
txn.put(db, b"key1", b"val1", WriteFlags::empty()).unwrap();
txn.put(db, b"key2", b"val2", WriteFlags::empty()).unwrap();
txn.put(db, b"key3", b"val3", WriteFlags::empty()).unwrap();
let cursor = txn.open_ro_cursor(db).unwrap();
assert_eq!((Some(&b"key1"[..]), &b"val1"[..]), cursor.get(None, None, MDB_FIRST).unwrap());
assert_eq!((Some(&b"key1"[..]), &b"val1"[..]), cursor.get(None, None, MDB_GET_CURRENT).unwrap());
assert_eq!((Some(&b"key2"[..]), &b"val2"[..]), cursor.get(None, None, MDB_NEXT).unwrap());
assert_eq!((Some(&b"key1"[..]), &b"val1"[..]), cursor.get(None, None, MDB_PREV).unwrap());
assert_eq!((Some(&b"key3"[..]), &b"val3"[..]), cursor.get(None, None, MDB_LAST).unwrap());
assert_eq!((None, &b"val2"[..]), cursor.get(Some(b"key2"), None, MDB_SET).unwrap());
assert_eq!((Some(&b"key3"[..]), &b"val3"[..]), cursor.get(Some(&b"key3"[..]), None, MDB_SET_KEY).unwrap());
assert_eq!((Some(&b"key3"[..]), &b"val3"[..]), cursor.get(Some(&b"key2\0"[..]), None, MDB_SET_RANGE).unwrap());
}
#[test]
fn test_get_dup() {
let dir = TempDir::new("test").unwrap();
let env = Environment::new().open(dir.path()).unwrap();
let db = env.create_db(None, DatabaseFlags::DUP_SORT).unwrap();
let mut txn = env.begin_rw_txn().unwrap();
txn.put(db, b"key1", b"val1", WriteFlags::empty()).unwrap();
txn.put(db, b"key1", b"val2", WriteFlags::empty()).unwrap();
txn.put(db, b"key1", b"val3", WriteFlags::empty()).unwrap();
txn.put(db, b"key2", b"val1", WriteFlags::empty()).unwrap();
txn.put(db, b"key2", b"val2", WriteFlags::empty()).unwrap();
txn.put(db, b"key2", b"val3", WriteFlags::empty()).unwrap();
let cursor = txn.open_ro_cursor(db).unwrap();
assert_eq!((Some(&b"key1"[..]), &b"val1"[..]), cursor.get(None, None, MDB_FIRST).unwrap());
assert_eq!((None, &b"val1"[..]), cursor.get(None, None, MDB_FIRST_DUP).unwrap());
assert_eq!((Some(&b"key1"[..]), &b"val1"[..]), cursor.get(None, None, MDB_GET_CURRENT).unwrap());
assert_eq!((Some(&b"key2"[..]), &b"val1"[..]), cursor.get(None, None, MDB_NEXT_NODUP).unwrap());
assert_eq!((Some(&b"key2"[..]), &b"val2"[..]), cursor.get(None, None, MDB_NEXT_DUP).unwrap());
assert_eq!((Some(&b"key2"[..]), &b"val3"[..]), cursor.get(None, None, MDB_NEXT_DUP).unwrap());
assert!(cursor.get(None, None, MDB_NEXT_DUP).is_err());
assert_eq!((Some(&b"key2"[..]), &b"val2"[..]), cursor.get(None, None, MDB_PREV_DUP).unwrap());
assert_eq!((None, &b"val3"[..]), cursor.get(None, None, MDB_LAST_DUP).unwrap());
assert_eq!((Some(&b"key1"[..]), &b"val3"[..]), cursor.get(None, None, MDB_PREV_NODUP).unwrap());
assert_eq!((None, &b"val1"[..]), cursor.get(Some(&b"key1"[..]), None, MDB_SET).unwrap());
assert_eq!((Some(&b"key2"[..]), &b"val1"[..]), cursor.get(Some(&b"key2"[..]), None, MDB_SET_KEY).unwrap());
assert_eq!((Some(&b"key2"[..]), &b"val1"[..]), cursor.get(Some(&b"key1\0"[..]), None, MDB_SET_RANGE).unwrap());
assert_eq!((None, &b"val3"[..]), cursor.get(Some(&b"key1"[..]), Some(&b"val3"[..]), MDB_GET_BOTH).unwrap());
assert_eq!(
(None, &b"val1"[..]),
cursor.get(Some(&b"key2"[..]), Some(&b"val"[..]), MDB_GET_BOTH_RANGE).unwrap()
);
}
#[test]
fn test_get_dupfixed() {
let dir = TempDir::new("test").unwrap();
let env = Environment::new().open(dir.path()).unwrap();
let db = env.create_db(None, DatabaseFlags::DUP_SORT | DatabaseFlags::DUP_FIXED).unwrap();
let mut txn = env.begin_rw_txn().unwrap();
txn.put(db, b"key1", b"val1", WriteFlags::empty()).unwrap();
txn.put(db, b"key1", b"val2", WriteFlags::empty()).unwrap();
txn.put(db, b"key1", b"val3", WriteFlags::empty()).unwrap();
txn.put(db, b"key2", b"val4", WriteFlags::empty()).unwrap();
txn.put(db, b"key2", b"val5", WriteFlags::empty()).unwrap();
txn.put(db, b"key2", b"val6", WriteFlags::empty()).unwrap();
let cursor = txn.open_ro_cursor(db).unwrap();
assert_eq!((Some(&b"key1"[..]), &b"val1"[..]), cursor.get(None, None, MDB_FIRST).unwrap());
assert_eq!((None, &b"val1val2val3"[..]), cursor.get(None, None, MDB_GET_MULTIPLE).unwrap());
assert!(cursor.get(None, None, MDB_NEXT_MULTIPLE).is_err());
}
#[test]
fn test_iter() {
let dir = TempDir::new("test").unwrap();
let env = Environment::new().open(dir.path()).unwrap();
let db = env.open_db(None).unwrap();
let items: Vec<(&[u8], &[u8])> =
vec![(b"key1", b"val1"), (b"key2", b"val2"), (b"key3", b"val3"), (b"key5", b"val5")];
{
let mut txn = env.begin_rw_txn().unwrap();
for &(ref key, ref data) in &items {
txn.put(db, key, data, WriteFlags::empty()).unwrap();
}
txn.commit().unwrap();
}
let txn = env.begin_ro_txn().unwrap();
let mut cursor = txn.open_ro_cursor(db).unwrap();
assert_eq!(items, cursor.iter().collect::<Result<Vec<_>>>().unwrap());
let retr: Result<Vec<_>> = cursor.iter_start().collect();
assert_eq!(items, retr.unwrap());
cursor.get(Some(b"key2"), None, MDB_SET).unwrap();
assert_eq!(
items.clone().into_iter().skip(2).collect::<Vec<_>>(),
cursor.iter().collect::<Result<Vec<_>>>().unwrap()
);
assert_eq!(items, cursor.iter_start().collect::<Result<Vec<_>>>().unwrap());
assert_eq!(
items.clone().into_iter().skip(1).collect::<Vec<_>>(),
cursor.iter_from(b"key2").collect::<Result<Vec<_>>>().unwrap()
);
assert_eq!(
items.clone().into_iter().skip(3).collect::<Vec<_>>(),
cursor.iter_from(b"key4").collect::<Result<Vec<_>>>().unwrap()
);
assert_eq!(
vec!().into_iter().collect::<Vec<(&[u8], &[u8])>>(),
cursor.iter_from(b"key6").collect::<Result<Vec<_>>>().unwrap()
);
}
#[test]
fn test_iter_empty_database() {
let dir = TempDir::new("test").unwrap();
let env = Environment::new().open(dir.path()).unwrap();
let db = env.open_db(None).unwrap();
let txn = env.begin_ro_txn().unwrap();
let mut cursor = txn.open_ro_cursor(db).unwrap();
assert_eq!(0, cursor.iter().count());
assert_eq!(0, cursor.iter_start().count());
assert_eq!(0, cursor.iter_from(b"foo").count());
}
#[test]
fn test_iter_empty_dup_database() {
let dir = TempDir::new("test").unwrap();
let env = Environment::new().open(dir.path()).unwrap();
let db = env.create_db(None, DatabaseFlags::DUP_SORT).unwrap();
let txn = env.begin_ro_txn().unwrap();
let mut cursor = txn.open_ro_cursor(db).unwrap();
assert_eq!(0, cursor.iter().count());
assert_eq!(0, cursor.iter_start().count());
assert_eq!(0, cursor.iter_from(b"foo").count());
assert_eq!(0, cursor.iter_dup().count());
assert_eq!(0, cursor.iter_dup_start().count());
assert_eq!(0, cursor.iter_dup_from(b"foo").count());
assert_eq!(0, cursor.iter_dup_of(b"foo").count());
}
#[test]
fn test_iter_dup() {
let dir = TempDir::new("test").unwrap();
let env = Environment::new().open(dir.path()).unwrap();
let db = env.create_db(None, DatabaseFlags::DUP_SORT).unwrap();
let items: Vec<(&[u8], &[u8])> = vec![
(b"a", b"1"),
(b"a", b"2"),
(b"a", b"3"),
(b"b", b"1"),
(b"b", b"2"),
(b"b", b"3"),
(b"c", b"1"),
(b"c", b"2"),
(b"c", b"3"),
(b"e", b"1"),
(b"e", b"2"),
(b"e", b"3"),
];
{
let mut txn = env.begin_rw_txn().unwrap();
for &(ref key, ref data) in &items {
txn.put(db, key, data, WriteFlags::empty()).unwrap();
}
txn.commit().unwrap();
}
let txn = env.begin_ro_txn().unwrap();
let mut cursor = txn.open_ro_cursor(db).unwrap();
assert_eq!(items, cursor.iter_dup().flatten().collect::<Result<Vec<_>>>().unwrap());
cursor.get(Some(b"b"), None, MDB_SET).unwrap();
assert_eq!(
items.clone().into_iter().skip(4).collect::<Vec<(&[u8], &[u8])>>(),
cursor.iter_dup().flatten().collect::<Result<Vec<_>>>().unwrap()
);
assert_eq!(items, cursor.iter_dup_start().flatten().collect::<Result<Vec<(&[u8], &[u8])>>>().unwrap());
assert_eq!(
items.clone().into_iter().skip(3).collect::<Vec<(&[u8], &[u8])>>(),
cursor.iter_dup_from(b"b").flatten().collect::<Result<Vec<_>>>().unwrap()
);
assert_eq!(
items.clone().into_iter().skip(3).collect::<Vec<(&[u8], &[u8])>>(),
cursor.iter_dup_from(b"ab").flatten().collect::<Result<Vec<_>>>().unwrap()
);
assert_eq!(
items.clone().into_iter().skip(9).collect::<Vec<(&[u8], &[u8])>>(),
cursor.iter_dup_from(b"d").flatten().collect::<Result<Vec<_>>>().unwrap()
);
assert_eq!(
vec!().into_iter().collect::<Vec<(&[u8], &[u8])>>(),
cursor.iter_dup_from(b"f").flatten().collect::<Result<Vec<_>>>().unwrap()
);
assert_eq!(
items.clone().into_iter().skip(3).take(3).collect::<Vec<(&[u8], &[u8])>>(),
cursor.iter_dup_of(b"b").collect::<Result<Vec<_>>>().unwrap()
);
assert_eq!(0, cursor.iter_dup_of(b"foo").count());
}
#[test]
fn test_put_del() {
let dir = TempDir::new("test").unwrap();
let env = Environment::new().open(dir.path()).unwrap();
let db = env.open_db(None).unwrap();
let mut txn = env.begin_rw_txn().unwrap();
let mut cursor = txn.open_rw_cursor(db).unwrap();
cursor.put(b"key1", b"val1", WriteFlags::empty()).unwrap();
cursor.put(b"key2", b"val2", WriteFlags::empty()).unwrap();
cursor.put(b"key3", b"val3", WriteFlags::empty()).unwrap();
assert_eq!((Some(&b"key3"[..]), &b"val3"[..]), cursor.get(None, None, MDB_GET_CURRENT).unwrap());
cursor.del(WriteFlags::empty()).unwrap();
assert_eq!((Some(&b"key2"[..]), &b"val2"[..]), cursor.get(None, None, MDB_LAST).unwrap());
}
}