database: Entry API

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

@@ -62,19 +62,12 @@ pub fn remove_output(
     pre_rct_output_id: &PreRctOutputId,
     tables: &mut impl TablesMut,
 ) -> DbResult<()> {
-    // Decrement the amount index by 1, or delete the entry out-right.
-    // FIXME: this would be much better expressed with a
-    // `btree_map::Entry`-like API, fix `trait DatabaseRw`.
+    // INVARIANT: `num_outputs` should never be 0.
     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)?
+        .remove_if(|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

@@ -6,6 +6,7 @@ use std::{cell::RefCell, ops::RangeBounds};
 use crate::{
     backend::heed::types::HeedDb,
     database::{DatabaseIter, DatabaseRo, DatabaseRw},
+    entry::{Entry, OccupiedEntry, VacantEntry},
     error::{DbResult, RuntimeError},
     table::Table,
 };
@@ -239,6 +240,19 @@ impl<T: Table> DatabaseRw<T> for HeedTableRw<'_, '_, T> {
             Ok(false) => unreachable!(),
         }
     }
+
+    #[inline]
+    fn entry<'a>(&'a mut self, key: &'a T::Key) -> DbResult<Entry<'a, T, Self>> {
+        match DatabaseRo::get(self, key) {
+            Ok(value) => Ok(Entry::Occupied(OccupiedEntry {
+                db: self,
+                key,
+                value,
+            })),
+            Err(RuntimeError::KeyNotFound) => Ok(Entry::Vacant(VacantEntry { db: self, key })),
+            Err(e) => Err(e),
+        }
+    }
 }
 
 //---------------------------------------------------------------------------------------------------- Tests

storage/database/src/backend/redb/database.rs

@@ -11,6 +11,7 @@ use crate::{
         types::{RedbTableRo, RedbTableRw},
     },
     database::{DatabaseIter, DatabaseRo, DatabaseRw},
+    entry::{Entry, OccupiedEntry, VacantEntry},
     error::{DbResult, RuntimeError},
     table::Table,
 };
@@ -162,7 +163,7 @@ unsafe impl<T: Table + 'static> DatabaseRo<T> for RedbTableRw<'_, T::Key, T::Val
     }
 }
 
-impl<T: Table + 'static> DatabaseRw<T> for RedbTableRw<'_, T::Key, T::Value> {
+impl<T: Table> DatabaseRw<T> for RedbTableRw<'_, T::Key, T::Value> {
     // `redb` returns the value after function calls so we end with Ok(()) instead.
 
     #[inline]
@@ -197,6 +198,19 @@ impl<T: Table + 'static> DatabaseRw<T> for RedbTableRw<'_, T::Key, T::Value> {
         let (key, value) = redb::Table::pop_last(self)?.ok_or(RuntimeError::KeyNotFound)?;
         Ok((key.value(), value.value()))
     }
+
+    #[inline]
+    fn entry<'a>(&'a mut self, key: &'a T::Key) -> DbResult<Entry<'a, T, Self>> {
+        match get::<T>(self, key) {
+            Ok(value) => Ok(Entry::Occupied(OccupiedEntry {
+                db: self,
+                key,
+                value,
+            })),
+            Err(RuntimeError::KeyNotFound) => Ok(Entry::Vacant(VacantEntry { db: self, key })),
+            Err(e) => Err(e),
+        }
+    }
 }
 
 //---------------------------------------------------------------------------------------------------- Tests

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();
 

storage/database/src/database.rs

@@ -4,6 +4,7 @@
 use std::ops::RangeBounds;
 
 use crate::{
+    entry::Entry,
     error::{DbResult, RuntimeError},
     table::Table,
 };
@@ -151,7 +152,7 @@ pub unsafe trait DatabaseRo<T: Table> {
 /// Database (key-value store) read/write abstraction.
 ///
 /// All [`DatabaseRo`] functions are also callable by [`DatabaseRw`].
-pub trait DatabaseRw<T: Table>: DatabaseRo<T> {
+pub trait DatabaseRw<T: Table>: DatabaseRo<T> + Sized {
     /// Insert a key-value pair into the database.
     ///
     /// This will overwrite any existing key-value pairs.
@@ -178,30 +179,6 @@ pub trait DatabaseRw<T: Table>: DatabaseRo<T> {
     #[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!()]
@@ -211,4 +188,9 @@ pub trait DatabaseRw<T: Table>: DatabaseRo<T> {
     ///
     #[doc = doc_database!()]
     fn pop_last(&mut self) -> DbResult<(T::Key, T::Value)>;
+
+    /// Gets the given key's corresponding entry for in-place manipulation.
+    ///
+    #[doc = doc_database!()]
+    fn entry<'a>(&'a mut self, key: &'a T::Key) -> DbResult<Entry<'a, T, Self>>;
 }