NextGraph
/
rust-rocksdb
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

More idiomatic iterators

Browse Source
master
Pete Hunt 9 years ago
parent f5d4d834b7
commit 09250315e8
3 changed files with 87 additions and 92 deletions
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Show Stats Download Patch File Download Diff File
  1. 2
      src/lib.rs
  2. 137
      src/rocksdb.rs
  3. 38
      test/test_iterator.rs

2
src/lib.rs
Unescape View File

@ -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;

137
src/rocksdb.rs
Unescape View File

@ -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());

38
test/test_iterator.rs
Unescape View File

@ -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);
}

Write Preview
Loading…
Cancel
Save