More idiomatic iterators

master
Pete Hunt 9 years ago
parent f5d4d834b7
commit 09250315e8
  1. 2
      src/lib.rs
  2. 137
      src/rocksdb.rs
  3. 38
      test/test_iterator.rs

@ -15,7 +15,7 @@
//
pub use ffi as rocksdb_ffi;
pub use ffi::{DBCompactionStyle, DBComparator, new_bloom_filter};
pub use rocksdb::{DB, DBIterator, DBVector, Direction, SubDBIterator, Writable, WriteBatch};
pub use rocksdb::{DB, DBIterator, DBVector, Direction, Writable, WriteBatch, IteratorMode};
pub use rocksdb_options::{BlockBasedOptions, Options};
pub use merge_operator::MergeOperands;
pub mod rocksdb;

@ -51,9 +51,8 @@ pub struct Snapshot<'a> {
inner: rocksdb_ffi::DBSnapshot,
}
pub struct DBIterator {
// TODO: should have a reference to DB to enforce scope, but it's trickier than I
// thought to add
pub struct DBIterator<'a> {
db: &'a DB,
inner: rocksdb_ffi::DBIterator,
direction: Direction,
just_seeked: bool,
@ -64,23 +63,18 @@ pub enum Direction {
reverse,
}
pub struct SubDBIterator<'a> {
iter: &'a mut DBIterator,
direction: Direction,
}
impl <'a> Iterator for SubDBIterator<'a> {
impl<'a> Iterator for DBIterator<'a> {
type Item = (Box<[u8]>, Box<[u8]>);
fn next(&mut self) -> Option<(Box<[u8]>, Box<[u8]>)> {
let native_iter = self.iter.inner;
if !self.iter.just_seeked {
let native_iter = self.inner;
if !self.just_seeked {
match self.direction {
Direction::forward => unsafe { rocksdb_ffi::rocksdb_iter_next(native_iter) },
Direction::reverse => unsafe { rocksdb_ffi::rocksdb_iter_prev(native_iter) },
}
} else {
self.iter.just_seeked = false;
self.just_seeked = false;
}
if unsafe { rocksdb_ffi::rocksdb_iter_valid(native_iter) } {
let mut key_len: size_t = 0;
@ -108,77 +102,80 @@ impl <'a> Iterator for SubDBIterator<'a> {
}
}
impl DBIterator {
// TODO alias db & opts to different lifetimes, and DBIterator to the db's
// lifetime
fn new(db: &DB, readopts: &ReadOptions) -> DBIterator {
pub enum IteratorMode<'a> {
Start,
End,
From(&'a [u8], Direction),
}
impl<'a> DBIterator<'a> {
fn new<'b>(db: &'a DB, readopts: &'b ReadOptions, mode: IteratorMode) -> DBIterator<'a> {
unsafe {
let iterator = rocksdb_ffi::rocksdb_create_iterator(db.inner,
readopts.inner);
rocksdb_ffi::rocksdb_iter_seek_to_first(iterator);
DBIterator {
let mut rv = DBIterator {
db: db,
inner: iterator,
direction: Direction::forward,
just_seeked: true,
direction: Direction::forward, // blown away by set_mode()
just_seeked: false,
};
rv.set_mode(mode);
rv
}
}
pub fn set_mode(&mut self, mode: IteratorMode) {
unsafe {
match mode {
IteratorMode::Start => {
rocksdb_ffi::rocksdb_iter_seek_to_first(self.inner);
self.direction = Direction::forward;
},
IteratorMode::End => {
rocksdb_ffi::rocksdb_iter_seek_to_last(self.inner);
self.direction = Direction::reverse;
},
IteratorMode::From(key, dir) => {
rocksdb_ffi::rocksdb_iter_seek(self.inner,
key.as_ptr(),
key.len() as size_t);
self.direction = dir;
}
};
self.just_seeked = true;
}
}
fn new_cf(db: &DB,
fn new_cf(db: &'a DB,
cf_handle: DBCFHandle,
readopts: &ReadOptions)
-> Result<DBIterator, String> {
readopts: &ReadOptions,
mode: IteratorMode)
-> Result<DBIterator<'a>, String> {
unsafe {
let iterator =
rocksdb_ffi::rocksdb_create_iterator_cf(db.inner,
readopts.inner,
cf_handle);
rocksdb_ffi::rocksdb_iter_seek_to_first(iterator);
Ok(DBIterator {
inner: iterator,
direction: Direction::forward,
just_seeked: true,
})
}
}
pub fn from_start(&mut self) -> SubDBIterator {
self.just_seeked = true;
unsafe {
rocksdb_ffi::rocksdb_iter_seek_to_first(self.inner);
}
SubDBIterator {
iter: self,
direction: Direction::forward,
}
}
let mut rv = DBIterator {
db: db,
inner: iterator,
direction: Direction::forward, // blown away by set_mode()
just_seeked: false,
};
pub fn from_end(&mut self) -> SubDBIterator {
self.just_seeked = true;
unsafe {
rocksdb_ffi::rocksdb_iter_seek_to_last(self.inner);
}
SubDBIterator {
iter: self,
direction: Direction::reverse,
}
}
rv.set_mode(mode);
pub fn from(&mut self, key: &[u8], dir: Direction) -> SubDBIterator {
self.just_seeked = true;
unsafe {
rocksdb_ffi::rocksdb_iter_seek(self.inner,
key.as_ptr(),
key.len() as size_t);
}
SubDBIterator {
iter: self,
direction: dir,
Ok(rv)
}
}
}
impl Drop for DBIterator {
impl<'a> Drop for DBIterator<'a> {
fn drop(&mut self) {
unsafe {
rocksdb_ffi::rocksdb_iter_destroy(self.inner);
@ -197,10 +194,10 @@ impl <'a> Snapshot<'a> {
}
}
pub fn iterator(&self) -> DBIterator {
pub fn iterator(&self, mode: IteratorMode) -> DBIterator {
let mut readopts = ReadOptions::new();
readopts.set_snapshot(self);
DBIterator::new(self.db, &readopts)
DBIterator::new(self.db, &readopts, mode)
}
}
@ -512,14 +509,14 @@ impl DB {
self.cfs.get(name)
}
pub fn iterator(&self) -> DBIterator {
pub fn iterator(&self, mode: IteratorMode) -> DBIterator {
let opts = ReadOptions::new();
DBIterator::new(&self, &opts)
DBIterator::new(&self, &opts, mode)
}
pub fn iterator_cf(&self, cf_handle: DBCFHandle) -> Result<DBIterator, String> {
pub fn iterator_cf(&self, cf_handle: DBCFHandle, mode: IteratorMode) -> Result<DBIterator, String> {
let opts = ReadOptions::new();
DBIterator::new_cf(&self, cf_handle, &opts)
DBIterator::new_cf(&self, cf_handle, &opts, mode)
}
pub fn snapshot(&self) -> Snapshot {
@ -891,8 +888,8 @@ fn iterator_test() {
assert!(p.is_ok());
let p = db.put(b"k3", b"v3333");
assert!(p.is_ok());
let mut iter = db.iterator();
for (k, v) in iter.from_start() {
let mut iter = db.iterator(IteratorMode::Start);
for (k, v) in iter {
println!("Hello {}: {}",
from_utf8(&*k).unwrap(),
from_utf8(&*v).unwrap());

@ -1,4 +1,4 @@
use rocksdb::{DB, Direction, Options, Writable};
use rocksdb::{DB, Direction, Options, Writable, IteratorMode};
fn cba(input: &Box<[u8]>) -> Box<[u8]> {
input.iter().cloned().collect::<Vec<_>>().into_boxed_slice()
@ -23,92 +23,90 @@ pub fn test_iterator() {
assert!(p.is_ok());
let p = db.put(&*k3, &*v3);
assert!(p.is_ok());
let mut view1 = db.iterator();
let expected = vec![(cba(&k1), cba(&v1)),
(cba(&k2), cba(&v2)),
(cba(&k3), cba(&v3))];
{
let iterator1 = view1.from_start();
let iterator1 = db.iterator(IteratorMode::Start);
assert_eq!(iterator1.collect::<Vec<_>>(), expected);
}
// Test that it's reusable a few times
// Test that it's idempotent
{
let iterator1 = view1.from_start();
let iterator1 = db.iterator(IteratorMode::Start);
assert_eq!(iterator1.collect::<Vec<_>>(), expected);
}
{
let iterator1 = view1.from_start();
let iterator1 = db.iterator(IteratorMode::Start);
assert_eq!(iterator1.collect::<Vec<_>>(), expected);
}
{
let iterator1 = view1.from_start();
let iterator1 = db.iterator(IteratorMode::Start);
assert_eq!(iterator1.collect::<Vec<_>>(), expected);
}
// Test it in reverse a few times
{
let iterator1 = view1.from_end();
let iterator1 = db.iterator(IteratorMode::End);
let mut tmp_vec = iterator1.collect::<Vec<_>>();
tmp_vec.reverse();
assert_eq!(tmp_vec, expected);
}
{
let iterator1 = view1.from_end();
let iterator1 = db.iterator(IteratorMode::End);
let mut tmp_vec = iterator1.collect::<Vec<_>>();
tmp_vec.reverse();
assert_eq!(tmp_vec, expected);
}
{
let iterator1 = view1.from_end();
let iterator1 = db.iterator(IteratorMode::End);
let mut tmp_vec = iterator1.collect::<Vec<_>>();
tmp_vec.reverse();
assert_eq!(tmp_vec, expected);
}
{
let iterator1 = view1.from_end();
let iterator1 = db.iterator(IteratorMode::End);
let mut tmp_vec = iterator1.collect::<Vec<_>>();
tmp_vec.reverse();
assert_eq!(tmp_vec, expected);
}
{
let iterator1 = view1.from_end();
let iterator1 = db.iterator(IteratorMode::End);
let mut tmp_vec = iterator1.collect::<Vec<_>>();
tmp_vec.reverse();
assert_eq!(tmp_vec, expected);
}
// Try it forward again
{
let iterator1 = view1.from_start();
let iterator1 = db.iterator(IteratorMode::Start);
assert_eq!(iterator1.collect::<Vec<_>>(), expected);
}
{
let iterator1 = view1.from_start();
let iterator1 = db.iterator(IteratorMode::Start);
assert_eq!(iterator1.collect::<Vec<_>>(), expected);
}
let old_iterator = db.iterator(IteratorMode::Start);
let p = db.put(&*k4, &*v4);
assert!(p.is_ok());
let mut view3 = db.iterator();
let expected2 = vec![(cba(&k1), cba(&v1)),
(cba(&k2), cba(&v2)),
(cba(&k3), cba(&v3)),
(cba(&k4), cba(&v4))];
{
let iterator1 = view1.from_start();
assert_eq!(iterator1.collect::<Vec<_>>(), expected);
assert_eq!(old_iterator.collect::<Vec<_>>(), expected);
}
{
let iterator1 = view3.from_start();
let iterator1 = db.iterator(IteratorMode::Start);
assert_eq!(iterator1.collect::<Vec<_>>(), expected2);
}
{
let iterator1 = view3.from(b"k2", Direction::forward);
let iterator1 = db.iterator(IteratorMode::From(b"k2", Direction::forward));
let expected = vec![(cba(&k2), cba(&v2)),
(cba(&k3), cba(&v3)),
(cba(&k4), cba(&v4))];
assert_eq!(iterator1.collect::<Vec<_>>(), expected);
}
{
let iterator1 = view3.from(b"k2", Direction::reverse);
let iterator1 = db.iterator(IteratorMode::From(b"k2", Direction::reverse));
let expected = vec![(cba(&k2), cba(&v2)), (cba(&k1), cba(&v1))];
assert_eq!(iterator1.collect::<Vec<_>>(), expected);
}

Loading…
Cancel
Save