diff --git a/Cargo.toml b/Cargo.toml index 1e5edd4..55670e0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,6 +16,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" parity-scale-codec = { version = "3.0.0", default-features = false, features = [ "derive", diff --git a/README.md b/README.md index e1d9c7b..d8b1d20 100644 --- a/README.md +++ b/README.md @@ -55,16 +55,19 @@ fn main() { // Create a simple incremental ID builder for commit IDs. // This is not necessary, you can use any kind of strictly monotonically increasing value to tag your commits. let mut id_builder = BasicIdBuilder::new(); + + // Define an idenfitier for a trie. All insert, get, remove and root hash will use this identifier. Define multiple identifier to use multiple tries that have separate root hash. + let identifier = vec![]; // Insert an item `pair1`. let pair1 = (vec![1, 2, 1], Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap()); let bitvec_1 = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec_1, &pair1.1).unwrap(); + bonsai_storage.insert(&identifier, &bitvec_1, &pair1.1).unwrap(); // Insert a second item `pair2`. let pair2 = (vec![1, 2, 2], Felt::from_hex("0x66342762FD54D033c195fec3ce2568b62052e").unwrap()); let bitvec = BitVec::from_vec(pair2.0.clone()); - bonsai_storage.insert(&bitvec, &pair2.1).unwrap(); + bonsai_storage.insert(&identifier, &bitvec, &pair2.1).unwrap(); // Commit the insertion of `pair1` and `pair2`. let id1 = id_builder.new_id() @@ -73,31 +76,31 @@ fn main() { // Insert a new item `pair3`. let pair3 = (vec![1, 2, 2], Felt::from_hex("0x664D033c195fec3ce2568b62052e").unwrap()); let bitvec = BitVec::from_vec(pair3.0.clone()); - bonsai_storage.insert(&bitvec, &pair3.1).unwrap(); + bonsai_storage.insert(&identifier, &bitvec, &pair3.1).unwrap(); // Commit the insertion of `pair3`. Save the commit ID to the `revert_to_id` variable. let revert_to_id = id_builder.new_id(); bonsai_storage.commit(revert_to_id); // Remove `pair3`. - bonsai_storage.remove(&bitvec).unwrap(); + bonsai_storage.remove(&identifier, &bitvec).unwrap(); // Commit the removal of `pair3`. bonsai_storage.commit(id_builder.new_id()); // Print the root hash and item `pair1`. - println!("root: {:#?}", bonsai_storage.root_hash()); + println!("root: {:#?}", bonsai_storage.root_hash(&identifier)); println!( "value: {:#?}", - bonsai_storage.get(&bitvec_1).unwrap() + bonsai_storage.get(&identifier, &bitvec_1).unwrap() ); // Revert the collection state back to the commit tagged by the `revert_to_id` variable. bonsai_storage.revert_to(revert_to_id).unwrap(); // Print the root hash and item `pair3`. - println!("root: {:#?}", bonsai_storage.root_hash()); - println!("value: {:#?}", bonsai_storage.get(&bitvec).unwrap()); + println!("root: {:#?}", bonsai_storage.root_hash(&identifier)); + println!("value: {:#?}", bonsai_storage.get(&identifier, &bitvec).unwrap()); // Launch two threads that will simultaneously take transactional states to the commit identified by `id1`, // asserting in both of them that the item `pair1` is present and has the right value. @@ -108,7 +111,7 @@ fn main() { .unwrap() .unwrap(); let bitvec = BitVec::from_vec(pair1.0.clone()); - assert_eq!(bonsai_at_txn.get(&bitvec).unwrap().unwrap(), pair1.1); + assert_eq!(bonsai_at_txn.get(&identifier, &bitvec).unwrap().unwrap(), pair1.1); }); s.spawn(|| { @@ -117,13 +120,13 @@ fn main() { .unwrap() .unwrap(); let bitvec = BitVec::from_vec(pair1.0.clone()); - assert_eq!(bonsai_at_txn.get(&bitvec).unwrap().unwrap(), pair1.1); + assert_eq!(bonsai_at_txn.get(&identifier, &bitvec).unwrap().unwrap(), pair1.1); }); }); // Read item `pair1`. let pair1_val = bonsai_storage - .get(&BitVec::from_vec(vec![1, 2, 2])) + .get(&identifier, &BitVec::from_vec(vec![1, 2, 2])) .unwrap(); // Insert a new item and commit. @@ -132,15 +135,15 @@ fn main() { Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(), ); bonsai_storage - .insert(&BitVec::from_vec(pair4.0.clone()), &pair4.1) + .insert(&identifier, &BitVec::from_vec(pair4.0.clone()), &pair4.1) .unwrap(); bonsai_storage.commit(id_builder.new_id()).unwrap(); let proof = bonsai_storage - .get_proof(&BitVec::from_vec(pair3.0.clone())) + .get_proof(&identifier, &BitVec::from_vec(pair3.0.clone())) .unwrap(); assert_eq!( BonsaiStorage::>::verify_proof( - bonsai_storage.root_hash().unwrap(), + bonsai_storage.root_hash(&identifier).unwrap(), &BitVec::from_vec(pair3.0.clone()), pair3.1, &proof diff --git a/ensure_no_std/Cargo.lock b/ensure_no_std/Cargo.lock index 2861af3..30a9d11 100644 --- a/ensure_no_std/Cargo.lock +++ b/ensure_no_std/Cargo.lock @@ -62,6 +62,7 @@ dependencies = [ "hashbrown", "log", "parity-scale-codec", + "rayon", "serde", "smallvec", "starknet-types-core", @@ -94,6 +95,31 @@ dependencies = [ "libc", ] +[[package]] +name = "crossbeam-deque" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613f8cc01fe9cf1a3eb3d7f488fd2fa8388403e97039e2f73692932e291a770d" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "248e3bacc7dc6baa3b21e405ee045c3047101a49145e7e9eca583ab4c2ca5345" + [[package]] name = "crypto-common" version = "0.1.6" @@ -125,6 +151,12 @@ dependencies = [ "crypto-common", ] +[[package]] +name = "either" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11157ac094ffbdde99aa67b23417ebdd801842852b500e395a45a9c0aac03e4a" + [[package]] name = "ensure_no_std" version = "0.1.0" @@ -339,6 +371,26 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" +[[package]] +name = "rayon" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b418a60154510ca1a002a752ca9714984e21e4241e804d32555251faf8b78ffa" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1465873a3dfdaa8ae7cb14b4383657caab0b3e8a0aa9ae8e04b044854c8dfce2" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "serde" version = "1.0.195" diff --git a/src/lib.rs b/src/lib.rs index 3de3f71..24d4f31 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -16,41 +16,42 @@ //! let db = create_rocks_db("./rocksdb").unwrap(); //! let config = BonsaiStorageConfig::default(); //! +//! let identifier = vec![]; //! let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); //! let mut id_builder = BasicIdBuilder::new(); //! //! let pair1 = (vec![1, 2, 1], Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap()); //! let bitvec_1 = BitVec::from_vec(pair1.0.clone()); -//! bonsai_storage.insert(&bitvec_1, &pair1.1).unwrap(); +//! bonsai_storage.insert(&identifier, &bitvec_1, &pair1.1).unwrap(); //! //! let pair2 = (vec![1, 2, 2], Felt::from_hex("0x66342762FD54D033c195fec3ce2568b62052e").unwrap()); //! let bitvec = BitVec::from_vec(pair2.0.clone()); -//! bonsai_storage.insert(&bitvec, &pair2.1).unwrap(); +//! bonsai_storage.insert(&identifier, &bitvec, &pair2.1).unwrap(); //! //! let id1 = id_builder.new_id(); //! bonsai_storage.commit(id1); //! //! let pair3 = (vec![1, 2, 2], Felt::from_hex("0x664D033c195fec3ce2568b62052e").unwrap()); //! let bitvec = BitVec::from_vec(pair3.0.clone()); -//! bonsai_storage.insert(&bitvec, &pair3.1).unwrap(); +//! bonsai_storage.insert(&identifier, &bitvec, &pair3.1).unwrap(); //! //! let revert_to_id = id_builder.new_id(); //! bonsai_storage.commit(revert_to_id); //! -//! bonsai_storage.remove(&bitvec).unwrap(); +//! bonsai_storage.remove(&identifier, &bitvec).unwrap(); //! //! bonsai_storage.commit(id_builder.new_id()); //! -//! println!("root: {:#?}", bonsai_storage.root_hash()); +//! println!("root: {:#?}", bonsai_storage.root_hash(&identifier)); //! println!( //! "value: {:#?}", -//! bonsai_storage.get(&bitvec_1).unwrap() +//! bonsai_storage.get(&identifier, &bitvec_1).unwrap() //! ); //! //! bonsai_storage.revert_to(revert_to_id).unwrap(); //! -//! println!("root: {:#?}", bonsai_storage.root_hash()); -//! println!("value: {:#?}", bonsai_storage.get(&bitvec).unwrap()); +//! println!("root: {:#?}", bonsai_storage.root_hash(&identifier)); +//! println!("value: {:#?}", bonsai_storage.get(&identifier, &bitvec).unwrap()); //! std::thread::scope(|s| { //! s.spawn(|| { //! let bonsai_at_txn: BonsaiStorage<_, _, Pedersen> = bonsai_storage @@ -58,7 +59,7 @@ //! .unwrap() //! .unwrap(); //! let bitvec = BitVec::from_vec(pair1.0.clone()); -//! assert_eq!(bonsai_at_txn.get(&bitvec).unwrap().unwrap(), pair1.1); +//! assert_eq!(bonsai_at_txn.get(&identifier, &bitvec).unwrap().unwrap(), pair1.1); //! }); //! //! s.spawn(|| { @@ -67,18 +68,18 @@ //! .unwrap() //! .unwrap(); //! let bitvec = BitVec::from_vec(pair1.0.clone()); -//! assert_eq!(bonsai_at_txn.get(&bitvec).unwrap().unwrap(), pair1.1); +//! assert_eq!(bonsai_at_txn.get(&identifier, &bitvec).unwrap().unwrap(), pair1.1); //! }); //! }); //! bonsai_storage -//! .get(&BitVec::from_vec(vec![1, 2, 2])) +//! .get(&identifier, &BitVec::from_vec(vec![1, 2, 2])) //! .unwrap(); //! let pair2 = ( //! vec![1, 2, 3], //! Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(), //! ); //! bonsai_storage -//! .insert(&BitVec::from_vec(pair2.0.clone()), &pair2.1) +//! .insert(&identifier, &BitVec::from_vec(pair2.0.clone()), &pair2.1) //! .unwrap(); //! bonsai_storage.commit(id_builder.new_id()).unwrap(); //! ``` @@ -111,6 +112,7 @@ pub mod id; pub use bonsai_database::{BonsaiDatabase, BonsaiPersistentDatabase, DBError, DatabaseKey}; pub use error::BonsaiStorageError; +use trie::merkle_tree::MerkleTrees; pub use trie::merkle_tree::{Membership, ProofNode}; #[cfg(test)] @@ -161,9 +163,9 @@ pub struct BonsaiStorage where DB: BonsaiDatabase, ChangeID: id::Id, - H: StarkHash, + H: StarkHash + Send + Sync, { - trie: MerkleTree, + tries: MerkleTrees, } /// Trie root hash type. @@ -173,7 +175,7 @@ impl BonsaiStorage where DB: BonsaiDatabase, ChangeID: id::Id, - H: StarkHash, + H: StarkHash + Send + Sync, { /// Create a new bonsai storage instance pub fn new( @@ -182,7 +184,7 @@ where ) -> Result> { let key_value_db = KeyValueDB::new(db, config.into(), None); Ok(Self { - trie: MerkleTree::new(key_value_db)?, + tries: MerkleTrees::new(key_value_db), }) } @@ -190,21 +192,36 @@ where db: DB, config: BonsaiStorageConfig, created_at: ChangeID, + identifiers: Vec>, ) -> Result> { let key_value_db = KeyValueDB::new(db, config.into(), Some(created_at)); - Ok(Self { - trie: MerkleTree::new(key_value_db)?, - }) + let mut tries = MerkleTrees::::new(key_value_db); + for identifier in identifiers { + tries.init_tree(&identifier)?; + } + Ok(Self { tries }) + } + + /// Initialize a new trie with the given identifier. + /// This function is useful when you want to create a new trie in the database without inserting any value. + /// If the trie already exists, it will do nothing. + /// When you insert a value in a trie, it will automatically create the trie if it doesn't exist. + pub fn init_tree( + &mut self, + identifier: &[u8], + ) -> Result<(), BonsaiStorageError> { + self.tries.init_tree(identifier) } /// Insert a new key/value in the trie, overwriting the previous value if it exists. /// If the value already exists it will overwrite it. pub fn insert( &mut self, + identifier: &[u8], key: &BitSlice, value: &Felt, ) -> Result<(), BonsaiStorageError> { - self.trie.set(key, *value)?; + self.tries.set(identifier, key, *value)?; Ok(()) } @@ -212,26 +229,29 @@ where /// If the value doesn't exist it will do nothing pub fn remove( &mut self, + identifier: &[u8], key: &BitSlice, ) -> Result<(), BonsaiStorageError> { - self.trie.set(key, Felt::ZERO)?; + self.tries.set(identifier, key, Felt::ZERO)?; Ok(()) } /// Get a value in the trie. pub fn get( &self, + identifier: &[u8], key: &BitSlice, ) -> Result, BonsaiStorageError> { - self.trie.get(key) + self.tries.get(identifier, key) } /// Checks if the key exists in the trie. pub fn contains( &self, + identifier: &[u8], key: &BitSlice, ) -> Result> { - self.trie.contains(key) + self.tries.contains(identifier, key) } /// Go to a specific commit ID. @@ -241,7 +261,7 @@ where &mut self, requested_id: ChangeID, ) -> Result<(), BonsaiStorageError> { - let kv = self.trie.db_mut(); + let kv = self.tries.db_mut(); // Clear current changes kv.changes_store.current_changes.0.clear(); @@ -266,7 +286,14 @@ where // Accumulate changes from requested to last recorded let mut full = Vec::new(); - for id in kv.changes_store.id_queue.iter().skip(id_position).rev() { + for id in kv + .changes_store + .id_queue + .iter() + .skip(id_position) + .rev() + .take_while(|id| *id != &requested_id) + { full.extend( ChangeBatch::deserialize( id, @@ -306,7 +333,7 @@ where // Write revert changes and trie logs truncation kv.db.write_batch(batch)?; - self.trie.reset_to_last_commit()?; + self.tries.reset_to_last_commit()?; Ok(()) } @@ -316,17 +343,20 @@ where &self, id: ChangeID, ) -> Result, Change>, BonsaiStorageError> { - self.trie.db_ref().get_changes(id) + self.tries.db_ref().get_changes(id) } #[cfg(test)] pub fn dump_database(&self) { - self.trie.db_ref().db.dump_database(); + self.tries.db_ref().db.dump_database(); } /// Get trie root hash at the latest commit - pub fn root_hash(&self) -> Result> { - Ok(self.trie.root_hash()) + pub fn root_hash( + &self, + identifier: &[u8], + ) -> Result> { + self.tries.root_hash(identifier) } /// This function must be used with transactional state only. @@ -335,8 +365,8 @@ where &mut self, id: ChangeID, ) -> Result<(), BonsaiStorageError> { - self.trie.commit()?; - self.trie.db_mut().commit(id)?; + self.tries.commit()?; + self.tries.db_mut().commit(id)?; Ok(()) } @@ -353,9 +383,18 @@ where /// 3. the root hash matches the known root pub fn get_proof( &self, + identifier: &[u8], key: &BitSlice, ) -> Result, BonsaiStorageError> { - self.trie.get_proof(key) + self.tries.get_proof(identifier, key) + } + + /// Get all the keys in a specific trie. + pub fn get_keys( + &self, + identifier: &[u8], + ) -> Result>, BonsaiStorageError> { + self.tries.get_keys(identifier) } /// Verifies a merkle-proof for a given `key` and `value`. @@ -365,7 +404,7 @@ where value: Felt, proofs: &[ProofNode], ) -> Option { - MerkleTree::::verify_proof(root, key, value, proofs) + MerkleTree::::verify_proof(root, key, value, proofs) } } @@ -373,16 +412,16 @@ impl BonsaiStorage where DB: BonsaiDatabase + BonsaiPersistentDatabase, ChangeID: id::Id, - H: StarkHash, + H: StarkHash + Send + Sync, { /// Update trie and database using all changes since the last commit. pub fn commit( &mut self, id: ChangeID, ) -> Result<(), BonsaiStorageError<::DatabaseError>> { - self.trie.commit()?; - self.trie.db_mut().commit(id)?; - self.trie.db_mut().create_snapshot(id); + self.tries.commit()?; + self.tries.db_mut().commit(id)?; + self.tries.db_mut().create_snapshot(id); Ok(()) } @@ -399,11 +438,12 @@ where Option>, BonsaiStorageError<::DatabaseError>, > { - if let Some(transaction) = self.trie.db_ref().get_transaction(change_id)? { + if let Some(transaction) = self.tries.db_ref().get_transaction(change_id)? { Ok(Some(BonsaiStorage::new_from_transactional_state( transaction, config, change_id, + self.tries.get_identifiers(), )?)) } else { Ok(None) @@ -412,7 +452,7 @@ where /// Get a copy of the config that can be used to create a transactional state or a new bonsai storage. pub fn get_config(&self) -> BonsaiStorageConfig { - self.trie.db_ref().get_config().into() + self.tries.db_ref().get_config().into() } /// Merge a transactional state into the main trie. @@ -421,8 +461,8 @@ where transactional_bonsai_storage: BonsaiStorage, ) -> Result<(), BonsaiStorageError<>::DatabaseError>> { - self.trie + self.tries .db_mut() - .merge(transactional_bonsai_storage.trie.db()) + .merge(transactional_bonsai_storage.tries.db()) } } diff --git a/src/tests/madara_comparison.rs b/src/tests/madara_comparison.rs index 48b1a82..870107d 100644 --- a/src/tests/madara_comparison.rs +++ b/src/tests/madara_comparison.rs @@ -10,6 +10,7 @@ use crate::{ #[test] fn trie_height_251() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -19,12 +20,14 @@ fn trie_height_251() { let mut key: BitVec = bits![u8, Msb0; 0; 251].to_bitvec(); key.set(i, true); let value = Felt::from_hex("0x01").unwrap(); - bonsai_storage.insert(key.as_bitslice(), &value).unwrap(); + bonsai_storage + .insert(&identifier, key.as_bitslice(), &value) + .unwrap(); } let mut id_builder = BasicIdBuilder::new(); let id = id_builder.new_id(); bonsai_storage.commit(id).unwrap(); - bonsai_storage.root_hash().unwrap(); + bonsai_storage.root_hash(&identifier).unwrap(); } // Test to add on Madara side to check with a tree of height 251 and see that we have same hash // #[test]// fn test_height_251() { diff --git a/src/tests/proof.rs b/src/tests/proof.rs index a315eda..bb85425 100644 --- a/src/tests/proof.rs +++ b/src/tests/proof.rs @@ -113,6 +113,7 @@ fn assert_eq_proof(bonsai_proof: &[ProofNode], pathfinder_proof: &[TrieNode]) { #[test] fn basic_proof() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -128,7 +129,9 @@ fn basic_proof() { Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(), ); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); pathfinder_merkle_tree .set( &storage, @@ -141,7 +144,9 @@ fn basic_proof() { Felt::from_hex("0x66342762FD54D033c195fec3ce2568b62052e").unwrap(), ); let bitvec = BitVec::from_vec(pair2.0.clone()); - bonsai_storage.insert(&bitvec, &pair2.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair2.1) + .unwrap(); pathfinder_merkle_tree .set( &storage, @@ -154,7 +159,9 @@ fn basic_proof() { Felt::from_hex("0x66342762FD54D033c195fec3ce2568b62052e").unwrap(), ); let bitvec = BitVec::from_vec(pair3.0.clone()); - bonsai_storage.insert(&bitvec, &pair3.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair3.1) + .unwrap(); pathfinder_merkle_tree .set( &storage, @@ -165,7 +172,7 @@ fn basic_proof() { bonsai_storage.commit(id_builder.new_id()).unwrap(); let (_, root_id) = commit_and_persist(pathfinder_merkle_tree.clone(), &mut storage); let bonsai_proof = bonsai_storage - .get_proof(&BitVec::from_vec(vec![1, 2, 1])) + .get_proof(&identifier, &BitVec::from_vec(vec![1, 2, 1])) .unwrap(); let pathfinder_proof = pathfinder_merkle_tree::tree::MerkleTree::::get_proof( @@ -177,7 +184,7 @@ fn basic_proof() { assert_eq_proof(&bonsai_proof, &pathfinder_proof); assert_eq!( BonsaiStorage::, Pedersen>::verify_proof( - bonsai_storage.root_hash().unwrap(), + bonsai_storage.root_hash(&identifier).unwrap(), &BitVec::from_vec(vec![1, 2, 1]), pair1.1, &bonsai_proof @@ -188,6 +195,7 @@ fn basic_proof() { #[test] fn multiple_proofs() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -211,7 +219,7 @@ fn multiple_proofs() { let value = Felt::from_hex(&element).unwrap(); let key = &value.to_bytes_be()[..31]; bonsai_storage - .insert(&BitVec::from_vec(key.to_vec()), &value) + .insert(&identifier, &BitVec::from_vec(key.to_vec()), &value) .unwrap(); pathfinder_merkle_tree .set( @@ -226,7 +234,7 @@ fn multiple_proofs() { let (_, root_id) = commit_and_persist(pathfinder_merkle_tree.clone(), &mut storage); for element in elements.iter() { let proof = bonsai_storage - .get_proof(&BitVec::from_vec(element.0.clone())) + .get_proof(&identifier, &BitVec::from_vec(element.0.clone())) .unwrap(); let pathfinder_proof = pathfinder_merkle_tree::tree::MerkleTree::::get_proof( @@ -238,7 +246,7 @@ fn multiple_proofs() { assert_eq_proof(&proof, &pathfinder_proof); assert_eq!( BonsaiStorage::, Pedersen>::verify_proof( - bonsai_storage.root_hash().unwrap(), + bonsai_storage.root_hash(&identifier).unwrap(), &BitVec::from_vec(element.0.clone()), element.1, &proof @@ -250,6 +258,7 @@ fn multiple_proofs() { #[test] fn one_element_proof() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -265,7 +274,9 @@ fn one_element_proof() { Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(), ); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); pathfinder_merkle_tree .set( &storage, @@ -276,7 +287,7 @@ fn one_element_proof() { bonsai_storage.commit(id_builder.new_id()).unwrap(); let (_, root_id) = commit_and_persist(pathfinder_merkle_tree.clone(), &mut storage); let bonsai_proof = bonsai_storage - .get_proof(&BitVec::from_vec(vec![1, 2, 1])) + .get_proof(&identifier, &BitVec::from_vec(vec![1, 2, 1])) .unwrap(); let pathfinder_proof = pathfinder_merkle_tree::tree::MerkleTree::::get_proof( @@ -288,7 +299,7 @@ fn one_element_proof() { assert_eq_proof(&bonsai_proof, &pathfinder_proof); assert_eq!( BonsaiStorage::, Pedersen>::verify_proof( - bonsai_storage.root_hash().unwrap(), + bonsai_storage.root_hash(&identifier).unwrap(), &BitVec::from_vec(vec![1, 2, 1]), pair1.1, &bonsai_proof @@ -299,6 +310,7 @@ fn one_element_proof() { #[test] fn zero_not_crashing() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -308,6 +320,6 @@ fn zero_not_crashing() { let mut id_builder = BasicIdBuilder::new(); bonsai_storage.commit(id_builder.new_id()).unwrap(); bonsai_storage - .get_proof(&BitVec::from_vec(vec![1, 2, 1])) - .unwrap(); + .get_proof(&identifier, &BitVec::from_vec(vec![1, 2, 1])) + .expect_err("Should error"); } diff --git a/src/tests/simple.rs b/src/tests/simple.rs index ef53de1..13483d6 100644 --- a/src/tests/simple.rs +++ b/src/tests/simple.rs @@ -4,11 +4,12 @@ use crate::{ id::{BasicId, BasicIdBuilder}, BonsaiStorage, BonsaiStorageConfig, Change, }; -use bitvec::{vec::BitVec, view::BitView}; +use bitvec::{order::Msb0, vec::BitVec, view::BitView}; use starknet_types_core::{felt::Felt, hash::Pedersen}; #[test] fn basics() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -20,41 +21,49 @@ fn basics() { Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(), ); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); bonsai_storage.commit(id_builder.new_id()).unwrap(); let pair2 = ( vec![1, 2, 2], Felt::from_hex("0x66342762FD54D033c195fec3ce2568b62052e").unwrap(), ); let bitvec = BitVec::from_vec(pair2.0.clone()); - bonsai_storage.insert(&bitvec, &pair2.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair2.1) + .unwrap(); bonsai_storage.commit(id_builder.new_id()).unwrap(); let pair3 = ( vec![1, 2, 3], Felt::from_hex("0x66342762FD54D033c195fec3ce2568b62052e").unwrap(), ); let bitvec = BitVec::from_vec(pair3.0.clone()); - bonsai_storage.insert(&bitvec, &pair3.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair3.1) + .unwrap(); bonsai_storage.commit(id_builder.new_id()).unwrap(); let bitvec = BitVec::from_vec(vec![1, 2, 1]); - bonsai_storage.remove(&bitvec).unwrap(); + bonsai_storage.remove(&identifier, &bitvec).unwrap(); assert_eq!( bonsai_storage - .get(&BitVec::from_vec(vec![1, 2, 1])) + .get(&identifier, &BitVec::from_vec(vec![1, 2, 1])) .unwrap(), None ); bonsai_storage.commit(id_builder.new_id()).unwrap(); assert_eq!( bonsai_storage - .get(&BitVec::from_vec(vec![1, 2, 1])) + .get(&identifier, &BitVec::from_vec(vec![1, 2, 1])) .unwrap(), None ); + assert_eq!(bonsai_storage.get_keys(&identifier).unwrap().len(), 2); } #[test] fn root_hash_similar_rocks_db() { + let identifier = vec![]; let root_hash_1 = { let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); @@ -68,14 +77,18 @@ fn root_hash_similar_rocks_db() { .unwrap(), ); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); let pair2 = ( vec![1, 2, 2], Felt::from_hex("0x100bd6fbfced88ded1b34bd1a55b747ce3a9fde9a914bca75571e4496b56443") .unwrap(), ); let bitvec = BitVec::from_vec(pair2.0.clone()); - bonsai_storage.insert(&bitvec, &pair2.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair2.1) + .unwrap(); bonsai_storage.commit(id_builder.new_id()).unwrap(); let pair3 = ( vec![1, 2, 3], @@ -83,16 +96,20 @@ fn root_hash_similar_rocks_db() { .unwrap(), ); let bitvec = BitVec::from_vec(pair3.0.clone()); - bonsai_storage.insert(&bitvec, &pair3.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair3.1) + .unwrap(); let pair4 = ( vec![1, 2, 4], Felt::from_hex("0x02808c7d8f3745e55655ad3f51f096d0c06a41f3d76caf96bad80f9be9ced171") .unwrap(), ); let bitvec = BitVec::from_vec(pair4.0.clone()); - bonsai_storage.insert(&bitvec, &pair4.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair4.1) + .unwrap(); bonsai_storage.commit(id_builder.new_id()).unwrap(); - bonsai_storage.root_hash().unwrap() + bonsai_storage.root_hash(&identifier).unwrap() }; let root_hash_2 = { let tempdir = tempfile::tempdir().unwrap(); @@ -107,16 +124,20 @@ fn root_hash_similar_rocks_db() { .unwrap(), ); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); let pair2 = ( vec![1, 2, 4], Felt::from_hex("0x02808c7d8f3745e55655ad3f51f096d0c06a41f3d76caf96bad80f9be9ced171") .unwrap(), ); let bitvec = BitVec::from_vec(pair2.0.clone()); - bonsai_storage.insert(&bitvec, &pair2.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair2.1) + .unwrap(); bonsai_storage.commit(id_builder.new_id()).unwrap(); - bonsai_storage.root_hash().unwrap() + bonsai_storage.root_hash(&identifier).unwrap() }; println!("root_hash_1: {:?}", root_hash_1.to_string()); println!("root_hash_2: {:?}", root_hash_2.to_string()); @@ -129,6 +150,7 @@ fn starknet_specific() { address: &'static str, state_hash: &'static str, } + let identifier = vec![]; let tempdir1 = tempfile::tempdir().unwrap(); let db1 = create_rocks_db(tempdir1.path()).unwrap(); @@ -160,10 +182,10 @@ fn starknet_specific() { let key = Felt::from_hex(key).unwrap().to_bytes_be().view_bits()[5..].to_bitvec(); let value = Felt::from_hex(value).unwrap(); bonsai_storage1 - .insert(&key, &value) + .insert(&identifier, &key, &value) .expect("Failed to insert storage update into trie"); bonsai_storage2 - .insert(&key, &value) + .insert(&identifier, &key, &value) .expect("Failed to insert storage update into trie"); } @@ -184,10 +206,10 @@ fn starknet_specific() { let value = Felt::from_hex(value).unwrap(); bonsai_storage1 - .insert(&key, &value) + .insert(&identifier, &key, &value) .expect("Failed to insert storage update into trie"); bonsai_storage2 - .insert(&key, &value) + .insert(&identifier, &key, &value) .expect("Failed to insert storage update into trie"); } @@ -196,20 +218,21 @@ fn starknet_specific() { .commit(id) .expect("Failed to commit to bonsai storage"); let root_hash1 = bonsai_storage1 - .root_hash() + .root_hash(&identifier) .expect("Failed to get root hash"); bonsai_storage2 .commit(id) .expect("Failed to commit to bonsai storage"); let root_hash2 = bonsai_storage2 - .root_hash() + .root_hash(&identifier) .expect("Failed to get root hash"); assert_eq!(root_hash1, root_hash2); } #[test] fn root_hash_similar_hashmap_db() { + let identifier = vec![]; let root_hash_1 = { let db = HashMapDb::::default(); let config = BonsaiStorageConfig::default(); @@ -222,14 +245,18 @@ fn root_hash_similar_hashmap_db() { .unwrap(), ); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); let pair2 = ( vec![1, 2, 2], Felt::from_hex("0x100bd6fbfced88ded1b34bd1a55b747ce3a9fde9a914bca75571e4496b56443") .unwrap(), ); let bitvec = BitVec::from_vec(pair2.0.clone()); - bonsai_storage.insert(&bitvec, &pair2.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair2.1) + .unwrap(); bonsai_storage.commit(id_builder.new_id()).unwrap(); let pair3 = ( vec![1, 2, 3], @@ -237,16 +264,20 @@ fn root_hash_similar_hashmap_db() { .unwrap(), ); let bitvec = BitVec::from_vec(pair3.0.clone()); - bonsai_storage.insert(&bitvec, &pair3.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair3.1) + .unwrap(); let pair4 = ( vec![1, 2, 4], Felt::from_hex("0x02808c7d8f3745e55655ad3f51f096d0c06a41f3d76caf96bad80f9be9ced171") .unwrap(), ); let bitvec = BitVec::from_vec(pair4.0.clone()); - bonsai_storage.insert(&bitvec, &pair4.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair4.1) + .unwrap(); bonsai_storage.commit(id_builder.new_id()).unwrap(); - bonsai_storage.root_hash().unwrap() + bonsai_storage.root_hash(&identifier).unwrap() }; let root_hash_2 = { let db = HashMapDb::::default(); @@ -260,16 +291,20 @@ fn root_hash_similar_hashmap_db() { .unwrap(), ); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); let pair2 = ( vec![1, 2, 4], Felt::from_hex("0x02808c7d8f3745e55655ad3f51f096d0c06a41f3d76caf96bad80f9be9ced171") .unwrap(), ); let bitvec = BitVec::from_vec(pair2.0.clone()); - bonsai_storage.insert(&bitvec, &pair2.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair2.1) + .unwrap(); bonsai_storage.commit(id_builder.new_id()).unwrap(); - bonsai_storage.root_hash().unwrap() + bonsai_storage.root_hash(&identifier).unwrap() }; println!("root_hash_1: {:?}", root_hash_1.to_string()); println!("root_hash_2: {:?}", root_hash_2.to_string()); @@ -278,6 +313,7 @@ fn root_hash_similar_hashmap_db() { #[test] fn double_insert() { + let identifier = vec![]; struct ContractState { address: &'static str, state_hash: &'static str, @@ -320,21 +356,96 @@ fn double_insert() { let bitkey = key.to_bytes_be().view_bits()[5..].to_bitvec(); let value = Felt::from_hex(value).unwrap(); bonsai_storage - .insert(&bitkey, &value) + .insert(&identifier, &bitkey, &value) .expect("Failed to insert storage update into trie"); // fails here for key 0x313ad57fdf765addc71329abf8d74ac2bce6d46da8c2b9b82255a5076620301 // and value 0x453ae0c9610197b18b13645c44d3d0a407083d96562e8752aab3fab616cecb0 - bonsai_storage.insert(&bitkey, &value).expect(&format!( - "Failed to insert storage update into trie for key {key:#x} & value {value:#x}" - )); + bonsai_storage + .insert(&identifier, &bitkey, &value) + .unwrap_or_else(|_| { + panic!( + "Failed to insert storage update into trie for key {key:#x} & value {value:#x}" + ) + }); } bonsai_storage.commit(id_builder.new_id()).unwrap(); - let root_hash = bonsai_storage.root_hash().unwrap(); + let root_hash = bonsai_storage.root_hash(&identifier).unwrap(); println!("root hash: {root_hash:#x}"); } +#[test] +fn double_identifier() { + let identifier = vec![2, 3]; + let identifier2 = vec![1, 3, 1]; + struct ContractState { + address: &'static str, + state_hash: &'static str, + } + let tempdir = tempfile::tempdir().unwrap(); + let db = create_rocks_db(tempdir.path()).unwrap(); + let config = BonsaiStorageConfig::default(); + let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + let mut id_builder = BasicIdBuilder::new(); + let contract_states = vec![ + ContractState { + address: "0x0000000000000000000000000000000000000000000000000000000000000005", + state_hash: "0x000000000000000000000000000000000000000000000000000000000000022b", + }, + ContractState { + address: "0x0313ad57fdf765addc71329abf8d74ac2bce6d46da8c2b9b82255a5076620300", + state_hash: "0x04e7e989d58a17cd279eca440c5eaa829efb6f9967aaad89022acbe644c39b36", + }, + // This seems to be what is causing the problem in case of double insertions. + // Other value are fine + ContractState { + address: "0x313ad57fdf765addc71329abf8d74ac2bce6d46da8c2b9b82255a5076620301", + state_hash: "0x453ae0c9610197b18b13645c44d3d0a407083d96562e8752aab3fab616cecb0", + }, + ContractState { + address: "0x05aee31408163292105d875070f98cb48275b8c87e80380b78d30647e05854d5", + state_hash: "0x00000000000000000000000000000000000000000000000000000000000007e5", + }, + ContractState { + address: "0x06cf6c2f36d36b08e591e4489e92ca882bb67b9c39a3afccf011972a8de467f0", + state_hash: "0x07ab344d88124307c07b56f6c59c12f4543e9c96398727854a322dea82c73240", + }, + ]; + for contract_state in contract_states { + let key = contract_state.address; + let value = contract_state.state_hash; + + let key = Felt::from_hex(key).unwrap(); + let bitkey = key.to_bytes_be().view_bits()[5..].to_bitvec(); + let value = Felt::from_hex(value).unwrap(); + bonsai_storage + .insert(&identifier, &bitkey, &value) + .expect("Failed to insert storage update into trie"); + // fails here for key 0x313ad57fdf765addc71329abf8d74ac2bce6d46da8c2b9b82255a5076620301 + // and value 0x453ae0c9610197b18b13645c44d3d0a407083d96562e8752aab3fab616cecb0 + bonsai_storage + .insert(&identifier2, &bitkey, &value) + .unwrap_or_else(|_| { + panic!( + "Failed to insert storage update into trie for key {key:#x} & value {value:#x}" + ) + }); + } + bonsai_storage.commit(id_builder.new_id()).unwrap(); + let root_hash = bonsai_storage.root_hash(&identifier).unwrap(); + println!("root hash: {root_hash:#x}"); + let root_hash2 = bonsai_storage.root_hash(&identifier2).unwrap(); + assert_eq!(root_hash, root_hash2); + assert_eq!( + bonsai_storage.get_keys(&identifier).unwrap(), + bonsai_storage.get_keys(&identifier2).unwrap() + ); + assert_eq!(bonsai_storage.get_keys(&identifier).unwrap().len(), 5); +} + #[test] fn get_changes() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -343,17 +454,25 @@ fn get_changes() { let mut id_builder = BasicIdBuilder::new(); let pair1 = (vec![1, 2, 1], Felt::from_hex("0x01").unwrap()); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); bonsai_storage.commit(id_builder.new_id()).unwrap(); let pair2 = (vec![1, 2, 2], Felt::from_hex("0x01").unwrap()); let bitvec = BitVec::from_vec(pair2.0.clone()); - bonsai_storage.insert(&bitvec, &pair2.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair2.1) + .unwrap(); let pair1_edited_1 = (vec![1, 2, 1], Felt::from_hex("0x02").unwrap()); let bitvec = BitVec::from_vec(pair1_edited_1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1_edited_1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1_edited_1.1) + .unwrap(); let pair1_edited_2 = (vec![1, 2, 1], Felt::from_hex("0x03").unwrap()); let bitvec = BitVec::from_vec(pair1_edited_2.0.clone()); - bonsai_storage.insert(&bitvec, &pair1_edited_2.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1_edited_2.1) + .unwrap(); let id = id_builder.new_id(); bonsai_storage.commit(id).unwrap(); let changes = bonsai_storage.get_changes(id).unwrap(); @@ -373,3 +492,560 @@ fn get_changes() { } ); } + +fn keyer(felt: Felt) -> BitVec { + felt.to_bytes_be().view_bits()[5..].to_bitvec() +} + +#[test] +fn test_insert_zero() { + let config = BonsaiStorageConfig::default(); + let bonsai_db = HashMapDb::::default(); + let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config) + .expect("Failed to create bonsai storage"); + let identifier = + "0x056e4fed965fccd7fb01fcadd827470338f35ced62275328929d0d725b5707ba".as_bytes(); + + // Insert Block 3 storage changes for contract `0x056e4fed965fccd7fb01fcadd827470338f35ced62275328929d0d725b5707ba` + let block_3 = [ + ("0x5", "0x456"), + ( + "0x378e096bb5e74b0f4ca78660a6b49b4a8035e571b024c018713c80b4b969735", + "0x205d119502a165dae3830f627fa93fbdf5bfb13edd8f00e4c72621d0cda24", + ), + ( + "0x41139bbf557d599fe8e96983251ecbfcb5bf4c4138c85946b0c4a6a68319f24", + "0x7eec291f712520293664c7e3a8bb39ab00babf51cb0d9c1fb543147f37b485f", + ), + ( + "0x77ae79c60260b3e48516a7da1aa173ac2765a5ced420f8ffd1539c394fbc03c", + "0x6025343ab6a7ac36acde4eba3b6fc21f53d5302ee26e6f28e8de5a62bbfd847", + ), + ( + "0x751901aac66fdc1f455c73022d02f1c085602cd0c9acda907cfca5418769e9c", + "0x3f23078d48a4bf1d5f8ca0348f9efe9300834603625a379cae5d6d81100adef", + ), + ( + "0x751901aac66fdc1f455c73022d02f1c085602cd0c9acda907cfca5418769e9d", + "0xbd858a06904cadc3787ecbad97409606dcee50ea6fc30b94930bcf3d8843d5", + ), + ]; + + for (key_hex, value_hex) in block_3.iter() { + let key: Felt = Felt::from_hex(key_hex).unwrap(); + let value = Felt::from_hex(value_hex).unwrap(); + bonsai_storage + .insert(identifier, keyer(key).as_bitslice(), &value) + .expect("Failed to insert storage update into trie"); + } + + let mut id_builder = BasicIdBuilder::new(); + let id = id_builder.new_id(); + bonsai_storage + .commit(id) + .expect("Failed to commit to bonsai storage"); + let root_hash = bonsai_storage + .root_hash(identifier) + .expect("Failed to get root hash"); + + println!( + "Expected: 0x069064A05C14A9A2B4ED81C479C14D30872A9AE9CE2DEA8E4B4509542C2DCC1F\nFound: {root_hash:#x}", + ); + assert_eq!( + root_hash, + Felt::from_hex("0x069064A05C14A9A2B4ED81C479C14D30872A9AE9CE2DEA8E4B4509542C2DCC1F") + .unwrap() + ); + + // Insert Block 4 storage changes for contract `0x056e4fed965fccd7fb01fcadd827470338f35ced62275328929d0d725b5707ba` + let block_4 = [ + ("0x5", "0x0"), // Inserting key = 0x0 + ( + "0x4b81c1bca2d1b7e08535a5abe231b2e94399674db5e8f1d851fd8f4af4abd34", + "0x7c7", + ), + ( + "0x6f8cf54aaec1f42d5f3868d597fcd7393da888264dc5a6e93c7bd528b6d6fee", + "0x7e5", + ), + ( + "0x2a315469199dfde4b05906db8c33f6962916d462d8f1cf5252b748dfa174a20", + "0xdae79d0308bb710af439eb36e82b405dc2bca23b351d08b4867d9525226e9d", + ), + ( + "0x2d1ed96c7561dd8e5919657790ffba8473b80872fea3f7ef8279a7253dc3b33", + "0x750387f4d66b0e9be1f2f330e8ad309733c46bb74e0be4df0a8c58fb4e89a25", + ), + ( + "0x6a93bcb89fc1f31fa544377c7de6de1dd3e726e1951abc95c4984995e84ad0d", + "0x7e5", + ), + ( + "0x6b3b4780013c33cdca6799e8aa3ef922b64f5a2d356573b33693d81504deccf", + "0x7c7", + ), + ]; + + for (key_hex, value_hex) in block_4.iter() { + let key: Felt = Felt::from_hex(key_hex).unwrap(); + let value = Felt::from_hex(value_hex).unwrap(); + bonsai_storage + .insert(identifier, keyer(key).as_bitslice(), &value) + .expect("Failed to insert storage update into trie"); + } + + let id = id_builder.new_id(); + bonsai_storage + .commit(id) + .expect("Failed to commit to bonsai storage"); + let root_hash = bonsai_storage + .root_hash(identifier) + .expect("Failed to get root hash"); + + println!( + "Expected: 0x0112998A41A3A2C720E758F82D184E4C39E9382620F12076B52C516D14622E57\nFound: {root_hash:#x}", + ); + assert_eq!( + root_hash, + Felt::from_hex("0x0112998A41A3A2C720E758F82D184E4C39E9382620F12076B52C516D14622E57") + .unwrap() + ); + + // Insert Block 5 storage changes for contract `0x056e4fed965fccd7fb01fcadd827470338f35ced62275328929d0d725b5707ba` + let block_5 = [("0x5", "0x456")]; + + for (key_hex, value_hex) in block_5.iter() { + let key: Felt = Felt::from_hex(key_hex).unwrap(); + let value = Felt::from_hex(value_hex).unwrap(); + bonsai_storage + .insert(identifier, keyer(key).as_bitslice(), &value) + .expect("Failed to insert storage update into trie"); + } + + let id = id_builder.new_id(); + bonsai_storage + .commit(id) + .expect("Failed to commit to bonsai storage"); + let root_hash = bonsai_storage + .root_hash(identifier) + .expect("Failed to get root hash"); + + println!( + "Expected: 0x072E79A6F71E3E63D7DE40EDF4322A22E64388D4D5BFE817C1271C78028B73BF\nFound: {root_hash:#x}" + ); + assert_eq!( + root_hash, + Felt::from_hex("0x072E79A6F71E3E63D7DE40EDF4322A22E64388D4D5BFE817C1271C78028B73BF") + .unwrap() + ); +} + +#[test] +fn test_block_7_starknet() { + let config = BonsaiStorageConfig::default(); + let bonsai_db = HashMapDb::::default(); + let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config) + .expect("Failed to create bonsai storage"); + let identifier = + "0x056e4fed965fccd7fb01fcadd827470338f35ced62275328929d0d725b5707ba".as_bytes(); + + // Insert Block 3 storage changes for contract `0x4d56b8ac0ed905936da10323328cba5def12957a2936920f043d8bf6a1e902d` + let block_3 = [ + ( + "0x67c2665fbdd32ded72c0665f9658c05a5f9233c8de2002b3eba8ae046174efd", + "0x2221def5413ed3e128051d5dff3ec816dbfb9db4454b98f4aa47804cb7a13d2", + ), + ("0x5", "0x66"), + ( + "0x101c2b102c8eb6bf091f5debcf97d8edde85983e23f9778e9cabbe0b5a4f997", + "0x99a58a9612fe930f39c4c399b6be14e8bb7c8229d06eab8d0a3a97877a6667", + ), + ( + "0x1aabd3b2e12959bab2c4ab530c1d8f0e675e0dc5ab29d1f10b7f1a154cabef9", + "0x41d4ae0ba9013f2f6e1551b62a9c9187053727e0e65217be97eae8922d5b2df", + ), + ( + "0x1aabd3b2e12959bab2c4ab530c1d8f0e675e0dc5ab29d1f10b7f1a154cabefa", + "0x6eda96627bd3de7af5b4f932ff1e858bd396c897229d64b6dd3f0f936f0ea17", + ), + ]; + + for (key_hex, value_hex) in block_3.iter() { + let key = Felt::from_hex(key_hex).unwrap(); + let value = Felt::from_hex(value_hex).unwrap(); + bonsai_storage + .insert(identifier, keyer(key).as_bitslice(), &value) + .expect("Failed to insert storage update into trie"); + } + + let mut id_builder = BasicIdBuilder::new(); + let id = id_builder.new_id(); + bonsai_storage + .commit(id) + .expect("Failed to commit to bonsai storage"); + let root_hash = bonsai_storage + .root_hash(identifier) + .expect("Failed to get root hash"); + + println!("Expected: 0x0297DE74ABD178CAF7EA2F1AE1B4588CA7433B1B11A98172B6F56E3E02739FD0\nFound: {root_hash:#x}"); + assert_eq!( + root_hash, + Felt::from_hex("0x0297DE74ABD178CAF7EA2F1AE1B4588CA7433B1B11A98172B6F56E3E02739FD0") + .unwrap() + ); + + // Insert Block 4 storage changes for contract `0x4d56b8ac0ed905936da10323328cba5def12957a2936920f043d8bf6a1e902d` + let block_4 = [ + ( + "0x3c14ddc99b06b00340bffd81ef1c4e10f74b800a911ee22c22bb28e4b516da5", + "0x7e5", + ), + ("0x5", "0x64"), + ( + "0x5201dd2a5f567a653e9a2b7a62816919d0d695d1e2f39d516f9befda30da720", + "0x29ed6ea046ebe50aaacb9cd6477ac368644c8f4242ee0687d31f6c2ac20c146", + ), + ( + "0x5b3856459ac954d3fd24d85924d978263709880e3ee4cafdfe0b7c95ee6b26a", + "0x4c90411b3376d5230a88496e58acf58c19431d52b89f1ab91924075f4b35ac1", + ), + ( + "0x5b3856459ac954d3fd24d85924d978263709880e3ee4cafdfe0b7c95ee6b26b", + "0x72a56d83fab34872a880dd35d936117a084b928fb9d47306abb2558472633c", + ), + ( + "0x6a93bcb89fc1f31fa544377c7de6de1dd3e726e1951abc95c4984995e84ad0d", + "0x7c7", + ), + ( + "0x6f8cf54aaec1f42d5f3868d597fcd7393da888264dc5a6e93c7bd528b6d6fee", + "0x7c7", + ), + ( + "0x6b30a5f1341c0c949f847afe7f761a6ea8cdc3337baa20e68a2891f62389052", + "0x7e5", + ), + ( + "0x6b3b4780013c33cdca6799e8aa3ef922b64f5a2d356573b33693d81504deccf", + "0x7e5", + ), + ( + "0x6f649e057570e0f3cc710d260c2067297542f8e18407a7e75008808e12e6099", + "0x61395ebfa1746f9449711a7e361254ddb90f642861807b7e5e05276c11033ec", + ), + ( + "0x6f649e057570e0f3cc710d260c2067297542f8e18407a7e75008808e12e609a", + "0x304d0ec8cc0ea6faf0f7ad67903bcffc6bc4474d25f93e1c961b239370b8c07", + ), + ]; + + for (key_hex, value_hex) in block_4.iter() { + let key = Felt::from_hex(key_hex).unwrap(); + let value = Felt::from_hex(value_hex).unwrap(); + bonsai_storage + .insert(identifier, keyer(key).as_bitslice(), &value) + .expect("Failed to insert storage update into trie"); + } + + let id = id_builder.new_id(); + bonsai_storage + .commit(id) + .expect("Failed to commit to bonsai storage"); + let root_hash = bonsai_storage + .root_hash(identifier) + .expect("Failed to get root hash"); + + println!("Expected: 0x07A4CA1440AF3858CEB11386BA7E2A0FC553BB73E741043218845D820009BCCB\nFound: {root_hash:#x}"); + assert_eq!( + root_hash, + Felt::from_hex("0x07A4CA1440AF3858CEB11386BA7E2A0FC553BB73E741043218845D820009BCCB") + .unwrap() + ); + + // Insert Block 5 storage changes for contract `0x4d56b8ac0ed905936da10323328cba5def12957a2936920f043d8bf6a1e902d` + let block_5 = [ + ( + "0x272cd29c23c7fd72ef13352ac037c6fabfee4c03056ea413c326be6501b4f31", + "0x7c7", + ), + ( + "0x2bb6a7dd9cbb9cec8fdad9c0557bd539683f7ea65d4f14d41fe4d72311775e3", + "0x7e5", + ), + ]; + + for (key_hex, value_hex) in block_5.iter() { + let key = Felt::from_hex(key_hex).unwrap(); + let value = Felt::from_hex(value_hex).unwrap(); + bonsai_storage + .insert(identifier, keyer(key).as_bitslice(), &value) + .expect("Failed to insert storage update into trie"); + } + + let id = id_builder.new_id(); + bonsai_storage + .commit(id) + .expect("Failed to commit to bonsai storage"); + let root_hash = bonsai_storage + .root_hash(identifier) + .expect("Failed to get root hash"); + + println!("Expected: 0x002363DCD04D065C6B50A4D46F930EBC91AC7F4B15DCF1B0A8D0165B0BA0F143\nFound: {root_hash:#x}"); + assert_eq!( + root_hash, + Felt::from_hex("0x002363DCD04D065C6B50A4D46F930EBC91AC7F4B15DCF1B0A8D0165B0BA0F143") + .unwrap() + ); + + // Insert Block 6 storage changes for contract `0x4d56b8ac0ed905936da10323328cba5def12957a2936920f043d8bf6a1e902d` + let block_6 = [("0x5", "0x22b")]; + + for (key_hex, value_hex) in block_6.iter() { + let key = Felt::from_hex(key_hex).unwrap(); + let value = Felt::from_hex(value_hex).unwrap(); + bonsai_storage + .insert(identifier, keyer(key).as_bitslice(), &value) + .expect("Failed to insert storage update into trie"); + } + + let id = id_builder.new_id(); + bonsai_storage + .commit(id) + .expect("Failed to commit to bonsai storage"); + let root_hash = bonsai_storage + .root_hash(identifier) + .expect("Failed to get root hash"); + + println!("Expected: 0x00C656C01BB43291BEA976CEACE3AFE89A5621045E3B6F23E4BCFFFBB4B66832\nFound: {root_hash:#x}"); + assert_eq!( + root_hash, + Felt::from_hex("0x00C656C01BB43291BEA976CEACE3AFE89A5621045E3B6F23E4BCFFFBB4B66832") + .unwrap() + ); + + // Insert Block 6 storage changes for contract `0x4d56b8ac0ed905936da10323328cba5def12957a2936920f043d8bf6a1e902d` + let block_7 = [("0x5", "0x0")]; + + for (key_hex, value_hex) in block_7.iter() { + let key = Felt::from_hex(key_hex).unwrap(); + let value = Felt::from_hex(value_hex).unwrap(); + bonsai_storage + .insert(identifier, keyer(key).as_bitslice(), &value) + .expect("Failed to insert storage update into trie"); + } + + let id = id_builder.new_id(); + bonsai_storage + .commit(id) + .expect("Failed to commit to bonsai storage"); + let root_hash = bonsai_storage + .root_hash(identifier) + .expect("Failed to get root hash"); + + println!("Expected: 0x032C61E78534A30DD005DB4B9136AA64893CC2F6E10C4535DD6F29BFB2ADC726\nFound: {root_hash:#x}"); + assert_eq!( + root_hash, + Felt::from_hex("0x032C61E78534A30DD005DB4B9136AA64893CC2F6E10C4535DD6F29BFB2ADC726") + .unwrap() + ); +} + +#[test] +fn test_block_7_starknet_2() { + let config = BonsaiStorageConfig::default(); + let bonsai_db = HashMapDb::::default(); + let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config) + .expect("Failed to create bonsai storage"); + let identifier = "0x421203c58e1b4a6c3675be26cfaa18d2b6b42695ca206be1f08ce29f7f1bc7c".as_bytes(); + + // Insert Block 5 storage changes for contract `0x421203c58e1b4a6c3675be26cfaa18d2b6b42695ca206be1f08ce29f7f1bc7c` + let block_5 = [ + ( + "0x2bb6a7dd9cbb9cec8fdad9c0557bd539683f7ea65d4f14d41fe4d72311775e3", + "0x7c7", + ), + ( + "0x584d53558c6731da8923f60f2d182027312ffa4e811e7eddc6401232d33400e", + "0x29bc2bad472c81f00b7873d7d27a68d63dc9ebd3a3661e2b4c3d6c90d732454", + ), + ( + "0x6c27ff92eab8802ca5141a60a5699e5075725d5526752c5fb368c12582af00c", + "0x645a108cc9b963369b91cad8a8b5c2ce774b79e871368d301d518012925abc6", + ), + ("0x5", "0x66"), + ( + "0x744f7d93c67c2ac6fbcdf632d530cebdbffa112d0cfacce28ed5773babfba60", + "0x2a49283d206395239d0c1d505a8ba2f446419e58a1fd40caccf796e810759d5", + ), + ( + "0x11f391c712bb4996774106b93766bc49f8bdb29b416cae0da0d981752c1a28b", + "0x43f3925b460d387343381e31e2f9299100609bc833f289bfd67316a0a06ce40", + ), + ( + "0x11f391c712bb4996774106b93766bc49f8bdb29b416cae0da0d981752c1a28c", + "0x2b72713e2fc2dec7cfe8e7c428f02728a031f17f876bb50841d4ee3eb12834", + ), + ( + "0x66631ce6af4e11972e05bed46e9b20a5480ffea4ae2a4d95e1d71fb37f25c0", + "0x1329ffd6765c348b5e7195b777241cf5eb84e438c0f5fa3acb5800ada846332", + ), + ]; + + for (key_hex, value_hex) in block_5.iter() { + let key = Felt::from_hex(key_hex).unwrap(); + let value = Felt::from_hex(value_hex).unwrap(); + bonsai_storage + .insert(identifier, keyer(key).as_bitslice(), &value) + .expect("Failed to insert storage update into trie"); + } + + let mut id_builder = BasicIdBuilder::new(); + let id = id_builder.new_id(); + bonsai_storage + .commit(id) + .expect("Failed to commit to bonsai storage"); + let root_hash = bonsai_storage + .root_hash(identifier) + .expect("Failed to get root hash"); + + println!("Expected: 0x03846F4AE281ADBCC68518766579DB77C27EF31955E9FC3183C397C2731A7627\nFound: {root_hash:#x}"); + assert_eq!( + root_hash, + Felt::from_hex("0x03846F4AE281ADBCC68518766579DB77C27EF31955E9FC3183C397C2731A7627") + .unwrap() + ); + + // Insert Block 6 storage changes for contract `0x421203c58e1b4a6c3675be26cfaa18d2b6b42695ca206be1f08ce29f7f1bc7c` + let block_6 = [ + ( + "0x591192c633e49a7e6ca0aae77da4e9a1df2c6db51cabb3cc929280a44745635", + "0x1b3479bec749469312a35a2001dc8cfaf38723c0a8763e01ad2abaefb2214e5", + ), + ( + "0x58bfc110ce09fc2bcff40dbb4887bfb32f5156f2195e8f6ea22e15784c01768", + "0x71cc8515287a6f5d8b81675bc7e41ca1fcd75afcc60984701033f0cdd05acd", + ), + ( + "0x58bfc110ce09fc2bcff40dbb4887bfb32f5156f2195e8f6ea22e15784c01769", + "0x6a8a49d797b80ef2be0ec8a72f71dccb655c07297f95e022a26a65787c3199c", + ), + ]; + + for (key_hex, value_hex) in block_6.iter() { + let key = Felt::from_hex(key_hex).unwrap(); + let value = Felt::from_hex(value_hex).unwrap(); + bonsai_storage + .insert(identifier, keyer(key).as_bitslice(), &value) + .expect("Failed to insert storage update into trie"); + } + + let id = id_builder.new_id(); + bonsai_storage + .commit(id) + .expect("Failed to commit to bonsai storage"); + let root_hash = bonsai_storage + .root_hash(identifier) + .expect("Failed to get root hash"); + + println!("Expected: 0x06E02FE529D3CBDCC5324D0981F991E777DAFC3F0C24E7CB56CE3D379BE9B510\nFound: {root_hash:#x}"); + assert_eq!( + root_hash, + Felt::from_hex("0x06E02FE529D3CBDCC5324D0981F991E777DAFC3F0C24E7CB56CE3D379BE9B510") + .unwrap() + ); + + // Insert Block 6 storage changes for contract `0x421203c58e1b4a6c3675be26cfaa18d2b6b42695ca206be1f08ce29f7f1bc7c` + let block_7 = [("0x5", "0x0")]; + + for (key_hex, value_hex) in block_7.iter() { + let key = Felt::from_hex(key_hex).unwrap(); + let value = Felt::from_hex(value_hex).unwrap(); + bonsai_storage + .insert(identifier, keyer(key).as_bitslice(), &value) + .expect("Failed to insert storage update into trie"); + } + + let id = id_builder.new_id(); + bonsai_storage + .commit(id) + .expect("Failed to commit to bonsai storage"); + let root_hash = bonsai_storage + .root_hash(identifier) + .expect("Failed to get root hash"); + + println!("Expected: 0x0528E360EA90E94F670451A76A7698900F0F7C1F2E88583F8B0162D486BF7947\nFound: {root_hash:#x}"); + assert_eq!( + root_hash, + Felt::from_hex("0x0528E360EA90E94F670451A76A7698900F0F7C1F2E88583F8B0162D486BF7947") + .unwrap() + ) +} + +#[test] +fn test_block_9() { + let config = BonsaiStorageConfig::default(); + let bonsai_db = HashMapDb::::default(); + let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config) + .expect("Failed to create bonsai storage"); + let identifier = + "0x06F3C934BA4EC49245CB9A42FC715E4D589AA502AF69BE13916127A538D525CE".as_bytes(); + + // Insert Block 8 storage changes for contract `0x06F3C934BA4EC49245CB9A42FC715E4D589AA502AF69BE13916127A538D525CE` + let block_8 = [ + ("0x5", "0x456"), + ( + "0x4b788ad12d2e47b2be358d61cc38d813aa79165ddbc0b29d4878ef0fbc18c15", + "0x612af3160e28962cb3dd6146a9c2f7bd7adeea1fddd39f767d936c7b5bcca97", + ), + ]; + + for (key_hex, value_hex) in block_8.iter() { + let key = Felt::from_hex(key_hex).unwrap(); + let value = Felt::from_hex(value_hex).unwrap(); + bonsai_storage + .insert(identifier, keyer(key).as_bitslice(), &value) + .expect("Failed to insert storage update into trie"); + } + + let mut id_builder = BasicIdBuilder::new(); + let id = id_builder.new_id(); + bonsai_storage + .commit(id) + .expect("Failed to commit to bonsai storage"); + let root_hash = bonsai_storage + .root_hash(identifier) + .expect("Failed to get root hash"); + + println!("Expected: 0x010AA5D1D36847AE64BA074B3A878BFD1A9AEAA952F6777C727EEA6AE6B2C99F\nFound: {root_hash:#x}"); + assert_eq!( + root_hash, + Felt::from_hex("0x010AA5D1D36847AE64BA074B3A878BFD1A9AEAA952F6777C727EEA6AE6B2C99F") + .unwrap() + ); + + // Insert Block 9 storage changes for contract `0x06F3C934BA4EC49245CB9A42FC715E4D589AA502AF69BE13916127A538D525CE` + let block_9 = [("0x5", "0x0")]; + + for (key_hex, value_hex) in block_9.iter() { + let key = Felt::from_hex(key_hex).unwrap(); + let value = Felt::from_hex(value_hex).unwrap(); + bonsai_storage + .insert(identifier, keyer(key).as_bitslice(), &value) + .expect("Failed to insert storage update into trie"); + } + + let id = id_builder.new_id(); + bonsai_storage + .commit(id) + .expect("Failed to commit to bonsai storage"); + let root_hash = bonsai_storage + .root_hash(identifier) + .expect("Failed to get root hash"); + + println!("Expected: 0x00072F7E2EC1A2F05342503B49AECD83E14884AE374A8570F2F6F7B868CF94AE\nFound: {root_hash:#x}"); + assert_eq!( + root_hash, + Felt::from_hex("0x00072F7E2EC1A2F05342503B49AECD83E14884AE374A8570F2F6F7B868CF94AE") + .unwrap() + ); +} diff --git a/src/tests/transactional_state.rs b/src/tests/transactional_state.rs index 84c1fd7..7271e8e 100644 --- a/src/tests/transactional_state.rs +++ b/src/tests/transactional_state.rs @@ -9,6 +9,7 @@ use starknet_types_core::{felt::Felt, hash::Pedersen}; #[test] fn basics() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -22,7 +23,9 @@ fn basics() { ); let id1 = id_builder.new_id(); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); bonsai_storage.commit(id1).unwrap(); let pair2 = ( @@ -31,7 +34,9 @@ fn basics() { ); let id2 = id_builder.new_id(); let bitvec = BitVec::from_vec(pair2.0.clone()); - bonsai_storage.insert(&bitvec, &pair2.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair2.1) + .unwrap(); bonsai_storage.commit(id2).unwrap(); let bonsai_at_txn: BonsaiStorage<_, _, Pedersen> = bonsai_storage @@ -39,13 +44,17 @@ fn basics() { .unwrap() .unwrap(); let bitvec = BitVec::from_vec(pair1.0.clone()); - assert_eq!(bonsai_at_txn.get(&bitvec).unwrap().unwrap(), pair1.1); + assert_eq!( + bonsai_at_txn.get(&identifier, &bitvec).unwrap().unwrap(), + pair1.1 + ); let bitvec = BitVec::from_vec(pair2.0.clone()); - assert!(bonsai_at_txn.get(&bitvec).unwrap().is_none()); + assert!(bonsai_at_txn.get(&identifier, &bitvec).unwrap().is_none()); } #[test] fn test_thread() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -59,7 +68,9 @@ fn test_thread() { ); let id1 = id_builder.new_id(); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); bonsai_storage.commit(id1).unwrap(); std::thread::scope(|s| { @@ -69,7 +80,10 @@ fn test_thread() { .unwrap() .unwrap(); let bitvec = BitVec::from_vec(pair1.0.clone()); - assert_eq!(bonsai_at_txn.get(&bitvec).unwrap().unwrap(), pair1.1); + assert_eq!( + bonsai_at_txn.get(&identifier, &bitvec).unwrap().unwrap(), + pair1.1 + ); }); s.spawn(|| { @@ -78,25 +92,29 @@ fn test_thread() { .unwrap() .unwrap(); let bitvec = BitVec::from_vec(pair1.0.clone()); - assert_eq!(bonsai_at_txn.get(&bitvec).unwrap().unwrap(), pair1.1); + assert_eq!( + bonsai_at_txn.get(&identifier, &bitvec).unwrap().unwrap(), + pair1.1 + ); }); }); bonsai_storage - .get(&BitVec::from_vec(vec![1, 2, 2])) + .get(&identifier, &BitVec::from_vec(vec![1, 2, 2])) .unwrap(); let pair2 = ( vec![1, 2, 3], Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(), ); bonsai_storage - .insert(&BitVec::from_vec(pair2.0.clone()), &pair2.1) + .insert(&identifier, &BitVec::from_vec(pair2.0.clone()), &pair2.1) .unwrap(); bonsai_storage.commit(id_builder.new_id()).unwrap(); } #[test] fn remove() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -110,12 +128,14 @@ fn remove() { ); let id1 = id_builder.new_id(); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); bonsai_storage.commit(id1).unwrap(); let id2 = id_builder.new_id(); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.remove(&bitvec).unwrap(); + bonsai_storage.remove(&identifier, &bitvec).unwrap(); bonsai_storage.commit(id2).unwrap(); bonsai_storage.dump_database(); @@ -124,18 +144,22 @@ fn remove() { .unwrap() .unwrap(); let bitvec = BitVec::from_vec(pair1.0.clone()); - assert_eq!(bonsai_at_txn.get(&bitvec).unwrap().unwrap(), pair1.1); + assert_eq!( + bonsai_at_txn.get(&identifier, &bitvec).unwrap().unwrap(), + pair1.1 + ); let bonsai_at_txn: BonsaiStorage<_, _, Pedersen> = bonsai_storage .get_transactional_state(id2, BonsaiStorageConfig::default()) .unwrap() .unwrap(); let bitvec = BitVec::from_vec(pair1.0.clone()); - assert!(bonsai_at_txn.get(&bitvec).unwrap().is_none()); + assert!(bonsai_at_txn.get(&identifier, &bitvec).unwrap().is_none()); } #[test] fn merge() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -149,7 +173,9 @@ fn merge() { ); let id1 = id_builder.new_id(); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); bonsai_storage.commit(id1).unwrap(); let mut bonsai_at_txn: BonsaiStorage<_, _, Pedersen> = bonsai_storage .get_transactional_state(id1, BonsaiStorageConfig::default()) @@ -160,7 +186,7 @@ fn merge() { Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(), ); bonsai_at_txn - .insert(&BitVec::from_vec(pair2.0.clone()), &pair2.1) + .insert(&identifier, &BitVec::from_vec(pair2.0.clone()), &pair2.1) .unwrap(); bonsai_at_txn .transactional_commit(id_builder.new_id()) @@ -168,7 +194,7 @@ fn merge() { bonsai_storage.merge(bonsai_at_txn).unwrap(); assert_eq!( bonsai_storage - .get(&BitVec::from_vec(vec![1, 2, 3])) + .get(&identifier, &BitVec::from_vec(vec![1, 2, 3])) .unwrap(), Some(pair2.1) ); @@ -176,6 +202,7 @@ fn merge() { #[test] fn merge_override() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -189,7 +216,9 @@ fn merge_override() { ); let id1 = id_builder.new_id(); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); bonsai_storage.commit(id1).unwrap(); let mut bonsai_at_txn: BonsaiStorage<_, _, Pedersen> = bonsai_storage .get_transactional_state(id1, BonsaiStorageConfig::default()) @@ -200,7 +229,7 @@ fn merge_override() { Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(), ); bonsai_at_txn - .insert(&BitVec::from_vec(pair2.0.clone()), &pair2.1) + .insert(&identifier, &BitVec::from_vec(pair2.0.clone()), &pair2.1) .unwrap(); bonsai_at_txn .transactional_commit(id_builder.new_id()) @@ -208,7 +237,7 @@ fn merge_override() { bonsai_storage.merge(bonsai_at_txn).unwrap(); assert_eq!( bonsai_storage - .get(&BitVec::from_vec(vec![1, 2, 2])) + .get(&identifier, &BitVec::from_vec(vec![1, 2, 2])) .unwrap(), Some(pair2.1) ); @@ -216,6 +245,7 @@ fn merge_override() { #[test] fn merge_remove() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -229,27 +259,32 @@ fn merge_remove() { ); let id1 = id_builder.new_id(); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); bonsai_storage.commit(id1).unwrap(); let mut bonsai_at_txn: BonsaiStorage<_, _, Pedersen> = bonsai_storage .get_transactional_state(id1, BonsaiStorageConfig::default()) .unwrap() .unwrap(); bonsai_at_txn - .remove(&BitVec::from_vec(pair1.0.clone())) + .remove(&identifier, &BitVec::from_vec(pair1.0.clone())) .unwrap(); bonsai_at_txn .transactional_commit(id_builder.new_id()) .unwrap(); bonsai_storage.merge(bonsai_at_txn).unwrap(); assert_eq!( - bonsai_storage.get(&BitVec::from_vec(pair1.0)).unwrap(), + bonsai_storage + .get(&identifier, &BitVec::from_vec(pair1.0)) + .unwrap(), None ); } #[test] fn merge_txn_revert() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -263,48 +298,53 @@ fn merge_txn_revert() { ); let id1 = id_builder.new_id(); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); bonsai_storage.commit(id1).unwrap(); - let root_hash1 = bonsai_storage.root_hash().unwrap(); + let root_hash1 = bonsai_storage.root_hash(&identifier).unwrap(); let mut bonsai_at_txn: BonsaiStorage<_, _, Pedersen> = bonsai_storage .get_transactional_state(id1, BonsaiStorageConfig::default()) .unwrap() .unwrap(); bonsai_at_txn - .remove(&BitVec::from_vec(pair1.0.clone())) + .remove(&identifier, &BitVec::from_vec(pair1.0.clone())) .unwrap(); let id2 = id_builder.new_id(); bonsai_at_txn.transactional_commit(id2).unwrap(); - let root_hash2 = bonsai_at_txn.root_hash().unwrap(); + let root_hash2 = bonsai_at_txn.root_hash(&identifier).unwrap(); let pair2 = ( vec![1, 2, 3], Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(), ); bonsai_at_txn - .insert(&BitVec::from_vec(pair2.0.clone()), &pair2.1) + .insert(&identifier, &BitVec::from_vec(pair2.0.clone()), &pair2.1) .unwrap(); let id3 = id_builder.new_id(); bonsai_at_txn.transactional_commit(id3).unwrap(); bonsai_at_txn.revert_to(id2).unwrap(); - let revert_hash2 = bonsai_at_txn.root_hash().unwrap(); + let revert_hash2 = bonsai_at_txn.root_hash(&identifier).unwrap(); bonsai_at_txn.revert_to(id1).unwrap(); - let revert_hash1 = bonsai_at_txn.root_hash().unwrap(); + let revert_hash1 = bonsai_at_txn.root_hash(&identifier).unwrap(); assert_eq!(root_hash2, revert_hash2); assert_eq!(root_hash1, revert_hash1); bonsai_storage.merge(bonsai_at_txn).unwrap(); assert_eq!( - bonsai_storage.get(&BitVec::from_vec(pair1.0)).unwrap(), + bonsai_storage + .get(&identifier, &BitVec::from_vec(pair1.0)) + .unwrap(), Some(pair1.1) ); } #[test] fn merge_invalid() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -318,7 +358,9 @@ fn merge_invalid() { ); let id1 = id_builder.new_id(); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); bonsai_storage.commit(id1).unwrap(); let mut bonsai_at_txn: BonsaiStorage<_, _, Pedersen> = bonsai_storage @@ -326,7 +368,7 @@ fn merge_invalid() { .unwrap() .unwrap(); bonsai_at_txn - .remove(&BitVec::from_vec(pair1.0.clone())) + .remove(&identifier, &BitVec::from_vec(pair1.0.clone())) .unwrap(); let id2 = id_builder.new_id(); bonsai_at_txn.transactional_commit(id2).unwrap(); @@ -336,7 +378,7 @@ fn merge_invalid() { Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(), ); bonsai_storage - .insert(&BitVec::from_vec(pair2.0.clone()), &pair2.1) + .insert(&identifier, &BitVec::from_vec(pair2.0.clone()), &pair2.1) .unwrap(); let id3 = id_builder.new_id(); bonsai_storage.commit(id3).unwrap(); @@ -346,6 +388,7 @@ fn merge_invalid() { #[test] fn many_snapshots() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig { @@ -362,7 +405,9 @@ fn many_snapshots() { ); let id1 = id_builder.new_id(); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); bonsai_storage.commit(id1).unwrap(); let pair2 = ( @@ -371,7 +416,9 @@ fn many_snapshots() { ); let id2 = id_builder.new_id(); let bitvec = BitVec::from_vec(pair2.0.clone()); - bonsai_storage.insert(&bitvec, &pair2.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair2.1) + .unwrap(); bonsai_storage.commit(id2).unwrap(); bonsai_storage diff --git a/src/tests/trie_log.rs b/src/tests/trie_log.rs index 56c6ecc..59be7b4 100644 --- a/src/tests/trie_log.rs +++ b/src/tests/trie_log.rs @@ -9,6 +9,7 @@ use starknet_types_core::{felt::Felt, hash::Pedersen}; #[test] fn basics() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -22,9 +23,11 @@ fn basics() { ); let id1 = id_builder.new_id(); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, pair1.1) + .unwrap(); bonsai_storage.commit(id1).unwrap(); - let root_hash1 = bonsai_storage.root_hash().unwrap(); + let root_hash1 = bonsai_storage.root_hash(&identifier).unwrap(); let id2 = id_builder.new_id(); let pair2 = ( @@ -32,20 +35,22 @@ fn basics() { &Felt::from_hex("0x66342762FDD54D3c195fec3ce2568b62052e").unwrap(), ); let bitvec = BitVec::from_vec(pair2.0.clone()); - bonsai_storage.insert(&bitvec, pair2.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, pair2.1) + .unwrap(); bonsai_storage.commit(id2).unwrap(); - let root_hash2 = bonsai_storage.root_hash().unwrap(); + let root_hash2 = bonsai_storage.root_hash(&identifier).unwrap(); let id3 = id_builder.new_id(); let bitvec = BitVec::from_vec(pair1.0); - bonsai_storage.remove(&bitvec).unwrap(); + bonsai_storage.remove(&identifier, &bitvec).unwrap(); bonsai_storage.commit(id3).unwrap(); bonsai_storage.revert_to(id2).unwrap(); - let revert_root_hash2 = bonsai_storage.root_hash().unwrap(); + let revert_root_hash2 = bonsai_storage.root_hash(&identifier).unwrap(); bonsai_storage.revert_to(id1).unwrap(); - let revert_root_hash1 = bonsai_storage.root_hash().unwrap(); + let revert_root_hash1 = bonsai_storage.root_hash(&identifier).unwrap(); assert_eq!(root_hash2, revert_root_hash2); assert_eq!(root_hash1, revert_root_hash1); @@ -53,6 +58,7 @@ fn basics() { #[test] fn unrecorded_revert() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -66,7 +72,9 @@ fn unrecorded_revert() { ); let id1 = id_builder.new_id(); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); bonsai_storage.commit(id1).unwrap(); let uncommited_id = id_builder.new_id(); @@ -75,6 +83,7 @@ fn unrecorded_revert() { #[test] fn in_place_revert() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -85,16 +94,19 @@ fn in_place_revert() { let pair1 = (vec![1, 2, 3], &BonsaiTrieHash::default()); let id1 = id_builder.new_id(); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, pair1.1) + .unwrap(); bonsai_storage.commit(id1).unwrap(); - let root_hash1 = bonsai_storage.root_hash().unwrap(); + let root_hash1 = bonsai_storage.root_hash(&identifier).unwrap(); bonsai_storage.revert_to(id1).unwrap(); - assert_eq!(root_hash1, bonsai_storage.root_hash().unwrap()); + assert_eq!(root_hash1, bonsai_storage.root_hash(&identifier).unwrap()); } #[test] fn truncated_revert() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -108,9 +120,11 @@ fn truncated_revert() { ); let id1 = id_builder.new_id(); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, pair1.1) + .unwrap(); bonsai_storage.commit(id1).unwrap(); - let root_hash1 = bonsai_storage.root_hash().unwrap(); + let root_hash1 = bonsai_storage.root_hash(&identifier).unwrap(); let id2 = id_builder.new_id(); let pair2 = ( @@ -118,11 +132,13 @@ fn truncated_revert() { &Felt::from_hex("0x66342762FDD54D3c195fec3ce2568b62052e").unwrap(), ); let bitvec = BitVec::from_vec(pair2.0.clone()); - bonsai_storage.insert(&bitvec, pair2.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, pair2.1) + .unwrap(); bonsai_storage.commit(id2).unwrap(); bonsai_storage.revert_to(id1).unwrap(); - let revert_root_hash1 = bonsai_storage.root_hash().unwrap(); + let revert_root_hash1 = bonsai_storage.root_hash(&identifier).unwrap(); bonsai_storage.revert_to(id2).unwrap_err(); assert_eq!(root_hash1, revert_root_hash1); @@ -130,6 +146,7 @@ fn truncated_revert() { #[test] fn double_revert() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -143,9 +160,11 @@ fn double_revert() { ); let id1 = id_builder.new_id(); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, pair1.1) + .unwrap(); bonsai_storage.commit(id1).unwrap(); - let root_hash1 = bonsai_storage.root_hash().unwrap(); + let root_hash1 = bonsai_storage.root_hash(&identifier).unwrap(); let id2 = id_builder.new_id(); let pair2 = ( @@ -153,13 +172,15 @@ fn double_revert() { &Felt::from_hex("0x66342762FDD54D3c195fec3ce2568b62052e").unwrap(), ); let bitvec = BitVec::from_vec(pair2.0.clone()); - bonsai_storage.insert(&bitvec, pair2.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, pair2.1) + .unwrap(); bonsai_storage.commit(id2).unwrap(); bonsai_storage.revert_to(id1).unwrap(); - let revert1 = bonsai_storage.root_hash().unwrap(); + let revert1 = bonsai_storage.root_hash(&identifier).unwrap(); bonsai_storage.revert_to(id1).unwrap(); - let revert2 = bonsai_storage.root_hash().unwrap(); + let revert2 = bonsai_storage.root_hash(&identifier).unwrap(); assert_eq!(root_hash1, revert1); assert_eq!(revert1, revert2); @@ -167,6 +188,7 @@ fn double_revert() { #[test] fn remove_and_reinsert() { + let identifier = vec![]; let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); @@ -180,14 +202,18 @@ fn remove_and_reinsert() { ); let id1 = id_builder.new_id(); let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); - bonsai_storage.remove(&bitvec).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); + bonsai_storage.remove(&identifier, &bitvec).unwrap(); bonsai_storage.commit(id1).unwrap(); - let root_hash1 = bonsai_storage.root_hash().unwrap(); + let root_hash1 = bonsai_storage.root_hash(&identifier).unwrap(); let id2 = id_builder.new_id(); - bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage + .insert(&identifier, &bitvec, &pair1.1) + .unwrap(); bonsai_storage.commit(id2).unwrap(); bonsai_storage.revert_to(id1).unwrap(); - assert_eq!(root_hash1, bonsai_storage.root_hash().unwrap()); + assert_eq!(root_hash1, bonsai_storage.root_hash(&identifier).unwrap()); } diff --git a/src/trie/merkle_tree.rs b/src/trie/merkle_tree.rs index 2d67881..2abae46 100644 --- a/src/trie/merkle_tree.rs +++ b/src/trie/merkle_tree.rs @@ -1,5 +1,5 @@ #[cfg(not(feature = "std"))] -use alloc::{format, string::ToString, vec, vec::Vec}; +use alloc::{format, string::ToString, vec::Vec}; use bitvec::{ prelude::{BitSlice, BitVec, Msb0}, view::BitView, @@ -11,6 +11,8 @@ 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; @@ -20,6 +22,7 @@ use crate::{error::BonsaiStorageError, id::Id, BonsaiDatabase, KeyValueDB}; use super::{ merkle_node::{BinaryNode, Direction, EdgeNode, Node, NodeHandle, NodeId}, path::Path, + trie_db::TrieKeyType, TrieKey, }; @@ -63,23 +66,200 @@ impl ProofNode { } } +pub(crate) struct MerkleTrees { + pub db: KeyValueDB, + _hasher: PhantomData, + pub trees: HashMap, MerkleTree>, +} + +impl MerkleTrees { + pub(crate) fn new(db: KeyValueDB) -> Self { + Self { + db, + _hasher: PhantomData, + trees: HashMap::new(), + } + } + + pub(crate) fn init_tree( + &mut self, + identifier: &[u8], + ) -> Result<(), BonsaiStorageError> { + let tree = MerkleTree::new(&mut self.db, identifier.to_vec())?; + self.trees.insert(identifier.to_vec(), tree); + Ok(()) + } + + pub(crate) fn set( + &mut self, + identifier: &[u8], + key: &BitSlice, + value: Felt, + ) -> Result<(), BonsaiStorageError> { + let tree = self.trees.get_mut(identifier); + if let Some(tree) = tree { + tree.set(&mut self.db, key, value) + } else { + let mut tree = MerkleTree::new(&mut self.db, identifier.to_vec())?; + tree.set(&mut self.db, key, value)?; + self.trees.insert(identifier.to_vec(), tree); + Ok(()) + } + } + + pub(crate) fn get( + &self, + identifier: &[u8], + key: &BitSlice, + ) -> Result, BonsaiStorageError> { + let tree = self.trees.get(identifier); + if let Some(tree) = tree { + tree.get(&self.db, key) + } else { + Ok(None) + } + } + + pub(crate) fn contains( + &self, + identifier: &[u8], + key: &BitSlice, + ) -> Result> { + let tree = self.trees.get(identifier); + if let Some(tree) = tree { + tree.contains(&self.db, key) + } else { + Ok(false) + } + } + + pub(crate) fn db_mut(&mut self) -> &mut KeyValueDB { + &mut self.db + } + + pub(crate) fn reset_to_last_commit( + &mut self, + ) -> Result<(), BonsaiStorageError> { + for tree in self.trees.values_mut() { + tree.reset_to_last_commit(&mut self.db)?; + } + Ok(()) + } + + pub(crate) fn db_ref(&self) -> &KeyValueDB { + &self.db + } + + pub(crate) fn db(self) -> KeyValueDB { + self.db + } + + pub(crate) fn root_hash( + &self, + identifier: &[u8], + ) -> Result> { + let tree = self.trees.get(identifier); + if let Some(tree) = tree { + Ok(tree.root_hash()) + } else { + Err(BonsaiStorageError::Trie("Tree not found".to_string())) + } + } + + pub(crate) fn get_keys( + &self, + identifier: &[u8], + ) -> Result>, BonsaiStorageError> { + self.db + .db + .get_by_prefix(&crate::DatabaseKey::Flat(identifier)) + .map(|key_value_pairs| { + // Remove the identifier from the key + key_value_pairs + .into_iter() + .map(|(key, _value)| key[identifier.len()..].to_vec()) + .collect() + }) + .map_err(|e| e.into()) + } + + pub(crate) fn commit(&mut self) -> Result<(), BonsaiStorageError> { + #[allow(clippy::type_complexity)] + #[cfg(not(feature = "std"))] + let db_changes: Vec< + Result< + HashMap>>, + BonsaiStorageError, + >, + > = self + .trees + .iter_mut() + .map(|(_, tree)| tree.get_updates::()) + .collect(); + #[allow(clippy::type_complexity)] + #[cfg(feature = "std")] + let db_changes: Vec< + Result< + HashMap>>, + BonsaiStorageError, + >, + > = self + .trees + .par_iter_mut() + .map(|(_, tree)| tree.get_updates::()) + .collect(); + let mut batch = self.db.create_batch(); + for changes in db_changes { + let changes = changes?; + for (key, value) in changes { + match value { + InsertOrRemove::Insert(value) => { + self.db.insert(&key, &value, Some(&mut batch))?; + } + InsertOrRemove::Remove => { + self.db.remove(&key, Some(&mut batch))?; + } + } + } + } + self.db.write_batch(batch)?; + Ok(()) + } + + pub(crate) fn get_proof( + &self, + identifier: &[u8], + key: &BitSlice, + ) -> Result, BonsaiStorageError> { + let tree = self.trees.get(identifier); + if let Some(tree) = tree { + tree.get_proof(&self.db, key) + } else { + Err(BonsaiStorageError::Trie("Tree not found".to_string())) + } + } + + pub(crate) fn get_identifiers(&self) -> Vec> { + self.trees.keys().cloned().collect() + } +} + /// A Starknet binary Merkle-Patricia tree with a specific root entry-point and storage. /// /// This is used to update, mutate and access global Starknet state as well as individual contract /// states. /// /// For more information on how this functions internally, see [here](super::merkle_node). -pub struct MerkleTree { +pub struct MerkleTree { /// The handle to the current root node could be hash if no modifications has been done /// since the last commit or in memory if there are some modifications. root_handle: NodeHandle, /// The last known root hash. Updated only each commit. (possibly outdated between two commits) root_hash: Felt, - /// Temporary storage used to store the nodes that are modified during a commit. + /// Identifier of the tree in the database. + identifier: Vec, /// This storage is used to avoid modifying the underlying database each time during a commit. storage_nodes: NodesMapping, - /// The underlying database used to store the nodes. - db: KeyValueDB, /// The id of the last node that has been added to the temporary storage. latest_node_id: NodeId, /// The list of nodes that should be removed from the underlying database during the next commit. @@ -91,23 +271,27 @@ pub struct MerkleTree { } #[derive(Debug, PartialEq, Eq)] -enum InsertOrRemove { +pub(crate) enum InsertOrRemove { Insert(T), Remove, } -impl MerkleTree { +impl MerkleTree { /// Less visible initialization for `MerkleTree` as the main entry points should be /// [`MerkleTree::::load`] for persistent trees and [`MerkleTree::empty`] for /// transient ones. - pub fn new(mut db: KeyValueDB) -> Result> { + + pub fn new( + db: &mut KeyValueDB, + identifier: Vec, + ) -> Result> { let nodes_mapping: HashMap = HashMap::new(); - let root_node = db.get(&TrieKey::Trie(vec![]))?; + let root_node = db.get(&TrieKey::new(&identifier, TrieKeyType::Trie, &[]))?; let node = if let Some(root_node) = root_node { Node::decode(&mut root_node.as_slice())? } else { db.insert( - &TrieKey::Trie(vec![]), + &TrieKey::new(&identifier, TrieKeyType::Trie, &[]), &Node::Unresolved(Felt::ZERO).encode(), None, )?; @@ -118,8 +302,8 @@ impl MerkleTree { Ok(Self { root_handle: NodeHandle::Hash(root), root_hash: root, + identifier, storage_nodes: NodesMapping(nodes_mapping), - db, latest_node_id: NodeId(0), death_row: Vec::new(), cache_leaf_modified: HashMap::new(), @@ -132,9 +316,12 @@ impl MerkleTree { } /// Remove all the modifications that have been done since the last commit. - pub fn reset_to_last_commit(&mut self) -> Result<(), BonsaiStorageError> { + pub fn reset_to_last_commit( + &mut self, + db: &mut KeyValueDB, + ) -> Result<(), BonsaiStorageError> { let node = self - .get_trie_branch_in_db_from_path(&Path(BitVec::::new()))? + .get_trie_branch_in_db_from_path(db, &Path(BitVec::::new()))? .ok_or(BonsaiStorageError::Trie( "root node doesn't exist in the storage".to_string(), ))?; @@ -149,29 +336,31 @@ impl MerkleTree { Ok(()) } - /// Persists all changes to storage and returns the new root hash. - pub fn commit(&mut self) -> Result> { - let mut batch = self.db.create_batch(); + /// Calculate all the new hashes and the root hash. + #[allow(clippy::type_complexity)] + pub(crate) fn get_updates( + &mut self, + ) -> Result>>, BonsaiStorageError> + { + let mut updates = HashMap::new(); for node_key in mem::take(&mut self.death_row) { - self.db.remove(&node_key, Some(&mut batch))?; + updates.insert(node_key, InsertOrRemove::Remove); } - let root_hash = self.commit_subtree(self.root_handle, Path(BitVec::new()), &mut batch)?; + let root_hash = + self.commit_subtree::(&mut updates, self.root_handle, Path(BitVec::new()))?; for (key, value) in mem::take(&mut self.cache_leaf_modified) { - match value { - InsertOrRemove::Insert(value) => { - self.db - .insert(&TrieKey::Flat(key), &value.encode(), Some(&mut batch))?; - } - InsertOrRemove::Remove => { - self.db.remove(&TrieKey::Flat(key), Some(&mut batch))?; - } - } + updates.insert( + TrieKey::new(&self.identifier, TrieKeyType::Flat, &key), + match value { + InsertOrRemove::Insert(value) => InsertOrRemove::Insert(value.encode()), + InsertOrRemove::Remove => InsertOrRemove::Remove, + }, + ); } - self.db.write_batch(batch)?; self.latest_node_id.reset(); self.root_hash = root_hash; self.root_handle = NodeHandle::Hash(root_hash); - Ok(root_hash) + Ok(updates) } /// Persists any changes in this subtree to storage. @@ -185,11 +374,11 @@ impl MerkleTree { /// # Arguments /// /// * `node` - The top node from the subtree to commit. - fn commit_subtree( + fn commit_subtree( &mut self, + updates: &mut HashMap>>, node_handle: NodeHandle, path: Path, - batch: &mut DB::Batch, ) -> Result> { use Node::*; let node_id = match node_handle { @@ -206,11 +395,10 @@ impl MerkleTree { ))? { Unresolved(hash) => { if path.0.is_empty() { - self.db.insert( - &TrieKey::Trie(vec![]), - &Node::Unresolved(hash).encode(), - Some(batch), - )?; + updates.insert( + TrieKey::new(&self.identifier, TrieKeyType::Trie, &[]), + InsertOrRemove::Insert(Node::Unresolved(hash).encode()), + ); Ok(hash) } else { Ok(hash) @@ -218,30 +406,25 @@ impl MerkleTree { } Binary(mut binary) => { let left_path = path.new_with_direction(Direction::Left); - let left_hash = self.commit_subtree(binary.left, left_path, batch)?; + let left_hash = self.commit_subtree::(updates, binary.left, left_path)?; let right_path = path.new_with_direction(Direction::Right); - let right_hash = self.commit_subtree(binary.right, right_path, batch)?; + let right_hash = self.commit_subtree::(updates, binary.right, right_path)?; let hash = H::hash(&left_hash, &right_hash); binary.hash = Some(hash); binary.left = NodeHandle::Hash(left_hash); binary.right = NodeHandle::Hash(right_hash); - let key_bytes = if path.0.is_empty() { - vec![] - } else { - [&[path.0.len() as u8], path.0.as_raw_slice()].concat() - }; - self.db.insert( - &TrieKey::Trie(key_bytes), - &Node::Binary(binary).encode(), - Some(batch), - )?; + let key_bytes: Vec = path.into(); + updates.insert( + TrieKey::new(&self.identifier, TrieKeyType::Trie, &key_bytes), + InsertOrRemove::Insert(Node::Binary(binary).encode()), + ); Ok(hash) } Edge(mut edge) => { let mut child_path = path.clone(); child_path.0.extend(&edge.path.0); - let child_hash = self.commit_subtree(edge.child, child_path, batch)?; + let child_hash = self.commit_subtree::(updates, edge.child, child_path)?; let mut bytes = [0u8; 32]; bytes.view_bits_mut::()[256 - edge.path.0.len()..] .copy_from_bitslice(&edge.path.0); @@ -255,16 +438,11 @@ impl MerkleTree { let hash = H::hash(&child_hash, &felt_path) + length; edge.hash = Some(hash); edge.child = NodeHandle::Hash(child_hash); - let key_bytes = if path.0.is_empty() { - vec![] - } else { - [&[path.0.len() as u8], path.0.as_raw_slice()].concat() - }; - self.db.insert( - &TrieKey::Trie(key_bytes), - &Node::Edge(edge).encode(), - Some(batch), - )?; + let key_bytes: Vec = path.into(); + updates.insert( + TrieKey::new(&self.identifier, TrieKeyType::Trie, &key_bytes), + InsertOrRemove::Insert(Node::Edge(edge).encode()), + ); Ok(hash) } } @@ -276,15 +454,31 @@ impl MerkleTree { /// /// * `key` - The key to set. /// * `value` - The value to set. - pub fn set( + pub fn set( &mut self, + db: &mut KeyValueDB, key: &BitSlice, value: Felt, ) -> Result<(), BonsaiStorageError> { if value == Felt::ZERO { - return self.delete_leaf(key); + return self.delete_leaf(db, key); + } + let key_bytes = bitslice_to_bytes(key); + if let Some(InsertOrRemove::Insert(value_db)) = self.cache_leaf_modified.get(&key_bytes) { + if &value == value_db { + return Ok(()); + } } - let path = self.preload_nodes(key)?; + if let Some(value_db) = db.get(&TrieKey::new( + &self.identifier, + TrieKeyType::Flat, + &key_bytes, + ))? { + if value == Felt::decode(&mut value_db.as_slice()).unwrap() { + return Ok(()); + } + } + let path = self.preload_nodes(db, key)?; // There are three possibilities. // // 1. The leaf exists, in which case we simply change its value. @@ -306,93 +500,104 @@ impl MerkleTree { Some(node_id) => { let mut nodes_to_add = Vec::new(); self.storage_nodes.0.entry(*node_id).and_modify(|node| { - if let Edge(edge) = node { - let common = edge.common_path(key); - // Height of the binary node - let branch_height = edge.height as usize + common.len(); - if branch_height == key.len() { - edge.child = NodeHandle::Hash(value); - // The leaf already exists, we simply change its value. - let key_bytes = bitslice_to_bytes(key); + match node { + Edge(edge) => { + let common = edge.common_path(key); + // Height of the binary node + let branch_height = edge.height as usize + common.len(); + if branch_height == key.len() { + edge.child = NodeHandle::Hash(value); + // The leaf already exists, we simply change its value. + self.cache_leaf_modified + .insert(key_bytes, InsertOrRemove::Insert(value)); + return; + } + // Height of the binary node's children + let child_height = branch_height + 1; + + // Path from binary node to new leaf + let new_path = key[child_height..].to_bitvec(); + // Path from binary node to existing child + let old_path = edge.path.0[common.len() + 1..].to_bitvec(); + + // The new leaf branch of the binary node. + // (this may be edge -> leaf, or just leaf depending). self.cache_leaf_modified .insert(key_bytes, InsertOrRemove::Insert(value)); - return; - } - // Height of the binary node's children - let child_height = branch_height + 1; - - // Path from binary node to new leaf - let new_path = key[child_height..].to_bitvec(); - // Path from binary node to existing child - let old_path = edge.path.0[common.len() + 1..].to_bitvec(); - - // The new leaf branch of the binary node. - // (this may be edge -> leaf, or just leaf depending). - let key_bytes = bitslice_to_bytes(key); - self.cache_leaf_modified - .insert(key_bytes, InsertOrRemove::Insert(value)); - - let new = if new_path.is_empty() { - NodeHandle::Hash(value) - } else { - let new_edge = Node::Edge(EdgeNode { - hash: None, - height: child_height as u64, - path: Path(new_path), - child: NodeHandle::Hash(value), - }); - let edge_id = self.latest_node_id.next_id(); - nodes_to_add.push((edge_id, new_edge)); - NodeHandle::InMemory(edge_id) - }; - // The existing child branch of the binary node. - let old = if old_path.is_empty() { - edge.child - } else { - let old_edge = Node::Edge(EdgeNode { - hash: None, - height: child_height as u64, - path: Path(old_path), - child: edge.child, - }); - let edge_id = self.latest_node_id.next_id(); - nodes_to_add.push((edge_id, old_edge)); - NodeHandle::InMemory(edge_id) - }; + let new = if new_path.is_empty() { + NodeHandle::Hash(value) + } else { + let new_edge = Node::Edge(EdgeNode { + hash: None, + height: child_height as u64, + path: Path(new_path), + child: NodeHandle::Hash(value), + }); + let edge_id = self.latest_node_id.next_id(); + nodes_to_add.push((edge_id, new_edge)); + NodeHandle::InMemory(edge_id) + }; - let new_direction = Direction::from(key[branch_height]); - let (left, right) = match new_direction { - Direction::Left => (new, old), - Direction::Right => (old, new), - }; + // The existing child branch of the binary node. + let old = if old_path.is_empty() { + edge.child + } else { + let old_edge = Node::Edge(EdgeNode { + hash: None, + height: child_height as u64, + path: Path(old_path), + child: edge.child, + }); + let edge_id = self.latest_node_id.next_id(); + nodes_to_add.push((edge_id, old_edge)); + NodeHandle::InMemory(edge_id) + }; - let branch = Node::Binary(BinaryNode { - hash: None, - height: branch_height as u64, - left, - right, - }); - - // We may require an edge leading to the binary node. - let new_node = if common.is_empty() { - branch - } else { - let branch_id = self.latest_node_id.next_id(); - nodes_to_add.push((branch_id, branch)); + let new_direction = Direction::from(key[branch_height]); + let (left, right) = match new_direction { + Direction::Left => (new, old), + Direction::Right => (old, new), + }; - Node::Edge(EdgeNode { + let branch = Node::Binary(BinaryNode { hash: None, - height: edge.height, - path: Path(common.to_bitvec()), - child: NodeHandle::InMemory(branch_id), - }) - }; - let path = key[..edge.height as usize].to_bitvec(); - let key_bytes = [&[path.len() as u8], path.into_vec().as_slice()].concat(); - self.death_row.push(TrieKey::Trie(key_bytes)); - *node = new_node; - }; + height: branch_height as u64, + left, + right, + }); + + // We may require an edge leading to the binary node. + let new_node = if common.is_empty() { + branch + } else { + let branch_id = self.latest_node_id.next_id(); + nodes_to_add.push((branch_id, branch)); + + Node::Edge(EdgeNode { + hash: None, + height: edge.height, + path: Path(common.to_bitvec()), + child: NodeHandle::InMemory(branch_id), + }) + }; + let path = key[..edge.height as usize].to_bitvec(); + let key_bytes = + [&[path.len() as u8], path.into_vec().as_slice()].concat(); + self.death_row.push(TrieKey::Trie(key_bytes)); + *node = new_node; + } + Binary(binary) => { + if (binary.height + 1) as usize == key.len() { + let direction = Direction::from(key[binary.height as usize]); + match direction { + Direction::Left => binary.left = NodeHandle::Hash(value), + Direction::Right => binary.right = NodeHandle::Hash(value), + }; + } + } + _ => {} + } }); for (id, node) in nodes_to_add { self.storage_nodes.0.insert(id, node); @@ -424,18 +629,6 @@ impl MerkleTree { } } - pub fn db_ref(&self) -> &KeyValueDB { - &self.db - } - - pub fn db(self) -> KeyValueDB { - self.db - } - - pub fn db_mut(&mut self) -> &mut KeyValueDB { - &mut self.db - } - /// Deletes a leaf node from the tree. /// /// This is not an external facing API; the functionality is instead accessed by calling @@ -444,8 +637,9 @@ impl MerkleTree { /// # Arguments /// /// * `key` - The key to delete. - fn delete_leaf( + fn delete_leaf( &mut self, + db: &mut KeyValueDB, key: &BitSlice, ) -> Result<(), BonsaiStorageError> { // Algorithm explanation: @@ -461,12 +655,22 @@ impl MerkleTree { // and other remaining child node -- if they're also edges. // // Then we are done. - let key_bytes = bitslice_to_bytes(key); + if db + .get(&TrieKey::new( + &self.identifier, + TrieKeyType::Flat, + &key_bytes, + ))? + .is_none() + && !self.cache_leaf_modified.contains_key(&key_bytes) + { + return Ok(()); + } self.cache_leaf_modified .insert(key_bytes.clone(), InsertOrRemove::Remove); - let path = self.preload_nodes(key)?; + let path = self.preload_nodes(db, key)?; let mut last_binary_path = Path(key.to_bitvec()); @@ -481,7 +685,14 @@ impl MerkleTree { for _ in 0..edge.path.0.len() { last_binary_path.0.pop(); } - self.death_row.push((&last_binary_path).into()); + let mut new_path = Path(BitVec::new()); + for i in last_binary_path.0.iter() { + new_path.0.push(*i); + } + last_binary_path = new_path; + let path: Vec = (&last_binary_path).into(); + self.death_row + .push(TrieKey::new(&self.identifier, TrieKeyType::Trie, &path)); } } !node.is_binary() @@ -495,30 +706,45 @@ impl MerkleTree { let node = self.storage_nodes.0.get_mut(&node_id).ok_or( BonsaiStorageError::Trie("Node not found in memory".to_string()), )?; - let (direction, height) = { - // SAFETY: This node must be a binary node due to the iteration condition. - let binary = node.as_binary().unwrap(); - (binary.direction(key).invert(), binary.height) - }; + // SAFETY: This node must be a binary node due to the iteration condition. + let binary = node.as_binary().unwrap(); + let (direction, height) = + { (binary.direction(key).invert(), binary.height) }; + last_binary_path.0.pop(); + last_binary_path.0.push(bool::from(direction)); // Create an edge node to replace the old binary node // i.e. with the remaining child (note the direction invert), // and a path of just a single bit. - last_binary_path.0.push(direction.into()); let path = Path(once(bool::from(direction)).collect::>()); let mut edge = EdgeNode { hash: None, height, path, - child: NodeHandle::InMemory(self.latest_node_id), + child: match direction { + Direction::Left => binary.left, + Direction::Right => binary.right, + }, }; // Merge the remaining child if it's an edge. - self.merge_edges(&mut edge)?; - + self.merge_edges::(&mut edge, db, &last_binary_path)?; edge }; - // Replace the old binary node with the new edge node. - self.storage_nodes.0.insert(node_id, Node::Edge(new_edge)); + // Check the parent of the new edge. If it is also an edge, then they must merge. + if let Some(parent_node_id) = parent_branch_node { + // Get a mutable reference to the parent node to merge them + let parent_node = self.storage_nodes.0.get_mut(&parent_node_id).ok_or( + BonsaiStorageError::Trie("Node not found in memory".to_string()), + )?; + if let Node::Edge(parent_edge) = parent_node { + parent_edge.path.0.extend_from_bitslice(&new_edge.path.0); + parent_edge.child = new_edge.child; + } else { + self.storage_nodes.0.insert(node_id, Node::Edge(new_edge)); + } + } else { + self.storage_nodes.0.insert(node_id, Node::Edge(new_edge)); + } } None => { // We reached the root without a hitting binary node. The new tree @@ -532,49 +758,6 @@ impl MerkleTree { return Ok(()); } }; - - // Check the parent of the new edge. If it is also an edge, then they must merge. - if let Some(node_id) = parent_branch_node { - let node = self - .storage_nodes - .0 - .get(&node_id) - .ok_or(BonsaiStorageError::Trie( - "Node not found in memory".to_string(), - ))?; - // If it's an edge node and the child is in memory and it's an edge too we - // return the child otherwise we leave - let child = - if let Node::Edge(edge) = node { - match edge.child { - NodeHandle::Hash(_) => return Ok(()), - NodeHandle::InMemory(child_id) => { - let child_node = self.storage_nodes.0.get(&child_id).ok_or( - BonsaiStorageError::Trie("Node not found in memory".to_string()), - )?; - if let Node::Edge(child_edge) = child_node { - child_edge.clone() - } else { - return Ok(()); - } - } - } - } else { - return Ok(()); - }; - // Get a mutable reference to the parent node to merge them - let edge = self - .storage_nodes - .0 - .get_mut(&node_id) - .ok_or(BonsaiStorageError::Trie( - "Node not found in memory".to_string(), - ))?; - if let Node::Edge(edge) = edge { - edge.path.0.extend_from_bitslice(&child.path.0); - edge.child = child.child; - } - } Ok(()) } @@ -587,8 +770,9 @@ impl MerkleTree { /// # Returns /// /// The value of the key. - pub fn get( + pub fn get( &self, + db: &KeyValueDB, key: &BitSlice, ) -> Result, BonsaiStorageError> { let key = bitslice_to_bytes(key); @@ -598,13 +782,13 @@ impl MerkleTree { InsertOrRemove::Insert(value) => return Ok(Some(*value)), } } - self.db - .get(&TrieKey::Flat(key.to_vec())) + db.get(&TrieKey::new(&self.identifier, TrieKeyType::Flat, &key)) .map(|r| r.map(|opt| Felt::decode(&mut opt.as_slice()).unwrap())) } - pub fn contains( + pub fn contains( &self, + db: &KeyValueDB, key: &BitSlice, ) -> Result> { let key = bitslice_to_bytes(key); @@ -614,7 +798,7 @@ impl MerkleTree { InsertOrRemove::Insert(_) => return Ok(true), } } - self.db.contains(&TrieKey::Flat(key.to_vec())) + db.contains(&TrieKey::new(&self.identifier, TrieKeyType::Flat, &key)) } /// Returns the list of nodes along the path. @@ -635,15 +819,16 @@ impl MerkleTree { /// # Returns /// /// The merkle proof and all the child nodes hashes. - pub fn get_proof( + pub fn get_proof( &self, + db: &KeyValueDB, key: &BitSlice, ) -> Result, BonsaiStorageError> { 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(&Path(BitVec::::new()))? + .get_trie_branch_in_db_from_path(db, &Path(BitVec::::new()))? .ok_or(BonsaiStorageError::Trie( "Couldn't fetch root node in db".to_string(), ))?; @@ -667,7 +852,8 @@ impl MerkleTree { 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(&Path(child_path))?; + let node = + self.get_trie_branch_in_db_from_path(db, &Path(child_path))?; if let Some(node) = node { node } else { @@ -711,7 +897,7 @@ impl MerkleTree { 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(&Path(next_path))? + .get_trie_branch_in_db_from_path(db, &Path(next_path))? .ok_or(BonsaiStorageError::Trie( "Couldn't fetch next node in db".to_string(), ))?, @@ -791,15 +977,16 @@ impl MerkleTree { /// # Returns /// /// The list of nodes along the path. - fn preload_nodes( + fn preload_nodes( &mut self, + db: &mut KeyValueDB, dst: &BitSlice, ) -> Result, BonsaiStorageError> { 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(&Path(BitVec::::new()))? + .get_trie_branch_in_db_from_path(db, &Path(BitVec::::new()))? .ok_or(BonsaiStorageError::Trie( "Couldn't fetch root node in db".to_string(), ))?; @@ -817,12 +1004,19 @@ impl MerkleTree { root_id } }; - self.preload_nodes_subtree(dst, node_id, Path(BitVec::::new()), &mut nodes)?; + self.preload_nodes_subtree( + db, + dst, + node_id, + Path(BitVec::::new()), + &mut nodes, + )?; Ok(nodes) } - fn preload_nodes_subtree( + fn preload_nodes_subtree( &mut self, + db: &mut KeyValueDB, dst: &BitSlice, root_id: NodeId, mut path: Path, @@ -851,7 +1045,7 @@ impl MerkleTree { let next = binary_node.get_child(next_direction); match next { NodeHandle::Hash(_) => { - let node = self.get_trie_branch_in_db_from_path(&path)?; + 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); @@ -867,14 +1061,14 @@ impl MerkleTree { self.storage_nodes .0 .insert(root_id, Node::Binary(binary_node)); - self.preload_nodes_subtree(dst, self.latest_node_id, path, nodes) + 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(dst, next_id, path, nodes) + self.preload_nodes_subtree(db, dst, next_id, path, nodes) } } } @@ -890,21 +1084,21 @@ impl MerkleTree { let next = edge_node.child; match next { NodeHandle::Hash(_) => { - let node = self.get_trie_branch_in_db_from_path(&path)?; + 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(dst, self.latest_node_id, path, nodes) + 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(dst, next_id, path, nodes) + self.preload_nodes_subtree(db, dst, next_id, path, nodes) } } } @@ -914,12 +1108,13 @@ impl MerkleTree { } /// Get the node of the trie that corresponds to the path. - fn get_trie_branch_in_db_from_path( + fn get_trie_branch_in_db_from_path( &self, + db: &KeyValueDB, path: &Path, ) -> Result, BonsaiStorageError> { - self.db - .get(&path.into())? + let path: Vec = path.into(); + db.get(&TrieKey::new(&self.identifier, TrieKeyType::Trie, &path))? .map(|node| { Node::decode(&mut node.as_slice()).map_err(|err| { BonsaiStorageError::Trie(format!("Couldn't decode node: {}", err)) @@ -939,20 +1134,29 @@ impl MerkleTree { /// # Arguments /// /// * `parent` - The parent node to merge the child with. - fn merge_edges( + fn merge_edges( &self, parent: &mut EdgeNode, + db: &KeyValueDB, + path: &Path, ) -> Result<(), BonsaiStorageError> { let child_node = match parent.child { - NodeHandle::Hash(_) => return Ok(()), - NodeHandle::InMemory(child_id) => { - self.storage_nodes - .0 - .get(&child_id) - .ok_or(BonsaiStorageError::Trie( - "Couldn't fetch node in memory".to_string(), - ))? + NodeHandle::Hash(_) => { + let node = self.get_trie_branch_in_db_from_path(db, path)?; + if let Some(node) = node { + node + } else { + return Ok(()); + } } + NodeHandle::InMemory(child_id) => self + .storage_nodes + .0 + .get(&child_id) + .ok_or(BonsaiStorageError::Trie( + "Couldn't fetch node in memory".to_string(), + ))? + .clone(), }; if let Node::Edge(child_edge) = child_node { parent.path.0.extend_from_bitslice(&child_edge.path.0); diff --git a/src/trie/mod.rs b/src/trie/mod.rs index 0d06bb8..dc6e304 100644 --- a/src/trie/mod.rs +++ b/src/trie/mod.rs @@ -3,4 +3,4 @@ pub mod merkle_tree; mod path; mod trie_db; -pub use trie_db::TrieKey; +pub(crate) use trie_db::TrieKey; diff --git a/src/trie/path.rs b/src/trie/path.rs index 78906e4..b5f8cc4 100644 --- a/src/trie/path.rs +++ b/src/trie/path.rs @@ -1,7 +1,7 @@ use bitvec::{order::Msb0, vec::BitVec}; use parity_scale_codec::{Decode, Encode, Error, Input, Output}; -use super::{merkle_node::Direction, TrieKey}; +use super::merkle_node::Direction; #[cfg(not(feature = "std"))] use alloc::vec::Vec; @@ -104,25 +104,26 @@ impl Path { } } -impl From for TrieKey { +/// Convert Path to Vec can be used, for example, to create keys for the database +impl From for Vec { fn from(path: Path) -> Self { let key = if path.0.is_empty() { Vec::new() } else { [&[path.0.len() as u8], path.0.as_raw_slice()].concat() }; - TrieKey::Trie(key) + key } } -impl From<&Path> for TrieKey { +impl From<&Path> for Vec { fn from(path: &Path) -> Self { let key = if path.0.is_empty() { Vec::new() } else { [&[path.0.len() as u8], path.0.as_raw_slice()].concat() }; - TrieKey::Trie(key) + key } } diff --git a/src/trie/trie_db.rs b/src/trie/trie_db.rs index b7debf2..97f3a84 100644 --- a/src/trie/trie_db.rs +++ b/src/trie/trie_db.rs @@ -4,13 +4,14 @@ use crate::bonsai_database::DatabaseKey; use alloc::vec::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. #[derive(Debug, Clone, Hash, PartialEq, Eq)] -pub enum TrieKey { +pub(crate) enum TrieKey { Trie(Vec), Flat(Vec), } -enum TrieKeyType { +pub(crate) enum TrieKeyType { Trie = 0, Flat = 1, } @@ -34,6 +35,15 @@ impl From<&TrieKey> for u8 { } impl TrieKey { + pub fn new(identifier: &[u8], key_type: TrieKeyType, key: &[u8]) -> Self { + let mut final_key = identifier.to_vec(); + final_key.extend_from_slice(key); + match key_type { + TrieKeyType::Trie => TrieKey::Trie(final_key), + TrieKeyType::Flat => TrieKey::Flat(final_key), + } + } + pub fn from_variant_and_bytes(variant: u8, bytes: Vec) -> Self { match variant { x if x == TrieKeyType::Trie as u8 => TrieKey::Trie(bytes),