From 0af1f8e6457bb2ee10edab9795cb559a78653b2d Mon Sep 17 00:00:00 2001 From: Michal Nazarewicz Date: Thu, 30 Jun 2022 11:07:34 +0200 Subject: [PATCH] Add `DBRawIteratorWithThreadMode::item` method returning (key, value) (#647) --- src/db_iterator.rs | 63 +++++++++++++++++++-------------- tests/test_raw_iterator.rs | 72 +++++++++++++++----------------------- 2 files changed, 65 insertions(+), 70 deletions(-) diff --git a/src/db_iterator.rs b/src/db_iterator.rs index f40c3bc..fa988e6 100644 --- a/src/db_iterator.rs +++ b/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 } diff --git a/tests/test_raw_iterator.rs b/tests/test_raw_iterator.rs index f48e5d3..311ad23 100644 --- a/tests/test_raw_iterator.rs +++ b/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(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(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"); } }