From 09250315e8812ad3ebb7e55a0faffd8eef5406ef Mon Sep 17 00:00:00 2001 From: Pete Hunt Date: Sat, 26 Dec 2015 18:26:10 -0500 Subject: [PATCH] More idiomatic iterators --- src/lib.rs | 2 +- src/rocksdb.rs | 139 +++++++++++++++++++++--------------------- test/test_iterator.rs | 38 ++++++------ 3 files changed, 87 insertions(+), 92 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 0c80cd6..b2233b3 100644 --- a/src/lib.rs +++ b/src/lib.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; diff --git a/src/rocksdb.rs b/src/rocksdb.rs index 1d309e0..90b2566 100644 --- a/src/rocksdb.rs +++ b/src/rocksdb.rs @@ -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 { + readopts: &ReadOptions, + mode: IteratorMode) + -> Result, 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 { + pub fn iterator_cf(&self, cf_handle: DBCFHandle, mode: IteratorMode) -> Result { 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()); diff --git a/test/test_iterator.rs b/test/test_iterator.rs index a411d14..ac8ea75 100644 --- a/test/test_iterator.rs +++ b/test/test_iterator.rs @@ -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::>().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::>(), 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::>(), expected); } { - let iterator1 = view1.from_start(); + let iterator1 = db.iterator(IteratorMode::Start); assert_eq!(iterator1.collect::>(), expected); } { - let iterator1 = view1.from_start(); + let iterator1 = db.iterator(IteratorMode::Start); assert_eq!(iterator1.collect::>(), 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::>(); 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::>(); 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::>(); 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::>(); 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::>(); 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::>(), expected); } { - let iterator1 = view1.from_start(); + let iterator1 = db.iterator(IteratorMode::Start); assert_eq!(iterator1.collect::>(), 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::>(), expected); + assert_eq!(old_iterator.collect::>(), expected); } { - let iterator1 = view3.from_start(); + let iterator1 = db.iterator(IteratorMode::Start); assert_eq!(iterator1.collect::>(), 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::>(), 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::>(), expected); }