Optimized the work of InMemoryTransaction for lookups and empty ins…

…ertion (#1974)

Just was walking around this code and noticed that we can remove not
needed `to_vec` and improve the case when we commit changes into empty
`InMemoryTransaction`.

### Before requesting review
- [x] I have reviewed the code myself

CHANGELOG.md

@@ -6,6 +6,9 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 ## [Unreleased]
 
+### Added
+- [#1974](https://github.com/FuelLabs/fuel-core/pull/1974): Optimized the work of `InMemoryTransaction` for lookups and empty insertion.
+
 ### Changed
 - [#1973](https://github.com/FuelLabs/fuel-core/pull/1973): Updated VM initialization benchmark to include many inputs and outputs.
 

crates/fuel-core/src/state.rs

@@ -78,7 +78,7 @@ where
         Ok(self
             .changes
             .get(&column.id())
-            .and_then(|tree| tree.get(&key.to_vec()))
+            .and_then(|tree| tree.get(key))
             .and_then(|operation| match operation {
                 WriteOperation::Insert(value) => Some(value.clone()),
                 WriteOperation::Remove => None,
@@ -100,7 +100,9 @@ where
         if let Some(tree) = self.changes.get(&column.id()) {
             fuel_core_storage::iter::iterator(tree, prefix, start, direction)
                 .filter_map(|(key, value)| match value {
-                    WriteOperation::Insert(value) => Some((key.clone(), value.clone())),
+                    WriteOperation::Insert(value) => {
+                        Some((key.clone().into(), value.clone()))
+                    }
                     WriteOperation::Remove => None,
                 })
                 .map(Ok)

crates/fuel-core/src/state/in_memory/memory_store.rs

@@ -22,7 +22,10 @@ use fuel_core_storage::{
         Value,
         WriteOperation,
     },
-    transactional::Changes,
+    transactional::{
+        Changes,
+        ReferenceBytesKey,
+    },
     Result as StorageResult,
 };
 use std::{
@@ -36,7 +39,7 @@ pub struct MemoryStore<Description = OnChain>
 where
     Description: DatabaseDescription,
 {
-    inner: Vec<Mutex<BTreeMap<Vec<u8>, Value>>>,
+    inner: Vec<Mutex<BTreeMap<ReferenceBytesKey, Value>>>,
     _marker: core::marker::PhantomData<Description>,
 }
 
@@ -68,12 +71,8 @@ where
     ) -> impl Iterator<Item = KVItem> {
         let lock = self.inner[column.as_usize()].lock().expect("poisoned");
 
-        fn clone<K: Clone, V: Clone>(kv: (&K, &V)) -> (K, V) {
-            (kv.0.clone(), kv.1.clone())
-        }
-
         let collection: Vec<_> = iterator(&lock, prefix, start, direction)
-            .map(clone)
+            .map(|(key, value)| (key.clone().into(), value.clone()))
             .collect();
 
         collection.into_iter().map(Ok)
@@ -90,7 +89,7 @@ where
         Ok(self.inner[column.as_usize()]
             .lock()
             .map_err(|e| anyhow::anyhow!("The lock is poisoned: {}", e))?
-            .get(&key.to_vec())
+            .get(key)
             .cloned())
     }
 }

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

@@ -770,7 +770,7 @@ mod tests {
         let ops = vec![(
             Column::Metadata.id(),
             BTreeMap::from_iter(vec![(
-                key.clone(),
+                key.clone().into(),
                 WriteOperation::Insert(value.clone()),
             )]),
         )];
@@ -789,7 +789,7 @@ mod tests {
 
         let ops = vec![(
             Column::Metadata.id(),
-            BTreeMap::from_iter(vec![(key.clone(), WriteOperation::Remove)]),
+            BTreeMap::from_iter(vec![(key.clone().into(), WriteOperation::Remove)]),
         )];
         db.commit_changes(None, HashMap::from_iter(ops)).unwrap();
 

crates/storage/src/iter.rs

@@ -12,6 +12,7 @@ use crate::{
         KeyValueInspect,
     },
     structured_storage::TableWithBlueprint,
+    transactional::ReferenceBytesKey,
 };
 use fuel_vm_private::fuel_storage::Mappable;
 use std::{
@@ -200,11 +201,11 @@ impl<S> IteratorOverTable for S {}
 
 /// Returns an iterator over the values in the `BTreeMap`.
 pub fn iterator<'a, V>(
-    tree: &'a BTreeMap<Vec<u8>, V>,
+    tree: &'a BTreeMap<ReferenceBytesKey, V>,
     prefix: Option<&[u8]>,
     start: Option<&[u8]>,
     direction: IterDirection,
-) -> impl Iterator<Item = (&'a Vec<u8>, &'a V)> + 'a {
+) -> impl Iterator<Item = (&'a ReferenceBytesKey, &'a V)> + 'a {
     match (prefix, start) {
         (None, None) => {
             if direction == IterDirection::Forward {

crates/storage/src/transactional.rs

@@ -13,8 +13,10 @@ use crate::{
     Result as StorageResult,
 };
 use std::{
+    borrow::Borrow,
     collections::{
         btree_map::Entry,
+        hash_map,
         BTreeMap,
         HashMap,
     },
@@ -117,8 +119,73 @@ pub trait Modifiable {
     fn commit_changes(&mut self, changes: Changes) -> StorageResult<()>;
 }
 
+/// The wrapper around the `Vec<u8>` that supports `Borrow<[u8]>`.
+/// It allows the use of bytes slices to do lookups in the collections.
+#[derive(
+    Debug,
+    Clone,
+    PartialEq,
+    Eq,
+    Ord,
+    PartialOrd,
+    Hash,
+    serde::Serialize,
+    serde::Deserialize,
+)]
+pub struct ReferenceBytesKey(Vec<u8>);
+
+impl From<Vec<u8>> for ReferenceBytesKey {
+    fn from(value: Vec<u8>) -> Self {
+        ReferenceBytesKey(value)
+    }
+}
+
+impl From<ReferenceBytesKey> for Vec<u8> {
+    fn from(value: ReferenceBytesKey) -> Self {
+        value.0
+    }
+}
+
+impl Borrow<[u8]> for ReferenceBytesKey {
+    fn borrow(&self) -> &[u8] {
+        self.0.as_slice()
+    }
+}
+
+impl Borrow<Vec<u8>> for ReferenceBytesKey {
+    fn borrow(&self) -> &Vec<u8> {
+        &self.0
+    }
+}
+
+impl AsRef<[u8]> for ReferenceBytesKey {
+    fn as_ref(&self) -> &[u8] {
+        self.0.as_slice()
+    }
+}
+
+impl AsMut<[u8]> for ReferenceBytesKey {
+    fn as_mut(&mut self) -> &mut [u8] {
+        self.0.as_mut_slice()
+    }
+}
+
+impl core::ops::Deref for ReferenceBytesKey {
+    type Target = Vec<u8>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl core::ops::DerefMut for ReferenceBytesKey {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
 /// The type describing the list of changes to the storage.
-pub type Changes = HashMap<u32, BTreeMap<Vec<u8>, WriteOperation>>;
+pub type Changes = HashMap<u32, BTreeMap<ReferenceBytesKey, WriteOperation>>;
 
 impl<Storage> From<StorageTransaction<Storage>> for Changes {
     fn from(transaction: StorageTransaction<Storage>) -> Self {
@@ -198,7 +265,14 @@ where
 impl<Storage> Modifiable for InMemoryTransaction<Storage> {
     fn commit_changes(&mut self, changes: Changes) -> StorageResult<()> {
         for (column, value) in changes.into_iter() {
-            let btree = self.changes.entry(column).or_default();
+            let btree = match self.changes.entry(column) {
+                hash_map::Entry::Vacant(vacant) => {
+                    vacant.insert(value);
+                    continue
+                }
+                hash_map::Entry::Occupied(occupied) => occupied.into_mut(),
+            };
+
             for (k, v) in value {
                 match &self.policy {
                     ConflictPolicy::Fail => {
@@ -233,10 +307,9 @@ where
     S: KeyValueInspect<Column = Column>,
 {
     fn get_from_changes(&self, key: &[u8], column: Column) -> Option<&WriteOperation> {
-        let k = key.to_vec();
         self.changes
             .get(&column.id())
-            .and_then(|btree| btree.get(&k))
+            .and_then(|btree| btree.get(key))
     }
 }
 
@@ -321,7 +394,7 @@ where
         column: Self::Column,
         value: Value,
     ) -> StorageResult<()> {
-        let k = key.to_vec();
+        let k = key.to_vec().into();
         self.changes
             .entry(column.id())
             .or_default()
@@ -335,7 +408,7 @@ where
         column: Self::Column,
         value: Value,
     ) -> StorageResult<Option<Value>> {
-        let k = key.to_vec();
+        let k = key.to_vec().into();
         let entry = self.changes.entry(column.id()).or_default().entry(k);
 
         match entry {
@@ -360,7 +433,7 @@ where
         column: Self::Column,
         buf: &[u8],
     ) -> StorageResult<usize> {
-        let k = key.to_vec();
+        let k = key.to_vec().into();
         self.changes
             .entry(column.id())
             .or_default()
@@ -369,7 +442,7 @@ where
     }
 
     fn take(&mut self, key: &[u8], column: Self::Column) -> StorageResult<Option<Value>> {
-        let k = key.to_vec();
+        let k = key.to_vec().into();
         let entry = self.changes.entry(column.id()).or_default().entry(k);
 
         match entry {
@@ -389,7 +462,7 @@ where
     }
 
     fn delete(&mut self, key: &[u8], column: Self::Column) -> StorageResult<()> {
-        let k = key.to_vec();
+        let k = key.to_vec().into();
         self.changes
             .entry(column.id())
             .or_default()
@@ -409,7 +482,7 @@ where
     {
         let btree = self.changes.entry(column.id()).or_default();
         entries.for_each(|(key, operation)| {
-            btree.insert(key, operation);
+            btree.insert(key.into(), operation);
         });
         Ok(())
     }
@@ -455,10 +528,10 @@ mod test {
                 for (key, value) in value {
                     match value {
                         WriteOperation::Insert(value) => {
-                            self.storage.insert((column, key), value);
+                            self.storage.insert((column, key.into()), value);
                         }
                         WriteOperation::Remove => {
-                            self.storage.remove(&(column, key));
+                            self.storage.remove(&(column, key.into()));
                         }
                     }
                 }