From b73346b2a4e645b28394432b22150a8aa56c031f Mon Sep 17 00:00:00 2001
From: green <xgreenx9999@gmail.com>
Date: Thu, 31 Oct 2024 11:41:26 +0100
Subject: [PATCH 1/2] Proposals to multi get PR

---
 crates/fuel-core/src/graphql_api/ports.rs     | 10 ++-----
 .../fuel-core/src/state/generic_database.rs   | 10 +++++--
 .../fuel-core/src/state/historical_rocksdb.rs | 18 +++++++++++
 crates/storage/src/blueprint.rs               | 30 ++++++++-----------
 crates/storage/src/codec.rs                   |  2 +-
 crates/storage/src/kv_store.rs                | 16 ++++++----
 crates/storage/src/lib.rs                     | 10 ++++---
 crates/storage/src/structured_storage.rs      | 10 +++++--
 crates/storage/src/transactional.rs           | 18 ++++++++++-
 9 files changed, 83 insertions(+), 41 deletions(-)

diff --git a/crates/fuel-core/src/graphql_api/ports.rs b/crates/fuel-core/src/graphql_api/ports.rs
index dbc2a200808..2ce9ecdf4b8 100644
--- a/crates/fuel-core/src/graphql_api/ports.rs
+++ b/crates/fuel-core/src/graphql_api/ports.rs
@@ -16,7 +16,6 @@ use fuel_core_storage::{
     },
     Error as StorageError,
     Result as StorageResult,
-    StorageBatchInspect,
     StorageInspect,
     StorageRead,
 };
@@ -174,9 +173,7 @@ pub trait DatabaseDaCompressedBlocks {
 }
 
 /// Trait that specifies all the getters required for messages.
-pub trait DatabaseMessages:
-    StorageInspect<Messages, Error = StorageError> + StorageBatchInspect<Messages>
-{
+pub trait DatabaseMessages: StorageInspect<Messages, Error = StorageError> {
     fn all_messages(
         &self,
         start_message_id: Option<Nonce>,
@@ -199,10 +196,9 @@ pub trait DatabaseRelayedTransactions {
 }
 
 /// Trait that specifies all the getters required for coins
-pub trait DatabaseCoins:
-    StorageInspect<Coins, Error = StorageError> + StorageBatchInspect<Coins>
-{
+pub trait DatabaseCoins: StorageInspect<Coins, Error = StorageError> {
     fn coin(&self, utxo_id: UtxoId) -> StorageResult<Coin>;
+
     fn coins<'a>(
         &'a self,
         utxo_ids: BoxedIter<'a, &'a UtxoId>,
diff --git a/crates/fuel-core/src/state/generic_database.rs b/crates/fuel-core/src/state/generic_database.rs
index 178d1066a22..6d0114c4508 100644
--- a/crates/fuel-core/src/state/generic_database.rs
+++ b/crates/fuel-core/src/state/generic_database.rs
@@ -70,10 +70,14 @@ where
     M: Mappable,
     StructuredStorage<Storage>: StorageBatchInspect<M>,
 {
-    fn get_batch<'a>(
+    fn get_batch<'a, Iter>(
         &'a self,
-        keys: BoxedIter<'a, &'a M::Key>,
-    ) -> BoxedIter<'a, StorageResult<Option<<M as Mappable>::OwnedValue>>> {
+        keys: Iter,
+    ) -> impl Iterator<Item = fuel_core_storage::Result<Option<M::OwnedValue>>> + 'a
+    where
+        Iter: 'a + Iterator<Item = &'a M::Key> + Send,
+        M::Key: 'a,
+    {
         self.storage.get_batch(keys)
     }
 }
diff --git a/crates/fuel-core/src/state/historical_rocksdb.rs b/crates/fuel-core/src/state/historical_rocksdb.rs
index 77c6e586434..559505ba67b 100644
--- a/crates/fuel-core/src/state/historical_rocksdb.rs
+++ b/crates/fuel-core/src/state/historical_rocksdb.rs
@@ -29,6 +29,7 @@ use crate::{
 use fuel_core_storage::{
     iter::{
         BoxedIter,
+        IntoBoxedIter,
         IterDirection,
         IterableStore,
         IteratorOverTable,
@@ -36,6 +37,7 @@ use fuel_core_storage::{
     kv_store::{
         KVItem,
         KeyValueInspect,
+        StorageColumn,
         Value,
         WriteOperation,
     },
@@ -56,6 +58,7 @@ use serde::{
     Serialize,
 };
 use std::{
+    borrow::Cow,
     num::NonZeroU64,
     path::Path,
 };
@@ -395,6 +398,21 @@ where
         self.db.get(key, Column::OriginalColumn(column))
     }
 
+    fn get_batch<'a>(
+        &'a self,
+        keys: BoxedIter<'a, Cow<'a, [u8]>>,
+        column: Self::Column,
+    ) -> BoxedIter<'a, StorageResult<Option<Value>>> {
+        self.db
+            .multi_get(Column::<Description>::OriginalColumn(column).id(), keys)
+            .map(|result| {
+                result
+                    .map(|value| value.map(Into::into))
+                    .map_err(Into::into)
+            })
+            .into_boxed()
+    }
+
     fn read(
         &self,
         key: &[u8],
diff --git a/crates/storage/src/blueprint.rs b/crates/storage/src/blueprint.rs
index a0e7d185144..656b76f615c 100644
--- a/crates/storage/src/blueprint.rs
+++ b/crates/storage/src/blueprint.rs
@@ -9,10 +9,7 @@ use crate::{
         Encode,
         Encoder,
     },
-    iter::{
-        BoxedIter,
-        IntoBoxedIter,
-    },
+    iter::IntoBoxedIter,
     kv_store::{
         BatchOperations,
         KeyValueInspect,
@@ -81,30 +78,29 @@ where
     }
 
     /// Returns multiple values from the storage.
-    fn get_batch<'a>(
+    fn get_batch<'a, Iter>(
         storage: &'a S,
-        keys: BoxedIter<'a, &'a M::Key>,
+        keys: Iter,
         column: S::Column,
-    ) -> BoxedIter<'a, StorageResult<Option<M::OwnedValue>>>
+    ) -> impl Iterator<Item = crate::Result<Option<M::OwnedValue>>> + 'a
     where
+        Iter: 'a + Iterator<Item = &'a M::Key> + Send,
+        M::Key: 'a,
         Self::KeyCodec: 'a,
     {
         let keys = keys
             .map(|key| Self::KeyCodec::encode(key).into_bytes())
             .into_boxed();
 
-        storage
-            .get_batch(keys, column)
-            .map(|result| {
-                result.and_then(|opt| {
-                    opt.map(|value| {
-                        Self::ValueCodec::decode_from_value(value)
-                            .map_err(crate::Error::Codec)
-                    })
-                    .transpose()
+        storage.get_batch(keys, column).map(|result| {
+            result.and_then(|opt| {
+                opt.map(|value| {
+                    Self::ValueCodec::decode_from_value(value)
+                        .map_err(crate::Error::Codec)
                 })
+                .transpose()
             })
-            .into_boxed()
+        })
     }
 }
 
diff --git a/crates/storage/src/codec.rs b/crates/storage/src/codec.rs
index f80dbfb6d1f..f77f91cc044 100644
--- a/crates/storage/src/codec.rs
+++ b/crates/storage/src/codec.rs
@@ -36,7 +36,7 @@ pub trait Encode<T: ?Sized> {
         T: 'a;
 
     /// Encodes the object to the bytes and passes it to the `Encoder`.
-    fn encode<'a>(t: &'a T) -> Self::Encoder<'a>;
+    fn encode(t: &T) -> Self::Encoder<'_>;
 
     /// Returns the serialized object as an [`Value`].
     fn encode_as_value(t: &T) -> Value {
diff --git a/crates/storage/src/kv_store.rs b/crates/storage/src/kv_store.rs
index 42fb1dda9b2..46f6b404f1c 100644
--- a/crates/storage/src/kv_store.rs
+++ b/crates/storage/src/kv_store.rs
@@ -1,10 +1,7 @@
 //! The module provides plain abstract definition of the key-value store.
 
 use crate::{
-    iter::{
-        BoxedIter,
-        IntoBoxedIter,
-    },
+    iter::BoxedIter,
     Error as StorageError,
     Result as StorageResult,
 };
@@ -17,6 +14,7 @@ use alloc::{
     vec::Vec,
 };
 
+use crate::iter::IntoBoxedIter;
 #[cfg(feature = "std")]
 use core::ops::Deref;
 
@@ -31,7 +29,7 @@ pub type KVItem = StorageResult<(Key, Value)>;
 pub type KeyItem = StorageResult<Key>;
 
 /// A column of the storage.
-pub trait StorageColumn: Copy + core::fmt::Debug + Send + Sync {
+pub trait StorageColumn: Copy + core::fmt::Debug + Send + Sync + 'static {
     /// Returns the name of the column.
     fn name(&self) -> String;
 
@@ -121,6 +119,14 @@ where
         self.deref().get(key, column)
     }
 
+    fn get_batch<'a>(
+        &'a self,
+        keys: BoxedIter<'a, Cow<'a, [u8]>>,
+        column: Self::Column,
+    ) -> BoxedIter<'a, StorageResult<Option<Value>>> {
+        self.deref().get_batch(keys, column)
+    }
+
     fn read(
         &self,
         key: &[u8],
diff --git a/crates/storage/src/lib.rs b/crates/storage/src/lib.rs
index 2def12e2d03..1977dacd1ef 100644
--- a/crates/storage/src/lib.rs
+++ b/crates/storage/src/lib.rs
@@ -17,7 +17,6 @@ extern crate alloc;
 use anyhow::anyhow;
 use core::array::TryFromSliceError;
 use fuel_core_types::services::executor::Error as ExecutorError;
-use iter::BoxedIter;
 
 #[cfg(feature = "alloc")]
 use alloc::{
@@ -182,10 +181,13 @@ pub trait StorageBatchMutate<Type: Mappable>: StorageMutate<Type> {
 /// getting values one by one.
 pub trait StorageBatchInspect<Type: Mappable> {
     /// Get a batch of values associated with the provided keys.
-    fn get_batch<'a>(
+    fn get_batch<'a, Iter>(
         &'a self,
-        keys: BoxedIter<'a, &'a Type::Key>,
-    ) -> BoxedIter<'a, Result<Option<Type::OwnedValue>>>;
+        keys: Iter,
+    ) -> impl Iterator<Item = Result<Option<Type::OwnedValue>>> + 'a
+    where
+        Iter: 'a + Iterator<Item = &'a Type::Key> + Send,
+        Type::Key: 'a;
 }
 
 /// Creates `StorageError::NotFound` error with file and line information inside.
diff --git a/crates/storage/src/structured_storage.rs b/crates/storage/src/structured_storage.rs
index b6a36028058..d0396a528f2 100644
--- a/crates/storage/src/structured_storage.rs
+++ b/crates/storage/src/structured_storage.rs
@@ -285,10 +285,14 @@ where
     M::Blueprint: BlueprintInspect<M, StructuredStorage<S>>,
     <M::Blueprint as BlueprintInspect<M, StructuredStorage<S>>>::KeyCodec: 'static,
 {
-    fn get_batch<'a>(
+    fn get_batch<'a, Iter>(
         &'a self,
-        keys: BoxedIter<'a, &'a M::Key>,
-    ) -> BoxedIter<'a, StorageResult<Option<M::OwnedValue>>> {
+        keys: Iter,
+    ) -> impl Iterator<Item = crate::Result<Option<M::OwnedValue>>> + 'a
+    where
+        Iter: 'a + Iterator<Item = &'a M::Key> + Send,
+        M::Key: 'a,
+    {
         M::Blueprint::get_batch(self, keys, M::column())
     }
 }
diff --git a/crates/storage/src/transactional.rs b/crates/storage/src/transactional.rs
index 14ec74159ed..14cb4cfe0c8 100644
--- a/crates/storage/src/transactional.rs
+++ b/crates/storage/src/transactional.rs
@@ -1,6 +1,10 @@
 //! The primitives to work with storage in transactional mode.
 
 use crate::{
+    iter::{
+        BoxedIter,
+        IntoBoxedIter,
+    },
     kv_store::{
         BatchOperations,
         KeyValueInspect,
@@ -14,6 +18,7 @@ use crate::{
 };
 use core::borrow::Borrow;
 
+use alloc::borrow::Cow;
 #[cfg(feature = "alloc")]
 use alloc::{
     boxed::Box,
@@ -27,7 +32,6 @@ use alloc::{
 #[cfg(feature = "test-helpers")]
 use crate::{
     iter::{
-        BoxedIter,
         IterDirection,
         IterableStore,
     },
@@ -384,6 +388,18 @@ where
         }
     }
 
+    fn get_batch<'a>(
+        &'a self,
+        keys: BoxedIter<'a, Cow<'a, [u8]>>,
+        column: Self::Column,
+    ) -> BoxedIter<'a, StorageResult<Option<Value>>> {
+        if !self.changes.contains_key(&column.id()) {
+            return self.storage.get_batch(keys, column);
+        } else {
+            keys.map(move |key| self.get(&key, column)).into_boxed()
+        }
+    }
+
     fn read(
         &self,
         key: &[u8],

From 9f5bcdfd5bc9796151dbe043f0c2baf432028e1a Mon Sep 17 00:00:00 2001
From: green <xgreenx9999@gmail.com>
Date: Thu, 31 Oct 2024 12:11:53 +0100
Subject: [PATCH 2/2] Proposals to multi get PR with coded and `to_vec`

---
 crates/storage/src/blueprint.rs          |  1 -
 crates/storage/src/blueprint/plain.rs    |  4 ++--
 crates/storage/src/blueprint/sparse.rs   | 10 ++++------
 crates/storage/src/codec.rs              |  2 +-
 crates/storage/src/kv_store.rs           |  8 ++++++--
 crates/storage/src/structured_storage.rs |  8 ++++++--
 crates/storage/src/transactional.rs      |  6 +++---
 7 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/crates/storage/src/blueprint.rs b/crates/storage/src/blueprint.rs
index 656b76f615c..9845c37b967 100644
--- a/crates/storage/src/blueprint.rs
+++ b/crates/storage/src/blueprint.rs
@@ -86,7 +86,6 @@ where
     where
         Iter: 'a + Iterator<Item = &'a M::Key> + Send,
         M::Key: 'a,
-        Self::KeyCodec: 'a,
     {
         let keys = keys
             .map(|key| Self::KeyCodec::encode(key).into_bytes())
diff --git a/crates/storage/src/blueprint/plain.rs b/crates/storage/src/blueprint/plain.rs
index 509fc1e9dab..9c41ac63b2f 100644
--- a/crates/storage/src/blueprint/plain.rs
+++ b/crates/storage/src/blueprint/plain.rs
@@ -134,7 +134,7 @@ where
             set.map(|(key, value)| {
                 let key_encoder =
                     <M::Blueprint as BlueprintInspect<M, S>>::KeyCodec::encode(key);
-                let key_bytes = key_encoder.as_bytes().to_vec();
+                let key_bytes = key_encoder.into_bytes();
                 let value =
                     <M::Blueprint as BlueprintInspect<M, S>>::ValueCodec::encode_as_value(
                         value,
@@ -158,7 +158,7 @@ where
             set.map(|key| {
                 let key_encoder =
                     <M::Blueprint as BlueprintInspect<M, S>>::KeyCodec::encode(key);
-                let key_bytes = key_encoder.as_bytes().to_vec();
+                let key_bytes = key_encoder.into_bytes();
                 (key_bytes, WriteOperation::Remove)
             }),
         )
diff --git a/crates/storage/src/blueprint/sparse.rs b/crates/storage/src/blueprint/sparse.rs
index 9d21db77d00..fdd54d266ce 100644
--- a/crates/storage/src/blueprint/sparse.rs
+++ b/crates/storage/src/blueprint/sparse.rs
@@ -326,7 +326,7 @@ where
 
         let encoded_set = set
             .map(|(key, value)| {
-                let key = KeyCodec::encode(key).as_bytes().into_owned();
+                let key = KeyCodec::encode(key).into_bytes();
                 let value = ValueCodec::encode(value).as_bytes().into_owned();
                 (key, value)
             })
@@ -346,9 +346,7 @@ where
         )?;
 
         let nodes = nodes.iter().map(|(key, value)| {
-            let key = NodeKeyCodec::<S, Nodes>::encode(key)
-                .as_bytes()
-                .into_owned();
+            let key = NodeKeyCodec::<S, Nodes>::encode(key).into_bytes();
             let value = NodeValueCodec::<S, Nodes>::encode_as_value(value);
             (key, WriteOperation::Insert(value))
         });
@@ -392,7 +390,7 @@ where
 
         let encoded_set = set
             .map(|(key, value)| {
-                let key = KeyCodec::encode(key).as_bytes().into_owned();
+                let key = KeyCodec::encode(key).into_bytes();
                 let value = ValueCodec::encode(value).as_bytes().into_owned();
                 (key, value)
             })
@@ -449,7 +447,7 @@ where
             .map_err(|err| StorageError::Other(anyhow::anyhow!("{err:?}")))?;
 
         let encoded_set = set
-            .map(|key| KeyCodec::encode(key).as_bytes().into_owned())
+            .map(|key| KeyCodec::encode(key).into_bytes())
             .collect_vec();
 
         for key_bytes in encoded_set.iter() {
diff --git a/crates/storage/src/codec.rs b/crates/storage/src/codec.rs
index f77f91cc044..373f5d298cf 100644
--- a/crates/storage/src/codec.rs
+++ b/crates/storage/src/codec.rs
@@ -29,7 +29,7 @@ pub trait Encoder<'a> {
 /// flexibility and more performant encoding, allowing the use of slices and arrays
 /// instead of vectors in some cases. Since the [`Encoder`] returns `Cow<[u8]>`,
 /// it is always possible to take ownership of the serialized value.
-pub trait Encode<T: ?Sized> {
+pub trait Encode<T: ?Sized>: 'static {
     /// The encoder type that stores serialized object.
     type Encoder<'a>: Encoder<'a>
     where
diff --git a/crates/storage/src/kv_store.rs b/crates/storage/src/kv_store.rs
index 46f6b404f1c..7d4564395d3 100644
--- a/crates/storage/src/kv_store.rs
+++ b/crates/storage/src/kv_store.rs
@@ -194,7 +194,11 @@ pub enum WriteOperation {
 #[impl_tools::autoimpl(for<T: trait> &mut T, Box<T>)]
 pub trait BatchOperations: KeyValueMutate {
     /// Writes the batch of the entries into the storage.
-    fn batch_write<I>(&mut self, column: Self::Column, entries: I) -> StorageResult<()>
+    fn batch_write<'a, I>(
+        &mut self,
+        column: Self::Column,
+        entries: I,
+    ) -> StorageResult<()>
     where
-        I: Iterator<Item = (Vec<u8>, WriteOperation)>;
+        I: Iterator<Item = (Cow<'a, [u8]>, WriteOperation)> + 'a;
 }
diff --git a/crates/storage/src/structured_storage.rs b/crates/storage/src/structured_storage.rs
index d0396a528f2..3bdb6f2f30b 100644
--- a/crates/storage/src/structured_storage.rs
+++ b/crates/storage/src/structured_storage.rs
@@ -218,9 +218,13 @@ impl<S> BatchOperations for StructuredStorage<S>
 where
     S: BatchOperations,
 {
-    fn batch_write<I>(&mut self, column: Self::Column, entries: I) -> StorageResult<()>
+    fn batch_write<'a, I>(
+        &mut self,
+        column: Self::Column,
+        entries: I,
+    ) -> StorageResult<()>
     where
-        I: Iterator<Item = (Vec<u8>, WriteOperation)>,
+        I: Iterator<Item = (Cow<'a, [u8]>, WriteOperation)> + 'a,
     {
         self.inner.batch_write(column, entries)
     }
diff --git a/crates/storage/src/transactional.rs b/crates/storage/src/transactional.rs
index 14cb4cfe0c8..9041b56e07e 100644
--- a/crates/storage/src/transactional.rs
+++ b/crates/storage/src/transactional.rs
@@ -519,13 +519,13 @@ where
     Column: StorageColumn,
     S: KeyValueInspect<Column = Column>,
 {
-    fn batch_write<I>(&mut self, column: Column, entries: I) -> StorageResult<()>
+    fn batch_write<'a, I>(&mut self, column: Column, entries: I) -> StorageResult<()>
     where
-        I: Iterator<Item = (Vec<u8>, WriteOperation)>,
+        I: Iterator<Item = (Cow<'a, [u8]>, WriteOperation)> + 'a,
     {
         let btree = self.changes.entry(column.id()).or_default();
         entries.for_each(|(key, operation)| {
-            btree.insert(key.into(), operation);
+            btree.insert(key.to_vec().into(), operation);
         });
         Ok(())
     }