refactor: fix iterator db loading, replace preload_nodes with iterator

src/bonsai_database.rs

@@ -27,7 +27,7 @@ pub trait DBError: Error + Send + Sync {}
 pub trait DBError: Send + Sync {}
 
 /// Trait to be implemented on any type that can be used as a database.
-pub trait BonsaiDatabase {
+pub trait BonsaiDatabase: core::fmt::Debug {
     type Batch: Default;
     #[cfg(feature = "std")]
     type DatabaseError: Error + DBError;

src/databases/rocks_db.rs

@@ -173,6 +173,12 @@ pub struct RocksDBTransaction<'a> {
     column_families: HashMap<String, ColumnFamilyRef<'a>>,
 }
 
+impl<'a> fmt::Debug for RocksDBTransaction<'a> {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.debug_struct("RocksDBTransaction").finish()
+    }
+}
+
 impl<'db, ID> BonsaiDatabase for RocksDB<'db, ID>
 where
     ID: Id,

src/tests/proptest.rs

@@ -4,9 +4,9 @@ use crate::id::BasicId;
 use crate::key_value_db::KeyValueDB;
 use crate::trie::tree::MerkleTree;
 use crate::{BitVec, HashMap};
-use core::fmt::{self, Debug};
 use bitvec::bitvec;
 use bitvec::order::Msb0;
+use core::fmt::{self, Debug};
 use proptest::prelude::*;
 use proptest_derive::Arbitrary;
 use smallvec::smallvec;
@@ -121,7 +121,7 @@ impl MerkleTreeInsertProblem {
 }
 
 proptest::proptest! {
-    #![proptest_config(ProptestConfig::with_cases(5))] // comment this when developing, this is mostly for faster ci & whole workspace `cargo test`
+    // #![proptest_config(ProptestConfig::with_cases(5))] // comment this when developing, this is mostly for faster ci & whole workspace `cargo test`
     #[test]
     fn proptest_inserts(pb in any::<MerkleTreeInsertProblem>()) {
         let _ = env_logger::builder().is_test(true).try_init();
@@ -275,3 +275,27 @@ fn test_merkle_pb_6() {
 
     pb.check();
 }
+
+#[test]
+fn test_merkle_pb_7() {
+    use Step::*;
+    let _ = env_logger::builder().is_test(true).try_init();
+    log::set_max_level(log::LevelFilter::Trace);
+    let pb = MerkleTreeInsertProblem(vec![
+        Insert(
+            Key(bitvec![u8, Msb0; 0,0,0,0,0]),
+            Value(Felt::from_hex("0x20").unwrap()),
+        ),
+        Insert(
+            Key(bitvec![u8, Msb0; 1,0,0,0,0]),
+            Value(Felt::from_hex("0x20").unwrap()),
+        ),
+        Commit,
+        Insert(
+            Key(bitvec![u8, Msb0; 0,0,0,0,0]),
+            Value(Felt::from_hex("0x40").unwrap()),
+        ),
+    ]);
+
+    pb.check();
+}

src/trie/iterator.rs

@@ -6,7 +6,7 @@ use super::{
     TrieKey,
 };
 use crate::{
-    id::Id, key_value_db::KeyValueDB, BitSlice, BonsaiDatabase, BonsaiStorageError, ByteVec,
+    id::Id, key_value_db::KeyValueDB, trie::tree::NodesMapping, BitSlice, BonsaiDatabase, BonsaiStorageError, ByteVec
 };
 use core::fmt;
 use starknet_types_core::hash::StarkHash;
@@ -42,6 +42,7 @@ impl<'a, H: StarkHash, DB: BonsaiDatabase, ID: Id> MerkleTreeIterator<'a, H, DB,
         }
     }
 
+    #[cfg(test)]
     pub fn cur_nodes(&self) -> Vec<NodeId> {
         self.cur_path_nodes_heights
             .iter()
@@ -164,7 +165,7 @@ impl<'a, H: StarkHash, DB: BonsaiDatabase, ID: Id> MerkleTreeIterator<'a, H, DB,
                         let path: ByteVec = self.current_path.clone().into();
                         log::trace!("Visiting db node {:?}", self.current_path);
                         let key = TrieKey::new(&self.tree.identifier, TrieKeyType::Trie, &path);
-                        let Some((node_id, node)) = self.tree.node_storage.nodes.load_db_node(
+                        let Some((node_id, node)) = NodesMapping::load_db_node_get_id(
                             &mut self.tree.node_storage.latest_node_id,
                             &self.tree.death_row,
                             &self.db,
@@ -176,6 +177,8 @@ impl<'a, H: StarkHash, DB: BonsaiDatabase, ID: Id> MerkleTreeIterator<'a, H, DB,
                                 "Could not get node from db".to_string(),
                             ));
                         };
+                        *next_to_visit = NodeHandle::InMemory(node_id);
+                        let node = self.tree.node_storage.nodes.load_db_node_to_id::<DB>(node_id, node)?;
                         (node_id, node)
                     }
                     NodeHandle::InMemory(node_id) => {
@@ -291,17 +294,26 @@ mod tests {
         // from scratch, should find the leaf
         iter.seek_to(&bits![u8, Msb0; 0,0,0,1,0,0,0,0]).unwrap();
         assert_eq!(iter.current_value, Some(NodeHandle::Hash(ONE)));
-        assert_eq!(iter.cur_nodes(), vec![NodeId(1), NodeId(7), NodeId(6), NodeId(4), NodeId(2)]);
+        assert_eq!(
+            iter.cur_nodes(),
+            vec![NodeId(1), NodeId(7), NodeId(6), NodeId(4), NodeId(2)]
+        );
         println!("{iter:?}");
         // from a closeby leaf, should backtrack and find the next one
         iter.seek_to(&bits![u8, Msb0; 0,0,0,1,0,0,0,1]).unwrap();
         assert_eq!(iter.current_value, Some(NodeHandle::Hash(TWO)));
-        assert_eq!(iter.cur_nodes(), vec![NodeId(1), NodeId(7), NodeId(6), NodeId(4), NodeId(2)]);
+        assert_eq!(
+            iter.cur_nodes(),
+            vec![NodeId(1), NodeId(7), NodeId(6), NodeId(4), NodeId(2)]
+        );
         println!("{iter:?}");
         // backtrack farther, should find the leaf
         iter.seek_to(&bits![u8, Msb0; 0,0,0,1,0,0,1,0]).unwrap();
         assert_eq!(iter.current_value, Some(NodeHandle::Hash(THREE)));
-        assert_eq!(iter.cur_nodes(), vec![NodeId(1), NodeId(7), NodeId(6), NodeId(4), NodeId(3)]);
+        assert_eq!(
+            iter.cur_nodes(),
+            vec![NodeId(1), NodeId(7), NodeId(6), NodeId(4), NodeId(3)]
+        );
         println!("{iter:?}");
         // backtrack farther, should find the leaf
         iter.seek_to(&bits![u8, Msb0; 0,1,0,0,0,0,0,0]).unwrap();
@@ -312,7 +324,10 @@ mod tests {
         // similar case
         iter.seek_to(&bits![u8, Msb0; 0,0,0,1,0,0,0,1]).unwrap();
         assert_eq!(iter.current_value, Some(NodeHandle::Hash(TWO)));
-        assert_eq!(iter.cur_nodes(), vec![NodeId(1), NodeId(7), NodeId(6), NodeId(4), NodeId(2)]);
+        assert_eq!(
+            iter.cur_nodes(),
+            vec![NodeId(1), NodeId(7), NodeId(6), NodeId(4), NodeId(2)]
+        );
         println!("{iter:?}");
 
         // SEEK MIDWAY INTO THE TREE
@@ -340,7 +355,10 @@ mod tests {
         // The current value should reflect the binary node
         assert_eq!(iter.current_value, Some(NodeHandle::InMemory(NodeId(2))));
         assert_eq!(iter.current_path.0, bits![u8, Msb0; 0,0,0,1,0,0,0]);
-        assert_eq!(iter.cur_nodes(), vec![NodeId(1), NodeId(7), NodeId(6), NodeId(4)]);
+        assert_eq!(
+            iter.cur_nodes(),
+            vec![NodeId(1), NodeId(7), NodeId(6), NodeId(4)]
+        );
         println!("{iter:?}");
 
         // jump to the end of an edge