diff --git a/Cargo.toml b/Cargo.toml index 47350aa..d9384cd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,46 +6,49 @@ version = "0.1.0" [features] default = ["std", "rocksdb"] rocksdb = ["dep:rocksdb"] -std = ["parity-scale-codec/std", "bitvec/std", "starknet-types-core/std", "rayon", "hashbrown/rayon"] +std = [ + "parity-scale-codec/std", + "bitvec/std", + "starknet-types-core/std", + "rayon", + "hashbrown/rayon", +] # internal bench = [] [dependencies] bitvec = { version = "1", default-features = false, features = ["alloc"] } derive_more = { version = "0.99.17", default-features = false, features = [ - "constructor", + "constructor", ] } hashbrown = "0.14.3" log = "0.4.20" rayon = { version = "1.9.0", optional = true } smallvec = { version = "1.11.2", features = ["serde"] } +slotmap = "1.0.7" parity-scale-codec = { version = "3.0.0", default-features = false, features = [ - "derive", + "derive", ] } serde = { version = "1.0.195", default-features = false, features = [ - "derive", - "alloc", + "derive", + "alloc", ] } starknet-types-core = { version = "0.1.5", default-features = false, features = [ - "hash", - "parity-scale-codec", - "alloc", + "hash", + "parity-scale-codec", + "alloc", ] } # Optionals rocksdb = { optional = true, version = "0.21.0", features = [ - "multi-threaded-cf", + "multi-threaded-cf", ] } [dev-dependencies] env_logger = "0.11.3" once_cell = "1.19.0" pprof = { version = "0.3", features = ["flamegraph"] } -pathfinder-common = { git = "https://github.com/massalabs/pathfinder.git", package = "pathfinder-common", rev = "b7b6d76a76ab0e10f92e5f84ce099b5f727cb4db" } -pathfinder-crypto = { git = "https://github.com/massalabs/pathfinder.git", package = "pathfinder-crypto", rev = "b7b6d76a76ab0e10f92e5f84ce099b5f727cb4db" } -pathfinder-merkle-tree = { git = "https://github.com/massalabs/pathfinder.git", package = "pathfinder-merkle-tree", rev = "b7b6d76a76ab0e10f92e5f84ce099b5f727cb4db" } -pathfinder-storage = { git = "https://github.com/massalabs/pathfinder.git", package = "pathfinder-storage", rev = "b7b6d76a76ab0e10f92e5f84ce099b5f727cb4db" } rand = { version = "0.8.5", features = ["small_rng"] } tempfile = "3.8.0" rstest = "0.18.2" diff --git a/benches/storage.rs b/benches/storage.rs index 687f135..d03f5a9 100644 --- a/benches/storage.rs +++ b/benches/storage.rs @@ -1,6 +1,5 @@ use std::hint::black_box; -use bitvec::vec::BitVec; use bonsai_trie::{ databases::HashMapDb, id::{BasicId, BasicIdBuilder}, @@ -10,7 +9,7 @@ use criterion::{criterion_group, criterion_main, BatchSize, Criterion}; use rand::{prelude::*, thread_rng}; use starknet_types_core::{ felt::Felt, - hash::{Pedersen, StarkHash}, + hash::{Pedersen, Poseidon, StarkHash}, }; mod flamegraph; @@ -244,7 +243,7 @@ fn multiple_contracts(c: &mut Criterion) { }); } -fn hash(c: &mut Criterion) { +fn pedersen_hash(c: &mut Criterion) { c.bench_function("pedersen hash", move |b| { let felt0 = Felt::from_hex("0x100bd6fbfced88ded1b34bd1a55b747ce3a9fde9a914bca75571e4496b56443") @@ -258,9 +257,23 @@ fn hash(c: &mut Criterion) { }); } +fn poseidon_hash(c: &mut Criterion) { + c.bench_function("poseidon hash", move |b| { + let felt0 = + Felt::from_hex("0x100bd6fbfced88ded1b34bd1a55b747ce3a9fde9a914bca75571e4496b56443") + .unwrap(); + let felt1 = + Felt::from_hex("0x00a038cda302fedbc4f6117648c6d3faca3cda924cb9c517b46232c6316b152f") + .unwrap(); + b.iter(|| { + black_box(Poseidon::hash(&felt0, &felt1)); + }) + }); +} + criterion_group! { name = benches; config = Criterion::default(); // .with_profiler(flamegraph::FlamegraphProfiler::new(100)); - targets = storage, one_update, five_updates, hash, drop_storage, storage_with_insert, multiple_contracts + targets = storage, one_update, five_updates, pedersen_hash, poseidon_hash, drop_storage, storage_with_insert, multiple_contracts } criterion_main!(benches); diff --git a/proptest-regressions/tests/proptest.txt b/proptest-regressions/tests/proptest.txt new file mode 100644 index 0000000..957ce40 --- /dev/null +++ b/proptest-regressions/tests/proptest.txt @@ -0,0 +1,7 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc 4e01e217c0cfdcb88c0fda9a6f6c5d2bac70c39f3cb9aed3afc4cbb2c0000ac9 # shrinks to pb = MerkleTreeInsertProblem([Insert([00110], 0x52889e3d0cd0a2a1c49d2f232fb0938760bb332ba356dbee7a698a5b05cebff), Commit, Commit, Commit, Commit, Insert([00100], 0x1f0a9ec83897033916a4417aa017cb8591c7374e9c53ea1b3895b8da022658b), Commit, Commit, Remove([00110]), Commit, Insert([00010], 0x31ba571cd01834c526a0a394ece58dcea478e34ee3af3f2fe9874beb8ac3a8c), Commit, Commit, Remove([00100]), Insert([00111], 0x3667307f1e8e8150930e6ef12390b595ac3e54aa68df53222bf25a9006b1298), Remove([11001]), Remove([10001]), Insert([00011], 0x32de2acb4350b53e3d87f4307dc81c180f1e3aa8c1177643e6f779c441abb92), Remove([11000]), Insert([10010], 0x3437d396d3992c99cbbf1d9aa492a6ea90d40d41897091a4fcb6f9bfa6e5bbc), Commit, Remove([11010]), Commit, Insert([01011], 0x4159a60c3014785b08904cfa2336b0c6168efe8b460350573b2d5bc9548cffa), Remove([01101]), Commit, Remove([00010]), Commit, Remove([10110]), Commit, Commit, Commit, Commit, Insert([01100], 0x3b88e95019dd205f2bfc89f2eb9594901c4690ba7de27124e43adbdbddafcfc), Commit, Remove([11110]), Remove([10010]), Commit, Insert([00011], 0x33778bf3f7cdeca9f1bfb1ccd857d5e57bae7ac0d93538155784214ef081b89), Insert([00111], 0x17d918366c70c47a80af7ffa5b27363affc93788f0ac06276394a46997c87c1), Commit, Commit, Remove([11110]), Insert([00001], 0x3c2616fce851e54a0b7f65a5816a77086cffdff98ef59075d9b72c641abaf1a), Remove([00011]), Remove([11011]), Remove([00001]), Commit, Commit, Remove([00001]), Commit, Insert([01001], 0x13d745e43bb4d094211e14c57b28505afa2519e8c30feee5fa4c6a3c96f728b), Insert([00010], 0xb980fa05a0424e41452c9c197235a415f726a202541d2ced481e0e04b8c6db), Remove([00010]), Commit, Remove([10100]), Remove([10100]), Insert([00000], 0x73fa95cefa03f0c723b6a595672e0cc830e49a364e69b6200106c7ff7a69610), Commit, Remove([11111]), Insert([10111], 0x342719e264ac5638ff3b14214a6c33851a5d100831f4588a482facba2040386), Commit, Commit, Insert([01011], 0x28b63e6950adc4dd378685c4f00e0c9df567db08c49924389b1af309adf9b8), Insert([10001], 0x17452a2bb9b207fdda9d116a52a72640738baefe326e0b6cb74ffe2a38147c8), Remove([01110]), Remove([10011])]) diff --git a/proptest-regressions/trie/iterator.txt b/proptest-regressions/trie/iterator.txt new file mode 100644 index 0000000..f893eb7 --- /dev/null +++ b/proptest-regressions/trie/iterator.txt @@ -0,0 +1,8 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc cdc544314a0c860fc660985c3ce379c8680d5d52b0bca8837b69b7ecfa301afc # shrinks to cases = [1, 4] +cc c41a719525322a5cc0c9799d0cfb090cc3d1a17a1acc4fd66122065c02cd48ba # shrinks to cases = [5, 5] diff --git a/src/bonsai_database.rs b/src/bonsai_database.rs index 7118774..592cb8d 100644 --- a/src/bonsai_database.rs +++ b/src/bonsai_database.rs @@ -27,7 +27,7 @@ pub trait DBError: Error + Send + Sync {} pub trait DBError: Send + Sync {} /// Trait to be implemented on any type that can be used as a database. -pub trait BonsaiDatabase { +pub trait BonsaiDatabase: core::fmt::Debug { type Batch: Default; #[cfg(feature = "std")] type DatabaseError: Error + DBError; diff --git a/src/databases/rocks_db.rs b/src/databases/rocks_db.rs index 1e2f176..92350c7 100644 --- a/src/databases/rocks_db.rs +++ b/src/databases/rocks_db.rs @@ -173,6 +173,12 @@ pub struct RocksDBTransaction<'a> { column_families: HashMap>, } +impl<'a> fmt::Debug for RocksDBTransaction<'a> { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("RocksDBTransaction").finish() + } +} + impl<'db, ID> BonsaiDatabase for RocksDB<'db, ID> where ID: Id, diff --git a/src/error.rs b/src/error.rs index 7cfea0c..d4245da 100644 --- a/src/error.rs +++ b/src/error.rs @@ -1,7 +1,7 @@ #[cfg(feature = "std")] use std::{error::Error, fmt::Display}; -use crate::{bonsai_database::DBError, String}; +use crate::{bonsai_database::DBError, BitVec, String}; /// All errors that can be returned by BonsaiStorage. #[derive(Debug)] @@ -21,6 +21,10 @@ where Database(DatabaseError), /// Error when decoding a node NodeDecodeError(parity_scale_codec::Error), + /// Error when creating a storage proof. + CreateProofKeyNotInTree { key: BitVec }, + /// Malformated trie key. + KeyLength { expected: usize, got: usize }, } impl core::convert::From @@ -52,6 +56,12 @@ where BonsaiStorageError::Merge(e) => write!(f, "Merge error: {}", e), BonsaiStorageError::Database(e) => write!(f, "Database error: {}", e), BonsaiStorageError::NodeDecodeError(e) => write!(f, "Node decode error: {}", e), + BonsaiStorageError::CreateProofKeyNotInTree { key } => { + write!(f, "Key not in tree: {key:b}") + } + BonsaiStorageError::KeyLength { expected, got } => { + write!(f, "Malformated key length: expected {expected}, got {got}") + } } } } diff --git a/src/key_value_db.rs b/src/key_value_db.rs index bde5402..8343608 100644 --- a/src/key_value_db.rs +++ b/src/key_value_db.rs @@ -1,8 +1,7 @@ use crate::{ - changes::key_new_value, format, trie::merkle_tree::bytes_to_bitvec, BTreeSet, ByteVec, + changes::key_new_value, format, trie::tree::bytes_to_bitvec, BTreeSet, BitVec, ByteVec, Change as ExternChange, ToString, }; -use bitvec::{order::Msb0, vec::BitVec}; use hashbrown::HashMap; use log::trace; use parity_scale_codec::Decode; @@ -92,8 +91,7 @@ where pub(crate) fn get_changes( &self, id: ID, - ) -> Result, ExternChange>, BonsaiStorageError> - { + ) -> Result, BonsaiStorageError> { if self.changes_store.id_queue.contains(&id) { let mut leaf_changes = HashMap::new(); let changes = ChangeBatch::deserialize( diff --git a/src/lib.rs b/src/lib.rs index 2d7b4c5..6662ae9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -111,12 +111,30 @@ pub(crate) use std::{ vec::Vec, }; -use crate::trie::merkle_tree::MerkleTree; - pub type ByteVec = smallvec::SmallVec<[u8; 32]>; +pub type BitVec = bitvec::vec::BitVec; +pub type BitSlice = bitvec::slice::BitSlice; + +mod changes; +mod key_value_db; +mod trie; + +mod bonsai_database; +/// All databases already implemented in this crate. +pub mod databases; +mod error; +/// Definition and basic implementation of an CommitID +pub mod id; + +pub use bonsai_database::{BonsaiDatabase, BonsaiPersistentDatabase, DBError, DatabaseKey}; +pub use error::BonsaiStorageError; +pub use trie::proof::{Membership, MultiProof, ProofNode}; + +#[cfg(test)] +mod tests; pub(crate) trait EncodeExt: parity_scale_codec::Encode { - fn encode_sbytevec(&self) -> ByteVec { + fn encode_bytevec(&self) -> ByteVec { struct Out(ByteVec); impl parity_scale_codec::Output for Out { #[inline] @@ -132,32 +150,10 @@ pub(crate) trait EncodeExt: parity_scale_codec::Encode { } impl EncodeExt for T {} -use bitvec::{order::Msb0, slice::BitSlice, vec::BitVec}; use changes::ChangeBatch; use key_value_db::KeyValueDB; -use starknet_types_core::{ - felt::Felt, - hash::{Pedersen, StarkHash}, -}; - -mod changes; -mod key_value_db; -mod trie; - -mod bonsai_database; -/// All databases already implemented in this crate. -pub mod databases; -mod error; -/// Definition and basic implementation of an CommitID -pub mod id; - -pub use bonsai_database::{BonsaiDatabase, BonsaiPersistentDatabase, DBError, DatabaseKey}; -pub use error::BonsaiStorageError; -use trie::merkle_tree::{bytes_to_bitvec, MerkleTrees}; -pub use trie::merkle_tree::{Membership, ProofNode}; - -#[cfg(test)] -mod tests; +use starknet_types_core::{felt::Felt, hash::StarkHash}; +use trie::{tree::bytes_to_bitvec, trees::MerkleTrees}; /// Structure that contains the configuration for the BonsaiStorage. /// A default implementation is provided with coherent values. @@ -241,21 +237,23 @@ where pub fn new( db: DB, config: BonsaiStorageConfig, + max_height: u8, ) -> Result> { let key_value_db = KeyValueDB::new(db, config.into(), None); Ok(Self { - tries: MerkleTrees::new(key_value_db), + tries: MerkleTrees::new(key_value_db, max_height), }) } pub fn new_from_transactional_state( db: DB, config: BonsaiStorageConfig, + max_height: u8, created_at: ChangeID, _identifiers: impl IntoIterator>, ) -> Result> { let key_value_db = KeyValueDB::new(db, config.into(), Some(created_at)); - let tries = MerkleTrees::::new(key_value_db); + let tries = MerkleTrees::::new(key_value_db, max_height); Ok(Self { tries }) } @@ -264,7 +262,7 @@ where pub fn insert( &mut self, identifier: &[u8], - key: &BitSlice, + key: &BitSlice, value: &Felt, ) -> Result<(), BonsaiStorageError> { self.tries.set(identifier, key, *value)?; @@ -276,7 +274,7 @@ where pub fn remove( &mut self, identifier: &[u8], - key: &BitSlice, + key: &BitSlice, ) -> Result<(), BonsaiStorageError> { self.tries.set(identifier, key, Felt::ZERO)?; Ok(()) @@ -286,7 +284,7 @@ where pub fn get( &self, identifier: &[u8], - key: &BitSlice, + key: &BitSlice, ) -> Result, BonsaiStorageError> { self.tries.get(identifier, key) } @@ -298,7 +296,7 @@ where pub fn get_at( &self, identifier: &[u8], - key: &BitSlice, + key: &BitSlice, id: ChangeID, ) -> Result, BonsaiStorageError> { self.tries.get_at(identifier, key, id) @@ -308,7 +306,7 @@ where pub fn contains( &self, identifier: &[u8], - key: &BitSlice, + key: &BitSlice, ) -> Result> { self.tries.contains(identifier, key) } @@ -402,7 +400,7 @@ where pub fn get_changes( &self, id: ChangeID, - ) -> Result, Change>, BonsaiStorageError> { + ) -> Result, BonsaiStorageError> { self.tries.db_ref().get_changes(id) } @@ -411,6 +409,11 @@ where self.tries.db_ref().db.dump_database(); } + #[cfg(test)] + pub fn dump(&self) { + self.tries.dump(); + } + /// Get trie root hash at the latest commit pub fn root_hash( &self, @@ -430,24 +433,25 @@ where Ok(()) } - /// Generates a merkle-proof for a given `key`. - /// - /// Returns vector of [`TrieNode`] which form a chain from the root to the key, - /// if it exists, or down to the node which proves that the key does not exist. - /// - /// The nodes are returned in order, root first. - /// - /// Verification is performed by confirming that: - /// 1. the chain follows the path of `key`, and - /// 2. the hashes are correct, and - /// 3. the root hash matches the known root - pub fn get_proof( - &self, - identifier: &[u8], - key: &BitSlice, - ) -> Result, BonsaiStorageError> { - self.tries.get_proof(identifier, key) - } + // /// Generates a merkle-proof for a given `key`. + // /// + // /// Returns vector of [`TrieNode`] which form a chain from the root to the key, + // /// if it exists, or down to the node which proves that the key does not exist. + // /// + // /// The nodes are returned in order, root first. + // /// + // /// Verification is performed by confirming that: + // /// 1. the chain follows the path of `key`, and + // /// 2. the hashes are correct, and + // /// 3. the root hash matches the known root + // pub fn get_proof( + // &self, + // identifier: &[u8], + // key: &BitSlice, + // ) -> Result, BonsaiStorageError> { + // todo!() + // // self.tries.get_proof(identifier, key) + // } /// Get all the keys in a specific trie. pub fn get_keys( @@ -471,14 +475,12 @@ where self.tries.db_ref().get_latest_id() } - /// Verifies a merkle-proof for a given `key` and `value`. - pub fn verify_proof( - root: Felt, - key: &BitSlice, - value: Felt, - proofs: &[ProofNode], - ) -> Option { - MerkleTree::::verify_proof(root, key, value, proofs) + pub fn get_multi_proof( + &mut self, + identifier: &[u8], + keys: impl IntoIterator>, + ) -> Result> { + self.tries.get_multi_proof(identifier, keys) } } @@ -516,6 +518,7 @@ where Ok(Some(BonsaiStorage::new_from_transactional_state( transaction, config, + self.tries.max_height, change_id, self.tries.get_identifiers(), )?)) @@ -546,7 +549,7 @@ where for (identifier, tree) in trees { for (k, op) in tree.cache_leaf_modified() { match op { - crate::trie::merkle_tree::InsertOrRemove::Insert(v) => { + crate::trie::tree::InsertOrRemove::Insert(v) => { self.insert(&identifier, &bytes_to_bitvec(k), v) .map_err(|e| { BonsaiStorageError::Merge(format!( @@ -555,7 +558,7 @@ where )) })?; } - crate::trie::merkle_tree::InsertOrRemove::Remove => { + crate::trie::tree::InsertOrRemove::Remove => { self.remove(&identifier, &bytes_to_bitvec(k)).map_err(|e| { BonsaiStorageError::Merge(format!( "While merging remove({:?}) faced error: {:?}", diff --git a/src/tests/madara_comparison.rs b/src/tests/madara_comparison.rs index eef4cc6..793faaf 100644 --- a/src/tests/madara_comparison.rs +++ b/src/tests/madara_comparison.rs @@ -1,11 +1,11 @@ #![cfg(all(feature = "std", feature = "rocksdb"))] -use bitvec::{bits, order::Msb0, vec::BitVec}; +use bitvec::{bits, order::Msb0}; use starknet_types_core::{felt::Felt, hash::Pedersen}; use crate::{ databases::{create_rocks_db, RocksDB, RocksDBConfig}, id::BasicIdBuilder, - BonsaiStorage, BonsaiStorageConfig, + BitVec, BonsaiStorage, BonsaiStorageConfig, }; #[test] @@ -15,9 +15,9 @@ fn trie_height_251() { 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(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); for i in 0..251 { - let mut key: BitVec = bits![u8, Msb0; 0; 251].to_bitvec(); + let mut key: BitVec = bits![u8, Msb0; 0; 251].to_bitvec(); key.set(i, true); let value = Felt::from_hex("0x01").unwrap(); bonsai_storage @@ -33,7 +33,7 @@ fn trie_height_251() { // #[test]// fn test_height_251() { // let mut tree = super::merkle_patricia_tree::merkle_tree::MerkleTree::::empty(); // for i in 0..251 { -// let mut key: BitVec = bits![u8, Msb0; 0; 251].to_bitvec(); +// let mut key: BitVec = bits![u8, Msb0; 0; 251].to_bitvec(); // key.set(i, true); // let value = Felt::from_hex_be("0x01").unwrap(); // tree.set(key.as_bitslice(), value); diff --git a/src/tests/merge.rs b/src/tests/merge.rs index 976d212..141671e 100644 --- a/src/tests/merge.rs +++ b/src/tests/merge.rs @@ -4,28 +4,27 @@ use crate::{ databases::{create_rocks_db, RocksDB, RocksDBConfig, RocksDBTransaction}, id::{BasicId, BasicIdBuilder}, - BonsaiStorage, BonsaiStorageConfig, + BitVec, BonsaiStorage, BonsaiStorageConfig, }; -use bitvec::vec::BitVec; use once_cell::sync::Lazy; use rocksdb::OptimisticTransactionDB; use starknet_types_core::{felt::Felt, hash::Pedersen}; -static PAIR1: Lazy<(BitVec, Felt)> = Lazy::new(|| { +static PAIR1: Lazy<(BitVec, Felt)> = Lazy::new(|| { ( BitVec::from_vec(vec![1, 2, 2]), Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(), ) }); -static PAIR2: Lazy<(BitVec, Felt)> = Lazy::new(|| { +static PAIR2: Lazy<(BitVec, Felt)> = Lazy::new(|| { ( BitVec::from_vec(vec![1, 2, 3]), Felt::from_hex("0x66342762FD54D033c195fec3ce2568b62052e").unwrap(), ) }); -static PAIR3: Lazy<(BitVec, Felt)> = Lazy::new(|| { +static PAIR3: Lazy<(BitVec, Felt)> = Lazy::new(|| { ( BitVec::from_vec(vec![1, 2, 4]), Felt::from_hex("0x66342762FD54D033c195fec3ce2568b62052e").unwrap(), @@ -61,8 +60,9 @@ fn init_test( let identifier = vec![]; let config = BonsaiStorageConfig::default(); - let mut bonsai_storage = BonsaiStorage::new(RocksDB::new(db, RocksDBConfig::default()), config) - .expect("Failed to create BonsaiStorage"); + let mut bonsai_storage = + BonsaiStorage::new(RocksDB::new(db, RocksDBConfig::default()), config, 24) + .expect("Failed to create BonsaiStorage"); let mut id_builder = BasicIdBuilder::new(); diff --git a/src/tests/merkle_tree.rs b/src/tests/merkle_tree.rs new file mode 100644 index 0000000..9fcd768 --- /dev/null +++ b/src/tests/merkle_tree.rs @@ -0,0 +1,732 @@ +#![cfg(all(test, feature = "std", feature = "rocksdb"))] +use bitvec::view::BitView; +use indexmap::IndexMap; +use starknet_types_core::{felt::Felt, hash::Pedersen}; + +use crate::{ + databases::{create_rocks_db, RocksDB, RocksDBConfig}, + id::BasicId, + BitVec, BonsaiStorage, BonsaiStorageConfig, ByteVec, +}; + +#[test_log::test] +// The whole point of this test is to make sure it is possible to reconstruct the original +// keys from the data present in the db. +fn test_key_retrieval() { + let tempdir = tempfile::tempdir().unwrap(); + let rocksdb = create_rocks_db(tempdir.path()).unwrap(); + let db = RocksDB::new(&rocksdb, RocksDBConfig::default()); + let mut bonsai = + BonsaiStorage::::new(db, BonsaiStorageConfig::default(), 251) + .unwrap(); + + let block_0 = vec![ + ( + str_to_felt_bytes("0x031c887d82502ceb218c06ebb46198da3f7b92864a8223746bc836dda3e34b52"), + vec![ + ( + str_to_felt_bytes( + "0x0000000000000000000000000000000000000000000000000000000000000005", + ), + str_to_felt_bytes( + "0x0000000000000000000000000000000000000000000000000000000000000065", + ), + ), + ( + str_to_felt_bytes( + "0x00cfc2e2866fd08bfb4ac73b70e0c136e326ae18fc797a2c090c8811c695577e", + ), + str_to_felt_bytes( + "0x05f1dd5a5aef88e0498eeca4e7b2ea0fa7110608c11531278742f0b5499af4b3", + ), + ), + ( + str_to_felt_bytes( + "0x05aee31408163292105d875070f98cb48275b8c87e80380b78d30647e05854d5", + ), + str_to_felt_bytes( + "0x00000000000000000000000000000000000000000000000000000000000007c7", + ), + ), + ( + str_to_felt_bytes( + "0x05fac6815fddf6af1ca5e592359862ede14f171e1544fd9e792288164097c35d", + ), + str_to_felt_bytes( + "0x00299e2f4b5a873e95e65eb03d31e532ea2cde43b498b50cd3161145db5542a5", + ), + ), + ( + str_to_felt_bytes( + "0x05fac6815fddf6af1ca5e592359862ede14f171e1544fd9e792288164097c35e", + ), + str_to_felt_bytes( + "0x03d6897cf23da3bf4fd35cc7a43ccaf7c5eaf8f7c5b9031ac9b09a929204175f", + ), + ), + ], + ), + ( + str_to_felt_bytes("0x06ee3440b08a9c805305449ec7f7003f27e9f7e287b83610952ec36bdc5a6bae"), + vec![ + ( + str_to_felt_bytes( + "0x01e2cd4b3588e8f6f9c4e89fb0e293bf92018c96d7a93ee367d29a284223b6ff", + ), + str_to_felt_bytes( + "0x071d1e9d188c784a0bde95c1d508877a0d93e9102b37213d1e13f3ebc54a7751", + ), + ), + ( + str_to_felt_bytes( + "0x0449908c349e90f81ab13042b1e49dc251eb6e3e51092d9a40f86859f7f415b0", + ), + str_to_felt_bytes( + "0x06cb6104279e754967a721b52bcf5be525fdc11fa6db6ef5c3a4db832acf7804", + ), + ), + ( + str_to_felt_bytes( + "0x048cba68d4e86764105adcdcf641ab67b581a55a4f367203647549c8bf1feea2", + ), + str_to_felt_bytes( + "0x0362d24a3b030998ac75e838955dfee19ec5b6eceb235b9bfbeccf51b6304d0b", + ), + ), + ( + str_to_felt_bytes( + "0x05bdaf1d47b176bfcd1114809af85a46b9c4376e87e361d86536f0288a284b65", + ), + str_to_felt_bytes( + "0x028dff6722aa73281b2cf84cac09950b71fa90512db294d2042119abdd9f4b87", + ), + ), + ( + str_to_felt_bytes( + "0x05bdaf1d47b176bfcd1114809af85a46b9c4376e87e361d86536f0288a284b66", + ), + str_to_felt_bytes( + "0x057a8f8a019ccab5bfc6ff86c96b1392257abb8d5d110c01d326b94247af161c", + ), + ), + ( + str_to_felt_bytes( + "0x05f750dc13ed239fa6fc43ff6e10ae9125a33bd05ec034fc3bb4dd168df3505f", + ), + str_to_felt_bytes( + "0x00000000000000000000000000000000000000000000000000000000000007e5", + ), + ), + ], + ), + ( + str_to_felt_bytes("0x0735596016a37ee972c42adef6a3cf628c19bb3794369c65d2c82ba034aecf2c"), + vec![ + ( + str_to_felt_bytes( + "0x0000000000000000000000000000000000000000000000000000000000000005", + ), + str_to_felt_bytes( + "0x0000000000000000000000000000000000000000000000000000000000000064", + ), + ), + ( + str_to_felt_bytes( + "0x002f50710449a06a9fa789b3c029a63bd0b1f722f46505828a9f815cf91b31d8", + ), + str_to_felt_bytes( + "0x02a222e62eabe91abdb6838fa8b267ffe81a6eb575f61e96ec9aa4460c0925a2", + ), + ), + ], + ), + ( + str_to_felt_bytes("0x020cfa74ee3564b4cd5435cdace0f9c4d43b939620e4a0bb5076105df0a626c6"), + vec![ + ( + str_to_felt_bytes( + "0x0000000000000000000000000000000000000000000000000000000000000005", + ), + str_to_felt_bytes( + "0x000000000000000000000000000000000000000000000000000000000000022b", + ), + ), + ( + str_to_felt_bytes( + "0x0313ad57fdf765addc71329abf8d74ac2bce6d46da8c2b9b82255a5076620300", + ), + str_to_felt_bytes( + "0x04e7e989d58a17cd279eca440c5eaa829efb6f9967aaad89022acbe644c39b36", + ), + ), + ( + str_to_felt_bytes( + "0x0313ad57fdf765addc71329abf8d74ac2bce6d46da8c2b9b82255a5076620301", + ), + str_to_felt_bytes( + "0x0453ae0c9610197b18b13645c44d3d0a407083d96562e8752aab3fab616cecb0", + ), + ), + ( + str_to_felt_bytes( + "0x05aee31408163292105d875070f98cb48275b8c87e80380b78d30647e05854d5", + ), + str_to_felt_bytes( + "0x00000000000000000000000000000000000000000000000000000000000007e5", + ), + ), + ( + str_to_felt_bytes( + "0x06cf6c2f36d36b08e591e4489e92ca882bb67b9c39a3afccf011972a8de467f0", + ), + str_to_felt_bytes( + "0x07ab344d88124307c07b56f6c59c12f4543e9c96398727854a322dea82c73240", + ), + ), + ], + ), + ( + str_to_felt_bytes("0x031c887d82502ceb218c06ebb46198da3f7b92864a8223746bc836dda3e34b52"), + vec![ + ( + str_to_felt_bytes( + "0x00df28e613c065616a2e79ca72f9c1908e17b8c913972a9993da77588dc9cae9", + ), + str_to_felt_bytes( + "0x01432126ac23c7028200e443169c2286f99cdb5a7bf22e607bcd724efa059040", + ), + ), + ( + str_to_felt_bytes( + "0x05f750dc13ed239fa6fc43ff6e10ae9125a33bd05ec034fc3bb4dd168df3505f", + ), + str_to_felt_bytes( + "0x00000000000000000000000000000000000000000000000000000000000007c7", + ), + ), + ], + ), + ]; + + let block_1 = [ + ( + str_to_felt_bytes("0x06538fdd3aa353af8a87f5fe77d1f533ea82815076e30a86d65b72d3eb4f0b80"), + vec![ + ( + str_to_felt_bytes( + "0x0000000000000000000000000000000000000000000000000000000000000005", + ), + str_to_felt_bytes( + "0x000000000000000000000000000000000000000000000000000000000000022b", + ), + ), + ( + str_to_felt_bytes( + "0x01aed933fd362faecd8ea54ee749092bd21f89901b7d1872312584ac5b636c6d", + ), + str_to_felt_bytes( + "0x00000000000000000000000000000000000000000000000000000000000007e5", + ), + ), + ( + str_to_felt_bytes( + "0x010212fa2be788e5d943714d6a9eac5e07d8b4b48ead96b8d0a0cbe7a6dc3832", + ), + str_to_felt_bytes( + "0x008a81230a7e3ffa40abe541786a9b69fbb601434cec9536d5d5b2ee4df90383", + ), + ), + ( + str_to_felt_bytes( + "0x00ffda4b5cf0dce9bc9b0d035210590c73375fdbb70cd94ec6949378bffc410c", + ), + str_to_felt_bytes( + "0x02b36318931915f71777f7e59246ecab3189db48408952cefda72f4b7977be51", + ), + ), + ( + str_to_felt_bytes( + "0x00ffda4b5cf0dce9bc9b0d035210590c73375fdbb70cd94ec6949378bffc410d", + ), + str_to_felt_bytes( + "0x07e928dcf189b05e4a3dae0bc2cb98e447f1843f7debbbf574151eb67cda8797", + ), + ), + ], + ), + ( + str_to_felt_bytes("0x0327d34747122d7a40f4670265b098757270a449ec80c4871450fffdab7c2fa8"), + vec![ + ( + str_to_felt_bytes( + "0x0000000000000000000000000000000000000000000000000000000000000005", + ), + str_to_felt_bytes( + "0x0000000000000000000000000000000000000000000000000000000000000065", + ), + ), + ( + str_to_felt_bytes( + "0x01aed933fd362faecd8ea54ee749092bd21f89901b7d1872312584ac5b636c6d", + ), + str_to_felt_bytes( + "0x00000000000000000000000000000000000000000000000000000000000007c7", + ), + ), + ( + str_to_felt_bytes( + "0x04184fa5a6d40f47a127b046ed6facfa3e6bc3437b393da65cc74afe47ca6c6e", + ), + str_to_felt_bytes( + "0x001ef78e458502cd457745885204a4ae89f3880ec24db2d8ca97979dce15fedc", + ), + ), + ( + str_to_felt_bytes( + "0x05591c8c3c8d154a30869b463421cd5933770a0241e1a6e8ebcbd91bdd69bec4", + ), + str_to_felt_bytes( + "0x026b5943d4a0c420607cee8030a8cdd859bf2814a06633d165820960a42c6aed", + ), + ), + ( + str_to_felt_bytes( + "0x05591c8c3c8d154a30869b463421cd5933770a0241e1a6e8ebcbd91bdd69bec5", + ), + str_to_felt_bytes( + "0x01518eec76afd5397cefd14eda48d01ad59981f9ce9e70c233ca67acd8754008", + ), + ), + ], + ), + ]; + + let block_2 = vec![ + ( + str_to_felt_bytes("0x001fb4457f3fe8a976bdb9c04dd21549beeeb87d3867b10effe0c4bd4064a8e4"), + vec![( + str_to_felt_bytes( + "0x056c060e7902b3d4ec5a327f1c6e083497e586937db00af37fe803025955678f", + ), + str_to_felt_bytes( + "0x075495b43f53bd4b9c9179db113626af7b335be5744d68c6552e3d36a16a747c", + ), + )], + ), + ( + str_to_felt_bytes("0x05790719f16afe1450b67a92461db7d0e36298d6a5f8bab4f7fd282050e02f4f"), + vec![( + str_to_felt_bytes( + "0x0772c29fae85f8321bb38c9c3f6edb0957379abedc75c17f32bcef4e9657911a", + ), + str_to_felt_bytes( + "0x06d4ca0f72b553f5338a95625782a939a49b98f82f449c20f49b42ec60ed891c", + ), + )], + ), + ( + str_to_felt_bytes("0x057b973bf2eb26ebb28af5d6184b4a044b24a8dcbf724feb95782c4d1aef1ca9"), + vec![( + str_to_felt_bytes( + "0x04f2c206f3f2f1380beeb9fe4302900701e1cb48b9b33cbe1a84a175d7ce8b50", + ), + str_to_felt_bytes( + "0x02a614ae71faa2bcdacc5fd66965429c57c4520e38ebc6344f7cf2e78b21bd2f", + ), + )], + ), + ( + str_to_felt_bytes("0x02d6c9569dea5f18628f1ef7c15978ee3093d2d3eec3b893aac08004e678ead3"), + vec![( + str_to_felt_bytes( + "0x07f93985c1baa5bd9b2200dd2151821bd90abb87186d0be295d7d4b9bc8ca41f", + ), + str_to_felt_bytes( + "0x0127cd00a078199381403a33d315061123ce246c8e5f19aa7f66391a9d3bf7c6", + ), + )], + ), + ]; + + let blocks = block_0.iter().chain(block_1.iter()).chain(block_2.iter()); + + // Inserts all storage updates into the bonsai + for (contract_address, storage) in blocks.clone() { + log::info!( + "contract address (write): {:#064x}", + Felt::from_bytes_be_slice(contract_address) + ); + + for (k, v) in storage { + // truncate only keeps the first 251 bits in a key + // so there should be no error during insertion + let ktrunc = &truncate(k); + let kfelt0 = Felt::from_bytes_be_slice(k); + let kfelt1 = Felt::from_bytes_be_slice(ktrunc.as_raw_slice()); + + // quick sanity check to make sure truncating a key does not remove any data + assert_eq!(kfelt0, kfelt1); + + let v = &Felt::from_bytes_be_slice(v); + assert!(bonsai.insert(contract_address, ktrunc, v).is_ok()); + } + } + assert!(bonsai.commit(BasicId::new(0)).is_ok()); + + // aggreates all storage changes to their latest state + // (replacements are takent into account) + let mut storage_map = IndexMap::>::new(); + for (contract_address, storage) in blocks.clone() { + let map = storage_map.entry((*contract_address).into()).or_default(); + + for (k, v) in storage { + let k = Felt::from_bytes_be_slice(k); + let v = Felt::from_bytes_be_slice(v); + map.insert(k, v); + } + } + + // checks for each contract if the original key can be reconstructed + // from the data stored in the db + for (contract_address, storage) in storage_map.iter() { + log::info!( + "contract address (read): {:#064x}", + Felt::from_bytes_be_slice(contract_address) + ); + + let keys = bonsai.get_keys(contract_address).unwrap(); + log::debug!("{keys:?}"); + for k in keys { + // if all has gone well, the db should contain the first 251 bits of the key, + // which should represent the entirety of the data + let k = Felt::from_bytes_be_slice(&k); + log::info!("looking for key: {k:#064x}"); + + assert!(storage.contains_key(&k)); + } + } + + // makes sure retrieving key-value pairs works for each contract + for (contract_address, storage) in storage_map.iter() { + log::info!( + "contract address (read): {:#064x}", + Felt::from_bytes_be_slice(contract_address) + ); + + let kv = bonsai.get_key_value_pairs(contract_address).unwrap(); + log::debug!("{kv:?}"); + for (k, v) in kv { + let k = Felt::from_bytes_be_slice(&k); + let v = Felt::from_bytes_be_slice(&v); + log::info!("checking for key-value pair:({k:#064x}, {v:#064x})"); + + assert_eq!(*storage.get(&k).unwrap(), v); + } + } +} + +fn str_to_felt_bytes(hex: &str) -> [u8; 32] { + Felt::from_hex(hex).unwrap().to_bytes_be() +} + +fn truncate(key: &[u8]) -> BitVec { + key.view_bits()[5..].to_owned() +} + +// use crate::{ +// databases::{create_rocks_db, RocksDB, RocksDBConfig}, +// id::BasicId, +// key_value_db::KeyValueDBConfig, +// KeyValueDB, +// }; +// use mp_felt::Felt252Wrapper; +// use mp_hashers::pedersen::PedersenHasher; +// use parity_scale_codec::{Decode, Encode}; +// use rand::prelude::*; +// use starknet_types_core::{felt::Felt, hash::Pedersen}; + +// // convert a Madara felt to a standard Felt +// fn felt_from_madara_felt(madara_felt: &Felt252Wrapper) -> Felt { +// let encoded = madara_felt.encode(); +// Felt::decode(&mut &encoded[..]).unwrap() +// } + +// // convert a standard Felt to a Madara felt +// fn madara_felt_from_felt(felt: &Felt) -> Felt252Wrapper { +// let encoded = felt.encode(); +// Felt252Wrapper::decode(&mut &encoded[..]).unwrap() +// } + +// #[test] +// fn one_commit_tree_compare() { +// let mut elements = vec![]; +// let tempdir = tempfile::tempdir().unwrap(); +// let mut rng = rand::thread_rng(); +// let tree_size = rng.gen_range(10..100); +// for _ in 0..tree_size { +// let mut element = String::from("0x"); +// let element_size = rng.gen_range(10..32); +// for _ in 0..element_size { +// let random_byte: u8 = rng.gen(); +// element.push_str(&format!("{:02x}", random_byte)); +// } +// elements.push(Felt::from_hex(&element).unwrap()); +// } +// let madara_elements = elements +// .iter() +// .map(madara_felt_from_felt) +// .collect::>(); +// let rocks_db = create_rocks_db(std::path::Path::new(tempdir.path())).unwrap(); +// let rocks_db = RocksDB::new(&rocks_db, RocksDBConfig::default()); +// let db = KeyValueDB::new(rocks_db, KeyValueDBConfig::default(), None); +// let mut bonsai_tree: super::MerkleTree, BasicId> = +// super::MerkleTree::new(db).unwrap(); +// let root_hash = mp_commitments::calculate_class_commitment_tree_root_hash::( +// &madara_elements, +// ); +// elements +// .iter() +// .zip(madara_elements.iter()) +// .for_each(|(element, madara_element)| { +// let final_hash = +// calculate_class_commitment_leaf_hash::(*madara_element); +// let key = &element.to_bytes_be()[..31]; +// bonsai_tree +// .set( +// &BitVec::from_vec(key.to_vec()), +// felt_from_madara_felt(&final_hash), +// ) +// .unwrap(); +// }); +// bonsai_tree.display(); +// assert_eq!( +// bonsai_tree.commit().unwrap(), +// felt_from_madara_felt(&root_hash) +// ); +// } + +// #[test] +// fn simple_commits() { +// let tempdir = tempfile::tempdir().unwrap(); +// let mut madara_tree = StateCommitmentTree::::default(); +// let rocks_db = create_rocks_db(std::path::Path::new(tempdir.path())).unwrap(); +// let rocks_db = RocksDB::new(&rocks_db, RocksDBConfig::default()); +// let db = KeyValueDB::new(rocks_db, KeyValueDBConfig::default(), None); +// let mut bonsai_tree: super::MerkleTree, BasicId> = +// super::MerkleTree::new(db).unwrap(); +// let elements = [ +// [Felt::from_hex("0x665342762FDD54D0303c195fec3ce2568b62052e").unwrap()], +// [Felt::from_hex("0x66342762FDD54D0303c195fec3ce2568b62052e").unwrap()], +// [Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap()], +// ]; +// for elem in elements { +// elem.iter().for_each(|class_hash| { +// let final_hash = +// felt_from_madara_felt(&calculate_class_commitment_leaf_hash::( +// madara_felt_from_felt(class_hash), +// )); +// madara_tree.set( +// madara_felt_from_felt(class_hash), +// madara_felt_from_felt(&final_hash), +// ); +// let key = &class_hash.to_bytes_be()[..31]; +// bonsai_tree +// .set(&BitVec::from_vec(key.to_vec()), final_hash) +// .unwrap(); +// }); +// } +// let madara_root_hash = madara_tree.commit(); +// let bonsai_root_hash = bonsai_tree.commit().unwrap(); +// assert_eq!(bonsai_root_hash, felt_from_madara_felt(&madara_root_hash)); +// } + +// #[test] +// fn simple_commits_and_delete() { +// let tempdir = tempfile::tempdir().unwrap(); +// let rocks_db = create_rocks_db(std::path::Path::new(tempdir.path())).unwrap(); +// let rocks_db = RocksDB::new(&rocks_db, RocksDBConfig::default()); +// let db = KeyValueDB::new(rocks_db, KeyValueDBConfig::default(), None); +// let mut bonsai_tree: super::MerkleTree, BasicId> = +// super::MerkleTree::new(db).unwrap(); +// let elements = [ +// [Felt::from_hex("0x665342762FDD54D0303c195fec3ce2568b62052e").unwrap()], +// [Felt::from_hex("0x66342762FDD54D0303c195fec3ce2568b62052e").unwrap()], +// [Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap()], +// ]; +// for elem in elements { +// elem.iter().for_each(|class_hash| { +// let final_hash = calculate_class_commitment_leaf_hash::( +// madara_felt_from_felt(class_hash), +// ); +// let key = &class_hash.to_bytes_be()[..31]; +// bonsai_tree +// .set( +// &BitVec::from_vec(key.to_vec()), +// felt_from_madara_felt(&final_hash), +// ) +// .unwrap(); +// }); +// } +// bonsai_tree.commit().unwrap(); +// for elem in elements { +// elem.iter().for_each(|class_hash| { +// let key = &class_hash.to_bytes_be()[..31]; +// bonsai_tree +// .set(&BitVec::from_vec(key.to_vec()), Felt::ZERO) +// .unwrap(); +// }); +// } +// bonsai_tree.commit().unwrap(); +// } + +// #[test] +// fn multiple_commits_tree_compare() { +// let mut rng = rand::thread_rng(); +// let tempdir = tempfile::tempdir().unwrap(); +// let mut madara_tree = StateCommitmentTree::::default(); +// let rocks_db = create_rocks_db(std::path::Path::new(tempdir.path())).unwrap(); +// let rocks_db = RocksDB::new(&rocks_db, RocksDBConfig::default()); +// let db = KeyValueDB::new(rocks_db, KeyValueDBConfig::default(), None); +// let mut bonsai_tree: super::MerkleTree, BasicId> = +// super::MerkleTree::new(db).unwrap(); +// let nb_commits = rng.gen_range(2..4); +// for _ in 0..nb_commits { +// let mut elements = vec![]; +// let tree_size = rng.gen_range(10..100); +// for _ in 0..tree_size { +// let mut element = String::from("0x"); +// let element_size = rng.gen_range(10..32); +// for _ in 0..element_size { +// let random_byte: u8 = rng.gen(); +// element.push_str(&format!("{:02x}", random_byte)); +// } +// elements.push(Felt::from_hex(&element).unwrap()); +// } +// elements.iter().for_each(|class_hash| { +// let final_hash = calculate_class_commitment_leaf_hash::( +// madara_felt_from_felt(class_hash), +// ); +// madara_tree.set(madara_felt_from_felt(class_hash), final_hash); +// let key = &class_hash.to_bytes_be()[..31]; +// bonsai_tree +// .set( +// &BitVec::from_vec(key.to_vec()), +// felt_from_madara_felt(&final_hash), +// ) +// .unwrap(); +// }); + +// let bonsai_root_hash = bonsai_tree.commit().unwrap(); +// let madara_root_hash = madara_tree.commit(); +// assert_eq!(bonsai_root_hash, felt_from_madara_felt(&madara_root_hash)); +// } +// } + +// #[test] // fn multiple_commits_tree_compare_with_deletes() { +// let mut rng = rand::thread_rng(); +// let mut madara_tree = StateCommitmentTree::::default(); +// let rocks_db = create_rocks_db(std::path::Path::new("test_db")).unwrap(); +// let mut db = RocksDB::new(&rocks_db, RocksDBConfig::default()); +// let mut bonsai_tree: super::MerkleTree = +// super::MerkleTree::empty(&mut db); +// let nb_commits = rng.gen_range(2..5); +// let mut elements_to_delete = vec![]; +// for _ in 0..nb_commits { +// let mut elements = vec![]; +// let tree_size = rng.gen_range(10..100); +// for _ in 0..tree_size { +// let mut element = String::from("0x"); +// let element_size = rng.gen_range(10..32); +// for _ in 0..element_size { +// let random_byte: u8 = rng.gen(); +// element.push_str(&format!("{:02x}", random_byte)); +// } +// if rng.gen_bool(0.1) { +// elements_to_delete.push(Felt::from_hex_be(&element).unwrap()); +// elements.push(Felt::from_hex_be(&element).unwrap()); +// } else { +// elements.push(Felt::from_hex_be(&element).unwrap()); +// } +// } +// elements.iter().for_each(|class_hash| { +// let final_hash = +// calculate_class_commitment_leaf_hash::(*class_hash); +// madara_tree.set(*class_hash, final_hash); +// let key = &class_hash.0.to_bytes_be()[..31]; +// bonsai_tree.set(&BitVec::from_vec(key.to_vec()), final_hash); +// }); + +// let bonsai_root_hash = bonsai_tree.commit(); +// let madara_root_hash = madara_tree.commit(); +// assert_eq!(bonsai_root_hash, madara_root_hash); +// } +// elements_to_delete.iter().for_each(|class_hash| { +// madara_tree.set(*class_hash, Felt::ZERO); +// let key = &class_hash.0.to_bytes_be()[..31]; +// bonsai_tree.set(&BitVec::from_vec(key.to_vec()), Felt::ZERO); +// }); + +// let bonsai_root_hash = bonsai_tree.commit(); +// let madara_root_hash = madara_tree.commit(); +// assert_eq!(bonsai_root_hash, madara_root_hash); +// } + +// #[test] +// fn test_proof() { +// let tempdir = tempfile::tempdir().unwrap(); +// let rocks_db = create_rocks_db(std::path::Path::new(tempdir.path())).unwrap(); +// let rocks_db = RocksDB::new(&rocks_db, RocksDBConfig::default()); +// let db = KeyValueDB::new(rocks_db, KeyValueDBConfig::default(), None); +// let mut bonsai_tree: super::MerkleTree, BasicId> = +// super::MerkleTree::new(db).unwrap(); +// let elements = [ +// [Felt252Wrapper::from_hex_be("0x665342762FDD54D0303c195fec3ce2568b62052e").unwrap()], +// [Felt252Wrapper::from_hex_be("0x66342762FDD54D0303c195fec3ce2568b62052e").unwrap()], +// [Felt252Wrapper::from_hex_be("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap()], +// ]; +// for elem in elements { +// elem.iter().for_each(|class_hash| { +// let final_hash = +// calculate_class_commitment_leaf_hash::(*class_hash); +// let key = &class_hash.0.to_bytes_be()[..31]; +// bonsai_tree +// .set( +// &BitVec::from_vec(key.to_vec()), +// Felt::from_bytes_be(&final_hash.0.to_bytes_be()), +// ) +// .unwrap(); +// }); +// } +// bonsai_tree.commit().unwrap(); +// let bonsai_proof = bonsai_tree +// .get_proof(&BitVec::from_vec( +// elements[0][0].0.to_bytes_be()[..31].to_vec(), +// )) +// .unwrap(); +// println!("bonsai_proof: {:?}", bonsai_proof); +// } + +// test in madara +// #[test] +// fn test_proof() { +// let mut tree = super::merkle_patricia_tree::merkle_tree::MerkleTree::::empty(); +// let elements = [ +// [Felt252Wrapper::from_hex_be("0x665342762FDD54D0303c195fec3ce2568b62052e").unwrap()], +// [Felt252Wrapper::from_hex_be("0x66342762FDD54D0303c195fec3ce2568b62052e").unwrap()], +// [Felt252Wrapper::from_hex_be("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap()], +// ]; +// for elem in elements { +// elem.iter().for_each(|class_hash| { +// let final_hash = +// calculate_class_commitment_leaf_hash::(*class_hash); +// let key = &class_hash.0.to_bytes_be()[..31]; +// tree +// .set(&BitVec::from_vec(key.to_vec()), final_hash) +// }); +// } +// tree.commit(); +// let bonsai_proof = tree.get_proof(&BitVec::from_vec( +// elements[0][0].0.to_bytes_be()[..31].to_vec(), +// )); +// println!("bonsai_proof: {:?}", bonsai_proof); +// } diff --git a/src/tests/mod.rs b/src/tests/mod.rs index 2c11aa9..da125d9 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -1,6 +1,6 @@ mod madara_comparison; mod merge; -mod proof; +mod merkle_tree; mod proptest; mod simple; mod transactional_state; diff --git a/src/tests/proof.rs b/src/tests/proof.rs deleted file mode 100644 index 844c3d6..0000000 --- a/src/tests/proof.rs +++ /dev/null @@ -1,372 +0,0 @@ -#![cfg(all(feature = "std", feature = "rocksdb"))] -use bitvec::{order::Msb0, vec::BitVec, view::BitView}; -use pathfinder_common::{hash::PedersenHash, trie::TrieNode}; -use pathfinder_crypto::Felt as PathfinderFelt; -use pathfinder_merkle_tree::tree::{MerkleTree, TestStorage}; -use pathfinder_storage::{Node, StoredNode}; -use rand::Rng; -use starknet_types_core::{felt::Felt, hash::Pedersen}; -use std::collections::HashMap; - -use crate::{ - databases::{create_rocks_db, RocksDB, RocksDBConfig}, - id::{BasicId, BasicIdBuilder}, - trie::merkle_tree::{Membership, ProofNode}, - BonsaiStorage, BonsaiStorageConfig, -}; - -/// Commits the tree changes and persists them to storage. -fn commit_and_persist( - tree: MerkleTree, - storage: &mut TestStorage, -) -> (PathfinderFelt, u64) { - use pathfinder_storage::Child; - - for (key, value) in &tree.leaves { - let key = PathfinderFelt::from_bits(key).unwrap(); - storage.leaves.insert(key, *value); - } - - let update = tree.commit(storage).unwrap(); - - let mut indices = HashMap::new(); - let mut idx = storage.nodes.len(); - for hash in update.nodes.keys() { - indices.insert(*hash, idx as u64); - idx += 1; - } - - for (hash, node) in update.nodes { - let node = match node { - Node::Binary { left, right } => { - let left = match left { - Child::Id(idx) => idx, - Child::Hash(hash) => { - *indices.get(&hash).expect("Left child should have an index") - } - }; - - let right = match right { - Child::Id(idx) => idx, - Child::Hash(hash) => *indices - .get(&hash) - .expect("Right child should have an index"), - }; - - StoredNode::Binary { left, right } - } - Node::Edge { child, path } => { - let child = match child { - Child::Id(idx) => idx, - Child::Hash(hash) => *indices.get(&hash).expect("Child should have an index"), - }; - - StoredNode::Edge { child, path } - } - Node::LeafBinary => StoredNode::LeafBinary, - Node::LeafEdge { path } => StoredNode::LeafEdge { path }, - }; - - storage - .nodes - .insert(*indices.get(&hash).unwrap(), (hash, node)); - } - - let index = *indices.get(&update.root).unwrap(); - - (update.root, index) -} - -fn assert_eq_proof(bonsai_proof: &[ProofNode], pathfinder_proof: &[TrieNode]) { - for (bonsai_node, pathfinder_node) in bonsai_proof.iter().zip(pathfinder_proof.iter()) { - match (bonsai_node, pathfinder_node) { - ( - ProofNode::Binary { left, right }, - pathfinder_common::trie::TrieNode::Binary { - left: pathfinder_left, - right: pathfinder_right, - }, - ) => { - let pathfinder_left_bits = pathfinder_left.to_hex_str(); - let pathfinder_felt = Felt::from_hex(&pathfinder_left_bits).unwrap(); - assert_eq!(left, &pathfinder_felt); - let pathfinder_right_bits = pathfinder_right.to_hex_str(); - let pathfinder_felt = Felt::from_hex(&pathfinder_right_bits).unwrap(); - assert_eq!(right, &pathfinder_felt); - } - ( - ProofNode::Edge { child, path }, - pathfinder_common::trie::TrieNode::Edge { - child: pathfinder_child, - path: pathfinder_path, - }, - ) => { - let pathfinder_child_bits = pathfinder_child.to_hex_str(); - let pathfinder_felt = Felt::from_hex(&pathfinder_child_bits).unwrap(); - assert_eq!(child, &pathfinder_felt); - assert_eq!(&path.0, pathfinder_path); - } - _ => panic!("Proofs are not the same"), - } - } -} - -#[test] -fn debug_deoxys() { - // Load storage_data.json file - let storage_data = include_str!("storage_data.json"); - let storage_data: Vec> = serde_json::from_str(storage_data).unwrap(); - let tempdir = tempfile::tempdir().unwrap(); - let db = create_rocks_db(tempdir.path()).unwrap(); - let config = BonsaiStorageConfig::default(); - let mut storage = pathfinder_merkle_tree::tree::TestStorage::default(); - let mut id_builder = BasicIdBuilder::new(); - let mut bonsai_storage = - BonsaiStorage::<_, _, Pedersen>::new(RocksDB::new(&db, RocksDBConfig::default()), config) - .unwrap(); - let mut pathfinder_merkle_tree: MerkleTree = - pathfinder_merkle_tree::tree::MerkleTree::empty(); - let identifier = - Felt::from_hex("0x04acd4b2a59eae7196f6a5c26ead8cb5f9d7ad3d911096418a23357bb2eac075") - .unwrap() - .to_bytes_be() - .to_vec(); - for block_changes in storage_data.iter() { - for pair in block_changes.iter() { - let key = keyer(Felt::from_hex(&pair.0).unwrap()); - let value = Felt::from_hex(&pair.1).unwrap(); - bonsai_storage.insert(&identifier, &key, &value).unwrap(); - pathfinder_merkle_tree - .set( - &storage, - key, - PathfinderFelt::from_hex_str(&pair.1).unwrap(), - ) - .unwrap(); - } - bonsai_storage.commit(id_builder.new_id()).unwrap(); - let (_, root_id) = commit_and_persist(pathfinder_merkle_tree.clone(), &mut storage); - let pathfinder_root = storage.nodes.get(&root_id).unwrap().0; - let bonsai_root = bonsai_storage.root_hash(&identifier).unwrap(); - println!("{:#02x}", bonsai_root); - println!("{:#02x}", pathfinder_root); - assert_eq!(pathfinder_root.to_be_bytes(), bonsai_root.to_bytes_be()); - } -} - -fn keyer(felt: Felt) -> BitVec { - felt.to_bytes_be().view_bits()[5..].to_bitvec() -} - -#[test] -fn basic_proof() { - let identifier = vec![]; - let tempdir = tempfile::tempdir().unwrap(); - let db = create_rocks_db(tempdir.path()).unwrap(); - let config = BonsaiStorageConfig::default(); - let mut storage = pathfinder_merkle_tree::tree::TestStorage::default(); - let mut bonsai_storage = - BonsaiStorage::<_, _, Pedersen>::new(RocksDB::new(&db, RocksDBConfig::default()), config) - .unwrap(); - let mut pathfinder_merkle_tree: MerkleTree = - pathfinder_merkle_tree::tree::MerkleTree::empty(); - let mut id_builder = BasicIdBuilder::new(); - let pair1 = ( - vec![1, 2, 1], - Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(), - ); - let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage - .insert(&identifier, &bitvec, &pair1.1) - .unwrap(); - pathfinder_merkle_tree - .set( - &storage, - bitvec, - PathfinderFelt::from_hex_str("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(), - ) - .unwrap(); - let pair2 = ( - vec![1, 2, 2], - Felt::from_hex("0x66342762FD54D033c195fec3ce2568b62052e").unwrap(), - ); - let bitvec = BitVec::from_vec(pair2.0.clone()); - bonsai_storage - .insert(&identifier, &bitvec, &pair2.1) - .unwrap(); - pathfinder_merkle_tree - .set( - &storage, - bitvec, - PathfinderFelt::from_hex_str("0x66342762FD54D033c195fec3ce2568b62052e").unwrap(), - ) - .unwrap(); - let pair3 = ( - vec![1, 2, 3], - Felt::from_hex("0x66342762FD54D033c195fec3ce2568b62052e").unwrap(), - ); - let bitvec = BitVec::from_vec(pair3.0.clone()); - bonsai_storage - .insert(&identifier, &bitvec, &pair3.1) - .unwrap(); - pathfinder_merkle_tree - .set( - &storage, - bitvec, - PathfinderFelt::from_hex_str("0x66342762FD54D033c195fec3ce2568b62052e").unwrap(), - ) - .unwrap(); - 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(&identifier, &BitVec::from_vec(vec![1, 2, 1])) - .unwrap(); - let pathfinder_proof = - pathfinder_merkle_tree::tree::MerkleTree::::get_proof( - root_id, - &storage, - &BitVec::from_vec(vec![1, 2, 1]), - ) - .unwrap(); - assert_eq_proof(&bonsai_proof, &pathfinder_proof); - assert_eq!( - BonsaiStorage::, Pedersen>::verify_proof( - bonsai_storage.root_hash(&identifier).unwrap(), - &BitVec::from_vec(vec![1, 2, 1]), - pair1.1, - &bonsai_proof - ), - Some(Membership::Member) - ); -} - -#[test] -fn multiple_proofs() { - let identifier = vec![]; - let tempdir = tempfile::tempdir().unwrap(); - let db = create_rocks_db(tempdir.path()).unwrap(); - let config = BonsaiStorageConfig::default(); - let mut storage = pathfinder_merkle_tree::tree::TestStorage::default(); - let mut bonsai_storage = - BonsaiStorage::<_, _, Pedersen>::new(RocksDB::new(&db, RocksDBConfig::default()), config) - .unwrap(); - let mut pathfinder_merkle_tree: MerkleTree = - pathfinder_merkle_tree::tree::MerkleTree::empty(); - let mut id_builder = BasicIdBuilder::new(); - let mut rng = rand::thread_rng(); - let tree_size = rng.gen_range(10..1000); - let mut elements = vec![]; - for _ in 0..tree_size { - let mut element = String::from("0x"); - let element_size = rng.gen_range(10..32); - for _ in 0..element_size { - let random_byte: u8 = rng.gen(); - element.push_str(&format!("{:02x}", random_byte)); - } - let value = Felt::from_hex(&element).unwrap(); - let key = &value.to_bytes_be()[..31]; - bonsai_storage - .insert(&identifier, &BitVec::from_vec(key.to_vec()), &value) - .unwrap(); - pathfinder_merkle_tree - .set( - &storage, - BitVec::from_vec(key.to_vec()), - PathfinderFelt::from_hex_str(&element).unwrap(), - ) - .unwrap(); - elements.push((key.to_vec(), value)); - } - bonsai_storage.commit(id_builder.new_id()).unwrap(); - let (_, root_id) = commit_and_persist(pathfinder_merkle_tree.clone(), &mut storage); - for element in elements.iter() { - let proof = bonsai_storage - .get_proof(&identifier, &BitVec::from_vec(element.0.clone())) - .unwrap(); - let pathfinder_proof = - pathfinder_merkle_tree::tree::MerkleTree::::get_proof( - root_id, - &storage, - &BitVec::from_vec(element.0.clone()), - ) - .unwrap(); - assert_eq_proof(&proof, &pathfinder_proof); - assert_eq!( - BonsaiStorage::, Pedersen>::verify_proof( - bonsai_storage.root_hash(&identifier).unwrap(), - &BitVec::from_vec(element.0.clone()), - element.1, - &proof - ), - Some(Membership::Member) - ); - } -} - -#[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(); - let mut storage = pathfinder_merkle_tree::tree::TestStorage::default(); - let mut bonsai_storage = - BonsaiStorage::<_, _, Pedersen>::new(RocksDB::new(&db, RocksDBConfig::default()), config) - .unwrap(); - let mut pathfinder_merkle_tree: MerkleTree = - pathfinder_merkle_tree::tree::MerkleTree::empty(); - let mut id_builder = BasicIdBuilder::new(); - let pair1 = ( - vec![1, 2, 1], - Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(), - ); - let bitvec = BitVec::from_vec(pair1.0.clone()); - bonsai_storage - .insert(&identifier, &bitvec, &pair1.1) - .unwrap(); - pathfinder_merkle_tree - .set( - &storage, - bitvec, - PathfinderFelt::from_hex_str("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(), - ) - .unwrap(); - 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(&identifier, &BitVec::from_vec(vec![1, 2, 1])) - .unwrap(); - let pathfinder_proof = - pathfinder_merkle_tree::tree::MerkleTree::::get_proof( - root_id, - &storage, - &BitVec::from_vec(vec![1, 2, 1]), - ) - .unwrap(); - assert_eq_proof(&bonsai_proof, &pathfinder_proof); - assert_eq!( - BonsaiStorage::, Pedersen>::verify_proof( - bonsai_storage.root_hash(&identifier).unwrap(), - &BitVec::from_vec(vec![1, 2, 1]), - pair1.1, - &bonsai_proof - ), - Some(Membership::Member) - ); -} - -#[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(); - let mut bonsai_storage = - BonsaiStorage::<_, _, Pedersen>::new(RocksDB::new(&db, RocksDBConfig::default()), config) - .unwrap(); - let mut id_builder = BasicIdBuilder::new(); - bonsai_storage.commit(id_builder.new_id()).unwrap(); - bonsai_storage - .get_proof(&identifier, &BitVec::from_vec(vec![1, 2, 1])) - .expect_err("Should error"); -} diff --git a/src/tests/proptest.rs b/src/tests/proptest.rs index 476c7f3..a623f82 100644 --- a/src/tests/proptest.rs +++ b/src/tests/proptest.rs @@ -1,17 +1,12 @@ #![cfg(feature = "std")] - -use crate::MerkleTree; - -use core::fmt::{self, Debug}; - use crate::databases::HashMapDb; use crate::id::BasicId; use crate::key_value_db::KeyValueDB; -use crate::HashMap; - +use crate::trie::tree::MerkleTree; +use crate::{BitVec, HashMap}; use bitvec::bitvec; use bitvec::order::Msb0; -use bitvec::vec::BitVec; +use core::fmt::{self, Debug}; use proptest::prelude::*; use proptest_derive::Arbitrary; use smallvec::smallvec; @@ -19,7 +14,7 @@ use starknet_types_core::felt::Felt; use starknet_types_core::hash::Pedersen; #[derive(PartialEq, Eq, Hash)] -struct Key(BitVec); +struct Key(BitVec); impl fmt::Debug for Key { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "{:b}", self.0) @@ -31,7 +26,7 @@ impl Arbitrary for Key { fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy { <[bool; 5]>::arbitrary() - .prop_map(|arr| arr.into_iter().collect::>()) + .prop_map(|arr| arr.into_iter().collect::()) .prop_map(Self) .boxed() } @@ -50,7 +45,7 @@ impl Arbitrary for Value { fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy { <[bool; 251]>::arbitrary() - .prop_map(|arr| arr.into_iter().collect::>()) + .prop_map(|arr| arr.into_iter().collect::()) .prop_map(|vec| Felt::from_bytes_be(vec.as_raw_slice().try_into().unwrap())) .prop_map(Self) .boxed() @@ -78,7 +73,7 @@ impl MerkleTreeInsertProblem { let mut ckv = HashMap::new(); // apply steps - let mut tree = MerkleTree::::new(smallvec![]); + let mut tree = MerkleTree::::new(smallvec![], 5); for step in &self.0 { match step { Step::Insert(k, v) => { @@ -97,7 +92,7 @@ impl MerkleTreeInsertProblem { } } log::trace!("TREE"); - tree.display(); + tree.dump(); } // check @@ -126,7 +121,7 @@ impl MerkleTreeInsertProblem { } proptest::proptest! { - #![proptest_config(ProptestConfig::with_cases(5))] // comment this when developing, this is mostly for faster ci & whole workspace `cargo test` + // #![proptest_config(ProptestConfig::with_cases(5))] // comment this when developing, this is mostly for faster ci & whole workspace `cargo test` #[test] fn proptest_inserts(pb in any::()) { let _ = env_logger::builder().is_test(true).try_init(); @@ -280,3 +275,27 @@ fn test_merkle_pb_6() { pb.check(); } + +#[test] +fn test_merkle_pb_7() { + use Step::*; + let _ = env_logger::builder().is_test(true).try_init(); + log::set_max_level(log::LevelFilter::Trace); + let pb = MerkleTreeInsertProblem(vec![ + Insert( + Key(bitvec![u8, Msb0; 0,0,0,0,0]), + Value(Felt::from_hex("0x20").unwrap()), + ), + Insert( + Key(bitvec![u8, Msb0; 1,0,0,0,0]), + Value(Felt::from_hex("0x20").unwrap()), + ), + Commit, + Insert( + Key(bitvec![u8, Msb0; 0,0,0,0,0]), + Value(Felt::from_hex("0x40").unwrap()), + ), + ]); + + pb.check(); +} diff --git a/src/tests/simple.rs b/src/tests/simple.rs index 7e8cd52..83dbcd0 100644 --- a/src/tests/simple.rs +++ b/src/tests/simple.rs @@ -2,9 +2,9 @@ use crate::{ databases::{create_rocks_db, HashMapDb, RocksDB, RocksDBConfig}, id::{BasicId, BasicIdBuilder}, - BonsaiStorage, BonsaiStorageConfig, Change, + BitVec, BonsaiStorage, BonsaiStorageConfig, Change, }; -use bitvec::{order::Msb0, vec::BitVec, view::BitView}; +use bitvec::view::BitView; use starknet_types_core::{felt::Felt, hash::Pedersen}; #[test] @@ -14,7 +14,7 @@ fn basics() { 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(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( vec![1, 2, 1], @@ -69,7 +69,7 @@ fn root_hash_similar_rocks_db() { 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(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( vec![1, 2, 1], @@ -116,7 +116,7 @@ fn root_hash_similar_rocks_db() { 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(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( vec![1, 2, 3], @@ -156,13 +156,13 @@ fn starknet_specific() { let db1 = create_rocks_db(tempdir1.path()).unwrap(); let config1 = BonsaiStorageConfig::default(); let mut bonsai_storage1: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(RocksDB::new(&db1, RocksDBConfig::default()), config1).unwrap(); + BonsaiStorage::new(RocksDB::new(&db1, RocksDBConfig::default()), config1, 251).unwrap(); let tempdir2 = tempfile::tempdir().unwrap(); let db2 = create_rocks_db(tempdir2.path()).unwrap(); let config2 = BonsaiStorageConfig::default(); let mut bonsai_storage2: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(RocksDB::new(&db2, RocksDBConfig::default()), config2).unwrap(); + BonsaiStorage::new(RocksDB::new(&db2, RocksDBConfig::default()), config2, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let contract_states = vec![ @@ -237,7 +237,7 @@ fn root_hash_similar_hashmap_db() { let db = HashMapDb::::default(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(db, config).unwrap(); + BonsaiStorage::new(db, config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( vec![1, 2, 1], @@ -283,7 +283,7 @@ fn root_hash_similar_hashmap_db() { let db = HashMapDb::::default(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(db, config).unwrap(); + BonsaiStorage::new(db, config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( vec![1, 2, 3], @@ -322,7 +322,7 @@ fn double_insert() { 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(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let contract_states = vec![ ContractState { @@ -385,7 +385,7 @@ fn double_identifier() { 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(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let contract_states = vec![ ContractState { @@ -450,7 +450,7 @@ fn get_changes() { 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(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); 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()); @@ -493,7 +493,7 @@ fn get_changes() { ); } -fn keyer(felt: Felt) -> BitVec { +fn keyer(felt: Felt) -> BitVec { felt.to_bytes_be().view_bits()[5..].to_bitvec() } @@ -501,7 +501,7 @@ fn keyer(felt: Felt) -> BitVec { fn test_insert_zero() { let config = BonsaiStorageConfig::default(); let bonsai_db = HashMapDb::::default(); - let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config) + let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config, 251) .expect("Failed to create bonsai storage"); let identifier = "0x056e4fed965fccd7fb01fcadd827470338f35ced62275328929d0d725b5707ba".as_bytes(); @@ -642,9 +642,10 @@ fn test_insert_zero() { #[test] fn test_block_7_starknet() { + let _ = env_logger::builder().is_test(true).try_init(); let config = BonsaiStorageConfig::default(); let bonsai_db = HashMapDb::::default(); - let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config) + let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config, 251) .expect("Failed to create bonsai storage"); let identifier = "0x056e4fed965fccd7fb01fcadd827470338f35ced62275328929d0d725b5707ba".as_bytes(); @@ -678,6 +679,7 @@ fn test_block_7_starknet() { .expect("Failed to insert storage update into trie"); } + bonsai_storage.dump(); let mut id_builder = BasicIdBuilder::new(); let id = id_builder.new_id(); bonsai_storage @@ -747,6 +749,7 @@ fn test_block_7_starknet() { .expect("Failed to insert storage update into trie"); } + bonsai_storage.dump(); let id = id_builder.new_id(); bonsai_storage .commit(id) @@ -854,7 +857,7 @@ fn test_block_7_starknet() { 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) + let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config, 251) .expect("Failed to create bonsai storage"); let identifier = "0x421203c58e1b4a6c3675be26cfaa18d2b6b42695ca206be1f08ce29f7f1bc7c".as_bytes(); @@ -985,7 +988,7 @@ fn test_block_7_starknet_2() { fn test_block_9() { let config = BonsaiStorageConfig::default(); let bonsai_db = HashMapDb::::default(); - let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config) + let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config, 251) .expect("Failed to create bonsai storage"); let identifier = "0x06F3C934BA4EC49245CB9A42FC715E4D589AA502AF69BE13916127A538D525CE".as_bytes(); diff --git a/src/tests/transactional_state.rs b/src/tests/transactional_state.rs index a1a1a25..664b933 100644 --- a/src/tests/transactional_state.rs +++ b/src/tests/transactional_state.rs @@ -2,9 +2,8 @@ use crate::{ databases::{create_rocks_db, RocksDB, RocksDBConfig}, id::BasicIdBuilder, - BonsaiStorage, BonsaiStorageConfig, + BitVec, BonsaiStorage, BonsaiStorageConfig, }; -use bitvec::vec::BitVec; use log::LevelFilter; use starknet_types_core::{felt::Felt, hash::Pedersen}; @@ -15,7 +14,7 @@ fn basics() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -59,8 +58,12 @@ fn test_thread() { let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); - let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config.clone()).unwrap(); + let mut bonsai_storage = BonsaiStorage::new( + RocksDB::new(&db, RocksDBConfig::default()), + config.clone(), + 24, + ) + .unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -120,7 +123,7 @@ fn remove() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -165,7 +168,7 @@ fn merge() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -208,7 +211,7 @@ fn merge_with_uncommitted_insert() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -270,7 +273,7 @@ fn merge_with_uncommitted_remove() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -328,7 +331,7 @@ fn transactional_state_after_uncommitted() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -370,7 +373,7 @@ fn merge_override() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -413,7 +416,7 @@ fn merge_remove() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -452,7 +455,7 @@ fn merge_txn_revert() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -512,7 +515,7 @@ fn merge_invalid() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -559,7 +562,7 @@ fn many_snapshots() { ..Default::default() }; let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( diff --git a/src/tests/trie_log.rs b/src/tests/trie_log.rs index 8146ad5..e07d692 100644 --- a/src/tests/trie_log.rs +++ b/src/tests/trie_log.rs @@ -2,9 +2,8 @@ use crate::{ databases::{create_rocks_db, RocksDB, RocksDBConfig}, id::BasicIdBuilder, - BonsaiStorage, BonsaiStorageConfig, BonsaiTrieHash, + BitVec, BonsaiStorage, BonsaiStorageConfig, BonsaiTrieHash, }; -use bitvec::vec::BitVec; use starknet_types_core::{felt::Felt, hash::Pedersen}; #[test] @@ -14,7 +13,7 @@ fn basics() { 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(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -63,7 +62,7 @@ fn unrecorded_revert() { 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(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -88,7 +87,7 @@ fn in_place_revert() { 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(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = (vec![1, 2, 3], &BonsaiTrieHash::default()); @@ -111,7 +110,7 @@ fn truncated_revert() { 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(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -151,7 +150,7 @@ fn double_revert() { 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(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -193,7 +192,7 @@ fn remove_and_reinsert() { 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(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( diff --git a/src/trie/iterator.rs b/src/trie/iterator.rs new file mode 100644 index 0000000..6306e65 --- /dev/null +++ b/src/trie/iterator.rs @@ -0,0 +1,465 @@ +use super::{ + merkle_node::{Direction, Node, NodeHandle}, + path::Path, + tree::{MerkleTree, NodeKey}, +}; +use crate::{id::Id, key_value_db::KeyValueDB, BitSlice, BonsaiDatabase, BonsaiStorageError, Vec}; +use core::{fmt, marker::PhantomData}; +use starknet_types_core::{felt::Felt, hash::StarkHash}; + +/// This trait's function will be called on every node visited during a seek operation. +pub trait NodeVisitor { + fn visit_node( + &mut self, + tree: &mut MerkleTree, + node_id: NodeKey, + prev_height: usize, + ) -> Result<(), BonsaiStorageError>; +} + +pub struct NoopVisitor(PhantomData); +impl NodeVisitor for NoopVisitor { + fn visit_node( + &mut self, + _tree: &mut MerkleTree, + _node_id: NodeKey, + _prev_height: usize, + ) -> Result<(), BonsaiStorageError> { + Ok(()) + } +} + +pub struct MerkleTreeIterator<'a, H: StarkHash, DB: BonsaiDatabase, ID: Id> { + pub(crate) tree: &'a mut MerkleTree, + pub(crate) db: &'a KeyValueDB, + /// Current iteration path. + pub(crate) current_path: Path, + /// The loaded nodes in the current path with their corresponding heights. Height is at the base of the node, meaning + /// the first node here will always have height 0. + pub(crate) current_nodes_heights: Vec<(NodeKey, usize)>, + /// Current leaf hash. Note that partial traversal (traversal that stops midway through the tree) will + /// also update this field if an exact match for the key is found, even though we may not have reached a leaf. + pub(crate) leaf_hash: Option, +} + +impl<'a, H: StarkHash, DB: BonsaiDatabase, ID: Id> fmt::Debug + for MerkleTreeIterator<'a, H, DB, ID> +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MerkleTreeIterator") + .field("cur_path", &self.current_path) + .field("current_nodes_heights", &self.current_nodes_heights) + .field("leaf_hash", &self.leaf_hash) + .finish() + } +} + +impl<'a, H: StarkHash + Send + Sync, DB: BonsaiDatabase, ID: Id> MerkleTreeIterator<'a, H, DB, ID> { + pub fn new(tree: &'a mut MerkleTree, db: &'a KeyValueDB) -> Self { + Self { + tree, + db, + current_path: Default::default(), + current_nodes_heights: Vec::with_capacity(251), + leaf_hash: None, + } + } + + #[cfg(test)] + /// For testing purposes. + pub fn cur_nodes_ids(&self) -> Vec { + use slotmap::Key; + self.current_nodes_heights + .iter() + .map(|n| n.0.data().as_ffi() & !(1 << 32)) + .collect::>() + } + + pub fn seek_to(&mut self, key: &BitSlice) -> Result<(), BonsaiStorageError> { + self.traverse_to(&mut NoopVisitor(PhantomData), key) + } + + fn traverse_one( + &mut self, + node_id: NodeKey, + height: usize, + key: &BitSlice, + ) -> Result, BonsaiStorageError> { + self.current_nodes_heights + .push((node_id, self.current_path.len())); + + let node = self.tree.get_node_mut::(node_id)?; + let (node_handle, path_matches) = match node { + Node::Binary(binary_node) => { + log::trace!( + "Continue from binary node current_path={:?} key={:b}", + self.current_path, + key, + ); + let next_direction = Direction::from(key[self.current_path.len()]); + self.current_path.push(bool::from(next_direction)); + (binary_node.get_child(next_direction), true) + } + Node::Edge(edge_node) => { + self.current_path.extend_from_bitslice(&edge_node.path); + (edge_node.child, edge_node.path_matches(key, height)) + } + }; + + // path_matches is false when the edge node doesn't match the path we want to preload so we return nothing. + log::trace!( + "Compare: path_matches={path_matches} {:?} ?= {:b} (node_handle {node_handle:?})", + self.current_path, + key + ); + if !path_matches || self.current_path.len() >= key.len() { + self.leaf_hash = if path_matches && self.current_path.len() == key.len() { + node_handle.as_hash() + } else { + None + }; + return Ok(None); // end of traversal + } + + let child_key = self + .tree + .load_node_handle(self.db, node_handle, &self.current_path)?; + + // update parent ref + match self.tree.get_node_mut::(node_id)? { + Node::Binary(binary_node) => { + *binary_node.get_child_mut(Direction::from( + *self + .current_path + .last() + .expect("current path should have a length > 0 at this point"), + )) = NodeHandle::InMemory(child_key); + } + Node::Edge(edge_node) => { + edge_node.child = NodeHandle::InMemory(child_key); + } + }; + + Ok(Some(child_key)) + } + + pub fn traverse_to>( + &mut self, + visitor: &mut V, + key: &BitSlice, + ) -> Result<(), BonsaiStorageError> { + // First, truncate the curent path and nodes list to match the new key. + log::trace!("Start traverse_to"); + + if key.is_empty() { + self.current_nodes_heights.clear(); + self.current_path.clear(); + self.leaf_hash = None; + return Ok(()); + } + + let shared_prefix_len = self + .current_path + .iter() + .zip(key) + .take_while(|(a, b)| *a == *b) + .count(); + + let nodes_new_len = if shared_prefix_len == 0 { + 0 + } else { + // partition point is a binary search under the hood + // TODO(perf): measure whether binary search is actually better than reverse iteration - the happy path may be that + // only the last few bits are different. + + self.current_nodes_heights + .partition_point(|(_node, height)| *height < shared_prefix_len) + }; + log::trace!( + "Truncate pre node id cache shared_prefix_len={:?}, nodes_new_len={:?}, cur_path_nodes_heights={:?}, current_path={:?}", + shared_prefix_len, nodes_new_len, + self.current_nodes_heights, + self.current_path, + ); + + self.current_nodes_heights.truncate(nodes_new_len); + self.current_path.truncate(key.len()); + + let mut next_to_visit = if let Some((node_id, height)) = self.current_nodes_heights.pop() { + self.current_path.truncate(height); + self.traverse_one(node_id, height, key)? + } else { + // Start from tree root. + self.current_path.clear(); + let Some(node_id) = self.tree.load_root_node(self.db)? else { + // empty tree, not found + self.leaf_hash = None; + return Ok(()); + }; + Some(node_id) + }; + + log::trace!( + "Starting traversal with path {:?}, next={:?}", + self.current_path, + next_to_visit + ); + + // Tree traversal :) + + loop { + log::trace!("Loop start cur={:?} key={:b}", self.current_path, key); + + let Some(node_id) = next_to_visit else { + return Ok(()); + }; + + visitor.visit_node::(self.tree, node_id, self.current_path.len())?; + next_to_visit = self.traverse_one(node_id, self.current_path.len(), key)?; + + log::trace!( + "Got nodeid={:?} height={}, cur path={:?}, next to visit={:?}", + node_id, + self.current_path.len(), + self.current_path, + next_to_visit + ); + } + } +} + +#[cfg(test)] +mod tests { + //! The tree used in this series of tests looks like this: + //! ``` + //! │ + //! ┌▼┐ + //! (1)│ │[0] + //! │ │ + //! └┬┘ + //! (7)┌▼┐ + //! ┌────┴─┴────────┐ + //! ┌▼┐ ┌▼┐ + //! (6)│ │[0100] (5)│ │[000000] + //! │ │ │ │ + //! └┬┘ │ │ + //! (4)┌▼┐ │ │ + //! ┌────┴─┴─────┐ │ │ + //! │ ┌▼┐ │ │ + //! (2)┌▼┐ (3)│ │[0] │ │ + //! ┌──┴─┴─┐ │ │ │ │ + //! │ │ └┬┘ └┬┘ + //! 0x1 0x2 0x3 0x4 + //! ``` + + use crate::{ + databases::{create_rocks_db, RocksDB, RocksDBConfig}, + id::{BasicId, Id}, + trie::iterator::MerkleTreeIterator, + BonsaiDatabase, BonsaiStorage, BonsaiStorageConfig, + }; + use bitvec::{bits, order::Msb0}; + use prop::{collection::vec, sample::size_range}; + use proptest::prelude::*; + use starknet_types_core::{ + felt::Felt, + hash::{Pedersen, StarkHash}, + }; + + const ONE: Felt = Felt::ONE; + const TWO: Felt = Felt::TWO; + const THREE: Felt = Felt::THREE; + const FOUR: Felt = Felt::from_hex_unchecked("0x4"); + + #[test] + fn test_iterator_seek_to() { + test_iterator_seek_to_inner((0..all_cases_len()).collect()); + } + fn test_iterator_seek_to_inner(cases: Vec) { + let _ = env_logger::builder().is_test(true).try_init(); + log::set_max_level(log::LevelFilter::Trace); + let tempdir = tempfile::tempdir().unwrap(); + let db = create_rocks_db(tempdir.path()).unwrap(); + let mut bonsai_storage: BonsaiStorage = BonsaiStorage::new( + RocksDB::::new(&db, RocksDBConfig::default()), + BonsaiStorageConfig::default(), + 8, + ) + .unwrap(); + + bonsai_storage + .insert(&[], bits![u8, Msb0; 0,0,0,1,0,0,0,0], &ONE) + .unwrap(); + bonsai_storage + .insert(&[], bits![u8, Msb0; 0,0,0,1,0,0,0,1], &TWO) + .unwrap(); + bonsai_storage + .insert(&[], bits![u8, Msb0; 0,0,0,1,0,0,1,0], &THREE) + .unwrap(); + bonsai_storage + .insert(&[], bits![u8, Msb0; 0,1,0,0,0,0,0,0], &FOUR) + .unwrap(); + + bonsai_storage.dump(); + + // Trie + + let tree = bonsai_storage + .tries + .trees + .get_mut(&smallvec::smallvec![]) + .unwrap(); + let mut iter = MerkleTreeIterator::new(tree, &bonsai_storage.tries.db); + + let cases_funcs = all_cases(); + for case in cases { + cases_funcs[case](&mut iter) + } + } + + #[allow(clippy::type_complexity)] + fn all_cases( + ) -> Vec)> { + vec![ + // SEEK TO LEAF + // case 0 + |iter| { + // from scratch, should find the leaf + iter.seek_to(bits![u8, Msb0; 0,0,0,1,0,0,0,0]).unwrap(); + assert_eq!(iter.leaf_hash, Some(ONE)); + assert_eq!(iter.cur_nodes_ids(), vec![1, 7, 6, 4, 2]); + println!("{iter:?}"); + }, + // case 1 + |iter| { + // from a closeby leaf, should backtrack and find the next one + iter.seek_to(bits![u8, Msb0; 0,0,0,1,0,0,0,1]).unwrap(); + assert_eq!(iter.leaf_hash, Some(TWO)); + assert_eq!(iter.cur_nodes_ids(), vec![1, 7, 6, 4, 2]); + println!("{iter:?}"); + }, + // case 2 + |iter| { + // backtrack farther, should find the leaf + iter.seek_to(bits![u8, Msb0; 0,0,0,1,0,0,1,0]).unwrap(); + assert_eq!(iter.leaf_hash, Some(THREE)); + assert_eq!(iter.cur_nodes_ids(), vec![1, 7, 6, 4, 3]); + println!("{iter:?}"); + }, + // case 3 + |iter| { + // backtrack farther, should find the leaf + iter.seek_to(bits![u8, Msb0; 0,1,0,0,0,0,0,0]).unwrap(); + assert_eq!(iter.leaf_hash, Some(FOUR)); + assert_eq!(iter.cur_nodes_ids(), vec![1, 7, 5]); + println!("{iter:?}"); + }, + // case 4 + |iter| { + // similar case + iter.seek_to(bits![u8, Msb0; 0,0,0,1,0,0,0,1]).unwrap(); + assert_eq!(iter.leaf_hash, Some(TWO)); + assert_eq!(iter.cur_nodes_ids(), vec![1, 7, 6, 4, 2]); + println!("{iter:?}"); + }, + // SEEK MIDWAY INTO THE TREE + // case 5 + |iter| { + // jump midway into an edge + iter.seek_to(bits![u8, Msb0; 0,1,0,0,0]).unwrap(); + // The current path should reflect the tip of the edge + assert_eq!(iter.current_path.0, bits![u8, Msb0; 0,1,0,0,0,0,0,0]); + assert_eq!(iter.leaf_hash, None); + assert_eq!(iter.cur_nodes_ids(), vec![1, 7, 5]); + println!("{iter:?}"); + }, + // case 6 + |iter| { + // jump midway into an edge, but its child is not a leaf + iter.seek_to(bits![u8, Msb0; 0,0,0]).unwrap(); + // The current path should reflect the edge + assert_eq!(iter.current_path.0, bits![u8, Msb0; 0,0,0,1,0,0]); + assert_eq!(iter.leaf_hash, None); + assert_eq!(iter.cur_nodes_ids(), vec![1, 7, 6]); + println!("{iter:?}"); + }, + // case 7 + |iter| { + // jump to a binary node + iter.seek_to(bits![u8, Msb0; 0,0,0,1,0,0,0]).unwrap(); + assert_eq!(iter.current_path.0, bits![u8, Msb0; 0,0,0,1,0,0,0]); + assert_eq!(iter.leaf_hash, None); + assert_eq!(iter.cur_nodes_ids(), vec![1, 7, 6, 4]); + println!("{iter:?}"); + }, + // case 8 + |iter| { + // jump to the end of an edge + iter.seek_to(bits![u8, Msb0; 0,0,0,1,0,0]).unwrap(); + // The current path should reflect the tip of the edge + assert_eq!(iter.current_path.0, bits![u8, Msb0; 0,0,0,1,0,0]); + assert_eq!(iter.leaf_hash, None); + assert_eq!(iter.cur_nodes_ids(), vec![1, 7, 6]); + println!("{iter:?}"); + }, + // case 9 + |iter| { + // jump to top + iter.seek_to(bits![u8, Msb0; ]).unwrap(); + assert_eq!(iter.leaf_hash, None); + assert_eq!(iter.current_path.0, bits![u8, Msb0; ]); + assert_eq!(iter.cur_nodes_ids(), vec![]); + println!("{iter:?}"); + }, + // case 10 + |iter| { + // jump to first node + iter.seek_to(bits![u8, Msb0; 0]).unwrap(); + assert_eq!(iter.current_path.0, bits![u8, Msb0; 0]); + assert_eq!(iter.leaf_hash, None); + assert_eq!(iter.cur_nodes_ids(), vec![1]); + println!("{iter:?}"); + }, + // case 11 + |iter| { + // jump to non existent node, returning same edge + iter.seek_to(bits![u8, Msb0; 0,1,0,1,0,0,0]).unwrap(); + assert_eq!(iter.current_path.0, bits![u8, Msb0; 0,1,0,0,0,0,0,0]); + assert_eq!(iter.leaf_hash, None); + assert_eq!(iter.cur_nodes_ids(), vec![1, 7, 5]); + println!("{iter:?}"); + }, + // case 12 + |iter| { + // jump to non existent node, deviating from edge, should not go into the children + iter.seek_to(bits![u8, Msb0; 1,0,0,1,0,0,0]).unwrap(); + assert_eq!(iter.current_path.0, bits![u8, Msb0; 0]); + assert_eq!(iter.leaf_hash, None); + assert_eq!(iter.cur_nodes_ids(), vec![1]); + println!("{iter:?}"); + }, + // case 13 + |iter| { + // jump to non existent node, deviating from first node + iter.seek_to(bits![u8, Msb0; 1]).unwrap(); + assert_eq!(iter.current_path.0, bits![u8, Msb0; 0]); + assert_eq!(iter.leaf_hash, None); + assert_eq!(iter.cur_nodes_ids(), vec![1]); + println!("{iter:?}"); + }, + ] + } + + fn all_cases_len() -> usize { + all_cases::, BasicId>().len() + } + + proptest::proptest! { + // #![proptest_config(ProptestConfig::with_cases(5))] // comment this when developing, this is mostly for faster ci & whole workspace `cargo test` + #[test] + /// This proptest will apply the above seek_to cases in a random order, and possibly with duplicates. + fn proptest_seek_to(cases in vec(0..all_cases_len(), size_range(0..20)).boxed()) { + test_iterator_seek_to_inner(cases) + } + } +} diff --git a/src/trie/merkle_node.rs b/src/trie/merkle_node.rs index 1fa1289..a98c8c4 100644 --- a/src/trie/merkle_node.rs +++ b/src/trie/merkle_node.rs @@ -4,30 +4,13 @@ //! For more information about how these Starknet trees are structured, see //! [`MerkleTree`](super::merkle_tree::MerkleTree). +use crate::BitSlice; +use bitvec::view::BitView; use core::fmt; - -use bitvec::order::Msb0; -use bitvec::slice::BitSlice; use parity_scale_codec::{Decode, Encode}; -use starknet_types_core::felt::Felt; - -use super::path::Path; - -/// Id of a Node within the tree -#[derive(Copy, Clone, Debug, PartialEq, Eq, Default, PartialOrd, Ord, Hash, Encode, Decode)] -pub struct NodeId(pub u64); +use starknet_types_core::{felt::Felt, hash::StarkHash}; -impl NodeId { - /// Mutates the given NodeId to be the next one and returns it. - pub fn next_id(&mut self) -> NodeId { - self.0 = self.0.checked_add(1).expect("Node id overflow"); - NodeId(self.0) - } - - pub fn reset(&mut self) { - self.0 = 0; - } -} +use super::{path::Path, tree::NodeKey}; /// A node in a Binary Merkle-Patricia Tree graph. #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Encode, Decode)] @@ -41,8 +24,17 @@ pub enum Node { #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Encode, Decode)] pub enum NodeHandle { Hash(Felt), - InMemory(NodeId), + InMemory(NodeKey), +} +impl NodeHandle { + pub fn as_hash(self) -> Option { + match self { + NodeHandle::Hash(felt) => Some(felt), + NodeHandle::InMemory(_) => None, + } + } } + impl fmt::Debug for NodeHandle { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { @@ -59,6 +51,7 @@ pub struct BinaryNode { /// has not yet been committed. pub hash: Option, /// The height of this node in the tree. + // TODO: this field will be removed in the future. pub height: u64, /// [Left](Direction::Left) child. pub left: NodeHandle, @@ -73,6 +66,7 @@ pub struct EdgeNode { /// has not yet been committed. pub hash: Option, /// The starting height of this node in the tree. + // TODO: this field will be removed in the future. pub height: u64, /// The path this edge takes. pub path: Path, @@ -138,7 +132,7 @@ impl BinaryNode { /// # Returns /// /// The direction of the key. - pub fn direction(&self, key: &BitSlice) -> Direction { + pub fn direction(&self, key: &BitSlice) -> Direction { key[self.height as usize].into() } @@ -185,12 +179,33 @@ impl Node { /// Convert to node to binary node type (returns None if it's not a binary node). pub fn as_binary(&self) -> Option<&BinaryNode> { match self { - Node::Binary(binary) => Some(binary), + Node::Binary(node) => Some(node), + _ => None, + } + } + /// Convert to node to binary node type (returns None if it's not a binary node). + pub fn as_binary_mut(&mut self) -> Option<&mut BinaryNode> { + match self { + Node::Binary(node) => Some(node), + _ => None, + } + } + /// Convert to node to edge node type (returns None if it's not an edge node). + pub fn as_edge(&self) -> Option<&EdgeNode> { + match self { + Node::Edge(node) => Some(node), + _ => None, + } + } + /// Convert to node to edge node type (returns None if it's not an edge node). + pub fn as_edge_mut(&mut self) -> Option<&mut EdgeNode> { + match self { + Node::Edge(node) => Some(node), _ => None, } } - pub fn hash(&self) -> Option { + pub fn get_hash(&self) -> Option { match self { Node::Binary(binary) => binary.hash, Node::Edge(edge) => edge.hash, @@ -204,9 +219,19 @@ impl EdgeNode { /// # Arguments /// /// * `key` - The key to check if the path matches with the edge node. - pub fn path_matches(&self, key: &BitSlice) -> bool { - self.path.0 - == key[(self.height as usize)..(self.height + self.path.0.len() as u64) as usize] + /// * `node_height` - The height of the edge node. + pub fn path_matches(&self, key: &BitSlice, node_height: usize) -> bool { + assert_eq!(self.height as usize, node_height); + let lower_bound = node_height.min(key.len()); + let upper_bound = (node_height + self.path.0.len()).min(key.len()); + log::trace!( + "path_matches {:b}{lower_bound}..{upper_bound} ({}) - {:b}0..{}", + &key[lower_bound..upper_bound], + upper_bound - lower_bound, + self.path.0, + self.path.len() + ); + self.path.starts_with(&key[lower_bound..upper_bound]) } /// Returns the common bit prefix between the edge node's path and the given key. @@ -216,7 +241,7 @@ impl EdgeNode { /// # Arguments /// /// * `key` - The key to get the common path from. - pub fn common_path(&self, key: &BitSlice) -> &BitSlice { + pub fn common_path(&self, key: &BitSlice) -> &BitSlice { let key_path = key.iter().skip(self.height as usize); let common_length = key_path .zip(self.path.0.iter()) @@ -227,12 +252,26 @@ impl EdgeNode { } } +pub fn hash_binary_node(left_hash: Felt, right_hash: Felt) -> Felt { + H::hash(&left_hash, &right_hash) +} +pub fn hash_edge_node(path: &Path, child_hash: Felt) -> Felt { + let mut bytes = [0u8; 32]; + bytes.view_bits_mut()[256 - path.len()..].copy_from_bitslice(path); + + let felt_path = Felt::from_bytes_be(&bytes); + let mut length = [0; 32]; + // Safe as len() is guaranteed to be <= 251 + length[31] = path.len() as u8; + + let length = Felt::from_bytes_be(&length); + H::hash(&child_hash, &felt_path) + length +} + #[test] fn test_path_matches_basic() { - let path = Path( - BitSlice::::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]) - .to_bitvec(), - ); + let path = + Path(BitSlice::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]).to_bitvec()); let edge = EdgeNode { hash: None, height: 0, @@ -240,16 +279,14 @@ fn test_path_matches_basic() { child: NodeHandle::Hash(Felt::ZERO), }; - let key = BitSlice::::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]); - assert!(edge.path_matches(key)); + let key = BitSlice::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]); + assert!(edge.path_matches(key, 0)); } #[test] fn test_path_matches_with_height() { - let path = Path( - BitSlice::::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]) - .to_bitvec(), - ); + let path = + Path(BitSlice::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]).to_bitvec()); let edge = EdgeNode { hash: None, height: 8, @@ -257,18 +294,14 @@ fn test_path_matches_with_height() { child: NodeHandle::Hash(Felt::ZERO), }; - let key = BitSlice::::from_slice(&[ - 0b10101010, 0b10101010, 0b01010101, 0b10101010, 0b01010101, - ]); - assert!(edge.path_matches(key)); + let key = BitSlice::from_slice(&[0b10101010, 0b10101010, 0b01010101, 0b10101010, 0b01010101]); + assert!(edge.path_matches(key, 8)); } #[test] fn test_path_matches_only_part_with_height() { - let path = Path( - BitSlice::::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]) - .to_bitvec(), - ); + let path = + Path(BitSlice::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]).to_bitvec()); let edge = EdgeNode { hash: None, height: 8, @@ -276,18 +309,16 @@ fn test_path_matches_only_part_with_height() { child: NodeHandle::Hash(Felt::ZERO), }; - let key = BitSlice::::from_slice(&[ + let key = BitSlice::from_slice(&[ 0b10101010, 0b10101010, 0b01010101, 0b10101010, 0b01010101, 0b10101010, ]); - assert!(edge.path_matches(key)); + assert!(edge.path_matches(key, 8)); } #[test] fn test_path_dont_match() { - let path = Path( - BitSlice::::from_slice(&[0b10111010, 0b01010101, 0b10101010, 0b01010101]) - .to_bitvec(), - ); + let path = + Path(BitSlice::from_slice(&[0b10111010, 0b01010101, 0b10101010, 0b01010101]).to_bitvec()); let edge = EdgeNode { hash: None, height: 0, @@ -295,18 +326,14 @@ fn test_path_dont_match() { child: NodeHandle::Hash(Felt::ZERO), }; - let key = BitSlice::::from_slice(&[ - 0b10101010, 0b01010101, 0b10101010, 0b01010101, 0b10101010, - ]); - assert!(!edge.path_matches(key)); + let key = BitSlice::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101, 0b10101010]); + assert!(!edge.path_matches(key, 0)); } #[test] fn test_common_path_basic() { - let path = Path( - BitSlice::::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]) - .to_bitvec(), - ); + let path = + Path(BitSlice::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]).to_bitvec()); let edge = EdgeNode { hash: None, height: 0, @@ -314,16 +341,14 @@ fn test_common_path_basic() { child: NodeHandle::Hash(Felt::ZERO), }; - let key = BitSlice::::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]); + let key = BitSlice::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]); assert_eq!(edge.common_path(key), &path.0); } #[test] fn test_common_path_only_part() { - let path = Path( - BitSlice::::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]) - .to_bitvec(), - ); + let path = + Path(BitSlice::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]).to_bitvec()); let edge = EdgeNode { hash: None, height: 0, @@ -331,19 +356,17 @@ fn test_common_path_only_part() { child: NodeHandle::Hash(Felt::ZERO), }; - let key = BitSlice::::from_slice(&[0b10101010, 0b01010101]); + let key = BitSlice::from_slice(&[0b10101010, 0b01010101]); assert_eq!( edge.common_path(key), - BitSlice::::from_slice(&[0b10101010, 0b01010101]) + BitSlice::from_slice(&[0b10101010, 0b01010101]) ); } #[test] fn test_common_path_part_with_height() { - let path = Path( - BitSlice::::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]) - .to_bitvec(), - ); + let path = + Path(BitSlice::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]).to_bitvec()); let edge = EdgeNode { hash: None, height: 8, @@ -351,19 +374,14 @@ fn test_common_path_part_with_height() { child: NodeHandle::Hash(Felt::ZERO), }; - let key = BitSlice::::from_slice(&[0b01010101, 0b10101010]); - assert_eq!( - edge.common_path(key), - BitSlice::::from_slice(&[0b10101010]) - ); + let key = BitSlice::from_slice(&[0b01010101, 0b10101010]); + assert_eq!(edge.common_path(key), BitSlice::from_slice(&[0b10101010])); } #[test] fn test_no_common_path() { - let path = Path( - BitSlice::::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]) - .to_bitvec(), - ); + let path = + Path(BitSlice::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]).to_bitvec()); let edge = EdgeNode { hash: None, height: 0, @@ -371,6 +389,6 @@ fn test_no_common_path() { child: NodeHandle::Hash(Felt::ZERO), }; - let key = BitSlice::::from_slice(&[0b01010101, 0b10101010]); - assert_eq!(edge.common_path(key), BitSlice::::empty()); + let key = BitSlice::from_slice(&[0b01010101, 0b10101010]); + assert_eq!(edge.common_path(key), BitSlice::empty()); } diff --git a/src/trie/merkle_tree.rs b/src/trie/merkle_tree.rs deleted file mode 100644 index 6776afe..0000000 --- a/src/trie/merkle_tree.rs +++ /dev/null @@ -1,2455 +0,0 @@ -use bitvec::{ - prelude::{BitSlice, BitVec, Msb0}, - view::BitView, -}; -use core::{fmt, marker::PhantomData}; -use core::{iter, mem}; -use derive_more::Constructor; -use parity_scale_codec::Decode; -#[cfg(feature = "std")] -use rayon::prelude::*; -use starknet_types_core::{felt::Felt, hash::StarkHash}; - -use crate::{ - error::BonsaiStorageError, format, hash_map, id::Id, vec, BonsaiDatabase, ByteVec, EncodeExt, - HashMap, HashSet, KeyValueDB, ToString, Vec, -}; - -use super::{ - merkle_node::{BinaryNode, Direction, EdgeNode, Node, NodeHandle, NodeId}, - path::Path, - trie_db::TrieKeyType, - TrieKey, -}; - -#[cfg(test)] -use log::trace; - -#[derive(Debug, PartialEq, Eq)] -pub enum Membership { - Member, - NonMember, -} - -/// Wrapper type for a [HashMap] object. (It's not really a wrapper it's a -/// copy of the type but we implement the necessary traits.) -#[derive(Clone, Debug, PartialEq, Eq, Default, Constructor)] -pub struct NodesMapping(pub(crate) HashMap); - -/// A node used in proof generated by the trie. -/// -/// See pathfinders merkle-tree crate for more information. -#[derive(Debug, Clone, PartialEq)] -pub enum ProofNode { - Binary { left: Felt, right: Felt }, - Edge { child: Felt, path: Path }, -} - -impl ProofNode { - pub fn hash(&self) -> Felt { - match self { - ProofNode::Binary { left, right } => H::hash(left, right), - ProofNode::Edge { child, path } => { - let mut bytes = [0u8; 32]; - bytes.view_bits_mut::()[256 - path.0.len()..].copy_from_bitslice(&path.0); - // SAFETY: path len is <= 251 - let path_hash = Felt::from_bytes_be(&bytes); - - let length = Felt::from(path.0.len() as u8); - H::hash(child, &path_hash) + length - } - } - } -} - -#[derive(Debug, Clone)] -pub(crate) enum RootHandle { - Empty, - Loaded(NodeId), -} - -pub(crate) struct MerkleTrees { - pub db: KeyValueDB, - pub trees: HashMap>, -} - -impl fmt::Debug - for MerkleTrees -{ - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("MerkleTrees") - .field("db", &self.db) - .field("trees", &self.trees) - .finish() - } -} - -#[cfg(feature = "bench")] -impl Clone - for MerkleTrees -{ - fn clone(&self) -> Self { - Self { - db: self.db.clone(), - trees: self.trees.clone(), - } - } -} - -impl MerkleTrees { - pub(crate) fn new(db: KeyValueDB) -> Self { - Self { - db, - trees: HashMap::new(), - } - } - - pub(crate) fn set( - &mut self, - identifier: &[u8], - key: &BitSlice, - value: Felt, - ) -> Result<(), BonsaiStorageError> { - let tree = self - .trees - .entry_ref(identifier) - .or_insert_with(|| MerkleTree::new(identifier.into())); - - tree.set(&self.db, key, value) - } - - pub(crate) fn get( - &self, - identifier: &[u8], - key: &BitSlice, - ) -> Result, BonsaiStorageError> { - if let Some(tree) = self.trees.get(identifier) { - tree.get(&self.db, key) - } else { - MerkleTree::::new(identifier.into()).get(&self.db, key) - } - } - - pub(crate) fn get_at( - &self, - identifier: &[u8], - key: &BitSlice, - id: CommitID, - ) -> Result, BonsaiStorageError> { - if let Some(tree) = self.trees.get(identifier) { - tree.get_at(&self.db, key, id) - } else { - MerkleTree::::new(identifier.into()).get_at(&self.db, key, id) - } - } - - pub(crate) fn contains( - &self, - identifier: &[u8], - key: &BitSlice, - ) -> Result> { - if let Some(tree) = self.trees.get(identifier) { - tree.contains(&self.db, key) - } else { - MerkleTree::::new(identifier.into()).contains(&self.db, key) - } - } - - pub(crate) fn db_mut(&mut self) -> &mut KeyValueDB { - &mut self.db - } - - pub(crate) fn reset_to_last_commit( - &mut self, - ) -> Result<(), BonsaiStorageError> { - self.trees.clear(); // just clear the map - Ok(()) - } - - pub(crate) fn db_ref(&self) -> &KeyValueDB { - &self.db - } - - pub(crate) fn root_hash( - &self, - identifier: &[u8], - ) -> Result> { - if let Some(tree) = self.trees.get(identifier) { - Ok(tree.root_hash(&self.db)?) - } else { - MerkleTree::::new(identifier.into()).root_hash(&self.db) - } - } - - 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() - // FIXME: this does not filter out keys values correctly for `HashMapDb` due - // to branches and leafs not being differenciated - .filter_map(|(key, _value)| { - if key.len() > identifier.len() { - Some(key[identifier.len() + 1..].into()) - } else { - None - } - }) - .collect() - }) - .map_err(|e| e.into()) - } - - #[allow(clippy::type_complexity)] - pub(crate) fn get_key_value_pairs( - &self, - identifier: &[u8], - ) -> Result, Vec)>, BonsaiStorageError> { - self.db - .db - .get_by_prefix(&crate::DatabaseKey::Flat(identifier)) - .map(|key_value_pairs| { - key_value_pairs - .into_iter() - // FIXME: this does not filter out keys values correctly for `HashMapDb` due - // to branches and leafs not being differenciated - .filter_map(|(key, value)| { - if key.len() > identifier.len() { - Some((key[identifier.len() + 1..].into(), value.into_vec())) - } else { - None - } - }) - .collect() - }) - .map_err(|e| e.into()) - } - - pub(crate) fn commit(&mut self) -> Result<(), BonsaiStorageError> { - #[cfg(not(feature = "std"))] - let db_changes = self - .trees - .iter_mut() - .map(|(_, tree)| tree.get_updates::()); - #[cfg(feature = "std")] - let db_changes = self - .trees - .par_iter_mut() - .map(|(_, tree)| tree.get_updates::()) - .collect_vec_list() - .into_iter() - .flatten(); - - let mut batch = self.db.create_batch(); - for changes in db_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> { - if let Some(tree) = self.trees.get(identifier) { - tree.get_proof(&self.db, key) - } else { - MerkleTree::::new(identifier.into()).get_proof(&self.db, key) - } - } - - pub(crate) fn get_identifiers(&self) -> Vec> { - self.trees.keys().cloned().map(ByteVec::into_vec).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 { - /// The root node. None means the node has not been loaded yet. - pub(crate) root_node: Option, - /// Identifier of the tree in the database. - pub(crate) identifier: ByteVec, - /// This storage is used to avoid modifying the underlying database each time during a commit. - pub(crate) storage_nodes: NodesMapping, - /// The id of the last node that has been added to the temporary storage. - pub(crate) latest_node_id: NodeId, - /// The list of nodes that should be removed from the underlying database during the next commit. - pub(crate) death_row: HashSet, - /// The list of leaves that have been modified during the current commit. - pub(crate) cache_leaf_modified: HashMap>, - /// The hasher used to hash the nodes. - _hasher: PhantomData, -} - -impl fmt::Debug for MerkleTree { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("MerkleTree") - .field("root_node", &self.root_node) - .field("identifier", &self.identifier) - .field("storage_nodes", &self.storage_nodes) - .field("latest_node_id", &self.latest_node_id) - .field("death_row", &self.death_row) - .field("cache_leaf_modified", &self.cache_leaf_modified) - .finish() - } -} - -// NB: #[derive(Clone)] does not work because it expands to an impl block which forces H: Clone, which Pedersen/Poseidon aren't. -#[cfg(feature = "bench")] -impl Clone for MerkleTree { - fn clone(&self) -> Self { - Self { - root_node: self.root_node.clone(), - identifier: self.identifier.clone(), - storage_nodes: self.storage_nodes.clone(), - latest_node_id: self.latest_node_id, - death_row: self.death_row.clone(), - cache_leaf_modified: self.cache_leaf_modified.clone(), - _hasher: PhantomData, - } - } -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub(crate) enum InsertOrRemove { - Insert(T), - Remove, -} -enum NodeOrFelt<'a> { - Node(&'a Node), - Felt(Felt), -} - -impl MerkleTree { - pub fn new(identifier: ByteVec) -> Self { - let nodes_mapping: HashMap = HashMap::new(); - Self { - // root_handle: NodeHandle::Hash(root), - // root_hash: root, - root_node: None, - identifier, - storage_nodes: NodesMapping(nodes_mapping), - latest_node_id: NodeId(0), - death_row: HashSet::new(), - cache_leaf_modified: HashMap::new(), - _hasher: PhantomData, - } - } - - // TODO: this does not accept &mut self because borrow splitting - // this needs to be moved into the nodes_storage api entirely and latest_node_id there too - fn load_db_node<'a, DB: BonsaiDatabase, ID: Id>( - storage_nodes: &'a mut NodesMapping, - latest_node_id: &mut NodeId, - death_row: &HashSet, - db: &KeyValueDB, - key: &TrieKey, - ) -> Result, BonsaiStorageError> { - if death_row.contains(key) { - return Ok(None); - } - - let node = db.get(key)?; - let Some(node) = node else { return Ok(None) }; - - let node = Node::decode(&mut node.as_slice())?; - let node_id = latest_node_id.next_id(); - // Insert and return reference at the same time. Entry occupied case should not be possible. - match storage_nodes.0.entry(node_id) { - hash_map::Entry::Occupied(_) => Err(BonsaiStorageError::Trie( - "duplicate node id in storage".to_string(), - )), - hash_map::Entry::Vacant(entry) => Ok(Some((node_id, entry.insert(node)))), - } - } - - /// Loads the root node or returns None if the tree is empty. - fn get_root_node<'a, DB: BonsaiDatabase, ID: Id>( - root_node: &mut Option, - storage_nodes: &'a mut NodesMapping, - latest_node_id: &mut NodeId, - death_row: &HashSet, - identifier: &[u8], - db: &KeyValueDB, - ) -> Result, BonsaiStorageError> { - match root_node { - Some(RootHandle::Loaded(id)) => { - let node = storage_nodes - .0 - .get_mut(&*id) - .ok_or(BonsaiStorageError::Trie( - "root node doesn't exist in the storage".to_string(), - ))?; - Ok(Some((*id, node))) - } - Some(RootHandle::Empty) => Ok(None), - None => { - // load the node - let node = Self::load_db_node( - storage_nodes, - latest_node_id, - death_row, - db, - &TrieKey::new(identifier, TrieKeyType::Trie, &[0]), - )?; - - match node { - Some((id, n)) => { - *root_node = Some(RootHandle::Loaded(id)); - Ok(Some((id, n))) - } - None => { - *root_node = Some(RootHandle::Empty); - Ok(None) - } - } - } - } - } - - /// # Panics - /// - /// Calling this function when the tree has uncommited changes is invalid as the hashes need to be recomputed. - pub fn root_hash( - &self, - db: &KeyValueDB, - ) -> Result> { - match self.root_node { - Some(RootHandle::Empty) => Ok(Felt::ZERO), - Some(RootHandle::Loaded(node_id)) => { - let node = self.storage_nodes.0.get(&node_id).ok_or_else(|| { - BonsaiStorageError::Trie("could not fetch root node from storage".into()) - })?; - node.hash().ok_or_else(|| { - BonsaiStorageError::Trie("the tree has uncommited changes".into()) - }) - } - None => { - let Some(node) = Self::get_trie_branch_in_db_from_path( - &self.death_row, - &self.identifier, - db, - &Path::default(), - )? - else { - return Ok(Felt::ZERO); - }; - // UNWRAP: the node has just been fetched - Ok(node.hash().unwrap()) - } - } - } - - pub fn cache_leaf_modified(&self) -> &HashMap> { - &self.cache_leaf_modified - } - - // /// Remove all the modifications that have been done since the last commit. - // pub fn reset_to_last_commit(&mut self) { - // self.death_row.clear(); - // self.latest_node_id.reset(); - // self.storage_nodes.0.clear(); - // self.cache_leaf_modified.clear(); - // self.root_node = None; - // } - - /// Calculate all the new hashes and the root hash. - #[allow(clippy::type_complexity)] - pub(crate) fn get_updates( - &mut self, - ) -> Result< - impl Iterator)>, - BonsaiStorageError, - > { - let mut updates = HashMap::new(); - for node_key in mem::take(&mut self.death_row) { - updates.insert(node_key, InsertOrRemove::Remove); - } - - if let Some(RootHandle::Loaded(node_id)) = &self.root_node { - // compute hashes - let mut hashes = vec![]; - self.compute_root_hash::(&mut hashes)?; - - // commit the tree - self.commit_subtree::( - &mut updates, - *node_id, - Path(BitVec::new()), - &mut hashes.drain(..), - )?; - } - - self.root_node = None; // unloaded - - for (key, value) in mem::take(&mut self.cache_leaf_modified) { - updates.insert( - TrieKey::new(&self.identifier, TrieKeyType::Flat, &key), - match value { - InsertOrRemove::Insert(value) => { - InsertOrRemove::Insert(value.encode_sbytevec()) - } - InsertOrRemove::Remove => InsertOrRemove::Remove, - }, - ); - } - self.latest_node_id.reset(); - - #[cfg(test)] - assert_eq!(self.storage_nodes.0, [].into()); // we should have visited the whole tree - - Ok(updates.into_iter()) - } - - // Commit a single merkle tree - #[cfg(test)] - pub(crate) fn commit( - &mut self, - db: &mut KeyValueDB, - ) -> Result<(), BonsaiStorageError> { - let db_changes = self.get_updates::()?; - - let mut batch = db.create_batch(); - for (key, value) in db_changes { - match value { - InsertOrRemove::Insert(value) => { - log::trace!("committing insert {:?} => {:?}", key, value); - db.insert(&key, &value, Some(&mut batch))?; - } - InsertOrRemove::Remove => { - log::trace!("committing remove {:?}", key); - db.remove(&key, Some(&mut batch))?; - } - } - } - db.write_batch(batch).unwrap(); - log::trace!("commit finished"); - - Ok(()) - } - - #[cfg(test)] - pub(crate) fn assert_empty(&self) { - assert_eq!(self.storage_nodes.0, [].into()); - } - - fn get_node_or_felt( - &self, - node_handle: &NodeHandle, - ) -> Result> { - let node_id = match node_handle { - NodeHandle::Hash(hash) => return Ok(NodeOrFelt::Felt(*hash)), - NodeHandle::InMemory(root_id) => root_id, - }; - let node = self - .storage_nodes - .0 - .get(node_id) - .ok_or(BonsaiStorageError::Trie( - "Couldn't fetch node in the temporary storage".to_string(), - ))?; - Ok(NodeOrFelt::Node(node)) - } - - fn compute_root_hash( - &self, - hashes: &mut Vec, - ) -> Result> { - let handle = match &self.root_node { - Some(RootHandle::Loaded(node_id)) => *node_id, - Some(RootHandle::Empty) => return Ok(Felt::ZERO), - None => { - return Err(BonsaiStorageError::Trie( - "root node is not loaded".to_string(), - )) - } - }; - let Some(node) = self.storage_nodes.0.get(&handle) else { - return Err(BonsaiStorageError::Trie( - "could not fetch root node from storage".to_string(), - )); - }; - self.compute_hashes::(node, Path(BitVec::new()), hashes) - } - - /// Compute the hashes of all of the updated nodes in the merkle tree. This step - /// is separate from [`commit_subtree`] as it is done in parallel using rayon. - /// Computed hashes are pushed to the `hashes` vector, depth first. - fn compute_hashes( - &self, - node: &Node, - path: Path, - hashes: &mut Vec, - ) -> Result> { - use Node::*; - - match node { - Binary(binary) => { - // we check if we have one or two changed children - - let left_path = path.new_with_direction(Direction::Left); - let node_left = self.get_node_or_felt::(&binary.left)?; - let right_path = path.new_with_direction(Direction::Right); - let node_right = self.get_node_or_felt::(&binary.right)?; - - let (left_hash, right_hash) = match (node_left, node_right) { - #[cfg(feature = "std")] - (NodeOrFelt::Node(left), NodeOrFelt::Node(right)) => { - // two children: use rayon - let (left, right) = rayon::join( - || self.compute_hashes::(left, left_path, hashes), - || { - let mut hashes = vec![]; - let felt = - self.compute_hashes::(right, right_path, &mut hashes)?; - Ok::<_, BonsaiStorageError>((felt, hashes)) - }, - ); - let (left_hash, (right_hash, hashes2)) = (left?, right?); - hashes.extend(hashes2); - - (left_hash, right_hash) - } - (left, right) => { - let left_hash = match left { - NodeOrFelt::Felt(felt) => felt, - NodeOrFelt::Node(node) => { - self.compute_hashes::(node, left_path, hashes)? - } - }; - let right_hash = match right { - NodeOrFelt::Felt(felt) => felt, - NodeOrFelt::Node(node) => { - self.compute_hashes::(node, right_path, hashes)? - } - }; - (left_hash, right_hash) - } - }; - - let hash = H::hash(&left_hash, &right_hash); - hashes.push(hash); - Ok(hash) - } - - Edge(edge) => { - let mut child_path = path.clone(); - child_path.0.extend(&edge.path.0); - let child_hash = match self.get_node_or_felt::(&edge.child)? { - NodeOrFelt::Felt(felt) => felt, - NodeOrFelt::Node(node) => { - self.compute_hashes::(node, child_path, hashes)? - } - }; - - let mut bytes = [0u8; 32]; - bytes.view_bits_mut::()[256 - edge.path.0.len()..] - .copy_from_bitslice(&edge.path.0); - - let felt_path = Felt::from_bytes_be(&bytes); - let mut length = [0; 32]; - // Safe as len() is guaranteed to be <= 251 - length[31] = edge.path.0.len() as u8; - - let length = Felt::from_bytes_be(&length); - let hash = H::hash(&child_hash, &felt_path) + length; - hashes.push(hash); - Ok(hash) - } - } - } - - /// Persists any changes in this subtree to storage. - /// - /// This necessitates recursively calculating the hash of, and - /// in turn persisting, any changed child nodes. This is necessary - /// as the parent node's hash relies on its children hashes. - /// Hash computation is done in parallel with [`compute_hashes`] beforehand. - /// - /// In effect, the entire tree gets persisted. - /// - /// # Arguments - /// - /// * `node_handle` - The top node from the subtree to commit. - /// * `hashes` - The precomputed hashes for the subtree as returned by [`compute_hashes`]. - /// The order is depth first, left to right. - /// - /// # Panics - /// - /// Panics if the precomputed `hashes` do not match the length of the modified subtree. - fn commit_subtree( - &mut self, - updates: &mut HashMap>, - node_id: NodeId, - path: Path, - hashes: &mut impl Iterator, - ) -> Result> { - match self - .storage_nodes - .0 - .remove(&node_id) - .ok_or(BonsaiStorageError::Trie( - "Couldn't fetch node in the temporary storage".to_string(), - ))? { - Node::Binary(mut binary) => { - let left_path = path.new_with_direction(Direction::Left); - let left_hash = match binary.left { - NodeHandle::Hash(left_hash) => left_hash, - NodeHandle::InMemory(node_id) => { - self.commit_subtree::(updates, node_id, left_path, hashes)? - } - }; - let right_path = path.new_with_direction(Direction::Right); - let right_hash = match binary.right { - NodeHandle::Hash(right_hash) => right_hash, - NodeHandle::InMemory(node_id) => { - self.commit_subtree::(updates, node_id, right_path, hashes)? - } - }; - - let hash = hashes.next().expect("mismatched hash state"); - - binary.hash = Some(hash); - binary.left = NodeHandle::Hash(left_hash); - binary.right = NodeHandle::Hash(right_hash); - let key_bytes: ByteVec = path.into(); - updates.insert( - TrieKey::new(&self.identifier, TrieKeyType::Trie, &key_bytes), - InsertOrRemove::Insert(Node::Binary(binary).encode_sbytevec()), - ); - Ok(hash) - } - Node::Edge(mut edge) => { - let mut child_path = path.clone(); - child_path.0.extend(&edge.path.0); - let child_hash = match edge.child { - NodeHandle::Hash(right_hash) => right_hash, - NodeHandle::InMemory(node_id) => { - self.commit_subtree::(updates, node_id, child_path, hashes)? - } - }; - let hash = hashes.next().expect("mismatched hash state"); - edge.hash = Some(hash); - edge.child = NodeHandle::Hash(child_hash); - let key_bytes: ByteVec = path.into(); - updates.insert( - TrieKey::new(&self.identifier, TrieKeyType::Trie, &key_bytes), - InsertOrRemove::Insert(Node::Edge(edge).encode_sbytevec()), - ); - Ok(hash) - } - } - } - - /// Sets the value of a key. To delete a key, set the value to [Felt::ZERO]. - /// - /// # Arguments - /// - /// * `key` - The key to set. - /// * `value` - The value to set. - pub fn set( - &mut self, - db: &KeyValueDB, - key: &BitSlice, - value: Felt, - ) -> Result<(), BonsaiStorageError> { - if value == Felt::ZERO { - return self.delete_leaf(db, key); - } - let key_bytes = bitslice_to_bytes(key); - log::trace!("key_bytes: {:?}", key_bytes); - - // TODO(perf): do not double lookup when changing the value later (borrow needs to be split for preload_nodes though) - let mut cache_leaf_entry = self.cache_leaf_modified.entry_ref(&key_bytes[..]); - - if let hash_map::EntryRef::Occupied(entry) = &mut cache_leaf_entry { - if matches!(entry.get(), InsertOrRemove::Insert(_)) { - entry.insert(InsertOrRemove::Insert(value)); - return Ok(()); - } - } - - 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_nodes, _path) = self.preload_nodes(db, key)?; - // There are three possibilities. - // - // 1. The leaf exists, in which case we simply change its value. - // - // 2. The tree is empty, we insert the new leaf and the root becomes an edge node connecting to it. - // - // 3. The leaf does not exist, and the tree is not empty. The final node in the traversal will be an - // edge node who's path diverges from our new leaf node's. - // - // This edge must be split into a new subtree containing both the existing edge's child and the - // new leaf. This requires an edge followed by a binary node and then further edges to both the - // current child and the new leaf. Any of these new edges may also end with an empty path in - // which case they should be elided. It depends on the common path length of the current edge - // and the new leaf i.e. the split may be at the first bit (in which case there is no leading - // edge), or the split may be in the middle (requires both leading and post edges), or the - // split may be the final bit (no post edge). - - log::trace!("preload nodes: {:?}", path_nodes); - use Node::*; - match path_nodes.last() { - Some(node_id) => { - let mut nodes_to_add = Vec::new(); - self.storage_nodes.0.entry(*node_id).and_modify(|node| { - 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. - log::trace!("change val: {:?} => {:#x}", key_bytes, 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). - log::trace!( - "cache_leaf_modified insert: {:?} => {:#x}", - key_bytes, - value - ); - 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_direction = Direction::from(key[branch_height]); - let (left, right) = match new_direction { - Direction::Left => (new, old), - Direction::Right => (old, new), - }; - - 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)); - - Node::Edge(EdgeNode { - hash: None, - height: edge.height, - path: Path(common.to_bitvec()), - child: NodeHandle::InMemory(branch_id), - }) - }; - let key_bytes = bitslice_to_bytes(&key[..edge.height as usize]); - log::trace!("2 death row add ({:?})", key_bytes); - self.death_row.insert(TrieKey::Trie(key_bytes)); - *node = new_node; - } - Binary(binary) => { - let child_height = binary.height + 1; - - if child_height 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), - }; - self.cache_leaf_modified - .insert(key_bytes, InsertOrRemove::Insert(value)); - } - } - } - }); - self.storage_nodes.0.extend(nodes_to_add); - Ok(()) - } - None => { - // Getting no travel nodes implies that the tree is empty. - // - // Create a new leaf node with the value, and the root becomes - // an edge node connecting to the leaf. - let edge = Node::Edge(EdgeNode { - hash: None, - height: 0, - path: Path(key.to_bitvec()), - child: NodeHandle::Hash(value), - }); - self.storage_nodes - .0 - .insert(self.latest_node_id.next_id(), edge); - - self.root_node = Some(RootHandle::Loaded(self.latest_node_id)); - - let key_bytes = bitslice_to_bytes(key); - self.cache_leaf_modified - .insert(key_bytes, InsertOrRemove::Insert(value)); - Ok(()) - } - } - } - - /// Deletes a leaf node from the tree. - /// - /// This is not an external facing API; the functionality is instead accessed by calling - /// [`MerkleTree::set`] with value set to [`Felt::ZERO`]. - /// - /// # Arguments - /// - /// * `key` - The key to delete. - fn delete_leaf( - &mut self, - db: &KeyValueDB, - key: &BitSlice, - ) -> Result<(), BonsaiStorageError> { - // Algorithm explanation: - // - // The leaf's parent node is either an edge, or a binary node. - // If it's an edge node, then it must also be deleted. And its parent - // must be a binary node. In either case we end up with a binary node - // who's one child is deleted. This changes the binary to an edge node. - // - // Note that its possible that there is no binary node -- if the resulting tree would be empty. - // - // This new edge node may need to merge with the old binary node's parent node - // and other remaining child node -- if they're also edges. - // - // Then we are done. - let key_bytes = bitslice_to_bytes(key); - let leaf_entry = self.cache_leaf_modified.entry(key_bytes.clone()); - - let tree_has_value = if let hash_map::Entry::Occupied(entry) = &leaf_entry { - !matches!(entry.get(), InsertOrRemove::Remove) - } else { - db.get(&TrieKey::new( - &self.identifier, - TrieKeyType::Flat, - &key_bytes, - ))? - .is_some() - }; - - if !tree_has_value { - return Ok(()); - } - leaf_entry.insert(InsertOrRemove::Remove); - - let (path_nodes, _path) = self.preload_nodes(db, key)?; - - let mut last_binary_path = Path(key.to_bitvec()); - - // Go backwards until we hit a branch node. - let mut node_iter = path_nodes.into_iter().rev().skip_while(|node| { - let node = match self.storage_nodes.0.entry(*node) { - hash_map::Entry::Occupied(entry) => entry, - // SAFETY: Has been populate by preload_nodes just above - hash_map::Entry::Vacant(_) => unreachable!(), - }; - - match node.get() { - Node::Binary(_) => false, - Node::Edge(edge) => { - for _ in 0..edge.path.0.len() { - last_binary_path.0.pop(); - } - 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.clone(); - let path: ByteVec = (&last_binary_path).into(); - log::trace!( - "iter leaf={:?} edge={edge:?}, new_path={new_path:?}", - TrieKey::new(&self.identifier, TrieKeyType::Trie, &path) - ); - - self.death_row - .insert(TrieKey::new(&self.identifier, TrieKeyType::Trie, &path)); - node.remove(); - - true - } - } - }); - let branch_node = node_iter.next(); - let parent_branch_node = node_iter.next(); - - log::trace!( - "remove leaf branch_node={branch_node:?} parent_branch_node={parent_branch_node:?}" - ); - - match branch_node { - Some(node_id) => { - let (new_edge, par_path) = - { - let node = self.storage_nodes.0.get_mut(&node_id).ok_or( - BonsaiStorageError::Trie("Node not found in memory".to_string()), - )?; - - // 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. - let path = - Path(iter::once(bool::from(direction)).collect::>()); - let mut edge = EdgeNode { - hash: None, - height, - path, - 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, db, &last_binary_path)?; - let cl = last_binary_path.clone(); - last_binary_path.0.pop(); - (edge, cl) - }; - // 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; - - let mut par_path = par_path; - par_path.0.pop(); - let path: ByteVec = par_path.into(); - self.death_row.insert(TrieKey::new( - &self.identifier, - TrieKeyType::Trie, - &path, - )); - self.storage_nodes.0.remove(&node_id); // very sad hashbrown doesn't have a get_many_entries api, we have to double-lookup - } 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 - // must therefore be empty. - - log::trace!("empty {:?}", self.root_node); - if let Some(RootHandle::Loaded(node_id)) = self.root_node { - self.storage_nodes.0.remove(&node_id); - } - self.death_row - .insert(TrieKey::new(&self.identifier, TrieKeyType::Trie, &[0])); - self.root_node = Some(RootHandle::Empty); - return Ok(()); - } - }; - Ok(()) - } - - /// Returns the value stored at key, or `None` if it does not exist. - /// - /// # Arguments - /// - /// * `key` - The key of the value to get. - /// - /// # Returns - /// - /// The value of the key. - pub fn get( - &self, - db: &KeyValueDB, - key: &BitSlice, - ) -> Result, BonsaiStorageError> { - log::trace!("get with key {:b}", key); - let key = bitslice_to_bytes(key); - log::trace!("get from cache with {:?}", key); - if let Some(value) = self.cache_leaf_modified.get(&key) { - log::trace!("get has cache_leaf_modified {:?} {:?}", key, value); - match value { - InsertOrRemove::Remove => return Ok(None), - InsertOrRemove::Insert(value) => return Ok(Some(*value)), - } - } - log::trace!( - "get from db with key {:?}", - &TrieKey::new(&self.identifier, TrieKeyType::Flat, &key) - ); - db.get(&TrieKey::new(&self.identifier, TrieKeyType::Flat, &key)) - .map(|r| r.map(|opt| Felt::decode(&mut opt.as_slice()).unwrap())) - } - - pub fn get_at( - &self, - db: &KeyValueDB, - key: &BitSlice, - id: ID, - ) -> Result, BonsaiStorageError> { - let key = bitslice_to_bytes(key); - db.get_at(&TrieKey::new(&self.identifier, TrieKeyType::Flat, &key), id) - .map(|r| r.map(|opt| Felt::decode(&mut opt.as_slice()).unwrap())) - } - - pub fn contains( - &self, - db: &KeyValueDB, - key: &BitSlice, - ) -> Result> { - let key = bitslice_to_bytes(key); - if let Some(value) = self.cache_leaf_modified.get(&key) { - match value { - InsertOrRemove::Remove => return Ok(false), - InsertOrRemove::Insert(_) => return Ok(true), - } - } - db.contains(&TrieKey::new(&self.identifier, TrieKeyType::Flat, &key)) - } - - /// Returns the list of nodes along the path. - /// - /// if it exists, or down to the node which proves that the key does not exist. - /// - /// The nodes are returned in order, root first. - /// - /// Verification is performed by confirming that: - /// 1. the chain follows the path of `key`, and - /// 2. the hashes are correct, and - /// 3. the root hash matches the known root - /// - /// # Arguments - /// - /// * `key` - The key to get the merkle proof of. - /// - /// # Returns - /// - /// The merkle proof and all the child nodes hashes. - pub fn get_proof( - &self, - db: &KeyValueDB, - key: &BitSlice, - ) -> Result, BonsaiStorageError> { - let mut nodes = Vec::with_capacity(251); - let mut node = match self.root_node { - Some(RootHandle::Empty) => { - return Ok(Vec::new()); - } - Some(RootHandle::Loaded(node_id)) => self - .storage_nodes - .0 - .get(&node_id) - .ok_or(BonsaiStorageError::Trie( - "Couldn't get root node from storage".to_string(), - ))? - .clone(), - None => Self::get_trie_branch_in_db_from_path( - &self.death_row, - &self.identifier, - db, - &Path(BitVec::::new()), - )? - .ok_or(BonsaiStorageError::Trie( - "Couldn't fetch root node in db".to_string(), - ))?, - }; - loop { - match node { - Node::Edge(edge) => { - 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( - &self.death_row, - &self.identifier, - db, - &Path(child_path), - )?; - if let Some(node) = node { - node - } else { - nodes.push(ProofNode::Edge { - child: hash, - path: edge.path.clone(), - }); - return Ok(nodes); - } - } - NodeHandle::InMemory(child_id) => self - .storage_nodes - .0 - .get(&child_id) - .ok_or(BonsaiStorageError::Trie( - "Couldn't fetch child node in the temporary storage".to_string(), - ))? - .clone(), - }; - nodes.push(ProofNode::Edge { - child: child_node.hash().ok_or(BonsaiStorageError::Trie( - "Couldn't fetch child node in the temporary storage".to_string(), - ))?, - path: edge.path.clone(), - }); - if edge.path_matches(key) { - node = child_node; - } else { - return Ok(nodes); - } - if edge.common_path(key) == key { - return Ok(nodes); - } - } - Node::Binary(binary) => { - let next_direction = key - .get(binary.height as usize) - .map(|b| Direction::from(*b)) - .ok_or(BonsaiStorageError::Trie("Key too short".to_string()))?; - let next = binary.get_child(next_direction); - let next_path = key[..binary.height as usize + 1].to_bitvec(); - let next_node = match next { - NodeHandle::Hash(_) => Self::get_trie_branch_in_db_from_path( - &self.death_row, - &self.identifier, - db, - &Path(next_path), - )? - .ok_or(BonsaiStorageError::Trie( - "Couldn't fetch next node in db".to_string(), - ))?, - NodeHandle::InMemory(next_id) => self - .storage_nodes - .0 - .get(&next_id) - .ok_or(BonsaiStorageError::Trie( - "Couldn't fetch next node in the temporary storage".to_string(), - ))? - .clone(), - }; - let other = binary.get_child(next_direction.invert()); - let other_hash = match other { - NodeHandle::Hash(hash) => hash, - NodeHandle::InMemory(other_id) => { - let other_node = self - .storage_nodes - .0 - .get(&other_id) - .ok_or(BonsaiStorageError::Trie( - "Couldn't fetch other node in the temporary storage" - .to_string(), - ))? - .clone(); - other_node.hash().ok_or(BonsaiStorageError::Trie( - "Couldn't fetch other node in the temporary storage".to_string(), - ))? - } - }; - match next_direction { - Direction::Left => { - nodes.push(ProofNode::Binary { - left: next_node.hash().ok_or(BonsaiStorageError::Trie( - "Couldn't fetch next node in the temporary storage".to_string(), - ))?, - right: other_hash, - }); - } - Direction::Right => { - nodes.push(ProofNode::Binary { - left: other_hash, - right: next_node.hash().ok_or(BonsaiStorageError::Trie( - "Couldn't fetch next node in the temporary storage".to_string(), - ))?, - }); - } - } - node = next_node; - } // Node::Unresolved(hash) => { - // nodes.push(ProofNode::Edge { - // child: hash, - // path: Path(BitVec::::new()), - // }); - // return Ok(nodes); - // } - } - } - } - - /// preload_nodes from the current root towards the destination [Leaf](Node::Leaf) node. - /// If the destination node exists, it will be the final node in the list. - /// - /// This means that the final node will always be either a the destination [Leaf](Node::Leaf) - /// node, or an [Edge](Node::Edge) node who's path suffix does not match the leaf's path. - /// - /// The final node can __not__ be a [Binary](Node::Binary) node since it would always be - /// possible to continue on towards the destination. Nor can it be an - /// [Unresolved](Node::Unresolved) node since this would be resolved to check if we can - /// travel further. - /// - /// # Arguments - /// - /// * `dst` - The node to get to. - /// - /// # Returns - /// - /// The list of nodes along the path. - fn preload_nodes( - &mut self, - db: &KeyValueDB, - dst: &BitSlice, - ) -> Result<(Vec, Path), BonsaiStorageError> { - let mut nodes = Vec::with_capacity(251); - let mut path = Path(BitVec::::with_capacity(251)); - - let mut prev_handle = None::<&mut NodeHandle>; // None signals tree root - - loop { - // get node from cache or database - let (node_id, node) = match prev_handle { - // tree root - None => { - match Self::get_root_node( - &mut self.root_node, - &mut self.storage_nodes, - &mut self.latest_node_id, - &self.death_row, - &self.identifier, - db, - )? { - Some((node_id, node)) => (node_id, node), - None => { - // empty tree - return Ok((nodes, path)); - } - } - } - // not tree root - Some(prev_handle) => match prev_handle { - NodeHandle::Hash(_) => { - // load from db - let Some(node) = Self::get_trie_branch_in_db_from_path( - &self.death_row, - &self.identifier, - db, - &path, - )? - else { - // end of path traversal - break; - }; - - // put it in inmemory storage - self.latest_node_id.next_id(); - *prev_handle = NodeHandle::InMemory(self.latest_node_id); - let node = self.storage_nodes.0.entry(self.latest_node_id).insert(node); - - (self.latest_node_id, node.into_mut()) - } - NodeHandle::InMemory(node_id) => { - let node_id = *node_id; - - let node = self.storage_nodes.0.get_mut(&node_id).ok_or( - BonsaiStorageError::Trie( - "Couldn't get node from temp storage".to_string(), - ), - )?; - - (node_id, node) - } - }, - }; - - nodes.push(node_id); - - // visit the child - match node { - Node::Binary(binary_node) => { - let next_direction = binary_node.direction(dst); - path.0.push(bool::from(next_direction)); - prev_handle = Some(binary_node.get_child_mut(next_direction)); - } - - Node::Edge(edge_node) if edge_node.path_matches(dst) => { - path.0.extend_from_bitslice(&edge_node.path.0); - if path.0 == dst { - break; // found it :) - } - - prev_handle = Some(&mut edge_node.child); - } - - // We are in a case where the edge node doesn't match the path we want to preload so we return nothing. - Node::Edge(_) => break, - } - } - - Ok((nodes, path)) - } - - /// Get the node of the trie that corresponds to the path. - fn get_trie_branch_in_db_from_path( - death_row: &HashSet, - identifier: &[u8], - db: &KeyValueDB, - path: &Path, - ) -> Result, BonsaiStorageError> { - log::trace!("getting: {:b}", path.0); - - let path: ByteVec = path.into(); - let key = TrieKey::new(identifier, TrieKeyType::Trie, &path); - - if death_row.contains(&key) { - return Ok(None); - } - - db.get(&key)? - .map(|node| { - log::trace!("got: {:?}", node); - Node::decode(&mut node.as_slice()).map_err(|err| { - BonsaiStorageError::Trie(format!("Couldn't decode node: {}", err)) - }) - }) - .map_or(Ok(None), |r| r.map(Some)) - } - - /// This is a convenience function which merges the edge node with its child __iff__ it is also - /// an edge. - /// - /// Does nothing if the child is not also an edge node. - /// - /// This can occur when mutating the tree (e.g. deleting a child of a binary node), and is an - /// illegal state (since edge nodes __must be__ maximal subtrees). - /// - /// # Arguments - /// - /// * `parent` - The parent node to merge the child with. - fn merge_edges( - &mut self, - parent: &mut EdgeNode, - db: &KeyValueDB, - path: &Path, - ) -> Result<(), BonsaiStorageError> { - match parent.child { - NodeHandle::Hash(_) => { - let node = Self::get_trie_branch_in_db_from_path( - &self.death_row, - &self.identifier, - db, - path, - )?; - log::trace!("case: Hash {:?}", node); - if let Some(Node::Edge(child_edge)) = node { - parent.path.0.extend_from_bitslice(&child_edge.path.0); - parent.child = child_edge.child; - // remove node from db - let path: ByteVec = path.into(); - log::trace!("4 death row {:?}", path); - self.death_row - .insert(TrieKey::new(&self.identifier, TrieKeyType::Trie, &path)); - } - } - NodeHandle::InMemory(child_id) => { - let node = match self.storage_nodes.0.entry(child_id) { - hash_map::Entry::Occupied(entry) => entry, - hash_map::Entry::Vacant(_) => { - return Err(BonsaiStorageError::Trie("getting node from memory".into())) - } - }; - log::trace!("case: InMemory {:?}", node.get()); - - if let Node::Edge(child_edge) = node.get() { - parent.path.0.extend_from_bitslice(&child_edge.path.0); - parent.child = child_edge.child; - - node.remove(); - - let path: ByteVec = path.into(); - log::trace!("3 death row {:?}", path); - self.death_row - .insert(TrieKey::new(&self.identifier, TrieKeyType::Trie, &path)); - } - } - }; - Ok(()) - } - - /// Function that come from pathfinder_merkle_tree::merkle_tree::MerkleTree - /// Verifies that the key `key` with value `value` is indeed part of the MPT that has root - /// `root`, given `proofs`. - /// Supports proofs of non-membership as well as proof of membership: this function returns - /// an enum corresponding to the membership of `value`, or returns `None` in case of a hash mismatch. - /// The algorithm follows this logic: - /// 1. init expected_hash <- root hash - /// 2. loop over nodes: current <- nodes[i] - /// 1. verify the current node's hash matches expected_hash (if not then we have a bad proof) - /// 2. move towards the target - if current is: - /// 1. binary node then choose the child that moves towards the target, else if - /// 2. edge node then check the path against the target bits - /// 1. If it matches then proceed with the child, else - /// 2. if it does not match then we now have a proof that the target does not exist - /// 3. nibble off target bits according to which child you got in (2). If all bits are gone then you - /// have reached the target and the child hash is the value you wanted and the proof is complete. - /// 4. set expected_hash <- to the child hash - /// 3. check that the expected_hash is `value` (we should've reached the leaf) - pub fn verify_proof( - root: Felt, - key: &BitSlice, - value: Felt, - proofs: &[ProofNode], - ) -> Option { - // Protect from ill-formed keys - if key.len() > 251 { - return None; - } - - let mut expected_hash = root; - let mut remaining_path: &BitSlice = key; - - for proof_node in proofs.iter() { - // Hash mismatch? Return None. - if proof_node.hash::() != expected_hash { - return None; - } - match proof_node { - ProofNode::Binary { left, right } => { - // Direction will always correspond to the 0th index - // because we're removing bits on every iteration. - let direction = Direction::from(remaining_path[0]); - - // Set the next hash to be the left or right hash, - // depending on the direction - expected_hash = match direction { - Direction::Left => *left, - Direction::Right => *right, - }; - - // Advance by a single bit - remaining_path = &remaining_path[1..]; - } - ProofNode::Edge { child, path } => { - if path.0 != remaining_path[..path.0.len()] { - // If paths don't match, we've found a proof of non membership because we: - // 1. Correctly moved towards the target insofar as is possible, and - // 2. hashing all the nodes along the path does result in the root hash, which means - // 3. the target definitely does not exist in this tree - return Some(Membership::NonMember); - } - - // Set the next hash to the child's hash - expected_hash = *child; - - // Advance by the whole edge path - remaining_path = &remaining_path[path.0.len()..]; - } - } - } - - // At this point, we should reach `value` ! - if expected_hash == value { - Some(Membership::Member) - } else { - // Hash mismatch. Return `None`. - None - } - } - - #[cfg(test)] - #[allow(dead_code)] - pub(crate) fn display(&self) { - match self.root_node { - Some(RootHandle::Empty) => { - trace!("tree is empty") - } - Some(RootHandle::Loaded(node)) => { - trace!("root is node {:?}", node); - self.print(&node); - } - None => trace!("root is not loaded"), - } - } - - #[cfg(test)] - #[allow(dead_code)] - fn print(&self, head: &NodeId) { - use Node::*; - - let current_tmp = self.storage_nodes.0.get(head).unwrap().clone(); - trace!("bonsai_node {:?} = {:?}", head, current_tmp); - - match current_tmp { - Binary(binary) => { - match &binary.get_child(Direction::Left) { - NodeHandle::Hash(hash) => { - trace!("left is hash {:#x}", hash); - } - NodeHandle::InMemory(left_id) => { - self.print(left_id); - } - } - match &binary.get_child(Direction::Right) { - NodeHandle::Hash(hash) => { - trace!("right is hash {:#x}", hash); - } - NodeHandle::InMemory(right_id) => { - self.print(right_id); - } - } - } - Edge(edge) => match &edge.child { - NodeHandle::Hash(hash) => { - trace!("child is hash {:#x}", hash); - } - NodeHandle::InMemory(child_id) => { - self.print(child_id); - } - }, - }; - } -} - -pub(crate) fn bitslice_to_bytes(bitslice: &BitSlice) -> ByteVec { - // TODO(perf): this should not copy to a bitvec :( - if bitslice.is_empty() { - return Default::default(); - } // special case: tree root - iter::once(bitslice.len() as u8) - .chain(bitslice.to_bitvec().as_raw_slice().iter().copied()) - .collect() -} - -pub(crate) fn bytes_to_bitvec(bytes: &[u8]) -> BitVec { - BitSlice::from_slice(&bytes[1..]).to_bitvec() -} - -#[cfg(test)] -#[cfg(all(test, feature = "std", feature = "rocksdb"))] -mod tests { - use bitvec::{order::Msb0, vec::BitVec, view::BitView}; - use indexmap::IndexMap; - use starknet_types_core::{felt::Felt, hash::Pedersen}; - - use crate::{ - databases::{create_rocks_db, RocksDB, RocksDBConfig}, - id::BasicId, - BonsaiStorage, BonsaiStorageConfig, ByteVec, - }; - - #[test_log::test] - // The whole point of this test is to make sure it is possible to reconstruct the original - // keys from the data present in the db. - fn test_key_retrieval() { - let tempdir = tempfile::tempdir().unwrap(); - let rocksdb = create_rocks_db(tempdir.path()).unwrap(); - let db = RocksDB::new(&rocksdb, RocksDBConfig::default()); - let mut bonsai = - BonsaiStorage::::new(db, BonsaiStorageConfig::default()).unwrap(); - - let block_0 = vec![ - ( - str_to_felt_bytes( - "0x031c887d82502ceb218c06ebb46198da3f7b92864a8223746bc836dda3e34b52", - ), - vec![ - ( - str_to_felt_bytes( - "0x0000000000000000000000000000000000000000000000000000000000000005", - ), - str_to_felt_bytes( - "0x0000000000000000000000000000000000000000000000000000000000000065", - ), - ), - ( - str_to_felt_bytes( - "0x00cfc2e2866fd08bfb4ac73b70e0c136e326ae18fc797a2c090c8811c695577e", - ), - str_to_felt_bytes( - "0x05f1dd5a5aef88e0498eeca4e7b2ea0fa7110608c11531278742f0b5499af4b3", - ), - ), - ( - str_to_felt_bytes( - "0x05aee31408163292105d875070f98cb48275b8c87e80380b78d30647e05854d5", - ), - str_to_felt_bytes( - "0x00000000000000000000000000000000000000000000000000000000000007c7", - ), - ), - ( - str_to_felt_bytes( - "0x05fac6815fddf6af1ca5e592359862ede14f171e1544fd9e792288164097c35d", - ), - str_to_felt_bytes( - "0x00299e2f4b5a873e95e65eb03d31e532ea2cde43b498b50cd3161145db5542a5", - ), - ), - ( - str_to_felt_bytes( - "0x05fac6815fddf6af1ca5e592359862ede14f171e1544fd9e792288164097c35e", - ), - str_to_felt_bytes( - "0x03d6897cf23da3bf4fd35cc7a43ccaf7c5eaf8f7c5b9031ac9b09a929204175f", - ), - ), - ], - ), - ( - str_to_felt_bytes( - "0x06ee3440b08a9c805305449ec7f7003f27e9f7e287b83610952ec36bdc5a6bae", - ), - vec![ - ( - str_to_felt_bytes( - "0x01e2cd4b3588e8f6f9c4e89fb0e293bf92018c96d7a93ee367d29a284223b6ff", - ), - str_to_felt_bytes( - "0x071d1e9d188c784a0bde95c1d508877a0d93e9102b37213d1e13f3ebc54a7751", - ), - ), - ( - str_to_felt_bytes( - "0x0449908c349e90f81ab13042b1e49dc251eb6e3e51092d9a40f86859f7f415b0", - ), - str_to_felt_bytes( - "0x06cb6104279e754967a721b52bcf5be525fdc11fa6db6ef5c3a4db832acf7804", - ), - ), - ( - str_to_felt_bytes( - "0x048cba68d4e86764105adcdcf641ab67b581a55a4f367203647549c8bf1feea2", - ), - str_to_felt_bytes( - "0x0362d24a3b030998ac75e838955dfee19ec5b6eceb235b9bfbeccf51b6304d0b", - ), - ), - ( - str_to_felt_bytes( - "0x05bdaf1d47b176bfcd1114809af85a46b9c4376e87e361d86536f0288a284b65", - ), - str_to_felt_bytes( - "0x028dff6722aa73281b2cf84cac09950b71fa90512db294d2042119abdd9f4b87", - ), - ), - ( - str_to_felt_bytes( - "0x05bdaf1d47b176bfcd1114809af85a46b9c4376e87e361d86536f0288a284b66", - ), - str_to_felt_bytes( - "0x057a8f8a019ccab5bfc6ff86c96b1392257abb8d5d110c01d326b94247af161c", - ), - ), - ( - str_to_felt_bytes( - "0x05f750dc13ed239fa6fc43ff6e10ae9125a33bd05ec034fc3bb4dd168df3505f", - ), - str_to_felt_bytes( - "0x00000000000000000000000000000000000000000000000000000000000007e5", - ), - ), - ], - ), - ( - str_to_felt_bytes( - "0x0735596016a37ee972c42adef6a3cf628c19bb3794369c65d2c82ba034aecf2c", - ), - vec![ - ( - str_to_felt_bytes( - "0x0000000000000000000000000000000000000000000000000000000000000005", - ), - str_to_felt_bytes( - "0x0000000000000000000000000000000000000000000000000000000000000064", - ), - ), - ( - str_to_felt_bytes( - "0x002f50710449a06a9fa789b3c029a63bd0b1f722f46505828a9f815cf91b31d8", - ), - str_to_felt_bytes( - "0x02a222e62eabe91abdb6838fa8b267ffe81a6eb575f61e96ec9aa4460c0925a2", - ), - ), - ], - ), - ( - str_to_felt_bytes( - "0x020cfa74ee3564b4cd5435cdace0f9c4d43b939620e4a0bb5076105df0a626c6", - ), - vec![ - ( - str_to_felt_bytes( - "0x0000000000000000000000000000000000000000000000000000000000000005", - ), - str_to_felt_bytes( - "0x000000000000000000000000000000000000000000000000000000000000022b", - ), - ), - ( - str_to_felt_bytes( - "0x0313ad57fdf765addc71329abf8d74ac2bce6d46da8c2b9b82255a5076620300", - ), - str_to_felt_bytes( - "0x04e7e989d58a17cd279eca440c5eaa829efb6f9967aaad89022acbe644c39b36", - ), - ), - ( - str_to_felt_bytes( - "0x0313ad57fdf765addc71329abf8d74ac2bce6d46da8c2b9b82255a5076620301", - ), - str_to_felt_bytes( - "0x0453ae0c9610197b18b13645c44d3d0a407083d96562e8752aab3fab616cecb0", - ), - ), - ( - str_to_felt_bytes( - "0x05aee31408163292105d875070f98cb48275b8c87e80380b78d30647e05854d5", - ), - str_to_felt_bytes( - "0x00000000000000000000000000000000000000000000000000000000000007e5", - ), - ), - ( - str_to_felt_bytes( - "0x06cf6c2f36d36b08e591e4489e92ca882bb67b9c39a3afccf011972a8de467f0", - ), - str_to_felt_bytes( - "0x07ab344d88124307c07b56f6c59c12f4543e9c96398727854a322dea82c73240", - ), - ), - ], - ), - ( - str_to_felt_bytes( - "0x031c887d82502ceb218c06ebb46198da3f7b92864a8223746bc836dda3e34b52", - ), - vec![ - ( - str_to_felt_bytes( - "0x00df28e613c065616a2e79ca72f9c1908e17b8c913972a9993da77588dc9cae9", - ), - str_to_felt_bytes( - "0x01432126ac23c7028200e443169c2286f99cdb5a7bf22e607bcd724efa059040", - ), - ), - ( - str_to_felt_bytes( - "0x05f750dc13ed239fa6fc43ff6e10ae9125a33bd05ec034fc3bb4dd168df3505f", - ), - str_to_felt_bytes( - "0x00000000000000000000000000000000000000000000000000000000000007c7", - ), - ), - ], - ), - ]; - - let block_1 = [ - ( - str_to_felt_bytes( - "0x06538fdd3aa353af8a87f5fe77d1f533ea82815076e30a86d65b72d3eb4f0b80", - ), - vec![ - ( - str_to_felt_bytes( - "0x0000000000000000000000000000000000000000000000000000000000000005", - ), - str_to_felt_bytes( - "0x000000000000000000000000000000000000000000000000000000000000022b", - ), - ), - ( - str_to_felt_bytes( - "0x01aed933fd362faecd8ea54ee749092bd21f89901b7d1872312584ac5b636c6d", - ), - str_to_felt_bytes( - "0x00000000000000000000000000000000000000000000000000000000000007e5", - ), - ), - ( - str_to_felt_bytes( - "0x010212fa2be788e5d943714d6a9eac5e07d8b4b48ead96b8d0a0cbe7a6dc3832", - ), - str_to_felt_bytes( - "0x008a81230a7e3ffa40abe541786a9b69fbb601434cec9536d5d5b2ee4df90383", - ), - ), - ( - str_to_felt_bytes( - "0x00ffda4b5cf0dce9bc9b0d035210590c73375fdbb70cd94ec6949378bffc410c", - ), - str_to_felt_bytes( - "0x02b36318931915f71777f7e59246ecab3189db48408952cefda72f4b7977be51", - ), - ), - ( - str_to_felt_bytes( - "0x00ffda4b5cf0dce9bc9b0d035210590c73375fdbb70cd94ec6949378bffc410d", - ), - str_to_felt_bytes( - "0x07e928dcf189b05e4a3dae0bc2cb98e447f1843f7debbbf574151eb67cda8797", - ), - ), - ], - ), - ( - str_to_felt_bytes( - "0x0327d34747122d7a40f4670265b098757270a449ec80c4871450fffdab7c2fa8", - ), - vec![ - ( - str_to_felt_bytes( - "0x0000000000000000000000000000000000000000000000000000000000000005", - ), - str_to_felt_bytes( - "0x0000000000000000000000000000000000000000000000000000000000000065", - ), - ), - ( - str_to_felt_bytes( - "0x01aed933fd362faecd8ea54ee749092bd21f89901b7d1872312584ac5b636c6d", - ), - str_to_felt_bytes( - "0x00000000000000000000000000000000000000000000000000000000000007c7", - ), - ), - ( - str_to_felt_bytes( - "0x04184fa5a6d40f47a127b046ed6facfa3e6bc3437b393da65cc74afe47ca6c6e", - ), - str_to_felt_bytes( - "0x001ef78e458502cd457745885204a4ae89f3880ec24db2d8ca97979dce15fedc", - ), - ), - ( - str_to_felt_bytes( - "0x05591c8c3c8d154a30869b463421cd5933770a0241e1a6e8ebcbd91bdd69bec4", - ), - str_to_felt_bytes( - "0x026b5943d4a0c420607cee8030a8cdd859bf2814a06633d165820960a42c6aed", - ), - ), - ( - str_to_felt_bytes( - "0x05591c8c3c8d154a30869b463421cd5933770a0241e1a6e8ebcbd91bdd69bec5", - ), - str_to_felt_bytes( - "0x01518eec76afd5397cefd14eda48d01ad59981f9ce9e70c233ca67acd8754008", - ), - ), - ], - ), - ]; - - let block_2 = vec![ - ( - str_to_felt_bytes( - "0x001fb4457f3fe8a976bdb9c04dd21549beeeb87d3867b10effe0c4bd4064a8e4", - ), - vec![( - str_to_felt_bytes( - "0x056c060e7902b3d4ec5a327f1c6e083497e586937db00af37fe803025955678f", - ), - str_to_felt_bytes( - "0x075495b43f53bd4b9c9179db113626af7b335be5744d68c6552e3d36a16a747c", - ), - )], - ), - ( - str_to_felt_bytes( - "0x05790719f16afe1450b67a92461db7d0e36298d6a5f8bab4f7fd282050e02f4f", - ), - vec![( - str_to_felt_bytes( - "0x0772c29fae85f8321bb38c9c3f6edb0957379abedc75c17f32bcef4e9657911a", - ), - str_to_felt_bytes( - "0x06d4ca0f72b553f5338a95625782a939a49b98f82f449c20f49b42ec60ed891c", - ), - )], - ), - ( - str_to_felt_bytes( - "0x057b973bf2eb26ebb28af5d6184b4a044b24a8dcbf724feb95782c4d1aef1ca9", - ), - vec![( - str_to_felt_bytes( - "0x04f2c206f3f2f1380beeb9fe4302900701e1cb48b9b33cbe1a84a175d7ce8b50", - ), - str_to_felt_bytes( - "0x02a614ae71faa2bcdacc5fd66965429c57c4520e38ebc6344f7cf2e78b21bd2f", - ), - )], - ), - ( - str_to_felt_bytes( - "0x02d6c9569dea5f18628f1ef7c15978ee3093d2d3eec3b893aac08004e678ead3", - ), - vec![( - str_to_felt_bytes( - "0x07f93985c1baa5bd9b2200dd2151821bd90abb87186d0be295d7d4b9bc8ca41f", - ), - str_to_felt_bytes( - "0x0127cd00a078199381403a33d315061123ce246c8e5f19aa7f66391a9d3bf7c6", - ), - )], - ), - ]; - - let blocks = block_0.iter().chain(block_1.iter()).chain(block_2.iter()); - - // Inserts all storage updates into the bonsai - for (contract_address, storage) in blocks.clone() { - log::info!( - "contract address (write): {:#064x}", - Felt::from_bytes_be_slice(contract_address) - ); - - for (k, v) in storage { - // truncate only keeps the first 251 bits in a key - // so there should be no error during insertion - let ktrunc = &truncate(k); - let kfelt0 = Felt::from_bytes_be_slice(k); - let kfelt1 = Felt::from_bytes_be_slice(ktrunc.as_raw_slice()); - - // quick sanity check to make sure truncating a key does not remove any data - assert_eq!(kfelt0, kfelt1); - - let v = &Felt::from_bytes_be_slice(v); - assert!(bonsai.insert(contract_address, ktrunc, v).is_ok()); - } - } - assert!(bonsai.commit(BasicId::new(0)).is_ok()); - - // aggreates all storage changes to their latest state - // (replacements are takent into account) - let mut storage_map = IndexMap::>::new(); - for (contract_address, storage) in blocks.clone() { - let map = storage_map.entry((*contract_address).into()).or_default(); - - for (k, v) in storage { - let k = Felt::from_bytes_be_slice(k); - let v = Felt::from_bytes_be_slice(v); - map.insert(k, v); - } - } - - // checks for each contract if the original key can be reconstructed - // from the data stored in the db - for (contract_address, storage) in storage_map.iter() { - log::info!( - "contract address (read): {:#064x}", - Felt::from_bytes_be_slice(contract_address) - ); - - let keys = bonsai.get_keys(contract_address).unwrap(); - log::debug!("{keys:?}"); - for k in keys { - // if all has gone well, the db should contain the first 251 bits of the key, - // which should represent the entirety of the data - let k = Felt::from_bytes_be_slice(&k); - log::info!("looking for key: {k:#064x}"); - - assert!(storage.contains_key(&k)); - } - } - - // makes sure retrieving key-value pairs works for each contract - for (contract_address, storage) in storage_map.iter() { - log::info!( - "contract address (read): {:#064x}", - Felt::from_bytes_be_slice(contract_address) - ); - - let kv = bonsai.get_key_value_pairs(contract_address).unwrap(); - log::debug!("{kv:?}"); - for (k, v) in kv { - let k = Felt::from_bytes_be_slice(&k); - let v = Felt::from_bytes_be_slice(&v); - log::info!("checking for key-value pair:({k:#064x}, {v:#064x})"); - - assert_eq!(*storage.get(&k).unwrap(), v); - } - } - } - - fn str_to_felt_bytes(hex: &str) -> [u8; 32] { - Felt::from_hex(hex).unwrap().to_bytes_be() - } - - fn truncate(key: &[u8]) -> BitVec { - key.view_bits()[5..].to_owned() - } - - // use crate::{ - // databases::{create_rocks_db, RocksDB, RocksDBConfig}, - // id::BasicId, - // key_value_db::KeyValueDBConfig, - // KeyValueDB, - // }; - // use bitvec::vec::BitVec; - // use mp_felt::Felt252Wrapper; - // use mp_hashers::pedersen::PedersenHasher; - // use parity_scale_codec::{Decode, Encode}; - // use rand::prelude::*; - // use starknet_types_core::{felt::Felt, hash::Pedersen}; - - // // convert a Madara felt to a standard Felt - // fn felt_from_madara_felt(madara_felt: &Felt252Wrapper) -> Felt { - // let encoded = madara_felt.encode(); - // Felt::decode(&mut &encoded[..]).unwrap() - // } - - // // convert a standard Felt to a Madara felt - // fn madara_felt_from_felt(felt: &Felt) -> Felt252Wrapper { - // let encoded = felt.encode(); - // Felt252Wrapper::decode(&mut &encoded[..]).unwrap() - // } - - // #[test] - // fn one_commit_tree_compare() { - // let mut elements = vec![]; - // let tempdir = tempfile::tempdir().unwrap(); - // let mut rng = rand::thread_rng(); - // let tree_size = rng.gen_range(10..100); - // for _ in 0..tree_size { - // let mut element = String::from("0x"); - // let element_size = rng.gen_range(10..32); - // for _ in 0..element_size { - // let random_byte: u8 = rng.gen(); - // element.push_str(&format!("{:02x}", random_byte)); - // } - // elements.push(Felt::from_hex(&element).unwrap()); - // } - // let madara_elements = elements - // .iter() - // .map(madara_felt_from_felt) - // .collect::>(); - // let rocks_db = create_rocks_db(std::path::Path::new(tempdir.path())).unwrap(); - // let rocks_db = RocksDB::new(&rocks_db, RocksDBConfig::default()); - // let db = KeyValueDB::new(rocks_db, KeyValueDBConfig::default(), None); - // let mut bonsai_tree: super::MerkleTree, BasicId> = - // super::MerkleTree::new(db).unwrap(); - // let root_hash = mp_commitments::calculate_class_commitment_tree_root_hash::( - // &madara_elements, - // ); - // elements - // .iter() - // .zip(madara_elements.iter()) - // .for_each(|(element, madara_element)| { - // let final_hash = - // calculate_class_commitment_leaf_hash::(*madara_element); - // let key = &element.to_bytes_be()[..31]; - // bonsai_tree - // .set( - // &BitVec::from_vec(key.to_vec()), - // felt_from_madara_felt(&final_hash), - // ) - // .unwrap(); - // }); - // bonsai_tree.display(); - // assert_eq!( - // bonsai_tree.commit().unwrap(), - // felt_from_madara_felt(&root_hash) - // ); - // } - - // #[test] - // fn simple_commits() { - // let tempdir = tempfile::tempdir().unwrap(); - // let mut madara_tree = StateCommitmentTree::::default(); - // let rocks_db = create_rocks_db(std::path::Path::new(tempdir.path())).unwrap(); - // let rocks_db = RocksDB::new(&rocks_db, RocksDBConfig::default()); - // let db = KeyValueDB::new(rocks_db, KeyValueDBConfig::default(), None); - // let mut bonsai_tree: super::MerkleTree, BasicId> = - // super::MerkleTree::new(db).unwrap(); - // let elements = [ - // [Felt::from_hex("0x665342762FDD54D0303c195fec3ce2568b62052e").unwrap()], - // [Felt::from_hex("0x66342762FDD54D0303c195fec3ce2568b62052e").unwrap()], - // [Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap()], - // ]; - // for elem in elements { - // elem.iter().for_each(|class_hash| { - // let final_hash = - // felt_from_madara_felt(&calculate_class_commitment_leaf_hash::( - // madara_felt_from_felt(class_hash), - // )); - // madara_tree.set( - // madara_felt_from_felt(class_hash), - // madara_felt_from_felt(&final_hash), - // ); - // let key = &class_hash.to_bytes_be()[..31]; - // bonsai_tree - // .set(&BitVec::from_vec(key.to_vec()), final_hash) - // .unwrap(); - // }); - // } - // let madara_root_hash = madara_tree.commit(); - // let bonsai_root_hash = bonsai_tree.commit().unwrap(); - // assert_eq!(bonsai_root_hash, felt_from_madara_felt(&madara_root_hash)); - // } - - // #[test] - // fn simple_commits_and_delete() { - // let tempdir = tempfile::tempdir().unwrap(); - // let rocks_db = create_rocks_db(std::path::Path::new(tempdir.path())).unwrap(); - // let rocks_db = RocksDB::new(&rocks_db, RocksDBConfig::default()); - // let db = KeyValueDB::new(rocks_db, KeyValueDBConfig::default(), None); - // let mut bonsai_tree: super::MerkleTree, BasicId> = - // super::MerkleTree::new(db).unwrap(); - // let elements = [ - // [Felt::from_hex("0x665342762FDD54D0303c195fec3ce2568b62052e").unwrap()], - // [Felt::from_hex("0x66342762FDD54D0303c195fec3ce2568b62052e").unwrap()], - // [Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap()], - // ]; - // for elem in elements { - // elem.iter().for_each(|class_hash| { - // let final_hash = calculate_class_commitment_leaf_hash::( - // madara_felt_from_felt(class_hash), - // ); - // let key = &class_hash.to_bytes_be()[..31]; - // bonsai_tree - // .set( - // &BitVec::from_vec(key.to_vec()), - // felt_from_madara_felt(&final_hash), - // ) - // .unwrap(); - // }); - // } - // bonsai_tree.commit().unwrap(); - // for elem in elements { - // elem.iter().for_each(|class_hash| { - // let key = &class_hash.to_bytes_be()[..31]; - // bonsai_tree - // .set(&BitVec::from_vec(key.to_vec()), Felt::ZERO) - // .unwrap(); - // }); - // } - // bonsai_tree.commit().unwrap(); - // } - - // #[test] - // fn multiple_commits_tree_compare() { - // let mut rng = rand::thread_rng(); - // let tempdir = tempfile::tempdir().unwrap(); - // let mut madara_tree = StateCommitmentTree::::default(); - // let rocks_db = create_rocks_db(std::path::Path::new(tempdir.path())).unwrap(); - // let rocks_db = RocksDB::new(&rocks_db, RocksDBConfig::default()); - // let db = KeyValueDB::new(rocks_db, KeyValueDBConfig::default(), None); - // let mut bonsai_tree: super::MerkleTree, BasicId> = - // super::MerkleTree::new(db).unwrap(); - // let nb_commits = rng.gen_range(2..4); - // for _ in 0..nb_commits { - // let mut elements = vec![]; - // let tree_size = rng.gen_range(10..100); - // for _ in 0..tree_size { - // let mut element = String::from("0x"); - // let element_size = rng.gen_range(10..32); - // for _ in 0..element_size { - // let random_byte: u8 = rng.gen(); - // element.push_str(&format!("{:02x}", random_byte)); - // } - // elements.push(Felt::from_hex(&element).unwrap()); - // } - // elements.iter().for_each(|class_hash| { - // let final_hash = calculate_class_commitment_leaf_hash::( - // madara_felt_from_felt(class_hash), - // ); - // madara_tree.set(madara_felt_from_felt(class_hash), final_hash); - // let key = &class_hash.to_bytes_be()[..31]; - // bonsai_tree - // .set( - // &BitVec::from_vec(key.to_vec()), - // felt_from_madara_felt(&final_hash), - // ) - // .unwrap(); - // }); - - // let bonsai_root_hash = bonsai_tree.commit().unwrap(); - // let madara_root_hash = madara_tree.commit(); - // assert_eq!(bonsai_root_hash, felt_from_madara_felt(&madara_root_hash)); - // } - // } - - // #[test] // fn multiple_commits_tree_compare_with_deletes() { - // let mut rng = rand::thread_rng(); - // let mut madara_tree = StateCommitmentTree::::default(); - // let rocks_db = create_rocks_db(std::path::Path::new("test_db")).unwrap(); - // let mut db = RocksDB::new(&rocks_db, RocksDBConfig::default()); - // let mut bonsai_tree: super::MerkleTree = - // super::MerkleTree::empty(&mut db); - // let nb_commits = rng.gen_range(2..5); - // let mut elements_to_delete = vec![]; - // for _ in 0..nb_commits { - // let mut elements = vec![]; - // let tree_size = rng.gen_range(10..100); - // for _ in 0..tree_size { - // let mut element = String::from("0x"); - // let element_size = rng.gen_range(10..32); - // for _ in 0..element_size { - // let random_byte: u8 = rng.gen(); - // element.push_str(&format!("{:02x}", random_byte)); - // } - // if rng.gen_bool(0.1) { - // elements_to_delete.push(Felt::from_hex_be(&element).unwrap()); - // elements.push(Felt::from_hex_be(&element).unwrap()); - // } else { - // elements.push(Felt::from_hex_be(&element).unwrap()); - // } - // } - // elements.iter().for_each(|class_hash| { - // let final_hash = - // calculate_class_commitment_leaf_hash::(*class_hash); - // madara_tree.set(*class_hash, final_hash); - // let key = &class_hash.0.to_bytes_be()[..31]; - // bonsai_tree.set(&BitVec::from_vec(key.to_vec()), final_hash); - // }); - - // let bonsai_root_hash = bonsai_tree.commit(); - // let madara_root_hash = madara_tree.commit(); - // assert_eq!(bonsai_root_hash, madara_root_hash); - // } - // elements_to_delete.iter().for_each(|class_hash| { - // madara_tree.set(*class_hash, Felt::ZERO); - // let key = &class_hash.0.to_bytes_be()[..31]; - // bonsai_tree.set(&BitVec::from_vec(key.to_vec()), Felt::ZERO); - // }); - - // let bonsai_root_hash = bonsai_tree.commit(); - // let madara_root_hash = madara_tree.commit(); - // assert_eq!(bonsai_root_hash, madara_root_hash); - // } - - // #[test] - // fn test_proof() { - // let tempdir = tempfile::tempdir().unwrap(); - // let rocks_db = create_rocks_db(std::path::Path::new(tempdir.path())).unwrap(); - // let rocks_db = RocksDB::new(&rocks_db, RocksDBConfig::default()); - // let db = KeyValueDB::new(rocks_db, KeyValueDBConfig::default(), None); - // let mut bonsai_tree: super::MerkleTree, BasicId> = - // super::MerkleTree::new(db).unwrap(); - // let elements = [ - // [Felt252Wrapper::from_hex_be("0x665342762FDD54D0303c195fec3ce2568b62052e").unwrap()], - // [Felt252Wrapper::from_hex_be("0x66342762FDD54D0303c195fec3ce2568b62052e").unwrap()], - // [Felt252Wrapper::from_hex_be("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap()], - // ]; - // for elem in elements { - // elem.iter().for_each(|class_hash| { - // let final_hash = - // calculate_class_commitment_leaf_hash::(*class_hash); - // let key = &class_hash.0.to_bytes_be()[..31]; - // bonsai_tree - // .set( - // &BitVec::from_vec(key.to_vec()), - // Felt::from_bytes_be(&final_hash.0.to_bytes_be()), - // ) - // .unwrap(); - // }); - // } - // bonsai_tree.commit().unwrap(); - // let bonsai_proof = bonsai_tree - // .get_proof(&BitVec::from_vec( - // elements[0][0].0.to_bytes_be()[..31].to_vec(), - // )) - // .unwrap(); - // println!("bonsai_proof: {:?}", bonsai_proof); - // } - - // test in madara - // #[test] - // fn test_proof() { - // let mut tree = super::merkle_patricia_tree::merkle_tree::MerkleTree::::empty(); - // let elements = [ - // [Felt252Wrapper::from_hex_be("0x665342762FDD54D0303c195fec3ce2568b62052e").unwrap()], - // [Felt252Wrapper::from_hex_be("0x66342762FDD54D0303c195fec3ce2568b62052e").unwrap()], - // [Felt252Wrapper::from_hex_be("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap()], - // ]; - // for elem in elements { - // elem.iter().for_each(|class_hash| { - // let final_hash = - // calculate_class_commitment_leaf_hash::(*class_hash); - // let key = &class_hash.0.to_bytes_be()[..31]; - // tree - // .set(&BitVec::from_vec(key.to_vec()), final_hash) - // }); - // } - // tree.commit(); - // let bonsai_proof = tree.get_proof(&BitVec::from_vec( - // elements[0][0].0.to_bytes_be()[..31].to_vec(), - // )); - // println!("bonsai_proof: {:?}", bonsai_proof); - // } -} diff --git a/src/trie/mod.rs b/src/trie/mod.rs index dc6e304..41e38d6 100644 --- a/src/trie/mod.rs +++ b/src/trie/mod.rs @@ -1,6 +1,9 @@ +pub(crate) mod iterator; mod merkle_node; -pub mod merkle_tree; mod path; -mod trie_db; +pub(crate) mod proof; +pub mod tree; +pub(crate) mod trees; +pub(crate) mod trie_db; pub(crate) use trie_db::TrieKey; diff --git a/src/trie/path.rs b/src/trie/path.rs index 62e7c4e..412475e 100644 --- a/src/trie/path.rs +++ b/src/trie/path.rs @@ -1,16 +1,34 @@ -use bitvec::{order::Msb0, vec::BitVec}; -use core::fmt; -use parity_scale_codec::{Decode, Encode, Error, Input, Output}; - use super::merkle_node::Direction; - -use crate::{ByteVec, EncodeExt}; +use crate::{BitVec, ByteVec, EncodeExt}; +use core::{ + fmt, + ops::{Deref, DerefMut}, +}; +use parity_scale_codec::{Decode, Encode, Error, Input, Output}; #[cfg(all(feature = "std", test))] use rstest::rstest; -#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Default)] -pub struct Path(pub BitVec); +#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct Path(pub BitVec); + +impl Default for Path { + fn default() -> Self { + Self(BitVec::with_capacity(251)) + } +} + +impl Path { + pub(crate) fn new_with_direction(&self, direction: Direction) -> Path { + let mut path = self.0.clone(); + path.push(direction.into()); + Path(path) + } + + pub fn len(&self) -> usize { + self.0.len() + } +} impl fmt::Debug for Path { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { @@ -66,7 +84,7 @@ impl Decode for Path { let mut remaining_bits = len as usize; let mut current_byte = None; let mut bit = 7; - let mut bits = BitVec::::new(); + let mut bits = BitVec::new(); // No bits left to decode; we're done. while remaining_bits != 0 { // Get the next store entry to pull from: @@ -102,24 +120,29 @@ impl Decode for Path { } } -impl Path { - pub(crate) fn new_with_direction(&self, direction: Direction) -> Path { - let mut path = self.0.clone(); - path.push(direction.into()); - Path(path) - } -} - /// Convert Path to SByteVec can be used, for example, to create keys for the database impl From for ByteVec { fn from(path: Path) -> Self { - path.encode_sbytevec() + path.encode_bytevec() } } impl From<&Path> for ByteVec { fn from(path: &Path) -> Self { - path.encode_sbytevec() + path.encode_bytevec() + } +} + +impl Deref for Path { + type Target = BitVec; + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl DerefMut for Path { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 } } @@ -132,7 +155,7 @@ impl From<&Path> for ByteVec { #[case(&[0b11111111])] #[case(&[0b11111111, 0b00000000, 0b10101010, 0b10101010, 0b11111111, 0b00000000, 0b10101010, 0b10101010, 0b11111111, 0b00000000, 0b10101010, 0b10101010])] fn test_shared_path_encode_decode(#[case] input: &[u8]) { - let path = Path(BitVec::::from_slice(input)); + let path = Path(BitVec::from_slice(input)); let mut encoded = Vec::new(); path.encode_to(&mut encoded); diff --git a/src/trie/proof.rs b/src/trie/proof.rs new file mode 100644 index 0000000..71b1daa --- /dev/null +++ b/src/trie/proof.rs @@ -0,0 +1,272 @@ +use super::{ + merkle_node::{hash_binary_node, hash_edge_node, Direction}, + path::Path, + tree::MerkleTree, +}; +use crate::{ + id::Id, + key_value_db::KeyValueDB, + trie::{ + iterator::NodeVisitor, + merkle_node::{Node, NodeHandle}, + tree::NodeKey, + }, + BitSlice, BitVec, BonsaiDatabase, BonsaiStorageError, HashMap, HashSet, +}; +use core::marker::PhantomData; +use hashbrown::hash_set; +use starknet_types_core::{felt::Felt, hash::StarkHash}; + +#[derive(Debug, PartialEq, Eq, Clone, Copy)] +pub enum Membership { + Member, + NonMember, +} + +impl From for bool { + fn from(value: Membership) -> Self { + value == Membership::Member + } +} + +impl From for Membership { + fn from(value: bool) -> Self { + match value { + true => Self::Member, + false => Self::NonMember, + } + } +} + +#[derive(Debug, Clone, PartialEq)] +pub enum ProofNode { + Binary { left: Felt, right: Felt }, + Edge { child: Felt, path: Path }, +} + +impl ProofNode { + pub fn hash(&self) -> Felt { + match self { + ProofNode::Binary { left, right } => hash_binary_node::(*left, *right), + ProofNode::Edge { child, path } => hash_edge_node::(path, *child), + } + } +} + +#[derive(Debug, Clone)] +pub struct MultiProof(pub HashMap); +impl MultiProof { + /// If the proof proves more than just the provided `key_values`, this function will not fail. + /// Not the most optimized way of doing it, but we don't actually need to verify proofs in madara. + /// As such, it has also not been properly proptested. + pub fn verify_proof<'a, 'b: 'a, H: StarkHash>( + &'b self, + root: Felt, + key_values: impl IntoIterator, Felt)> + 'a, + tree_height: u8, + ) -> impl Iterator + 'a { + let mut checked_cache: HashSet = Default::default(); + let mut current_path = BitVec::with_capacity(251); + key_values.into_iter().map(move |(k, v)| { + let k = k.as_ref(); + + if k.len() != tree_height as usize { + return Membership::NonMember; + } + + // Go down the tree, starting from the root. + current_path.clear(); // hoisted alloc + let mut current_felt = root; + + loop { + log::trace!("Start verify loop: {current_path:b} => {current_felt:#x}"); + if current_path.len() == k.len() { + // End of traversal, check if value is correct + log::trace!("End of traversal"); + break (v == current_felt).into(); + } + if current_path.len() > k.len() { + // We overshot. + log::trace!("Overshot"); + break Membership::NonMember; + } + let Some(node) = self.0.get(¤t_felt) else { + // Missing node. + log::trace!("Missing"); + break Membership::NonMember; + }; + + // Check hash and save to verification cache. + if let hash_set::Entry::Vacant(entry) = checked_cache.entry(v) { + let computed_hash = node.hash::(); + if computed_hash != current_felt { + // Hash mismatch. + log::trace!("Hash mismatch: {computed_hash:#x} {current_felt:#x}"); + break Membership::NonMember; + } + entry.insert(); + } + + match node { + ProofNode::Binary { left, right } => { + // PANIC: We checked above that current_path.len() < k.len(). + let direction = Direction::from(k[current_path.len()]); + log::trace!("Binary {direction:?}"); + current_path.push(direction.into()); + current_felt = match direction { + Direction::Left => *left, + Direction::Right => *right, + } + } + ProofNode::Edge { child, path } => { + log::trace!("Edge"); + if k.get(current_path.len()..(current_path.len() + path.len())) + != Some(&path.0) + { + log::trace!("Wrong edge: {path:?}"); + // Wrong edge path. + break Membership::NonMember; + } + current_path.extend_from_bitslice(&path.0); + current_felt = *child; + } + } + } + }) + } +} + +impl MerkleTree { + /// This function is designed to be very efficient if the `keys` are sorted - this allows for + /// the minimal amount of backtracking when switching from one key to the next. + pub fn get_multi_proof( + &mut self, + db: &KeyValueDB, + keys: impl IntoIterator>, + ) -> Result> { + let max_height = self.max_height; + + struct ProofVisitor(MultiProof, PhantomData); + impl NodeVisitor for ProofVisitor { + fn visit_node( + &mut self, + tree: &mut MerkleTree, + node_id: NodeKey, + _prev_height: usize, + ) -> Result<(), BonsaiStorageError> { + let proof_node = match tree.get_node_mut::(node_id)? { + Node::Binary(binary_node) => { + let (left, right) = (binary_node.left, binary_node.right); + ProofNode::Binary { + left: tree.get_or_compute_node_hash::(left)?, + right: tree.get_or_compute_node_hash::(right)?, + } + } + Node::Edge(edge_node) => { + let (child, path) = (edge_node.child, edge_node.path.clone()); + ProofNode::Edge { + child: tree.get_or_compute_node_hash::(child)?, + path, + } + } + }; + let hash = tree.get_or_compute_node_hash::(NodeHandle::InMemory(node_id))?; + self.0 .0.insert(hash, proof_node); + Ok(()) + } + } + let mut visitor = ProofVisitor::(MultiProof(Default::default()), PhantomData); + + let mut iter = self.iter(db); + for key in keys { + let key = key.as_ref(); + if key.len() != max_height as _ { + return Err(BonsaiStorageError::KeyLength { + expected: self.max_height as _, + got: key.len(), + }); + } + log::debug!("go to = {key:b}"); + iter.traverse_to(&mut visitor, key)?; + + log::debug!("iter = {iter:?}"); + // We should have found a leaf here. + iter.leaf_hash + .ok_or(BonsaiStorageError::CreateProofKeyNotInTree { + key: key.to_bitvec(), + })?; + } + + Ok(visitor.0) + } +} + +#[cfg(test)] +mod tests { + use crate::{ + databases::{create_rocks_db, RocksDB, RocksDBConfig}, + id::BasicId, + BonsaiStorage, BonsaiStorageConfig, + }; + use bitvec::{bits, order::Msb0}; + use starknet_types_core::{felt::Felt, hash::Pedersen}; + + const ONE: Felt = Felt::ONE; + const TWO: Felt = Felt::TWO; + const THREE: Felt = Felt::THREE; + const FOUR: Felt = Felt::from_hex_unchecked("0x4"); + + #[test] + fn test_multiproof() { + let _ = env_logger::builder().is_test(true).try_init(); + log::set_max_level(log::LevelFilter::Trace); + let tempdir = tempfile::tempdir().unwrap(); + let db = create_rocks_db(tempdir.path()).unwrap(); + let mut bonsai_storage: BonsaiStorage = BonsaiStorage::new( + RocksDB::::new(&db, RocksDBConfig::default()), + BonsaiStorageConfig::default(), + 8, + ) + .unwrap(); + + bonsai_storage + .insert(&[], bits![u8, Msb0; 0,0,0,1,0,0,0,0], &ONE) + .unwrap(); + bonsai_storage + .insert(&[], bits![u8, Msb0; 0,0,0,1,0,0,0,1], &TWO) + .unwrap(); + bonsai_storage + .insert(&[], bits![u8, Msb0; 0,0,0,1,0,0,1,0], &THREE) + .unwrap(); + bonsai_storage + .insert(&[], bits![u8, Msb0; 0,1,0,0,0,0,0,0], &FOUR) + .unwrap(); + + bonsai_storage.dump(); + + let tree = bonsai_storage + .tries + .trees + .get_mut(&smallvec::smallvec![]) + .unwrap(); + + let proof = tree + .get_multi_proof( + &bonsai_storage.tries.db, + [ + bits![u8, Msb0; 0,0,0,1,0,0,0,1], + bits![u8, Msb0; 0,1,0,0,0,0,0,0], + ], + ) + .unwrap(); + + log::trace!("proof: {proof:?}"); + assert!(proof + .verify_proof::( + tree.root_hash(&bonsai_storage.tries.db).unwrap(), + [(bits![u8, Msb0; 0,0,0,1,0,0,0,0], ONE)], + 8 + ) + .all(|v| v.into())); + } +} diff --git a/src/trie/tree.rs b/src/trie/tree.rs new file mode 100644 index 0000000..e824658 --- /dev/null +++ b/src/trie/tree.rs @@ -0,0 +1,1133 @@ +use core::{fmt, marker::PhantomData}; +use core::{iter, mem}; +use parity_scale_codec::Decode; +use slotmap::SlotMap; +use starknet_types_core::{felt::Felt, hash::StarkHash}; + +use crate::trie::merkle_node::{hash_binary_node, hash_edge_node}; +use crate::BitVec; +use crate::{ + error::BonsaiStorageError, format, hash_map, id::Id, vec, BitSlice, BonsaiDatabase, ByteVec, + EncodeExt, HashMap, HashSet, KeyValueDB, ToString, Vec, +}; + +use super::iterator::MerkleTreeIterator; +use super::{ + merkle_node::{BinaryNode, Direction, EdgeNode, Node, NodeHandle}, + path::Path, + trie_db::TrieKeyType, + TrieKey, +}; + +#[cfg(test)] +use log::trace; + +slotmap::new_key_type! { + /// Key for an inmemory node. + pub struct NodeKey; +} + +// TODO: implement encode and decode by hand in Node +// these cases should never happen, otherwise that would mean we are saving in-memory node keys to the db, which would be very bad. +impl parity_scale_codec::Encode for NodeKey { + fn using_encoded R>(&self, _f: F) -> R { + unreachable!("Cannot encode NodeKey") + } +} +impl parity_scale_codec::Decode for NodeKey { + fn decode( + _input: &mut I, + ) -> Result { + unreachable!("Cannot decode NodeKey") + } +} + +#[derive(Debug, Clone, Copy)] +pub(crate) enum RootHandle { + Empty, + Loaded(NodeKey), +} + +/// 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 { + /// The root node. None means the node has not been loaded yet. + pub(crate) root_node: Option, + /// In-memory nodes. + pub(crate) nodes: SlotMap, + /// Identifier of the tree in the database. + pub(crate) identifier: ByteVec, + /// The list of nodes that should be removed from the underlying database during the next commit. + pub(crate) death_row: HashSet, + /// The list of leaves that have been modified during the current commit. + pub(crate) cache_leaf_modified: HashMap>, + /// The maximum height of the tree. This is an u8 because we may rely on the fact that it's less than 256 in the future for optimizations. + pub(crate) max_height: u8, + /// The hasher used to hash the nodes. + _hasher: PhantomData, +} + +impl fmt::Debug for MerkleTree { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MerkleTree") + .field("root_node", &self.root_node) + .field("nodes", &self.nodes) + .field("identifier", &self.identifier) + .field("death_row", &self.death_row) + .field("cache_leaf_modified", &self.cache_leaf_modified) + .finish() + } +} + +// NB: #[derive(Clone)] does not work because it expands to an impl block which forces H: Clone, which Pedersen/Poseidon aren't. +#[cfg(feature = "bench")] +impl Clone for MerkleTree { + fn clone(&self) -> Self { + Self { + max_height: self.max_height, + root_node: self.root_node, + nodes: self.nodes.clone(), + identifier: self.identifier.clone(), + death_row: self.death_row.clone(), + cache_leaf_modified: self.cache_leaf_modified.clone(), + _hasher: PhantomData, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) enum InsertOrRemove { + Insert(T), + Remove, +} +enum NodeOrFelt<'a> { + Node(&'a Node), + Felt(Felt), +} + +impl MerkleTree { + pub fn new(identifier: ByteVec, max_height: u8) -> Self { + Self { + root_node: None, + nodes: Default::default(), + identifier, + death_row: HashSet::new(), + cache_leaf_modified: HashMap::new(), + max_height, + _hasher: PhantomData, + } + } + + /// Loads the root node or returns None if the tree is empty. + pub(crate) fn load_root_node( + &mut self, + db: &KeyValueDB, + ) -> Result, BonsaiStorageError> { + // try_get_or_insert + match self.root_node { + Some(RootHandle::Loaded(id)) => Ok(Some(id)), + Some(RootHandle::Empty) => Ok(None), + None => { + // load the node + let id = self + .load_db_node(db, &TrieKey::new(&self.identifier, TrieKeyType::Trie, &[0]))?; + + match id { + Some(id) => { + self.root_node = Some(RootHandle::Loaded(id)); + Ok(Some(id)) + } + None => { + self.root_node = Some(RootHandle::Empty); + Ok(None) + } + } + } + } + } + + /// First step of two phase init. + pub(crate) fn load_db_node( + &mut self, + db: &KeyValueDB, + key: &TrieKey, + ) -> Result, BonsaiStorageError> { + if self.death_row.contains(key) { + return Ok(None); + } + let node = db.get(key)?; + let Some(node) = node else { return Ok(None) }; + + let node = Node::decode(&mut node.as_slice())?; + let key = self.nodes.insert(node); + + Ok(Some(key)) + } + + pub(crate) fn get_node_mut( + &mut self, + node_key: NodeKey, + ) -> Result<&mut Node, BonsaiStorageError> { + self.nodes.get_mut(node_key).ok_or_else(|| { + BonsaiStorageError::Trie(format!("Dangling in-memory node key: {node_key:?}")) + }) + } + + pub(crate) fn load_node_handle( + &mut self, + db: &KeyValueDB, + handle: NodeHandle, + path: &Path, + ) -> Result> { + match handle { + NodeHandle::Hash(_) => { + // TODO(perf): useless allocs everywhere here... + let path: ByteVec = path.clone().into(); + log::trace!("Visiting db node {:?}", path); + let key = TrieKey::new(&self.identifier, TrieKeyType::Trie, &path); + let Some(node_key) = self.load_db_node(db, &key)? else { + // Dangling node id in db + return Err(BonsaiStorageError::Trie( + "Could not get node from db".to_string(), + )); + }; + Ok(node_key) + } + NodeHandle::InMemory(node_key) => Ok(node_key), + } + } + + /// Get or compute the hash of a node. + pub(crate) fn get_or_compute_node_hash( + &mut self, + node: NodeHandle, + ) -> Result> { + match node { + NodeHandle::Hash(felt) => Ok(felt), + NodeHandle::InMemory(node_key) => { + let computed_hash = match self.get_node_mut::(node_key)? { + Node::Binary(binary_node) => { + if let Some(hash) = binary_node.hash { + return Ok(hash); + } + let (left, right) = (binary_node.left, binary_node.right); + let left_hash = self.get_or_compute_node_hash::(left)?; + let right_hash = self.get_or_compute_node_hash::(right)?; + hash_binary_node::(left_hash, right_hash) + } + Node::Edge(edge_node) => { + if let Some(hash) = edge_node.hash { + return Ok(hash); + } + let (path, child) = (edge_node.path.clone(), edge_node.child); + // edge_node borrow ends here + let child_hash = self.get_or_compute_node_hash::(child)?; + hash_edge_node::(&path, child_hash) + } + }; + + // reborrow, for lifetime reasons (can't go into children if a borrow is alive) + match self.get_node_mut::(node_key)? { + Node::Binary(binary_node) => binary_node.hash = Some(computed_hash), + Node::Edge(edge_node) => edge_node.hash = Some(computed_hash), + } + + Ok(computed_hash) + } + } + } + + /// Note: as iterators load nodes from the database, this takes an &mut self. However, + /// note that it will not modify anything in the database - hence the &db. + pub fn iter<'a, DB: BonsaiDatabase, ID: Id>( + &'a mut self, + db: &'a KeyValueDB, + ) -> MerkleTreeIterator<'a, H, DB, ID> { + MerkleTreeIterator::new(self, db) + } + + /// # Panics + /// + /// Calling this function when the tree has uncommited changes is invalid as the hashes need to be recomputed. + pub fn root_hash( + &self, + db: &KeyValueDB, + ) -> Result> { + match self.root_node { + Some(RootHandle::Empty) => Ok(Felt::ZERO), + Some(RootHandle::Loaded(node_id)) => { + let node = self.nodes.get(node_id).ok_or_else(|| { + BonsaiStorageError::Trie("Could not fetch root node from storage".into()) + })?; + node.get_hash().ok_or_else(|| { + BonsaiStorageError::Trie("The tree has uncommited changes".into()) + }) + } + None => { + let Some(node) = Self::get_trie_branch_in_db_from_path( + &self.death_row, + &self.identifier, + db, + &Path::default(), + )? + else { + return Ok(Felt::ZERO); + }; + Ok(node + .get_hash() + .expect("The fetched node has no computed hash")) + } + } + } + + pub fn cache_leaf_modified(&self) -> &HashMap> { + &self.cache_leaf_modified + } + + /// Calculate all the new hashes and the root hash. + #[allow(clippy::type_complexity)] + pub(crate) fn get_updates( + &mut self, + ) -> Result< + impl Iterator)>, + BonsaiStorageError, + > { + let mut updates = HashMap::new(); + for node_key in mem::take(&mut self.death_row) { + updates.insert(node_key, InsertOrRemove::Remove); + } + + if let Some(RootHandle::Loaded(node_id)) = &self.root_node { + // compute hashes + let mut hashes = vec![]; + self.compute_root_hash::(&mut hashes)?; + + // commit the tree + self.commit_subtree::( + &mut updates, + *node_id, + Path::default(), + &mut hashes.into_iter(), + )?; + } + + self.root_node = None; // unloaded + + for (key, value) in mem::take(&mut self.cache_leaf_modified) { + updates.insert( + TrieKey::new(&self.identifier, TrieKeyType::Flat, &key), + match value { + InsertOrRemove::Insert(value) => InsertOrRemove::Insert(value.encode_bytevec()), + InsertOrRemove::Remove => InsertOrRemove::Remove, + }, + ); + } + #[cfg(test)] + self.assert_empty(); // we should have visited the whole tree + + Ok(updates.into_iter()) + } + + // Commit a single merkle tree + #[cfg(test)] + pub(crate) fn commit( + &mut self, + db: &mut KeyValueDB, + ) -> Result<(), BonsaiStorageError> { + let db_changes = self.get_updates::()?; + + let mut batch = db.create_batch(); + for (key, value) in db_changes { + match value { + InsertOrRemove::Insert(value) => { + log::trace!("committing insert {:?} => {:?}", key, value); + db.insert(&key, &value, Some(&mut batch))?; + } + InsertOrRemove::Remove => { + log::trace!("committing remove {:?}", key); + db.remove(&key, Some(&mut batch))?; + } + } + } + db.write_batch(batch).unwrap(); + log::trace!("commit finished"); + + Ok(()) + } + + #[cfg(test)] + pub(crate) fn assert_empty(&self) { + // we don't use is_empty here for better error messages :) + assert_eq!(self.nodes.iter().collect::>(), vec![]); + } + + fn get_node_or_felt( + &self, + node_handle: &NodeHandle, + ) -> Result> { + let node_id = match node_handle { + NodeHandle::Hash(hash) => return Ok(NodeOrFelt::Felt(*hash)), + NodeHandle::InMemory(node_id) => *node_id, + }; + let node = self.nodes.get(node_id).ok_or(BonsaiStorageError::Trie( + "Couldn't fetch node in the temporary storage".to_string(), + ))?; + Ok(NodeOrFelt::Node(node)) + } + + fn compute_root_hash( + &self, + hashes: &mut Vec, + ) -> Result> { + let handle = match &self.root_node { + Some(RootHandle::Loaded(node_id)) => *node_id, + Some(RootHandle::Empty) => return Ok(Felt::ZERO), + None => { + return Err(BonsaiStorageError::Trie( + "Root node is not loaded".to_string(), + )) + } + }; + let Some(node) = self.nodes.get(handle) else { + return Err(BonsaiStorageError::Trie( + "Could not fetch root node from storage".to_string(), + )); + }; + self.compute_hashes::(node, Path::default(), hashes) + } + + /// Compute the hashes of all of the updated nodes in the merkle tree. This step + /// is separate from [`commit_subtree`] as it is done in parallel using rayon. + /// Computed hashes are pushed to the `hashes` vector, depth first. + fn compute_hashes( + &self, + node: &Node, + path: Path, + hashes: &mut Vec, + ) -> Result> { + use Node::*; + + match node { + Binary(binary) => { + // we check if we have one or two changed children + + let left_path = path.new_with_direction(Direction::Left); + let node_left = self.get_node_or_felt::(&binary.left)?; + let right_path = path.new_with_direction(Direction::Right); + let node_right = self.get_node_or_felt::(&binary.right)?; + + let (left_hash, right_hash) = match (node_left, node_right) { + #[cfg(feature = "std")] + (NodeOrFelt::Node(left), NodeOrFelt::Node(right)) => { + // two children: use rayon + let (left, right) = rayon::join( + || self.compute_hashes::(left, left_path, hashes), + || { + let mut hashes = vec![]; + let felt = + self.compute_hashes::(right, right_path, &mut hashes)?; + Ok::<_, BonsaiStorageError>((felt, hashes)) + }, + ); + let (left_hash, (right_hash, hashes2)) = (left?, right?); + hashes.extend(hashes2); + + (left_hash, right_hash) + } + (left, right) => { + let left_hash = match left { + NodeOrFelt::Felt(felt) => felt, + NodeOrFelt::Node(node) => { + self.compute_hashes::(node, left_path, hashes)? + } + }; + let right_hash = match right { + NodeOrFelt::Felt(felt) => felt, + NodeOrFelt::Node(node) => { + self.compute_hashes::(node, right_path, hashes)? + } + }; + (left_hash, right_hash) + } + }; + + let hash = hash_binary_node::(left_hash, right_hash); + + hashes.push(hash); + Ok(hash) + } + + Edge(edge) => { + let mut child_path = path.clone(); + child_path.0.extend(&edge.path.0); + let child_hash = match self.get_node_or_felt::(&edge.child)? { + NodeOrFelt::Felt(felt) => felt, + NodeOrFelt::Node(node) => { + self.compute_hashes::(node, child_path, hashes)? + } + }; + + let hash = hash_edge_node::(&edge.path, child_hash); + hashes.push(hash); + + Ok(hash) + } + } + } + + /// Persists any changes in this subtree to storage. + /// + /// This necessitates recursively calculating the hash of, and + /// in turn persisting, any changed child nodes. This is necessary + /// as the parent node's hash relies on its children hashes. + /// Hash computation is done in parallel with [`compute_hashes`] beforehand. + /// + /// In effect, the entire tree gets persisted. + /// + /// # Arguments + /// + /// * `node_handle` - The top node from the subtree to commit. + /// * `hashes` - The precomputed hashes for the subtree as returned by [`compute_hashes`]. + /// The order is depth first, left to right. + /// + /// # Panics + /// + /// Panics if the precomputed `hashes` do not match the length of the modified subtree. + fn commit_subtree( + &mut self, + updates: &mut HashMap>, + node_id: NodeKey, + path: Path, + hashes: &mut impl Iterator, + ) -> Result> { + match self.nodes.remove(node_id).ok_or(BonsaiStorageError::Trie( + "Couldn't fetch node in the temporary storage".to_string(), + ))? { + Node::Binary(mut binary) => { + let left_path = path.new_with_direction(Direction::Left); + let left_hash = match binary.left { + NodeHandle::Hash(left_hash) => left_hash, + NodeHandle::InMemory(node_id) => { + self.commit_subtree::(updates, node_id, left_path, hashes)? + } + }; + let right_path = path.new_with_direction(Direction::Right); + let right_hash = match binary.right { + NodeHandle::Hash(right_hash) => right_hash, + NodeHandle::InMemory(node_id) => { + self.commit_subtree::(updates, node_id, right_path, hashes)? + } + }; + + let hash = hashes.next().expect("mismatched hash state"); + + binary.hash = Some(hash); + binary.left = NodeHandle::Hash(left_hash); + binary.right = NodeHandle::Hash(right_hash); + let key_bytes: ByteVec = path.into(); + updates.insert( + TrieKey::new(&self.identifier, TrieKeyType::Trie, &key_bytes), + InsertOrRemove::Insert(Node::Binary(binary).encode_bytevec()), + ); + Ok(hash) + } + Node::Edge(mut edge) => { + let mut child_path = path.clone(); + child_path.0.extend(&edge.path.0); + let child_hash = match edge.child { + NodeHandle::Hash(right_hash) => right_hash, + NodeHandle::InMemory(node_id) => { + self.commit_subtree::(updates, node_id, child_path, hashes)? + } + }; + let hash = hashes.next().expect("mismatched hash state"); + edge.hash = Some(hash); + edge.child = NodeHandle::Hash(child_hash); + let key_bytes: ByteVec = path.into(); + updates.insert( + TrieKey::new(&self.identifier, TrieKeyType::Trie, &key_bytes), + InsertOrRemove::Insert(Node::Edge(edge).encode_bytevec()), + ); + Ok(hash) + } + } + } + + /// Sets the value of a key. To delete a key, set the value to [Felt::ZERO]. + /// + /// # Arguments + /// + /// * `key` - The key to set. + /// * `value` - The value to set. + pub fn set( + &mut self, + db: &KeyValueDB, + key: &BitSlice, + value: Felt, + ) -> Result<(), BonsaiStorageError> { + if value == Felt::ZERO { + return self.delete_leaf(db, key); + } + if key.len() != self.max_height as _ { + return Err(BonsaiStorageError::KeyLength { + expected: self.max_height as _, + got: key.len(), + }); + } + let key_bytes = bitslice_to_bytes(key); + log::trace!("key_bytes: {:?}", key_bytes); + + // TODO(perf): do not double lookup when changing the value later (borrow needs to be split for preload_nodes though) + let mut cache_leaf_entry = self.cache_leaf_modified.entry_ref(&key_bytes[..]); + + if let hash_map::EntryRef::Occupied(entry) = &mut cache_leaf_entry { + if matches!(entry.get(), InsertOrRemove::Insert(_)) { + entry.insert(InsertOrRemove::Insert(value)); + return Ok(()); + } + } + + 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 mut iter = self.iter(db); + iter.seek_to(key)?; + log::trace!("Iter is {:?}", iter); + let path_nodes = iter.current_nodes_heights; + + // There are three possibilities. + // + // 1. The leaf exists, in which case we simply change its value. + // + // 2. The tree is empty, we insert the new leaf and the root becomes an edge node connecting to it. + // + // 3. The leaf does not exist, and the tree is not empty. The final node in the traversal will be an + // edge node who's path diverges from our new leaf node's. + // + // This edge must be split into a new subtree containing both the existing edge's child and the + // new leaf. This requires an edge followed by a binary node and then further edges to both the + // current child and the new leaf. Any of these new edges may also end with an empty path in + // which case they should be elided. It depends on the common path length of the current edge + // and the new leaf i.e. the split may be at the first bit (in which case there is no leading + // edge), or the split may be in the middle (requires both leading and post edges), or the + // split may be the final bit (no post edge). + + log::trace!("preload nodes: {:?}", path_nodes); + use Node::*; + match path_nodes.last() { + Some((node_id, _)) => { + let mut node = self.get_node_mut::(*node_id)?.clone(); + match &mut 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. + log::trace!("change val: {:?} => {:#x}", key_bytes, value); + self.cache_leaf_modified + .insert(key_bytes, InsertOrRemove::Insert(value)); + self.nodes[*node_id] = node; + return Ok(()); + } + // 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[common.len() + 1..].to_bitvec(); + + // The new leaf branch of the binary node. + // (this may be edge -> leaf, or just leaf depending). + log::trace!( + "cache_leaf_modified insert: {:?} => {:#x}", + key_bytes, + value + ); + self.cache_leaf_modified + .insert(key_bytes, InsertOrRemove::Insert(value)); + + let new = if new_path.is_empty() { + NodeHandle::Hash(value) + } else { + let edge_id = self.nodes.insert(Node::Edge(EdgeNode { + hash: None, + height: child_height as u64, + path: Path(new_path), + child: NodeHandle::Hash(value), + })); + NodeHandle::InMemory(edge_id) + }; + + // The existing child branch of the binary node. + let old = if old_path.is_empty() { + edge.child + } else { + let edge_id = self.nodes.insert(Node::Edge(EdgeNode { + hash: None, + height: child_height as u64, + path: Path(old_path), + child: edge.child, + })); + 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), + }; + + 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.nodes.insert(branch); + Node::Edge(EdgeNode { + hash: None, + height: edge.height, + path: Path(common.to_bitvec()), + child: NodeHandle::InMemory(branch_id), + }) + }; + let key_bytes = bitslice_to_bytes(&key[..edge.height as usize]); + log::trace!("2 death row add ({:?})", key_bytes); + self.death_row.insert(TrieKey::Trie(key_bytes)); + node = new_node; + } + Binary(binary) => { + let child_height = binary.height + 1; + + if child_height 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), + }; + self.cache_leaf_modified + .insert(key_bytes, InsertOrRemove::Insert(value)); + } + } + }; + + // Update the node + self.nodes[*node_id] = node; + Ok(()) + } + None => { + // Getting no travel nodes implies that the tree is empty. + // + // Create a new leaf node with the value, and the root becomes + // an edge node connecting to the leaf. + let edge = Node::Edge(EdgeNode { + hash: None, + height: 0, + path: Path(key.to_bitvec()), + child: NodeHandle::Hash(value), + }); + let node_id = self.nodes.insert(edge); + self.root_node = Some(RootHandle::Loaded(node_id)); + + let key_bytes = bitslice_to_bytes(key); + self.cache_leaf_modified + .insert(key_bytes, InsertOrRemove::Insert(value)); + Ok(()) + } + } + } + + /// Deletes a leaf node from the tree. + /// + /// This is not an external facing API; the functionality is instead accessed by calling + /// [`MerkleTree::set`] with value set to [`Felt::ZERO`]. + /// + /// # Arguments + /// + /// * `key` - The key to delete. + fn delete_leaf( + &mut self, + db: &KeyValueDB, + key: &BitSlice, + ) -> Result<(), BonsaiStorageError> { + if key.len() != self.max_height as _ { + return Err(BonsaiStorageError::KeyLength { + expected: self.max_height as _, + got: key.len(), + }); + } + log::trace!("delete leaf"); + // Algorithm explanation: + // + // The leaf's parent node is either an edge, or a binary node. + // If it's an edge node, then it must also be deleted. And its parent + // must be a binary node. In either case we end up with a binary node + // who's one child is deleted. This changes the binary to an edge node. + // + // Note that its possible that there is no binary node -- if the resulting tree would be empty. + // + // This new edge node may need to merge with the old binary node's parent node + // and other remaining child node -- if they're also edges. + // + // Then we are done. + let key_bytes = bitslice_to_bytes(key); + let leaf_entry = self.cache_leaf_modified.entry(key_bytes.clone()); + + let tree_has_value = if let hash_map::Entry::Occupied(entry) = &leaf_entry { + !matches!(entry.get(), InsertOrRemove::Remove) + } else { + db.get(&TrieKey::new( + &self.identifier, + TrieKeyType::Flat, + &key_bytes, + ))? + .is_some() + }; + + if !tree_has_value { + return Ok(()); + } + leaf_entry.insert(InsertOrRemove::Remove); + + let mut iter = self.iter(db); + iter.seek_to(key)?; + log::trace!("Iter is {:?}", iter); + let mut path_nodes = iter.current_nodes_heights; + + let mut last_binary_path = Path(key.to_bitvec()); + + // Remove the final edge if present, we are starting from the closest binary node. + if let Some((node_key, _height)) = path_nodes.last() { + match self.get_node_mut::(*node_key)? { + Node::Binary(_) => {} + Node::Edge(edge) => { + // todo(perf) this is kinda dumb isnt it + for _ in 0..edge.path.len() { + last_binary_path.pop(); + } + let mut new_path = Path(BitVec::new()); + for i in last_binary_path.iter() { + new_path.push(*i); + } + last_binary_path = new_path.clone(); + let path: ByteVec = (&last_binary_path).into(); + log::trace!( + "iter leaf= edge={edge:?}, new_path={new_path:?}", + // TrieKey::new(self.identifier.clone(), TrieKeyType::Trie, &path) + ); + + self.death_row + .insert(TrieKey::new(&self.identifier, TrieKeyType::Trie, &path)); + self.nodes.remove(*node_key); + path_nodes.pop(); + } + } + } + + let mut node_iter = path_nodes.into_iter().rev().peekable(); + + let branch_node = node_iter.next(); + let parent_branch_node = node_iter.next(); + + log::trace!( + "remove leaf branch_node={branch_node:?} parent_branch_node={parent_branch_node:?}" + ); + + match branch_node { + Some((node_id, _)) => { + let (new_edge, par_path) = { + let node = self.get_node_mut::(node_id)?; + + let binary = node + .as_binary() + .expect("The node must be a binary node due to the iteration condition"); + let (direction, height) = { (binary.direction(key).invert(), binary.height) }; + last_binary_path.pop(); + last_binary_path.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. + let path = Path(iter::once(bool::from(direction)).collect::()); + let mut edge = EdgeNode { + hash: None, + height, + path, + 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, db, &last_binary_path)?; + let cl = last_binary_path.clone(); + last_binary_path.pop(); + (edge, cl) + }; + // 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.get_node_mut::(parent_node_id)?; + if let Node::Edge(parent_edge) = parent_node { + parent_edge.path.extend_from_bitslice(&new_edge.path.0); + parent_edge.child = new_edge.child; + + let mut par_path = par_path; + par_path.pop(); + let path: ByteVec = par_path.into(); + self.death_row.insert(TrieKey::new( + &self.identifier, + TrieKeyType::Trie, + &path, + )); + self.nodes.remove(node_id); + } else { + self.nodes[node_id] = Node::Edge(new_edge); + } + } else { + self.nodes[node_id] = Node::Edge(new_edge); + } + } + None => { + // We reached the root without a hitting binary node. The new tree + // must therefore be empty. + + log::trace!("empty {:?}", self.root_node); + if let Some(RootHandle::Loaded(node_id)) = self.root_node { + self.nodes.remove(node_id); + } + self.death_row + .insert(TrieKey::new(&self.identifier, TrieKeyType::Trie, &[0])); + self.root_node = Some(RootHandle::Empty); + return Ok(()); + } + }; + Ok(()) + } + + /// Returns the value stored at key, or `None` if it does not exist. + /// + /// # Arguments + /// + /// * `key` - The key of the value to get. + /// + /// # Returns + /// + /// The value of the key. + pub fn get( + &self, + db: &KeyValueDB, + key: &BitSlice, + ) -> Result, BonsaiStorageError> { + log::trace!("get with key {:b}", key); + let key = bitslice_to_bytes(key); + log::trace!("get from cache with {:?}", key); + if let Some(value) = self.cache_leaf_modified.get(&key) { + log::trace!("get has cache_leaf_modified {:?} {:?}", key, value); + match value { + InsertOrRemove::Remove => return Ok(None), + InsertOrRemove::Insert(value) => return Ok(Some(*value)), + } + } + log::trace!( + "get from db with key {:?}", + &TrieKey::new(&self.identifier, TrieKeyType::Flat, &key) + ); + db.get(&TrieKey::new(&self.identifier, TrieKeyType::Flat, &key)) + .map(|r| r.map(|opt| Felt::decode(&mut opt.as_slice()).unwrap())) + } + + pub fn get_at( + &self, + db: &KeyValueDB, + key: &BitSlice, + id: ID, + ) -> Result, BonsaiStorageError> { + let key = bitslice_to_bytes(key); + db.get_at(&TrieKey::new(&self.identifier, TrieKeyType::Flat, &key), id) + .map(|r| r.map(|opt| Felt::decode(&mut opt.as_slice()).unwrap())) + } + + pub fn contains( + &self, + db: &KeyValueDB, + key: &BitSlice, + ) -> Result> { + let key = bitslice_to_bytes(key); + if let Some(value) = self.cache_leaf_modified.get(&key) { + match value { + InsertOrRemove::Remove => return Ok(false), + InsertOrRemove::Insert(_) => return Ok(true), + } + } + db.contains(&TrieKey::new(&self.identifier, TrieKeyType::Flat, &key)) + } + + /// Get the node of the trie that corresponds to the path. + fn get_trie_branch_in_db_from_path( + death_row: &HashSet, + identifier: &[u8], + db: &KeyValueDB, + path: &Path, + ) -> Result, BonsaiStorageError> { + log::trace!("getting: {:b}", path.0); + + let path: ByteVec = path.into(); + let key = TrieKey::new(identifier, TrieKeyType::Trie, &path); + + if death_row.contains(&key) { + return Ok(None); + } + + db.get(&key)? + .map(|node| { + log::trace!("got: {:?}", node); + Node::decode(&mut node.as_slice()).map_err(|err| { + BonsaiStorageError::Trie(format!("Couldn't decode node: {}", err)) + }) + }) + .map_or(Ok(None), |r| r.map(Some)) + } + + /// This is a convenience function which merges the edge node with its child __iff__ it is also + /// an edge. + /// + /// Does nothing if the child is not also an edge node. + /// + /// This can occur when mutating the tree (e.g. deleting a child of a binary node), and is an + /// illegal state (since edge nodes __must be__ maximal subtrees). + /// + /// # Arguments + /// + /// * `parent` - The parent node to merge the child with. + fn merge_edges( + &mut self, + parent: &mut EdgeNode, + db: &KeyValueDB, + path: &Path, + ) -> Result<(), BonsaiStorageError> { + match parent.child { + NodeHandle::Hash(_) => { + let node = Self::get_trie_branch_in_db_from_path( + &self.death_row, + &self.identifier, + db, + path, + )?; + log::trace!("case: Hash {:?}", node); + if let Some(Node::Edge(child_edge)) = node { + parent.path.0.extend_from_bitslice(&child_edge.path.0); + parent.child = child_edge.child; + // remove node from db + let path: ByteVec = path.into(); + log::trace!("4 death row {:?}", path); + self.death_row + .insert(TrieKey::new(&self.identifier, TrieKeyType::Trie, &path)); + } + } + NodeHandle::InMemory(child_id) => { + let node = self.get_node_mut::(child_id)?; + log::trace!("case: InMemory {:?}", node); + + if let Node::Edge(child_edge) = node { + parent.path.0.extend_from_bitslice(&child_edge.path.0); + parent.child = child_edge.child; + + self.nodes.remove(child_id); + + let path: ByteVec = path.into(); + log::trace!("3 death row {:?}", path); + self.death_row + .insert(TrieKey::new(&self.identifier, TrieKeyType::Trie, &path)); + } + } + }; + Ok(()) + } + + #[cfg(test)] + #[allow(dead_code)] + pub(crate) fn dump(&self) { + match self.root_node { + Some(RootHandle::Empty) => { + trace!("tree is empty") + } + Some(RootHandle::Loaded(node)) => { + trace!("root is node {:?}", node); + self.dump_node(node); + } + None => trace!("root is not loaded"), + } + } + + #[cfg(test)] + #[allow(dead_code)] + fn dump_node(&self, head: NodeKey) { + use Node::*; + + let current_tmp = self.nodes[head].clone(); + trace!("bonsai_node {:?} = {:?}", head, current_tmp); + + match current_tmp { + Binary(binary) => { + match &binary.get_child(Direction::Left) { + NodeHandle::Hash(hash) => { + trace!("left is hash {:#x}", hash); + } + NodeHandle::InMemory(left_id) => { + self.dump_node(*left_id); + } + } + match &binary.get_child(Direction::Right) { + NodeHandle::Hash(hash) => { + trace!("right is hash {:#x}", hash); + } + NodeHandle::InMemory(right_id) => { + self.dump_node(*right_id); + } + } + } + Edge(edge) => match &edge.child { + NodeHandle::Hash(hash) => { + trace!("child is hash {:#x}", hash); + } + NodeHandle::InMemory(child_id) => { + self.dump_node(*child_id); + } + }, + }; + } +} + +pub(crate) fn bitslice_to_bytes(bitslice: &BitSlice) -> ByteVec { + // TODO(perf): this should not copy to a bitvec :( + if bitslice.is_empty() { + return Default::default(); + } // special case: tree root + iter::once(bitslice.len() as u8) + .chain(bitslice.to_bitvec().as_raw_slice().iter().copied()) + .collect() +} + +pub(crate) fn bytes_to_bitvec(bytes: &[u8]) -> BitVec { + BitSlice::from_slice(&bytes[1..]).to_bitvec() +} diff --git a/src/trie/trees.rs b/src/trie/trees.rs new file mode 100644 index 0000000..b6b5a81 --- /dev/null +++ b/src/trie/trees.rs @@ -0,0 +1,247 @@ +use super::{proof::MultiProof, tree::MerkleTree}; +use crate::{ + id::Id, key_value_db::KeyValueDB, trie::tree::InsertOrRemove, BitSlice, BonsaiDatabase, + BonsaiStorageError, ByteVec, HashMap, Vec, +}; +use core::fmt; +use starknet_types_core::{felt::Felt, hash::StarkHash}; + +pub(crate) struct MerkleTrees { + pub db: KeyValueDB, + pub trees: HashMap>, + pub max_height: u8, +} + +impl fmt::Debug + for MerkleTrees +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MerkleTrees") + .field("db", &self.db) + .field("trees", &self.trees) + .finish() + } +} + +#[cfg(feature = "bench")] +impl Clone + for MerkleTrees +{ + fn clone(&self) -> Self { + Self { + db: self.db.clone(), + trees: self.trees.clone(), + max_height: self.max_height, + } + } +} + +impl MerkleTrees { + pub(crate) fn new(db: KeyValueDB, tree_height: u8) -> Self { + Self { + db, + trees: HashMap::new(), + max_height: tree_height, + } + } + + pub(crate) fn set( + &mut self, + identifier: &[u8], + key: &BitSlice, + value: Felt, + ) -> Result<(), BonsaiStorageError> { + let tree = self + .trees + .entry_ref(identifier) + .or_insert_with(|| MerkleTree::new(identifier.into(), self.max_height)); + + tree.set(&self.db, key, value) + } + + pub(crate) fn get( + &self, + identifier: &[u8], + key: &BitSlice, + ) -> Result, BonsaiStorageError> { + if let Some(tree) = self.trees.get(identifier) { + tree.get(&self.db, key) + } else { + MerkleTree::::new(identifier.into(), self.max_height).get(&self.db, key) + } + } + + pub(crate) fn get_at( + &self, + identifier: &[u8], + key: &BitSlice, + id: CommitID, + ) -> Result, BonsaiStorageError> { + if let Some(tree) = self.trees.get(identifier) { + tree.get_at(&self.db, key, id) + } else { + MerkleTree::::new(identifier.into(), self.max_height).get_at(&self.db, key, id) + } + } + + pub(crate) fn contains( + &self, + identifier: &[u8], + key: &BitSlice, + ) -> Result> { + if let Some(tree) = self.trees.get(identifier) { + tree.contains(&self.db, key) + } else { + MerkleTree::::new(identifier.into(), self.max_height).contains(&self.db, key) + } + } + + pub(crate) fn db_mut(&mut self) -> &mut KeyValueDB { + &mut self.db + } + + pub(crate) fn reset_to_last_commit( + &mut self, + ) -> Result<(), BonsaiStorageError> { + self.trees.clear(); // just clear the map + Ok(()) + } + + pub(crate) fn db_ref(&self) -> &KeyValueDB { + &self.db + } + + #[cfg(test)] + pub fn dump(&self) { + log::trace!("====== NUMBER OF TREES: {} ======", self.trees.len()); + self.trees.iter().for_each(|(k, tree)| { + log::trace!("TREE identifier={:?}:", k); + tree.dump(); + }); + } + + pub(crate) fn root_hash( + &self, + identifier: &[u8], + ) -> Result> { + if let Some(tree) = self.trees.get(identifier) { + Ok(tree.root_hash(&self.db)?) + } else { + MerkleTree::::new(identifier.into(), self.max_height).root_hash(&self.db) + } + } + + 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() + // FIXME: this does not filter out keys values correctly for `HashMapDb` due + // to branches and leafs not being differenciated + .filter_map(|(key, _value)| { + if key.len() > identifier.len() { + Some(key[identifier.len() + 1..].into()) + } else { + None + } + }) + .collect() + }) + .map_err(|e| e.into()) + } + + #[allow(clippy::type_complexity)] + pub(crate) fn get_key_value_pairs( + &self, + identifier: &[u8], + ) -> Result, Vec)>, BonsaiStorageError> { + self.db + .db + .get_by_prefix(&crate::DatabaseKey::Flat(identifier)) + .map(|key_value_pairs| { + key_value_pairs + .into_iter() + // FIXME: this does not filter out keys values correctly for `HashMapDb` due + // to branches and leafs not being differenciated + .filter_map(|(key, value)| { + if key.len() > identifier.len() { + Some((key[identifier.len() + 1..].into(), value.into_vec())) + } else { + None + } + }) + .collect() + }) + .map_err(|e| e.into()) + } + + pub(crate) fn commit(&mut self) -> Result<(), BonsaiStorageError> { + #[cfg(feature = "std")] + use rayon::prelude::*; + + #[cfg(not(feature = "std"))] + let db_changes = self + .trees + .iter_mut() + .map(|(_, tree)| tree.get_updates::()); + #[cfg(feature = "std")] + let db_changes = self + .trees + .par_iter_mut() + .map(|(_, tree)| tree.get_updates::()) + .collect_vec_list() + .into_iter() + .flatten(); + + let mut batch = self.db.create_batch(); + for changes in db_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> { + // if let Some(tree) = self.trees.get(identifier) { + // tree.get_proof(&self.db, key) + // } else { + // MerkleTree::::new(identifier.into()).get_proof(&self.db, key) + // } + // } + + pub(crate) fn get_identifiers(&self) -> Vec> { + self.trees.keys().cloned().map(ByteVec::into_vec).collect() + } + + pub fn get_multi_proof( + &mut self, + identifier: &[u8], + keys: impl IntoIterator>, + ) -> Result> { + let tree = self + .trees + .entry_ref(identifier) + .or_insert_with(|| MerkleTree::new(identifier.into(), self.max_height)); + + tree.get_multi_proof(&self.db, keys) + } +}