Add `DBRawIteratorWithThreadMode::item` method returning (key, value) (#647)

src/db_iterator.rs

@@ -299,16 +299,7 @@ impl<'a, D: DBAccess> DBRawIteratorWithThreadMode<'a, D> {
     /// Returns a slice of the current key.
     pub fn key(&self) -> Option<&[u8]> {
         if self.valid() {
-            // Safety Note: This is safe as all methods that may invalidate the buffer returned
-            // take `&mut self`, so borrow checker will prevent use of buffer after seek.
-            unsafe {
-                let mut key_len: size_t = 0;
-                let key_len_ptr: *mut size_t = &mut key_len;
-                let key_ptr =
-                    ffi::rocksdb_iter_key(self.inner.as_ptr(), key_len_ptr) as *const c_uchar;
-
-                Some(slice::from_raw_parts(key_ptr, key_len as usize))
-            }
+            Some(self.key_impl())
         } else {
             None
         }
@@ -317,20 +308,44 @@ impl<'a, D: DBAccess> DBRawIteratorWithThreadMode<'a, D> {
     /// Returns a slice of the current value.
     pub fn value(&self) -> Option<&[u8]> {
         if self.valid() {
-            // Safety Note: This is safe as all methods that may invalidate the buffer returned
-            // take `&mut self`, so borrow checker will prevent use of buffer after seek.
-            unsafe {
-                let mut val_len: size_t = 0;
-                let val_len_ptr: *mut size_t = &mut val_len;
-                let val_ptr =
-                    ffi::rocksdb_iter_value(self.inner.as_ptr(), val_len_ptr) as *const c_uchar;
-
-                Some(slice::from_raw_parts(val_ptr, val_len as usize))
-            }
+            Some(self.value_impl())
+        } else {
+            None
+        }
+    }
+
+    /// Returns pair with slice of the current key and current value.
+    pub fn item(&self) -> Option<(&[u8], &[u8])> {
+        if self.valid() {
+            Some((self.key_impl(), self.value_impl()))
         } else {
             None
         }
     }
+
+    /// Returns a slice of the current key; assumes the iterator is valid.
+    fn key_impl(&self) -> &[u8] {
+        // Safety Note: This is safe as all methods that may invalidate the buffer returned
+        // take `&mut self`, so borrow checker will prevent use of buffer after seek.
+        unsafe {
+            let mut key_len: size_t = 0;
+            let key_len_ptr: *mut size_t = &mut key_len;
+            let key_ptr = ffi::rocksdb_iter_key(self.inner.as_ptr(), key_len_ptr);
+            slice::from_raw_parts(key_ptr as *const c_uchar, key_len as usize)
+        }
+    }
+
+    /// Returns a slice of the current value; assumes the iterator is valid.
+    fn value_impl(&self) -> &[u8] {
+        // Safety Note: This is safe as all methods that may invalidate the buffer returned
+        // take `&mut self`, so borrow checker will prevent use of buffer after seek.
+        unsafe {
+            let mut val_len: size_t = 0;
+            let val_len_ptr: *mut size_t = &mut val_len;
+            let val_ptr = ffi::rocksdb_iter_value(self.inner.as_ptr(), val_len_ptr);
+            slice::from_raw_parts(val_ptr as *const c_uchar, val_len as usize)
+        }
+    }
 }
 
 impl<'a, D: DBAccess> Drop for DBRawIteratorWithThreadMode<'a, D> {
@@ -476,12 +491,8 @@ impl<'a, D: DBAccess> Iterator for DBIteratorWithThreadMode<'a, D> {
             }
         }
 
-        if self.raw.valid() {
-            // .key() and .value() only ever return None if valid == false, which we've just checked
-            Some((
-                Box::from(self.raw.key().unwrap()),
-                Box::from(self.raw.value().unwrap()),
-            ))
+        if let Some((key, value)) = self.raw.item() {
+            Some((Box::from(key), Box::from(value)))
         } else {
             None
         }

tests/test_raw_iterator.rs

@@ -16,9 +16,23 @@ mod util;
 
 use pretty_assertions::assert_eq;
 
-use rocksdb::DB;
+use rocksdb::{DBAccess, DBRawIteratorWithThreadMode, DB};
 use util::DBPath;
 
+fn assert_item<D: DBAccess>(iter: &DBRawIteratorWithThreadMode<'_, D>, key: &[u8], value: &[u8]) {
+    assert!(iter.valid());
+    assert_eq!(iter.key(), Some(key));
+    assert_eq!(iter.value(), Some(value));
+    assert_eq!(iter.item(), Some((key, value)));
+}
+
+fn assert_no_item<D: DBAccess>(iter: &DBRawIteratorWithThreadMode<'_, D>) {
+    assert!(!iter.valid());
+    assert_eq!(iter.key(), None);
+    assert_eq!(iter.value(), None);
+    assert_eq!(iter.item(), None);
+}
+
 #[test]
 pub fn test_forwards_iteration() {
     let n = DBPath::new("forwards_iteration");
@@ -31,24 +45,16 @@ pub fn test_forwards_iteration() {
 
         let mut iter = db.raw_iterator();
         iter.seek_to_first();
-
-        assert!(iter.valid());
-        assert_eq!(iter.key(), Some(b"k1".as_ref()));
-        assert_eq!(iter.value(), Some(b"v1".as_ref()));
+        assert_item(&iter, b"k1", b"v1");
 
         iter.next();
-
-        assert!(iter.valid());
-        assert_eq!(iter.key(), Some(b"k2".as_ref()));
-        assert_eq!(iter.value(), Some(b"v2".as_ref()));
+        assert_item(&iter, b"k2", b"v2");
 
         iter.next(); // k3
         iter.next(); // k4
-        iter.next(); // invalid!
 
-        assert!(!iter.valid());
-        assert_eq!(iter.key(), None);
-        assert_eq!(iter.value(), None);
+        iter.next(); // invalid!
+        assert_no_item(&iter);
     }
 }
 
@@ -64,24 +70,16 @@ pub fn test_seek_last() {
 
         let mut iter = db.raw_iterator();
         iter.seek_to_last();
-
-        assert!(iter.valid());
-        assert_eq!(iter.key(), Some(b"k4".as_ref()));
-        assert_eq!(iter.value(), Some(b"v4".as_ref()));
+        assert_item(&iter, b"k4", b"v4");
 
         iter.prev();
-
-        assert!(iter.valid());
-        assert_eq!(iter.key(), Some(b"k3".as_ref()));
-        assert_eq!(iter.value(), Some(b"v3".as_ref()));
+        assert_item(&iter, b"k3", b"v3");
 
         iter.prev(); // k2
         iter.prev(); // k1
-        iter.prev(); // invalid!
 
-        assert!(!iter.valid());
-        assert_eq!(iter.key(), None);
-        assert_eq!(iter.value(), None);
+        iter.prev(); // invalid!
+        assert_no_item(&iter);
     }
 }
 
@@ -97,16 +95,11 @@ pub fn test_seek() {
         let mut iter = db.raw_iterator();
         iter.seek(b"k2");
 
-        assert!(iter.valid());
-        assert_eq!(iter.key(), Some(b"k2".as_ref()));
-        assert_eq!(iter.value(), Some(b"v2".as_ref()));
+        assert_item(&iter, b"k2", b"v2");
 
         // Check it gets the next key when the key doesn't exist
         iter.seek(b"k3");
-
-        assert!(iter.valid());
-        assert_eq!(iter.key(), Some(b"k4".as_ref()));
-        assert_eq!(iter.value(), Some(b"v4".as_ref()));
+        assert_item(&iter, b"k4", b"v4");
     }
 }
 
@@ -121,10 +114,7 @@ pub fn test_seek_to_nonexistant() {
 
         let mut iter = db.raw_iterator();
         iter.seek(b"k2");
-
-        assert!(iter.valid());
-        assert_eq!(iter.key(), Some(b"k3".as_ref()));
-        assert_eq!(iter.value(), Some(b"v3".as_ref()));
+        assert_item(&iter, b"k3", b"v3");
     }
 }
 
@@ -139,16 +129,10 @@ pub fn test_seek_for_prev() {
 
         let mut iter = db.raw_iterator();
         iter.seek(b"k2");
-
-        assert!(iter.valid());
-        assert_eq!(iter.key(), Some(b"k2".as_ref()));
-        assert_eq!(iter.value(), Some(b"v2".as_ref()));
+        assert_item(&iter, b"k2", b"v2");
 
         // Check it gets the previous key when the key doesn't exist
         iter.seek_for_prev(b"k3");
-
-        assert!(iter.valid());
-        assert_eq!(iter.key(), Some(b"k2".as_ref()));
-        assert_eq!(iter.value(), Some(b"v2".as_ref()));
+        assert_item(&iter, b"k2", b"v2");
     }
 }