make iterator don't copy (#9)

src/lib.rs

@@ -25,9 +25,9 @@ pub mod rocksdb_options;
 pub mod merge_operator;
 pub mod comparator;
 
-pub use librocksdb_sys::{DBCompactionStyle, DBComparator, new_bloom_filter,
-                         self as rocksdb_ffi, DBCompressionType};
-pub use rocksdb::{DB, DBIterator, DBVector, Direction, IteratorMode, Writable,
-                  WriteBatch};
+pub use librocksdb_sys::{DBCompactionStyle, DBComparator, DBCompressionType,
+                         new_bloom_filter, self as rocksdb_ffi};
+pub use rocksdb::{DB, DBIterator, DBVector, Direction, IteratorMode, Kv,
+                  Writable, WriteBatch};
 pub use rocksdb_options::{BlockBasedOptions, Options, WriteOptions};
 pub use merge_operator::MergeOperands;

src/rocksdb.rs

@@ -48,7 +48,10 @@ pub struct Snapshot<'a> {
     inner: rocksdb_ffi::DBSnapshot,
 }
 
-pub struct DBIterator {
+// We need to find a better way to add a lifetime in here.
+#[allow(dead_code)]
+pub struct DBIterator<'a> {
+    db: &'a DB,
     inner: rocksdb_ffi::DBIterator,
     direction: Direction,
     just_seeked: bool,
@@ -59,13 +62,12 @@ pub enum Direction {
     Reverse,
 }
 
-// TODO: should we use Vec<u8> instead?
-pub type Kv = (Box<[u8]>, Box<[u8]>);
+pub type Kv<'a> = (&'a [u8], &'a [u8]);
 
-impl Iterator for DBIterator {
-    type Item = Kv;
+impl<'a> Iterator for DBIterator<'a> {
+    type Item = Kv<'a>;
 
-    fn next(&mut self) -> Option<Kv> {
+    fn next(&mut self) -> Option<Kv<'a>> {
         let native_iter = self.inner;
         if self.just_seeked {
             self.just_seeked = false;
@@ -97,8 +99,7 @@ impl Iterator for DBIterator {
                 slice::from_raw_parts(val_ptr, val_len as usize)
             };
 
-            Some((key.to_vec().into_boxed_slice(),
-                  val.to_vec().into_boxed_slice()))
+            Some((key, val))
         } else {
             None
         }
@@ -112,13 +113,17 @@ pub enum IteratorMode<'a> {
 }
 
 
-impl DBIterator {
-    fn new(db: &DB, readopts: &ReadOptions, mode: IteratorMode) -> DBIterator {
+impl<'a> DBIterator<'a> {
+    fn new(db: &'a DB,
+           readopts: &ReadOptions,
+           mode: IteratorMode)
+           -> DBIterator<'a> {
         unsafe {
             let iterator = rocksdb_ffi::rocksdb_create_iterator(db.inner,
                                                                 readopts.inner);
 
             let mut rv = DBIterator {
+                db: db,
                 inner: iterator,
                 direction: Direction::Forward, // blown away by set_mode()
                 just_seeked: false,
@@ -156,11 +161,11 @@ impl DBIterator {
         unsafe { rocksdb_ffi::rocksdb_iter_valid(self.inner) }
     }
 
-    fn new_cf(db: &DB,
+    fn new_cf(db: &'a DB,
               cf_handle: DBCFHandle,
               readopts: &ReadOptions,
               mode: IteratorMode)
-              -> Result<DBIterator, String> {
+              -> Result<DBIterator<'a>, String> {
         unsafe {
             let iterator =
                 rocksdb_ffi::rocksdb_create_iterator_cf(db.inner,
@@ -168,6 +173,7 @@ impl DBIterator {
                                                         cf_handle);
 
             let mut rv = DBIterator {
+                db: db,
                 inner: iterator,
                 direction: Direction::Forward, // blown away by set_mode()
                 just_seeked: false,
@@ -180,7 +186,7 @@ impl DBIterator {
     }
 }
 
-impl Drop for DBIterator {
+impl<'a> Drop for DBIterator<'a> {
     fn drop(&mut self) {
         unsafe {
             rocksdb_ffi::rocksdb_iter_destroy(self.inner);

test/test_iterator.rs

-use rocksdb::{DB, Direction, IteratorMode, Writable};
+use rocksdb::{DB, Direction, IteratorMode, Writable, Kv};
 use tempdir::TempDir;
 
-fn cba(input: &Box<[u8]>) -> Box<[u8]> {
-    input.iter().cloned().collect::<Vec<_>>().into_boxed_slice()
+fn collect<'a, T: Iterator<Item=Kv<'a>>>(iter: T) -> Vec<(Vec<u8>, Vec<u8>)> {
+    iter.map(|(k, v)| (k.to_vec(), v.to_vec())).collect()
 }
 
 #[test]
 pub fn test_iterator() {
     let path = TempDir::new("_rust_rocksdb_iteratortest").expect("");
 
-    let k1: Box<[u8]> = b"k1".to_vec().into_boxed_slice();
-    let k2: Box<[u8]> = b"k2".to_vec().into_boxed_slice();
-    let k3: Box<[u8]> = b"k3".to_vec().into_boxed_slice();
-    let k4: Box<[u8]> = b"k4".to_vec().into_boxed_slice();
-    let v1: Box<[u8]> = b"v1111".to_vec().into_boxed_slice();
-    let v2: Box<[u8]> = b"v2222".to_vec().into_boxed_slice();
-    let v3: Box<[u8]> = b"v3333".to_vec().into_boxed_slice();
-    let v4: Box<[u8]> = b"v4444".to_vec().into_boxed_slice();
+    let k1 = b"k1";
+    let k2 = b"k2";
+    let k3 = b"k3";
+    let k4 = b"k4";
+    let v1 = b"v1111";
+    let v2 = b"v2222";
+    let v3 = b"v3333";
+    let v4 = b"v4444";
     let db = DB::open_default(path.path().to_str().unwrap()).unwrap();
-    let p = db.put(&*k1, &*v1);
+    let p = db.put(k1, v1);
     assert!(p.is_ok());
-    let p = db.put(&*k2, &*v2);
+    let p = db.put(k2, v2);
     assert!(p.is_ok());
-    let p = db.put(&*k3, &*v3);
+    let p = db.put(k3, v3);
     assert!(p.is_ok());
-    let expected = vec![(cba(&k1), cba(&v1)),
-                        (cba(&k2), cba(&v2)),
-                        (cba(&k3), cba(&v3))];
+    let expected = vec![(k1.to_vec(), v1.to_vec()),
+                        (k2.to_vec(), v2.to_vec()),
+                        (k3.to_vec(), v3.to_vec())];
 
     let iterator1 = db.iterator(IteratorMode::Start);
-    assert_eq!(iterator1.collect::<Vec<_>>(), expected);
+    assert_eq!(collect(iterator1), expected);
 
     // Test that it's idempotent
     let iterator1 = db.iterator(IteratorMode::Start);
-    assert_eq!(iterator1.collect::<Vec<_>>(), expected);
+    assert_eq!(collect(iterator1), expected);
 
     let iterator1 = db.iterator(IteratorMode::Start);
-    assert_eq!(iterator1.collect::<Vec<_>>(), expected);
+    assert_eq!(collect(iterator1), expected);
 
     let iterator1 = db.iterator(IteratorMode::Start);
-    assert_eq!(iterator1.collect::<Vec<_>>(), expected);
+    assert_eq!(collect(iterator1), expected);
 
     // Test it in reverse a few times
     let iterator1 = db.iterator(IteratorMode::End);
-    let mut tmp_vec = iterator1.collect::<Vec<_>>();
+    let mut tmp_vec = collect(iterator1);
     tmp_vec.reverse();
     assert_eq!(tmp_vec, expected);
 
     let iterator1 = db.iterator(IteratorMode::End);
-    let mut tmp_vec = iterator1.collect::<Vec<_>>();
+    let mut tmp_vec = collect(iterator1);
     tmp_vec.reverse();
     assert_eq!(tmp_vec, expected);
 
     let iterator1 = db.iterator(IteratorMode::End);
-    let mut tmp_vec = iterator1.collect::<Vec<_>>();
+    let mut tmp_vec = collect(iterator1);
     tmp_vec.reverse();
     assert_eq!(tmp_vec, expected);
 
     let iterator1 = db.iterator(IteratorMode::End);
-    let mut tmp_vec = iterator1.collect::<Vec<_>>();
+    let mut tmp_vec = collect(iterator1);
     tmp_vec.reverse();
     assert_eq!(tmp_vec, expected);
 
     let iterator1 = db.iterator(IteratorMode::End);
-    let mut tmp_vec = iterator1.collect::<Vec<_>>();
+    let mut tmp_vec = collect(iterator1);
     tmp_vec.reverse();
     assert_eq!(tmp_vec, expected);
 
     // Try it forward again
     let iterator1 = db.iterator(IteratorMode::Start);
-    assert_eq!(iterator1.collect::<Vec<_>>(), expected);
+    assert_eq!(collect(iterator1), expected);
 
     let iterator1 = db.iterator(IteratorMode::Start);
-    assert_eq!(iterator1.collect::<Vec<_>>(), expected);
+    assert_eq!(collect(iterator1), expected);
 
     let old_iterator = db.iterator(IteratorMode::Start);
     let p = db.put(&*k4, &*v4);
     assert!(p.is_ok());
-    let expected2 = vec![(cba(&k1), cba(&v1)),
-                         (cba(&k2), cba(&v2)),
-                         (cba(&k3), cba(&v3)),
-                         (cba(&k4), cba(&v4))];
-    assert_eq!(old_iterator.collect::<Vec<_>>(), expected);
+    let expected2 = vec![(k1.to_vec(), v1.to_vec()),
+                         (k2.to_vec(), v2.to_vec()),
+                         (k3.to_vec(), v3.to_vec()),
+                         (k4.to_vec(), v4.to_vec())];
+    assert_eq!(collect(old_iterator), expected);
 
     let iterator1 = db.iterator(IteratorMode::Start);
-    assert_eq!(iterator1.collect::<Vec<_>>(), expected2);
+    assert_eq!(collect(iterator1), expected2);
 
     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 expected = vec![(k2.to_vec(), v2.to_vec()),
+                        (k3.to_vec(), v3.to_vec()),
+                        (k4.to_vec(), v4.to_vec())];
+    assert_eq!(collect(iterator1), expected);
 
     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);
+    let expected = vec![(k2.to_vec(), v2.to_vec()), (k1.to_vec(), v1.to_vec())];
+    assert_eq!(collect(iterator1), expected);
 
     let iterator1 = db.iterator(IteratorMode::From(b"k0", Direction::Forward));
     assert!(iterator1.valid());