Remove duplicating logic in the KeyValueStore trait (#1559)

Preparation before start work on
#1548.

The `KeyValueStore` trait has some duplicated logic. This PR removes it,
minimizing the number of methods that we need to implement. Also I
applied the original ordering of the method as in the trait.

bin/fuel-core/src/cli/run.rs

@@ -425,16 +425,12 @@ async fn shutdown_signal() -> anyhow::Result<()> {
 
         let mut sigint =
             tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt())?;
-        loop {
-            tokio::select! {
-                _ = sigterm.recv() => {
-                    tracing::info!("sigterm received");
-                    break;
-                }
-                _ = sigint.recv() => {
-                    tracing::info!("sigint received");
-                    break;
-                }
+        tokio::select! {
+            _ = sigterm.recv() => {
+                tracing::info!("sigterm received");
+            }
+            _ = sigint.recv() => {
+                tracing::info!("sigint received");
             }
         }
     }

crates/fuel-core/src/database.rs

@@ -56,6 +56,7 @@ type DatabaseResult<T> = Result<T>;
 // TODO: Extract `Database` and all belongs into `fuel-core-database`.
 #[cfg(feature = "rocksdb")]
 use crate::state::rocks_db::RocksDb;
+use crate::state::Value;
 #[cfg(feature = "rocksdb")]
 use std::path::Path;
 #[cfg(feature = "rocksdb")]
@@ -84,7 +85,14 @@ pub mod transactions;
 /// Database tables column ids to the corresponding [`fuel_core_storage::Mappable`] table.
 #[repr(u32)]
 #[derive(
-    Copy, Clone, Debug, strum_macros::EnumCount, PartialEq, Eq, enum_iterator::Sequence,
+    Copy,
+    Clone,
+    Debug,
+    strum_macros::EnumCount,
+    strum_macros::IntoStaticStr,
+    PartialEq,
+    Eq,
+    enum_iterator::Sequence,
 )]
 pub enum Column {
     /// The column id of metadata about the blockchain
@@ -152,6 +160,16 @@ impl Column {
     }
 }
 
+impl crate::state::StorageColumn for Column {
+    fn name(&self) -> &'static str {
+        self.into()
+    }
+
+    fn id(&self) -> u32 {
+        *self as u32
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct Database {
     data: DataSource,
@@ -253,13 +271,13 @@ impl Database {
 /// Mutable methods.
 // TODO: Add `&mut self` to them.
 impl Database {
-    fn insert<K: AsRef<[u8]>, V: Serialize, R: DeserializeOwned>(
+    fn insert<K: AsRef<[u8]>, V: Serialize + ?Sized, R: DeserializeOwned>(
         &self,
         key: K,
         column: Column,
         value: &V,
     ) -> DatabaseResult<Option<R>> {
-        let result = self.data.put(
+        let result = self.data.replace(
             key.as_ref(),
             column,
             Arc::new(postcard::to_stdvec(value).map_err(|_| DatabaseError::Codec)?),
@@ -273,6 +291,16 @@ impl Database {
         }
     }
 
+    fn insert_raw<K: AsRef<[u8]>, V: AsRef<[u8]>>(
+        &self,
+        key: K,
+        column: Column,
+        value: V,
+    ) -> DatabaseResult<Option<Value>> {
+        self.data
+            .replace(key.as_ref(), column, Arc::new(value.as_ref().to_vec()))
+    }
+
     fn batch_insert<K: AsRef<[u8]>, V: Serialize, S>(
         &self,
         column: Column,
@@ -299,36 +327,19 @@ impl Database {
         self.data.batch_write(&mut set.into_iter())
     }
 
-    fn remove<V: DeserializeOwned>(
+    fn take<V: DeserializeOwned>(
         &self,
         key: &[u8],
         column: Column,
     ) -> DatabaseResult<Option<V>> {
         self.data
-            .delete(key, column)?
+            .take(key, column)?
             .map(|val| postcard::from_bytes(&val).map_err(|_| DatabaseError::Codec))
             .transpose()
     }
 
-    fn write(&self, key: &[u8], column: Column, buf: &[u8]) -> DatabaseResult<usize> {
-        self.data.write(key, column, buf)
-    }
-
-    fn replace(
-        &self,
-        key: &[u8],
-        column: Column,
-        buf: &[u8],
-    ) -> DatabaseResult<(usize, Option<Vec<u8>>)> {
-        self.data
-            .replace(key, column, buf)
-            .map(|(size, value)| (size, value.map(|value| value.deref().clone())))
-    }
-
-    fn take(&self, key: &[u8], column: Column) -> DatabaseResult<Option<Vec<u8>>> {
-        self.data
-            .take(key, column)
-            .map(|value| value.map(|value| value.deref().clone()))
+    fn take_raw(&self, key: &[u8], column: Column) -> DatabaseResult<Option<Value>> {
+        self.data.take(key, column)
     }
 }
 
@@ -353,7 +364,7 @@ impl Database {
 
     fn read_alloc(&self, key: &[u8], column: Column) -> DatabaseResult<Option<Vec<u8>>> {
         self.data
-            .read_alloc(key, column)
+            .get(key, column)
             .map(|value| value.map(|value| value.deref().clone()))
     }
 

crates/fuel-core/src/database/balances.rs

@@ -100,7 +100,7 @@ impl StorageMutate<ContractsAssets> for Database {
         &mut self,
         key: &<ContractsAssets as Mappable>::Key,
     ) -> Result<Option<<ContractsAssets as Mappable>::OwnedValue>, Self::Error> {
-        let prev = Database::remove(self, key.as_ref(), Column::ContractsAssets)
+        let prev = Database::take(self, key.as_ref(), Column::ContractsAssets)
             .map_err(Into::into);
 
         // Get latest metadata entry for this contract id

crates/fuel-core/src/database/block.rs

@@ -106,7 +106,7 @@ impl StorageMutate<FuelBlocks> for Database {
 
     fn remove(&mut self, key: &BlockId) -> Result<Option<CompressedBlock>, Self::Error> {
         let prev: Option<CompressedBlock> =
-            Database::remove(self, key.as_slice(), Column::FuelBlocks)?;
+            Database::take(self, key.as_slice(), Column::FuelBlocks)?;
 
         if let Some(block) = &prev {
             let height = block.header().height();

crates/fuel-core/src/database/coin.rs

@@ -92,7 +92,7 @@ impl StorageMutate<Coins> for Database {
 
     fn remove(&mut self, key: &UtxoId) -> Result<Option<CompressedCoin>, Self::Error> {
         let coin: Option<CompressedCoin> =
-            Database::remove(self, &utxo_id_to_bytes(key), Column::Coins)?;
+            Database::take(self, &utxo_id_to_bytes(key), Column::Coins)?;
 
         // cleanup secondary index
         if let Some(coin) = &coin {

crates/fuel-core/src/database/contracts.rs

@@ -22,7 +22,6 @@ use fuel_core_storage::{
     StorageMutate,
     StorageRead,
     StorageSize,
-    StorageWrite,
 };
 use fuel_core_types::{
     entities::contract::ContractUtxoInfo,
@@ -74,19 +73,18 @@ impl StorageMutate<ContractsRawCode> for Database {
         key: &<ContractsRawCode as Mappable>::Key,
         value: &<ContractsRawCode as Mappable>::Value,
     ) -> Result<Option<<ContractsRawCode as Mappable>::OwnedValue>, Self::Error> {
-        let existing =
-            Database::replace(self, key.as_ref(), Column::ContractsRawCode, value)?;
-        Ok(existing.1.map(Contract::from))
+        let result = Database::insert_raw(self, key, Column::ContractsRawCode, value)?;
+
+        Ok(result.map(|v| Contract::from(v.as_ref().clone())))
     }
 
     fn remove(
         &mut self,
         key: &<ContractsRawCode as Mappable>::Key,
     ) -> Result<Option<<ContractsRawCode as Mappable>::OwnedValue>, Self::Error> {
-        Ok(
-            <Self as StorageWrite<ContractsRawCode>>::take(self, key)?
-                .map(Contract::from),
-        )
+        let result = Database::take_raw(self, key.as_ref(), Column::ContractsRawCode)?;
+
+        Ok(result.map(|v| Contract::from(v.as_ref().clone())))
     }
 }
 
@@ -110,44 +108,6 @@ impl StorageRead<ContractsRawCode> for Database {
     }
 }
 
-impl StorageWrite<ContractsRawCode> for Database {
-    fn write(&mut self, key: &ContractId, buf: Vec<u8>) -> Result<usize, Self::Error> {
-        Ok(Database::write(
-            self,
-            key.as_ref(),
-            Column::ContractsRawCode,
-            &buf,
-        )?)
-    }
-
-    fn replace(
-        &mut self,
-        key: &<ContractsRawCode as Mappable>::Key,
-        buf: Vec<u8>,
-    ) -> Result<(usize, Option<Vec<u8>>), <Self as StorageInspect<ContractsRawCode>>::Error>
-    where
-        Self: StorageSize<ContractsRawCode>,
-    {
-        Ok(Database::replace(
-            self,
-            key.as_ref(),
-            Column::ContractsRawCode,
-            &buf,
-        )?)
-    }
-
-    fn take(
-        &mut self,
-        key: &<ContractsRawCode as Mappable>::Key,
-    ) -> Result<Option<Vec<u8>>, Self::Error> {
-        Ok(Database::take(
-            self,
-            key.as_ref(),
-            Column::ContractsRawCode,
-        )?)
-    }
-}
-
 impl Database {
     pub fn get_contract_config_by_id(
         &self,

crates/fuel-core/src/database/message.rs

@@ -67,10 +67,10 @@ impl StorageMutate<Messages> for Database {
 
     fn remove(&mut self, key: &Nonce) -> Result<Option<Message>, Self::Error> {
         let result: Option<Message> =
-            Database::remove(self, key.database_key().as_ref(), Column::Messages)?;
+            Database::take(self, key.database_key().as_ref(), Column::Messages)?;
 
         if let Some(message) = &result {
-            Database::remove::<bool>(
+            Database::take::<bool>(
                 self,
                 &owner_msg_id_key(&message.recipient, key),
                 Column::OwnedMessageIds,

crates/fuel-core/src/database/state.rs

@@ -99,7 +99,7 @@ impl StorageMutate<ContractsState> for Database {
         &mut self,
         key: &<ContractsState as Mappable>::Key,
     ) -> Result<Option<<ContractsState as Mappable>::OwnedValue>, Self::Error> {
-        let prev = Database::remove(self, key.as_ref(), Column::ContractsState)
+        let prev = Database::take(self, key.as_ref(), Column::ContractsState)
             .map_err(Into::into);
 
         // Get latest metadata entry for this contract id

crates/fuel-core/src/database/storage.rs

@@ -232,8 +232,7 @@ where
     }
 
     fn remove(&mut self, key: &T::Key) -> StorageResult<Option<T::OwnedValue>> {
-        Database::remove(self, key.database_key().as_ref(), T::column())
-            .map_err(Into::into)
+        Database::take(self, key.database_key().as_ref(), T::column()).map_err(Into::into)
     }
 }