diff --git a/heed/src/databases/database.rs b/heed/src/databases/database.rs index 46ccc902..f8361c4e 100644 --- a/heed/src/databases/database.rs +++ b/heed/src/databases/database.rs @@ -1069,16 +1069,16 @@ impl Database { /// db.put(&mut wtxn, &13, "i-am-thirteen")?; /// /// let mut iter = db.iter_mut(&mut wtxn)?; - /// assert_eq!(iter.next().transpose()?, Some((13, "i-am-thirteen"))); + /// assert_eq!(Iterator::next(&mut iter).transpose()?, Some((13, "i-am-thirteen"))); /// let ret = unsafe { iter.del_current()? }; /// assert!(ret); /// - /// assert_eq!(iter.next().transpose()?, Some((27, "i-am-twenty-seven"))); - /// assert_eq!(iter.next().transpose()?, Some((42, "i-am-forty-two"))); + /// assert_eq!(Iterator::next(&mut iter).transpose()?, Some((27, "i-am-twenty-seven"))); + /// assert_eq!(Iterator::next(&mut iter).transpose()?, Some((42, "i-am-forty-two"))); /// let ret = unsafe { iter.put_current(&42, "i-am-the-new-forty-two")? }; /// assert!(ret); /// - /// assert_eq!(iter.next().transpose()?, None); + /// assert_eq!(Iterator::next(&mut iter).transpose()?, None); /// /// drop(iter); /// diff --git a/heed/src/databases/encrypted_database.rs b/heed/src/databases/encrypted_database.rs index 5020877d..376e5fb8 100644 --- a/heed/src/databases/encrypted_database.rs +++ b/heed/src/databases/encrypted_database.rs @@ -914,16 +914,16 @@ impl EncryptedDatabase { /// db.put(&mut wtxn, &13, "i-am-thirteen")?; /// /// let mut iter = db.iter_mut(&mut wtxn)?; - /// assert_eq!(iter.next().transpose()?, Some((13, "i-am-thirteen"))); + /// assert_eq!(Iterator::next(&mut iter).transpose()?, Some((13, "i-am-thirteen"))); /// let ret = unsafe { iter.del_current()? }; /// assert!(ret); /// - /// assert_eq!(iter.next().transpose()?, Some((27, "i-am-twenty-seven"))); - /// assert_eq!(iter.next().transpose()?, Some((42, "i-am-forty-two"))); + /// assert_eq!(Iterator::next(&mut iter).transpose()?, Some((27, "i-am-twenty-seven"))); + /// assert_eq!(Iterator::next(&mut iter).transpose()?, Some((42, "i-am-forty-two"))); /// let ret = unsafe { iter.put_current(&42, "i-am-the-new-forty-two")? }; /// assert!(ret); /// - /// assert_eq!(iter.next().transpose()?, None); + /// assert_eq!(Iterator::next(&mut iter).transpose()?, None); /// /// drop(iter); /// diff --git a/heed/src/iterator/iter.rs b/heed/src/iterator/iter.rs index 2f85fca5..d2c13755 100644 --- a/heed/src/iterator/iter.rs +++ b/heed/src/iterator/iter.rs @@ -1,7 +1,7 @@ use std::borrow::Cow; use std::marker; -use types::LazyDecode; +use types::{Lazy, LazyDecode}; use crate::iteration_method::{IterationMethod, MoveBetweenKeys, MoveThroughDuplicateValues}; use crate::*; @@ -437,6 +437,185 @@ impl fmt::Debug for RwIter<'_, KC, DC, IM> { } } +pub struct RwIterItem<'a, 'txn, KC, DC, IM = MoveThroughDuplicateValues> { + iter: &'a mut RwIter<'txn, KC, DC, IM>, + key: &'txn [u8], + data: &'txn [u8], +} + +impl<'a, 'txn, KC, DC, IM> RwIterItem<'a, 'txn, KC, DC, IM> { + /// Get the key at the current cursor, which is valid for at least as long + /// as the `RwIterItem` exists. + pub fn key<'b>(&'b self) -> Lazy<'b, KC> + where + KC: 'static, + { + LazyDecode::::bytes_decode(self.key).unwrap() + } + + /// Get the data at the current cursor, which is valid for at least as long + /// as the `RwIterItem` exists. + pub fn data<'b>(&'b self) -> Lazy<'b, DC> + where + DC: 'static, + { + LazyDecode::::bytes_decode(self.data).unwrap() + } + + /// Keep the entry the cursor is currently pointing to. + /// Returns the key and data at the current cursor, which are valid until + /// the end of the transaction. + pub fn keep(self) -> (Lazy<'txn, KC>, Lazy<'txn, DC>) + where + KC: 'static, + DC: 'static, + { + let key = LazyDecode::::bytes_decode(self.key).unwrap(); + let data = LazyDecode::::bytes_decode(self.data).unwrap(); + (key, data) + } + + /// Delete the entry the cursor is currently pointing to. + /// + /// Returns `true`` if the entry was successfully deleted. + pub fn delete(self) -> Result { + // SAFETY: It is not possible to keep reference to the current value + // after deleting, because this method takes `self` by value, and + // the current key / value are only available for the lifetime of + // `self`. + unsafe { self.iter.del_current() } + } + + /// Write a new value to the current entry. The entry is written with the specified flags. + /// + /// > This is intended to be used when the new data is the same size as the old. + /// > Otherwise it will simply perform a delete of the old record followed by an insert. + pub fn put_reserved_with_flags( + mut self, + flags: PutFlags, + data_size: usize, + write_func: F, + ) -> Result + where + F: FnOnce(&mut ReservedSpace<'_>) -> io::Result<()>, + { + let key_owned: Vec = self.key.to_owned(); + // SAFETY: The key is owned, and `write_func` cannot use a reference + // from this entry while modifying it, because this method takes `self` + // by value, and the current key / value are only available for the + // lifetime of `self`. + if unsafe { + self.iter.cursor.put_current_reserved_with_flags( + flags, + key_owned.as_slice(), + data_size, + write_func, + ) + }? { + let (key, data) = self.iter.cursor.current()?.ok_or(Error::Mdb(MdbError::NotFound))?; + self.key = key; + self.data = data; + Ok(self) + } else { + Err(Error::Mdb(MdbError::NotFound)) + } + } + + /// Insert a key-value pair in this database. + /// The entry is written with the specified flags and data codec. + pub fn put_with_options<'b, NDC>( + mut self, + flags: PutFlags, + data: &'b NDC::EItem, + ) -> Result + where + NDC: BytesEncode<'b>, + { + let key_owned: Vec = self.key.to_owned(); + let data_bytes: Cow<[u8]> = NDC::bytes_encode(data).map_err(Error::Encoding)?; + // SAFETY: The key is owned, and `data` cannot use a reference + // from this entry while modifying it, because this method takes `self` + // by value, and the current key / value are only available for the + // lifetime of `self`. + let () = unsafe { + self.iter.cursor.put_current_with_flags(flags, key_owned.as_slice(), &data_bytes) + }?; + let (key, data) = self.iter.cursor.current()?.ok_or(Error::Mdb(MdbError::NotFound))?; + self.key = key; + self.data = data; + Ok(self) + } + + /// Write a new value to the current entry. + /// + /// > This is intended to be used when the new data is the same size as the old. + /// > Otherwise it will simply perform a delete of the old record followed by an insert. + pub fn put<'b>(mut self, data: &'b DC::EItem) -> Result + where + DC: BytesEncode<'b>, + { + let key_owned: Vec = self.key.to_owned(); + let data_bytes: Cow<[u8]> = DC::bytes_encode(data).map_err(Error::Encoding)?; + // SAFETY: The key is owned, and `data` cannot use a reference + // from this entry while modifying it, because this method takes `self` + // by value, and the current key / value are only available for the + // lifetime of `self`. + if unsafe { self.iter.cursor.put_current(key_owned.as_slice(), &data_bytes) }? { + let (key, data) = self.iter.cursor.current()?.ok_or(Error::Mdb(MdbError::NotFound))?; + self.key = key; + self.data = data; + Ok(self) + } else { + Err(Error::Mdb(MdbError::NotFound)) + } + } +} + +impl<'txn, KC, DC, IM> RwIter<'txn, KC, DC, IM> { + /// Advance the iterator to the next item. + pub fn next<'a>(&'a mut self) -> Option>> + where + IM: IterationMethod, + { + let result = if self.move_on_first { + self.move_on_first = false; + self.cursor.move_on_first(IM::MOVE_OPERATION) + } else { + self.cursor.move_on_next(IM::MOVE_OPERATION) + }; + + match result { + Ok(Some((key, data))) => Some(Ok(RwIterItem { iter: self, key, data })), + Ok(None) => None, + Err(e) => Some(Err(e)), + } + } + + /// Advance the iterator to the last item. + pub fn last<'a>(&'a mut self) -> Option>> + where + IM: IterationMethod, + { + let result = if self.move_on_first { + self.cursor.move_on_last(IM::MOVE_OPERATION) + } else { + match (self.cursor.current(), self.cursor.move_on_last(IM::MOVE_OPERATION)) { + (Ok(Some((ckey, _))), Ok(Some((key, data)))) if ckey != key => { + Ok(Some((key, data))) + } + (Ok(_), Ok(_)) => Ok(None), + (Err(e), _) | (_, Err(e)) => Err(e), + } + }; + + match result { + Ok(Some((key, data))) => Some(Ok(RwIterItem { iter: self, key, data })), + Ok(None) => None, + Err(e) => Some(Err(e)), + } + } +} + /// A reverse read-only iterator structure. pub struct RoRevIter<'txn, KC, DC, IM = MoveThroughDuplicateValues> { cursor: RoCursor<'txn>,