fix tests

storage/database/src/backend/heed/database.rs

@@ -189,20 +189,6 @@ impl<T: Table> DatabaseRw<T> for HeedTableRw<'_, '_, T> {
         Ok(())
     }
 
-    #[inline]
-    fn take(&mut self, key: &T::Key) -> DbResult<T::Value> {
-        // LMDB/heed does not return the value on deletion.
-        // So, fetch it first - then delete.
-        let value = get::<T>(&self.db, &self.tx_rw.borrow(), key)?;
-        match self.db.delete(&mut self.tx_rw.borrow_mut(), key) {
-            Ok(true) => Ok(value),
-            Err(e) => Err(e.into()),
-            // We just `get()`'ed the value - it is
-            // incorrect for it to suddenly not exist.
-            Ok(false) => unreachable!(),
-        }
-    }
-
     #[inline]
     fn pop_first(&mut self) -> DbResult<(T::Key, T::Value)> {
         let tx_rw = &mut self.tx_rw.borrow_mut();

storage/database/src/backend/tests.rs

@@ -263,17 +263,17 @@ fn db_read_write() {
         }
     }
 
-    // Assert `update()` works.
+    // Assert `Entry` works.
     {
         const NEW_VALUE: u64 = 999;
 
         assert_ne!(table.get(&KEY).unwrap(), NEW_VALUE);
 
-        #[expect(unused_assignments)]
         table
-            .update(&KEY, |mut value| {
-                value = NEW_VALUE;
-                Some(value)
+            .entry(&KEY)
+            .unwrap()
+            .and_update(|value| {
+                *value = NEW_VALUE;
             })
             .unwrap();
 
@@ -293,11 +293,11 @@ fn db_read_write() {
         assert_value(value);
     }
 
-    // Assert `take()` works.
+    // Assert `Entry` returns the correct value.
     {
         let mut key = KEY;
         key += 1;
-        let value = table.take(&key).unwrap();
+        let value = table.entry(&key).unwrap().and_remove().unwrap();
         assert_eq!(value, VALUE);
 
         let get = table.get(&KEY);

storage/database/src/database.rs

@@ -171,38 +171,6 @@ pub trait DatabaseRw<T: Table>: DatabaseRo<T> + Sized {
     /// This will never [`RuntimeError::KeyExists`].
     fn delete(&mut self, key: &T::Key) -> DbResult<()>;
 
-    /// Delete and return a key-value pair in the database.
-    ///
-    /// This is the same as [`DatabaseRw::delete`], however,
-    /// it will serialize the `T::Value` and return it.
-    ///
-    #[doc = doc_database!()]
-    fn take(&mut self, key: &T::Key) -> DbResult<T::Value>;
-
-    /// Fetch the value, and apply a function to it - or delete the entry.
-    ///
-    /// This will call [`DatabaseRo::get`] and call your provided function `f` on it.
-    ///
-    /// The [`Option`] `f` returns will dictate whether `update()`:
-    /// - Updates the current value OR
-    /// - Deletes the `(key, value)` pair
-    ///
-    /// - If `f` returns `Some(value)`, that will be [`DatabaseRw::put`] as the new value
-    /// - If `f` returns `None`, the entry will be [`DatabaseRw::delete`]d
-    ///
-    #[doc = doc_database!()]
-    fn update<F>(&mut self, key: &T::Key, mut f: F) -> DbResult<()>
-    where
-        F: FnMut(T::Value) -> Option<T::Value>,
-    {
-        let value = DatabaseRo::get(self, key)?;
-
-        match f(value) {
-            Some(value) => DatabaseRw::put(self, key, &value),
-            None => DatabaseRw::delete(self, key),
-        }
-    }
-
     /// Removes and returns the first `(key, value)` pair in the database.
     ///
     #[doc = doc_database!()]

storage/database/src/entry/entry.rs

@@ -2,7 +2,7 @@
 
 use crate::{
     entry::{OccupiedEntry, VacantEntry},
-    DatabaseRw, DbResult, Table,
+    DatabaseRw, DbResult, RuntimeError, Table,
 };
 
 /// A view into a single entry in a [`DatabaseRw`], which may either be vacant or occupied.
@@ -74,15 +74,42 @@ where
         }
     }
 
-    /// Returns a reference to this entry's key (if the entry is [`OccupiedEntry`]).
-    pub const fn value(&self) -> Option<&T::Value> {
+    /// Returns a reference to this entry's value (if the entry is [`OccupiedEntry`]).
+    ///
+    /// # Errors
+    /// This returns [`RuntimeError::KeyNotFound`] if the entry is [`VacantEntry`].
+    pub const fn value(&self) -> DbResult<&T::Value> {
         match self {
-            Self::Occupied(entry) => Some(entry.value()),
-            Self::Vacant(_) => None,
+            Self::Occupied(entry) => Ok(entry.value()),
+            Self::Vacant(_) => Err(RuntimeError::KeyNotFound),
         }
     }
 
-    /// Provides in-place mutable access to an occupied entry before any potential inserts.
+    /// Take ownership of entry's value (if the entry is [`OccupiedEntry`]).
+    ///
+    /// # Errors
+    /// This returns [`RuntimeError::KeyNotFound`] if the entry is [`VacantEntry`].
+    pub fn into_value(self) -> DbResult<T::Value> {
+        match self {
+            Self::Occupied(entry) => Ok(entry.into_value()),
+            Self::Vacant(_) => Err(RuntimeError::KeyNotFound),
+        }
+    }
+
+    /// [`OccupiedEntry::remove`] the value if it already exists, else do nothing.
+    ///
+    /// # Errors
+    /// This returns [`RuntimeError::KeyNotFound`] if the entry is [`VacantEntry`].
+    pub fn and_remove(self) -> DbResult<T::Value> {
+        match self {
+            Self::Occupied(entry) => entry.remove(),
+            Self::Vacant(_) => Err(RuntimeError::KeyNotFound),
+        }
+    }
+
+    /// [`OccupiedEntry::update`] the value if it already exists
+    ///
+    /// This functions does nothing if the entry is [`VacantEntry`].
     pub fn and_update<F>(self, f: F) -> DbResult<Self>
     where
         F: FnOnce(&mut T::Value),