diff --git a/Cargo.toml b/Cargo.toml
index 0dfc007..16f1cd9 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -6,7 +6,7 @@ version = "0.1.0"
 [features]
 default = ["std", "rocksdb"]
 rocksdb = ["dep:rocksdb"]
-std = ["parity-scale-codec/std", "bitvec/std", "starknet-types-core/std"]
+std = ["parity-scale-codec/std", "bitvec/std", "starknet-types-core/std", "rayon", "hashbrown/rayon"]
 # internal
 bench = []
 
@@ -18,7 +18,7 @@ derive_more = { version = "0.99.17", default-features = false, features = [
 hashbrown = "0.14.3"
 log = "0.4.20"
 smallvec = "1.11.2"
-rayon = "1.9.0"
+rayon = { version = "1.9.0", optional = true }
 
 parity-scale-codec = { version = "3.0.0", default-features = false, features = [
     "derive",
diff --git a/benches/storage.rs b/benches/storage.rs
index 373e92f..03c8c75 100644
--- a/benches/storage.rs
+++ b/benches/storage.rs
@@ -6,7 +6,7 @@ use bonsai_trie::{
     id::{BasicId, BasicIdBuilder},
     BonsaiStorage, BonsaiStorageConfig,
 };
-use criterion::{criterion_group, criterion_main, Criterion};
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
 use rand::{prelude::*, thread_rng};
 use starknet_types_core::{
     felt::Felt,
@@ -15,6 +15,40 @@ use starknet_types_core::{
 
 mod flamegraph;
 
+fn storage_with_insert(c: &mut Criterion) {
+    c.bench_function("storage commit with insert", move |b| {
+        let mut rng = thread_rng();
+        b.iter_batched_ref(
+            || {
+                let bonsai_storage: BonsaiStorage<BasicId, _, Pedersen> = BonsaiStorage::new(
+                    HashMapDb::<BasicId>::default(),
+                    BonsaiStorageConfig::default(),
+                )
+                .unwrap();
+                bonsai_storage
+            },
+            |bonsai_storage| {
+                let felt = Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap();
+                for _ in 0..40000 {
+                    let bitvec = BitVec::from_vec(vec![
+                        rng.gen(),
+                        rng.gen(),
+                        rng.gen(),
+                        rng.gen(),
+                        rng.gen(),
+                        rng.gen(),
+                    ]);
+                    bonsai_storage.insert(&[], &bitvec, &felt).unwrap();
+                }
+
+                // let mut id_builder = BasicIdBuilder::new();
+                // bonsai_storage.commit(id_builder.new_id()).unwrap();
+            },
+            BatchSize::LargeInput,
+        );
+    });
+}
+
 fn storage(c: &mut Criterion) {
     c.bench_function("storage commit", move |b| {
         let mut bonsai_storage: BonsaiStorage<BasicId, _, Pedersen> = BonsaiStorage::new(
@@ -38,9 +72,9 @@ fn storage(c: &mut Criterion) {
         }
 
         let mut id_builder = BasicIdBuilder::new();
-        b.iter_batched(
+        b.iter_batched_ref(
             || bonsai_storage.clone(),
-            |mut bonsai_storage| {
+            |bonsai_storage| {
                 bonsai_storage.commit(id_builder.new_id()).unwrap();
             },
             criterion::BatchSize::LargeInput,
@@ -73,9 +107,9 @@ fn one_update(c: &mut Criterion) {
         let mut id_builder = BasicIdBuilder::new();
         bonsai_storage.commit(id_builder.new_id()).unwrap();
 
-        b.iter_batched(
+        b.iter_batched_ref(
             || bonsai_storage.clone(),
-            |mut bonsai_storage| {
+            |bonsai_storage| {
                 let bitvec = BitVec::from_vec(vec![0, 1, 2, 3, 4, 5]);
                 bonsai_storage.insert(&[], &bitvec, &felt).unwrap();
                 bonsai_storage.commit(id_builder.new_id()).unwrap();
@@ -110,9 +144,9 @@ fn five_updates(c: &mut Criterion) {
         let mut id_builder = BasicIdBuilder::new();
         bonsai_storage.commit(id_builder.new_id()).unwrap();
 
-        b.iter_batched(
+        b.iter_batched_ref(
             || bonsai_storage.clone(),
-            |mut bonsai_storage| {
+            |bonsai_storage| {
                 bonsai_storage
                     .insert(&[], &BitVec::from_vec(vec![0, 1, 2, 3, 4, 5]), &felt)
                     .unwrap();
@@ -135,6 +169,46 @@ fn five_updates(c: &mut Criterion) {
     });
 }
 
+fn multiple_contracts(c: &mut Criterion) {
+    c.bench_function("multiple contracts", move |b| {
+        let mut bonsai_storage: BonsaiStorage<BasicId, _, Pedersen> = BonsaiStorage::new(
+            HashMapDb::<BasicId>::default(),
+            BonsaiStorageConfig::default(),
+        )
+        .unwrap();
+        let mut rng = thread_rng();
+
+        let felt = Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap();
+        for _ in 0..1000 {
+            let bitvec = BitVec::from_vec(vec![rng.gen(), rng.gen(), rng.gen(), rng.gen()]);
+            bonsai_storage
+                .insert(
+                    &[
+                        rng.gen(),
+                        rng.gen(),
+                        rng.gen(),
+                        rng.gen(),
+                        rng.gen(),
+                        rng.gen(),
+                    ],
+                    &bitvec,
+                    &felt,
+                )
+                .unwrap();
+        }
+
+        let mut id_builder = BasicIdBuilder::new();
+
+        b.iter_batched_ref(
+            || bonsai_storage.clone(),
+            |bonsai_storage| {
+                bonsai_storage.commit(id_builder.new_id()).unwrap();
+            },
+            criterion::BatchSize::LargeInput,
+        );
+    });
+}
+
 fn hash(c: &mut Criterion) {
     c.bench_function("pedersen hash", move |b| {
         let felt0 =
@@ -152,6 +226,6 @@ fn hash(c: &mut Criterion) {
 criterion_group! {
     name = benches;
     config = Criterion::default(); // .with_profiler(flamegraph::FlamegraphProfiler::new(100));
-    targets = storage, one_update, five_updates, hash
+    targets = storage, one_update, five_updates, hash, storage_with_insert, multiple_contracts
 }
 criterion_main!(benches);
diff --git a/src/bonsai_database.rs b/src/bonsai_database.rs
index 39b07ae..6c0b3ae 100644
--- a/src/bonsai_database.rs
+++ b/src/bonsai_database.rs
@@ -1,6 +1,4 @@
-use crate::id::Id;
-#[cfg(not(feature = "std"))]
-use alloc::vec::Vec;
+use crate::{id::Id, Vec};
 #[cfg(feature = "std")]
 use std::error::Error;
 
diff --git a/src/changes.rs b/src/changes.rs
index 5de582b..6413fc7 100644
--- a/src/changes.rs
+++ b/src/changes.rs
@@ -1,12 +1,5 @@
-use crate::{id::Id, trie::TrieKey};
+use crate::{hash_map::Entry, id::Id, trie::TrieKey, HashMap, Vec, VecDeque};
 use serde::{Deserialize, Serialize};
-#[cfg(feature = "std")]
-use std::collections::{hash_map::Entry, HashMap, VecDeque};
-#[cfg(not(feature = "std"))]
-use {
-    alloc::{collections::VecDeque, vec::Vec},
-    hashbrown::{hash_map::Entry, HashMap},
-};
 
 #[derive(Debug, Clone, Serialize, Deserialize, Default)]
 pub struct Change {
diff --git a/src/databases/hashmap_db.rs b/src/databases/hashmap_db.rs
index f11914e..65249fe 100644
--- a/src/databases/hashmap_db.rs
+++ b/src/databases/hashmap_db.rs
@@ -1,15 +1,9 @@
 use crate::{
     bonsai_database::{BonsaiPersistentDatabase, DBError},
     id::Id,
-    BonsaiDatabase,
+    BTreeMap, BonsaiDatabase, HashMap, Vec,
 };
-#[cfg(not(feature = "std"))]
-use alloc::{collections::BTreeMap, vec::Vec};
 use core::{fmt, fmt::Display};
-#[cfg(not(feature = "std"))]
-use hashbrown::HashMap;
-#[cfg(feature = "std")]
-use std::collections::{BTreeMap, HashMap};
 
 #[derive(Debug)]
 pub struct HashMapDbError {}
diff --git a/src/error.rs b/src/error.rs
index 92e8ff4..7cfea0c 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -1,10 +1,8 @@
 #[cfg(feature = "std")]
 use std::{error::Error, fmt::Display};
 
-use crate::bonsai_database::DBError;
+use crate::{bonsai_database::DBError, String};
 
-#[cfg(not(feature = "std"))]
-use alloc::string::String;
 /// All errors that can be returned by BonsaiStorage.
 #[derive(Debug)]
 pub enum BonsaiStorageError<DatabaseError>
diff --git a/src/id.rs b/src/id.rs
index f0e4bd4..aa54311 100644
--- a/src/id.rs
+++ b/src/id.rs
@@ -1,5 +1,4 @@
-#[cfg(not(feature = "std"))]
-use alloc::vec::Vec;
+use crate::Vec;
 use core::{fmt::Debug, hash};
 
 /// Trait to be implemented on any type that can be used as an ID.
diff --git a/src/key_value_db.rs b/src/key_value_db.rs
index c5e5f71..f9316c8 100644
--- a/src/key_value_db.rs
+++ b/src/key_value_db.rs
@@ -1,13 +1,12 @@
-use crate::{changes::key_new_value, trie::merkle_tree::bytes_to_bitvec, Change as ExternChange};
-#[cfg(not(feature = "std"))]
-use alloc::{collections::BTreeSet, format, string::ToString, vec::Vec};
+use crate::{
+    changes::key_new_value, format, trie::merkle_tree::bytes_to_bitvec, BTreeSet,
+    Change as ExternChange, ToString, Vec,
+};
 use bitvec::{order::Msb0, vec::BitVec};
 use hashbrown::HashMap;
 use log::trace;
 use parity_scale_codec::Decode;
 use starknet_types_core::felt::Felt;
-#[cfg(feature = "std")]
-use std::collections::BTreeSet;
 
 use crate::{
     bonsai_database::{BonsaiDatabase, BonsaiPersistentDatabase, DatabaseKey},
diff --git a/src/lib.rs b/src/lib.rs
index b51edc8..7595f01 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -84,15 +84,33 @@
 //! bonsai_storage.commit(id_builder.new_id()).unwrap();
 //! ```
 #![cfg_attr(not(feature = "std"), no_std)]
+
+// hashbrown uses ahash by default instead of siphash
+pub(crate) type HashMap<K, V> = hashbrown::HashMap<K, V>;
+pub(crate) use hashbrown::hash_map;
+
 #[cfg(not(feature = "std"))]
 extern crate alloc;
-
-use crate::trie::merkle_tree::{bytes_to_bitvec, MerkleTree};
 #[cfg(not(feature = "std"))]
-use alloc::{format, vec::Vec};
+pub(crate) use alloc::{
+    collections::{BTreeMap, BTreeSet, VecDeque},
+    format,
+    string::{String, ToString},
+    vec,
+    vec::Vec,
+};
+#[cfg(feature = "std")]
+pub(crate) use std::{
+    collections::{BTreeMap, BTreeSet, VecDeque},
+    format,
+    string::{String, ToString},
+    vec,
+    vec::Vec,
+};
+
+use crate::trie::merkle_tree::MerkleTree;
 use bitvec::{order::Msb0, slice::BitSlice, vec::BitVec};
 use changes::ChangeBatch;
-use hashbrown::HashMap;
 use key_value_db::KeyValueDB;
 use starknet_types_core::{
     felt::Felt,
@@ -112,7 +130,7 @@ pub mod id;
 
 pub use bonsai_database::{BonsaiDatabase, BonsaiPersistentDatabase, DBError, DatabaseKey};
 pub use error::BonsaiStorageError;
-use trie::merkle_tree::MerkleTrees;
+use trie::merkle_tree::{bytes_to_bitvec, MerkleTrees};
 pub use trie::merkle_tree::{Membership, ProofNode};
 
 #[cfg(test)]
diff --git a/src/tests/madara_comparison.rs b/src/tests/madara_comparison.rs
index 870107d..eef4cc6 100644
--- a/src/tests/madara_comparison.rs
+++ b/src/tests/madara_comparison.rs
@@ -1,4 +1,4 @@
-#![cfg(feature = "std")]
+#![cfg(all(feature = "std", feature = "rocksdb"))]
 use bitvec::{bits, order::Msb0, vec::BitVec};
 use starknet_types_core::{felt::Felt, hash::Pedersen};
 
diff --git a/src/tests/proof.rs b/src/tests/proof.rs
index bb85425..9c942ff 100644
--- a/src/tests/proof.rs
+++ b/src/tests/proof.rs
@@ -1,4 +1,4 @@
-#![cfg(feature = "std")]
+#![cfg(all(feature = "std", feature = "rocksdb"))]
 use bitvec::vec::BitVec;
 use pathfinder_common::{hash::PedersenHash, trie::TrieNode};
 use pathfinder_crypto::Felt as PathfinderFelt;
diff --git a/src/tests/simple.rs b/src/tests/simple.rs
index 13483d6..7e8cd52 100644
--- a/src/tests/simple.rs
+++ b/src/tests/simple.rs
@@ -1,4 +1,4 @@
-#![cfg(feature = "std")]
+#![cfg(all(feature = "std", feature = "rocksdb"))]
 use crate::{
     databases::{create_rocks_db, HashMapDb, RocksDB, RocksDBConfig},
     id::{BasicId, BasicIdBuilder},
diff --git a/src/tests/transactional_state.rs b/src/tests/transactional_state.rs
index 1a4d38a..afb8dc6 100644
--- a/src/tests/transactional_state.rs
+++ b/src/tests/transactional_state.rs
@@ -1,4 +1,4 @@
-#![cfg(feature = "std")]
+#![cfg(all(feature = "std", feature = "rocksdb"))]
 use crate::{
     databases::{create_rocks_db, RocksDB, RocksDBConfig},
     id::BasicIdBuilder,
diff --git a/src/tests/trie_log.rs b/src/tests/trie_log.rs
index 59be7b4..8146ad5 100644
--- a/src/tests/trie_log.rs
+++ b/src/tests/trie_log.rs
@@ -1,4 +1,4 @@
-#![cfg(feature = "std")]
+#![cfg(all(feature = "std", feature = "rocksdb"))]
 use crate::{
     databases::{create_rocks_db, RocksDB, RocksDBConfig},
     id::BasicIdBuilder,
diff --git a/src/trie/merkle_node.rs b/src/trie/merkle_node.rs
index fa02128..fca4293 100644
--- a/src/trie/merkle_node.rs
+++ b/src/trie/merkle_node.rs
@@ -154,6 +154,25 @@ impl BinaryNode {
             Direction::Right => self.right,
         }
     }
+
+    /// Returns the [Left] or [Right] child.
+    ///
+    /// [Left]: Direction::Left
+    /// [Right]: Direction::Right
+    ///
+    /// # Arguments
+    ///
+    /// `direction` - The direction where to get the child from.
+    ///
+    /// # Returns
+    ///
+    /// The child in the specified direction.
+    pub fn get_child_mut(&mut self, direction: Direction) -> &mut NodeHandle {
+        match direction {
+            Direction::Left => &mut self.left,
+            Direction::Right => &mut self.right,
+        }
+    }
 }
 
 impl Node {
diff --git a/src/trie/merkle_tree.rs b/src/trie/merkle_tree.rs
index 77452d0..05f366c 100644
--- a/src/trie/merkle_tree.rs
+++ b/src/trie/merkle_tree.rs
@@ -1,5 +1,3 @@
-#[cfg(not(feature = "std"))]
-use alloc::{format, string::ToString, vec, vec::Vec};
 use bitvec::{
     prelude::{BitSlice, BitVec, Msb0},
     view::BitView,
@@ -8,16 +6,15 @@ use core::iter::once;
 use core::marker::PhantomData;
 use core::mem;
 use derive_more::Constructor;
-#[cfg(not(feature = "std"))]
-use hashbrown::HashMap;
 use parity_scale_codec::{Decode, Encode};
 #[cfg(feature = "std")]
 use rayon::prelude::*;
 use starknet_types_core::{felt::Felt, hash::StarkHash};
-#[cfg(feature = "std")]
-use std::collections::HashMap;
 
-use crate::{error::BonsaiStorageError, id::Id, BonsaiDatabase, KeyValueDB};
+use crate::{
+    error::BonsaiStorageError, format, id::Id, vec, BonsaiDatabase, HashMap, KeyValueDB, ToString,
+    Vec,
+};
 
 use super::{
     merkle_node::{BinaryNode, Direction, EdgeNode, Node, NodeHandle, NodeId},
@@ -241,18 +238,19 @@ impl<H: StarkHash + Send + Sync, DB: BonsaiDatabase, CommitID: Id> MerkleTrees<H
         let db_changes = self
             .trees
             .iter_mut()
-            .map(|(_, tree)| tree.get_updates::<DB>())
-            .collect::<Result<Vec<_>, BonsaiStorageError<DB::DatabaseError>>>()?;
+            .map(|(_, tree)| tree.get_updates::<DB>());
         #[cfg(feature = "std")]
         let db_changes = self
             .trees
             .par_iter_mut()
             .map(|(_, tree)| tree.get_updates::<DB>())
-            .collect::<Result<Vec<_>, BonsaiStorageError<DB::DatabaseError>>>()?;
+            .collect_vec_list()
+            .into_iter()
+            .flatten();
 
         let mut batch = self.db.create_batch();
         for changes in db_changes {
-            for (key, value) in changes {
+            for (key, value) in changes? {
                 match value {
                     InsertOrRemove::Insert(value) => {
                         self.db.insert(&key, &value, Some(&mut batch))?;
@@ -386,11 +384,14 @@ impl<H: StarkHash + Send + Sync> MerkleTree<H> {
         &mut self,
         db: &mut KeyValueDB<DB, ID>,
     ) -> Result<(), BonsaiStorageError<DB::DatabaseError>> {
-        let node = self
-            .get_trie_branch_in_db_from_path(db, &Path(BitVec::<u8, Msb0>::new()))?
-            .ok_or(BonsaiStorageError::Trie(
-                "root node doesn't exist in the storage".to_string(),
-            ))?;
+        let node = Self::get_trie_branch_in_db_from_path(
+            &self.identifier,
+            db,
+            &Path(BitVec::<u8, Msb0>::new()),
+        )?
+        .ok_or(BonsaiStorageError::Trie(
+            "root node doesn't exist in the storage".to_string(),
+        ))?;
         let node_hash = node.hash().ok_or(BonsaiStorageError::Trie(
             "Root doesn't exist in the storage".to_string(),
         ))?;
@@ -788,9 +789,7 @@ impl<H: StarkHash + Send + Sync> MerkleTree<H> {
                         _ => {}
                     }
                 });
-                for (id, node) in nodes_to_add {
-                    self.storage_nodes.0.insert(id, node);
-                }
+                self.storage_nodes.0.extend(nodes_to_add);
                 Ok(())
             }
             None => {
@@ -1027,11 +1026,14 @@ impl<H: StarkHash + Send + Sync> MerkleTree<H> {
         let mut nodes = Vec::with_capacity(251);
         let mut node = match self.root_handle {
             NodeHandle::Hash(_) => {
-                let node = self
-                    .get_trie_branch_in_db_from_path(db, &Path(BitVec::<u8, Msb0>::new()))?
-                    .ok_or(BonsaiStorageError::Trie(
-                        "Couldn't fetch root node in db".to_string(),
-                    ))?;
+                let node = Self::get_trie_branch_in_db_from_path(
+                    &self.identifier,
+                    db,
+                    &Path(BitVec::<u8, Msb0>::new()),
+                )?
+                .ok_or(BonsaiStorageError::Trie(
+                    "Couldn't fetch root node in db".to_string(),
+                ))?;
                 if node.is_empty() {
                     return Ok(Vec::new());
                 }
@@ -1052,8 +1054,11 @@ impl<H: StarkHash + Send + Sync> MerkleTree<H> {
                     let child_path = key[..edge.height as usize + edge.path.0.len()].to_bitvec();
                     let child_node = match edge.child {
                         NodeHandle::Hash(hash) => {
-                            let node =
-                                self.get_trie_branch_in_db_from_path(db, &Path(child_path))?;
+                            let node = Self::get_trie_branch_in_db_from_path(
+                                &self.identifier,
+                                db,
+                                &Path(child_path),
+                            )?;
                             if let Some(node) = node {
                                 node
                             } else {
@@ -1096,11 +1101,14 @@ impl<H: StarkHash + Send + Sync> MerkleTree<H> {
                     let next = binary.get_child(next_direction);
                     let next_path = key[..binary.height as usize + 1].to_bitvec();
                     let next_node = match next {
-                        NodeHandle::Hash(_) => self
-                            .get_trie_branch_in_db_from_path(db, &Path(next_path))?
-                            .ok_or(BonsaiStorageError::Trie(
-                                "Couldn't fetch next node in db".to_string(),
-                            ))?,
+                        NodeHandle::Hash(_) => Self::get_trie_branch_in_db_from_path(
+                            &self.identifier,
+                            db,
+                            &Path(next_path),
+                        )?
+                        .ok_or(BonsaiStorageError::Trie(
+                            "Couldn't fetch next node in db".to_string(),
+                        ))?,
                         NodeHandle::InMemory(next_id) => self
                             .storage_nodes
                             .0
@@ -1183,138 +1191,85 @@ impl<H: StarkHash + Send + Sync> MerkleTree<H> {
         dst: &BitSlice<u8, Msb0>,
     ) -> Result<Vec<NodeId>, BonsaiStorageError<DB::DatabaseError>> {
         let mut nodes = Vec::with_capacity(251);
-        let node_id = match self.root_handle {
-            NodeHandle::Hash(_) => {
-                let node = self
-                    .get_trie_branch_in_db_from_path(db, &Path(BitVec::<u8, Msb0>::new()))?
-                    .ok_or(BonsaiStorageError::Trie(
-                        "Couldn't fetch root node in db".to_string(),
-                    ))?;
-                if node.is_empty() {
-                    return Ok(Vec::new());
-                }
-                self.latest_node_id.next_id();
-                self.root_handle = NodeHandle::InMemory(self.latest_node_id);
-                self.storage_nodes.0.insert(self.latest_node_id, node);
-                nodes.push(self.latest_node_id);
-                self.latest_node_id
-            }
-            NodeHandle::InMemory(root_id) => {
-                nodes.push(root_id);
-                root_id
-            }
-        };
-        self.preload_nodes_subtree(
-            db,
-            dst,
-            node_id,
-            Path(BitVec::<u8, Msb0>::new()),
-            &mut nodes,
-        )?;
-        Ok(nodes)
-    }
 
-    fn preload_nodes_subtree<DB: BonsaiDatabase, ID: Id>(
-        &mut self,
-        db: &mut KeyValueDB<DB, ID>,
-        dst: &BitSlice<u8, Msb0>,
-        root_id: NodeId,
-        mut path: Path,
-        nodes: &mut Vec<NodeId>,
-    ) -> Result<(), BonsaiStorageError<DB::DatabaseError>> {
-        let node = self
-            .storage_nodes
-            .0
-            .get(&root_id)
-            .ok_or(BonsaiStorageError::Trie(
-                "Couldn't fetch node in the temporary storage".to_string(),
-            ))?
-            .clone();
-        match node {
-            // We are in a case where the trie is empty and so there is nothing to preload.
-            Node::Unresolved(_hash) => Ok(()),
-            // We are checking which side of the binary we should load in memory (if we don't have it already)
-            // We load this "child-side" node in the memory and refer his memory handle in the binary node.
-            // We also add the "child-side" node in the list that accumulate all the nodes we want to preload.
-            // We override the binary node in the memory with this new version that has the "child-side" memory handle
-            // instead of the hash.
-            // We call recursively the function with the "child-side" node.
-            Node::Binary(mut binary_node) => {
-                let next_direction = binary_node.direction(dst);
-                path.0.push(bool::from(next_direction));
-                let next = binary_node.get_child(next_direction);
-                match next {
-                    NodeHandle::Hash(_) => {
-                        let node = self.get_trie_branch_in_db_from_path(db, &path)?;
-                        if let Some(node) = node {
-                            self.latest_node_id.next_id();
-                            self.storage_nodes.0.insert(self.latest_node_id, node);
-                            nodes.push(self.latest_node_id);
-                            match next_direction {
-                                Direction::Left => {
-                                    binary_node.left = NodeHandle::InMemory(self.latest_node_id)
-                                }
-                                Direction::Right => {
-                                    binary_node.right = NodeHandle::InMemory(self.latest_node_id)
-                                }
-                            };
-                            self.storage_nodes
-                                .0
-                                .insert(root_id, Node::Binary(binary_node));
-                            self.preload_nodes_subtree(db, dst, self.latest_node_id, path, nodes)
-                        } else {
-                            Ok(())
-                        }
-                    }
-                    NodeHandle::InMemory(next_id) => {
-                        nodes.push(next_id);
-                        self.preload_nodes_subtree(db, dst, next_id, path, nodes)
+        let mut prev_handle = &mut self.root_handle;
+        let mut path = Path(BitVec::<u8, Msb0>::with_capacity(251));
+
+        loop {
+            // get node from cache or database
+            let (node_id, node) = match prev_handle {
+                NodeHandle::Hash(_) => {
+                    // load from db
+                    let node = Self::get_trie_branch_in_db_from_path(&self.identifier, db, &path)?
+                        .ok_or(BonsaiStorageError::Trie(
+                            "Couldn't fetch node from db".to_string(),
+                        ))?;
+
+                    if node.is_empty() {
+                        // empty tree
+                        break;
                     }
+
+                    // put it in inmemory storage
+                    self.latest_node_id.next_id();
+                    *prev_handle = NodeHandle::InMemory(self.latest_node_id);
+                    let node = self.storage_nodes.0.entry(self.latest_node_id).insert(node);
+
+                    (self.latest_node_id, node.into_mut())
                 }
-            }
-            // If the edge node match the path we want to preload then we load the child node in memory (if we don't have it already)
-            // and we override the edge node in the memory with this new version that has the child memory handle instead of the hash.
-            // We also add the child node in the list that accumulate all the nodes we want to preload.
-            // We call recursively the function with the child node.
-            Node::Edge(mut edge_node) if edge_node.path_matches(dst) => {
-                path.0.extend_from_bitslice(&edge_node.path.0);
-                if path.0 == dst {
-                    return Ok(());
+                NodeHandle::InMemory(node_id) => {
+                    let node_id = *node_id;
+
+                    let node =
+                        self.storage_nodes
+                            .0
+                            .get_mut(&node_id)
+                            .ok_or(BonsaiStorageError::Trie(
+                                "Couldn't get node from temp storage".to_string(),
+                            ))?;
+
+                    (node_id.clone(), node)
                 }
-                let next = edge_node.child;
-                match next {
-                    NodeHandle::Hash(_) => {
-                        let node = self.get_trie_branch_in_db_from_path(db, &path)?;
-                        if let Some(node) = node {
-                            self.latest_node_id.next_id();
-                            self.storage_nodes.0.insert(self.latest_node_id, node);
-                            nodes.push(self.latest_node_id);
-                            edge_node.child = NodeHandle::InMemory(self.latest_node_id);
-                            self.storage_nodes.0.insert(root_id, Node::Edge(edge_node));
-                            self.preload_nodes_subtree(db, dst, self.latest_node_id, path, nodes)
-                        } else {
-                            Ok(())
-                        }
-                    }
-                    NodeHandle::InMemory(next_id) => {
-                        nodes.push(next_id);
-                        self.preload_nodes_subtree(db, dst, next_id, path, nodes)
+            };
+
+            nodes.push(node_id);
+
+            // visit the child
+            match node {
+                // We are in a case where the trie is empty and so there is nothing to preload.
+                Node::Unresolved(_felt) => break,
+
+                Node::Binary(binary_node) => {
+                    let next_direction = binary_node.direction(dst);
+                    path.0.push(bool::from(next_direction));
+                    prev_handle = binary_node.get_child_mut(next_direction);
+                }
+
+                Node::Edge(edge_node) if edge_node.path_matches(dst) => {
+                    path.0.extend_from_bitslice(&edge_node.path.0);
+                    if path.0 == dst {
+                        break; // found it :)
                     }
+
+                    prev_handle = &mut edge_node.child;
                 }
+
+                // We are in a case where the edge node doesn't match the path we want to preload so we return nothing.
+                Node::Edge(_) => break,
             }
-            // We are in a case where the edge node doesn't match the path we want to preload so we return nothing.
-            Node::Edge(_) => Ok(()),
         }
+
+        Ok(nodes)
     }
 
     /// Get the node of the trie that corresponds to the path.
     fn get_trie_branch_in_db_from_path<DB: BonsaiDatabase, ID: Id>(
-        &self,
+        identifier: &[u8],
         db: &KeyValueDB<DB, ID>,
         path: &Path,
     ) -> Result<Option<Node>, BonsaiStorageError<DB::DatabaseError>> {
         let path: Vec<u8> = path.into();
-        db.get(&TrieKey::new(&self.identifier, TrieKeyType::Trie, &path))?
+        db.get(&TrieKey::new(&identifier, TrieKeyType::Trie, &path))?
             .map(|node| {
                 Node::decode(&mut node.as_slice()).map_err(|err| {
                     BonsaiStorageError::Trie(format!("Couldn't decode node: {}", err))
@@ -1342,7 +1297,7 @@ impl<H: StarkHash + Send + Sync> MerkleTree<H> {
     ) -> Result<(), BonsaiStorageError<DB::DatabaseError>> {
         let child_node = match parent.child {
             NodeHandle::Hash(_) => {
-                let node = self.get_trie_branch_in_db_from_path(db, path)?;
+                let node = Self::get_trie_branch_in_db_from_path(&self.identifier, db, path)?;
                 if let Some(node) = node {
                     node
                 } else {
@@ -1510,7 +1465,7 @@ pub(crate) fn bytes_to_bitvec(bytes: &[u8]) -> BitVec<u8, Msb0> {
 }
 
 #[cfg(test)]
-#[cfg(all(test, feature = "std"))]
+#[cfg(all(test, feature = "std", feature = "rocksdb"))]
 mod tests {
     use bitvec::{order::Msb0, vec::BitVec, view::BitView};
     use indexmap::IndexMap;
diff --git a/src/trie/path.rs b/src/trie/path.rs
index b5f8cc4..de8716d 100644
--- a/src/trie/path.rs
+++ b/src/trie/path.rs
@@ -3,8 +3,7 @@ use parity_scale_codec::{Decode, Encode, Error, Input, Output};
 
 use super::merkle_node::Direction;
 
-#[cfg(not(feature = "std"))]
-use alloc::vec::Vec;
+use crate::Vec;
 
 #[cfg(all(feature = "std", test))]
 use rstest::rstest;
diff --git a/src/trie/trie_db.rs b/src/trie/trie_db.rs
index 97f3a84..857548c 100644
--- a/src/trie/trie_db.rs
+++ b/src/trie/trie_db.rs
@@ -1,7 +1,5 @@
 use crate::bonsai_database::DatabaseKey;
-
-#[cfg(not(feature = "std"))]
-use alloc::vec::Vec;
+use crate::Vec;
 
 /// Key in the database of the different elements that are used in the storage of the trie data.
 /// Use `new` function to create a new key.