update call-sites

storage/blockchain/src/ops/alt_block/chain.rs

@@ -28,10 +28,10 @@ pub fn update_alt_chain_info(
         Err(e) => return Err(e),
     };
 
-    // try update the info if one exists for this chain.
-    let update = tables
+    tables
         .alt_chain_infos_mut()
-        .update(&alt_block_height.chain_id, |mut info| {
+        .entry(&alt_block_height.chain_id)?
+        .and_update(|info| {
             if info.chain_height < alt_block_height.height + 1 {
                 // If the chain height is increasing we only need to update the chain height.
                 info.chain_height = alt_block_height.height + 1;
@@ -43,25 +43,12 @@ pub fn update_alt_chain_info(
             }
 
             info.chain_height = alt_block_height.height + 1;
-            Some(info)
-        });
-
-    match update {
-        Ok(()) => return Ok(()),
-        Err(RuntimeError::KeyNotFound) => (),
-        Err(e) => return Err(e),
-    }
-
-    // If one doesn't already exist add it.
-
-    tables.alt_chain_infos_mut().put(
-        &alt_block_height.chain_id,
-        &AltChainInfo {
+        })?
+        .or_insert_with(|| AltChainInfo {
             parent_chain: parent_chain.into(),
             common_ancestor_height: alt_block_height.height.checked_sub(1).unwrap(),
             chain_height: alt_block_height.height + 1,
-        },
-    )
+        })
 }
 
 /// Get the height history of an alt-chain in reverse chronological order.

storage/blockchain/src/ops/output.rs

@@ -67,14 +67,9 @@ pub fn remove_output(
     // `btree_map::Entry`-like API, fix `trait DatabaseRw`.
     tables
         .num_outputs_mut()
-        .update(&pre_rct_output_id.amount, |num_outputs| {
-            // INVARIANT: Should never be 0.
-            if num_outputs == 1 {
-                None
-            } else {
-                Some(num_outputs - 1)
-            }
-        })?;
+        .entry(&pre_rct_output_id.amount)?
+        .and_remove(|num_outputs| *num_outputs == 1)?
+        .and_update(|num_outputs| *num_outputs -= 1)?;
 
     // Delete the output data itself.
     tables.outputs_mut().delete(pre_rct_output_id)

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

@@ -189,6 +189,20 @@ 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

@@ -293,11 +293,11 @@ fn db_read_write() {
         assert_value(value);
     }
 
-    // Assert `Entry` returns the correct value.
+    // Assert `take()` works.
     {
         let mut key = KEY;
         key += 1;
-        let value = table.entry(&key).unwrap().and_remove().unwrap();
+        let value = table.take(&key).unwrap();
         assert_eq!(value, VALUE);
 
         let get = table.get(&KEY);

storage/database/src/database.rs

@@ -171,6 +171,14 @@ 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>;
+
     /// Removes and returns the first `(key, value)` pair in the database.
     ///
     #[doc = doc_database!()]

storage/database/src/entry/entry.rs

@@ -29,6 +29,16 @@ where
     T: Table,
     D: DatabaseRw<T>,
 {
+    /// TODO
+    pub const fn is_occupied(&self) -> bool {
+        matches!(self, Self::Occupied(_))
+    }
+
+    /// TODO
+    pub const fn is_vacant(&self) -> bool {
+        matches!(self, Self::Vacant(_))
+    }
+
     /// Ensures a value is in the entry by inserting the `default` if empty.
     ///
     /// This only inserts if the entry is [`VacantEntry`].
@@ -44,24 +54,24 @@ where
     /// This only inserts if the entry is [`VacantEntry`].
     pub fn or_insert_with<F>(self, default: F) -> DbResult<()>
     where
-        F: FnOnce() -> &'a T::Value,
+        F: FnOnce() -> T::Value,
     {
         match self {
             Self::Occupied(_) => Ok(()),
-            Self::Vacant(entry) => entry.insert(default()),
+            Self::Vacant(entry) => entry.insert(&default()),
         }
     }
 
     /// Same as [`Self::or_insert_with`] but gives access to the key.
     pub fn or_insert_with_key<F>(self, default: F) -> DbResult<()>
     where
-        F: FnOnce(&'a T::Key) -> &'a T::Value,
+        F: FnOnce(&'a T::Key) -> T::Value,
     {
         match self {
             Self::Occupied(_) => Ok(()),
             Self::Vacant(entry) => {
                 let key = entry.key;
-                entry.insert(default(key))
+                entry.insert(&default(key))
             }
         }
     }
@@ -96,15 +106,23 @@ where
         }
     }
 
-    /// [`OccupiedEntry::remove`] the value if it already exists, else do nothing.
+    /// Conditionally [`OccupiedEntry::remove`] the value if it already exists.
     ///
-    /// # 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),
-        }
+    /// This functions does nothing if the entry is [`VacantEntry`].
+    pub fn and_remove<F>(self, f: F) -> DbResult<Self>
+    where
+        F: FnOnce(&T::Value) -> bool,
+    {
+        Ok(match self {
+            Self::Occupied(entry) => {
+                if f(&entry.value) {
+                    entry.remove()?.0
+                } else {
+                    Self::Occupied(entry)
+                }
+            }
+            Self::Vacant(entry) => Self::Vacant(entry),
+        })
     }
 
     /// [`OccupiedEntry::update`] the value if it already exists

storage/database/src/entry/occupied_entry.rs

@@ -2,6 +2,8 @@
 
 use crate::{DatabaseRw, DbResult, Table};
 
+use super::{Entry, VacantEntry};
+
 /// A view into an occupied entry in a [`DatabaseRw`]. It is part of [`crate::entry::Entry`].
 pub struct OccupiedEntry<'a, T, D>
 where
@@ -13,7 +15,7 @@ where
     pub(crate) value: T::Value,
 }
 
-impl<T, D> OccupiedEntry<'_, T, D>
+impl<'a, T, D> OccupiedEntry<'a, T, D>
 where
     T: Table,
     D: DatabaseRw<T>,
@@ -52,7 +54,18 @@ where
     }
 
     /// Remove this entry.
-    pub fn remove(self) -> DbResult<T::Value> {
-        DatabaseRw::delete(self.db, self.key).map(|()| Ok(self.value))?
+    ///
+    /// The returns values are:
+    /// - An [`Entry::VacantEntry`]
+    /// - The value that was removed
+    pub fn remove(self) -> DbResult<(Entry<'a, T, D>, T::Value)> {
+        DatabaseRw::delete(self.db, self.key)?;
+        Ok((
+            Entry::Vacant(VacantEntry {
+                db: self.db,
+                key: self.key,
+            }),
+            self.value,
+        ))
     }
 }

storage/txpool/src/service/write.rs

@@ -114,10 +114,11 @@ fn promote(env: &ConcreteEnv, tx_hash: &TransactionHash) -> DbResult<TxpoolWrite
     let res = || {
         let mut tx_infos = env_inner.open_db_rw::<TransactionInfos>(&tx_rw)?;
 
-        tx_infos.update(tx_hash, |mut info| {
+        tx_infos.entry(tx_hash)?.and_update(|info| {
             info.flags.remove(TxStateFlags::STATE_STEM);
-            Some(info)
-        })
+        })?;
+
+        Ok(())
     };
 
     if let Err(e) = res() {