add safe methods to get key / value from iterator

rim99 · Oct 13, 2024 · eeccde1 · eeccde1
1 parent 68414f9
commit eeccde1
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 16 deletions.
diff --git a/src/database/iterator.rs b/src/database/iterator.rs
@@ -156,32 +156,56 @@ pub trait LevelDBIterator<'a, K: Serializable + Ord> {
                 self.start();
             }
             if let Some(end) = self.to_key() {
-                if !self.valid() || end <= &self.key() {
+                if !self.valid() || end <= &self.unsafe_key() {
                     self.stop();
                 }
             }
         }
         self.valid()
     }
 
-    fn key(&self) -> K {
+    fn unsafe_key(&self) -> K {
         unsafe {
             let length: size_t = 0;
             let value = leveldb_iter_key(self.raw_iterator(), &length) as *const u8;
             from_u8(from_raw_parts(value, length as usize))
         }
     }
 
-    fn value(&self) -> Vec<u8> {
+    fn unsafe_value(&self) -> Vec<u8> {
         unsafe {
             let length: size_t = 0;
             let value = leveldb_iter_value(self.raw_iterator(), &length) as *const u8;
             from_raw_parts(value, length as usize).to_vec()
         }
     }
 
-    fn entry(&self) -> (K, Vec<u8>) {
-        (self.key(), self.value())
+    fn unsafe_entry(&self) -> (K, Vec<u8>) {
+        (self.unsafe_key(), self.unsafe_value())
+    }
+
+    fn key(&self) -> Option<K> {
+        if self.valid() {
+            Some(self.unsafe_key())
+        } else {
+            None
+        }
+    }
+
+    fn value(&self) -> Option<Vec<u8>> {
+        if self.valid() {
+            Some(self.unsafe_value())
+        } else {
+            None
+        }
+    }
+
+    fn entry(&self) -> Option<(K, Vec<u8>)> {
+        if self.valid() {
+            Some(self.unsafe_entry())
+        } else {
+            None
+        }
     }
 
     fn seek_to_first(&self) {
@@ -232,7 +256,7 @@ impl<'a, K: Serializable + Ord> Iterator<'a, K> {
     /// return the last element of the iterator
     pub fn last(self) -> Option<(K, Vec<u8>)> {
         self.seek_to_last();
-        Some((self.key(), self.value()))
+        Some((self.unsafe_key(), self.unsafe_value()))
     }
 }
 
@@ -384,7 +408,7 @@ impl<'a, K: Serializable + Ord> KeyIterator<'a, K> {
     /// return the last element of the iterator
     pub fn last(self) -> Option<K> {
         self.seek_to_last();
-        Some(self.key())
+        Some(self.unsafe_key())
     }
 }
 
@@ -398,7 +422,7 @@ impl<'a, K: Serializable + Ord> ValueIterator<'a, K> {
     /// return the last element of the iterator
     pub fn last(self) -> Option<Vec<u8>> {
         self.seek_to_last();
-        Some(self.value())
+        Some(self.unsafe_value())
     }
 }
 
@@ -487,9 +511,9 @@ macro_rules! impl_iterator {
     };
 }
 
-impl_iterator!(Iterator<'a, K>, (K, Vec<u8>), entry);
-impl_iterator!(RevIterator<'a, K>, (K, Vec<u8>), entry);
-impl_iterator!(KeyIterator<'a, K>, K, key);
-impl_iterator!(RevKeyIterator<'a, K>, K, key);
-impl_iterator!(ValueIterator<'a, K>, Vec<u8>, value);
-impl_iterator!(RevValueIterator<'a, K>, K, key);
+impl_iterator!(Iterator<'a, K>, (K, Vec<u8>), unsafe_entry);
+impl_iterator!(RevIterator<'a, K>, (K, Vec<u8>), unsafe_entry);
+impl_iterator!(KeyIterator<'a, K>, K, unsafe_key);
+impl_iterator!(RevKeyIterator<'a, K>, K, unsafe_key);
+impl_iterator!(ValueIterator<'a, K>, Vec<u8>, unsafe_value);
+impl_iterator!(RevValueIterator<'a, K>, K, unsafe_key);
diff --git a/tests/iterator.rs b/tests/iterator.rs
@@ -260,7 +260,8 @@ fn test_seek_before_inserted() {
     let iter = database.iter(read_opts);
 
     iter.seek(&0);
-    assert!(iter.valid())
+    assert!(iter.valid());
+    assert_eq!(iter.key(), Some(1));
 }
 
 #[test]
@@ -274,5 +275,6 @@ fn test_seek_after_inserted() {
     let iter = database.iter(read_opts);
 
     iter.seek(&3);
-    assert!(!iter.valid())
+    assert!(!iter.valid());
+    assert_eq!(iter.key(), None);
 }