diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 4426b0f..fc24173 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -37,7 +37,7 @@ jobs: key: ${{ runner.os }}-build-${{ env.cache-name }}-${{ hashFiles('Cargo.lock') }} - name: Generate test result and coverage report run: | - cargo test $CARGO_OPTIONS --no-default-features + cargo test $CARGO_OPTIONS # - name: Upload test results # uses: EnricoMi/publish-unit-test-result-action@v1 # with: diff --git a/Cargo.toml b/Cargo.toml index 3883221..47350aa 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,8 +17,8 @@ derive_more = { version = "0.99.17", default-features = false, features = [ ] } hashbrown = "0.14.3" log = "0.4.20" -smallvec = "1.11.2" rayon = { version = "1.9.0", optional = true } +smallvec = { version = "1.11.2", features = ["serde"] } parity-scale-codec = { version = "3.0.0", default-features = false, features = [ "derive", @@ -27,9 +27,10 @@ serde = { version = "1.0.195", default-features = false, features = [ "derive", "alloc", ] } -starknet-types-core = { version = "0.1", default-features = false, features = [ +starknet-types-core = { version = "0.1.5", default-features = false, features = [ "hash", "parity-scale-codec", + "alloc", ] } # Optionals @@ -45,12 +46,14 @@ pathfinder-common = { git = "https://github.com/massalabs/pathfinder.git", packa 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 = "0.8.5" +rand = { version = "0.8.5", features = ["small_rng"] } tempfile = "3.8.0" rstest = "0.18.2" test-log = "0.2.15" indexmap = "2.2.6" criterion = "0.5.1" +proptest = "1.4.0" +proptest-derive = "0.4.0" serde_json = "1.0.68" [[bench]] diff --git a/benches/storage.rs b/benches/storage.rs index 03c8c75..687f135 100644 --- a/benches/storage.rs +++ b/benches/storage.rs @@ -15,6 +15,42 @@ use starknet_types_core::{ mod flamegraph; +fn drop_storage(c: &mut Criterion) { + c.bench_function("drop storage", move |b| { + b.iter_batched( + || { + let mut bonsai_storage: BonsaiStorage = BonsaiStorage::new( + HashMapDb::::default(), + BonsaiStorageConfig::default(), + ) + .unwrap(); + + let mut rng = SmallRng::seed_from_u64(42); + let felt = Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(); + for _ in 0..4000 { + let bitvec = BitVec::from_vec(vec![ + rng.gen(), + rng.gen(), + rng.gen(), + rng.gen(), + rng.gen(), + rng.gen(), + ]); + bonsai_storage.insert(&[], &bitvec, &felt).unwrap(); + } + + let mut id_builder = BasicIdBuilder::new(); + let id1 = id_builder.new_id(); + bonsai_storage.commit(id1).unwrap(); + + bonsai_storage + }, + std::mem::drop, + BatchSize::LargeInput, + ); + }); +} + fn storage_with_insert(c: &mut Criterion) { c.bench_function("storage commit with insert", move |b| { let mut rng = thread_rng(); @@ -40,7 +76,6 @@ fn storage_with_insert(c: &mut Criterion) { ]); bonsai_storage.insert(&[], &bitvec, &felt).unwrap(); } - // let mut id_builder = BasicIdBuilder::new(); // bonsai_storage.commit(id_builder.new_id()).unwrap(); }, @@ -56,7 +91,7 @@ fn storage(c: &mut Criterion) { BonsaiStorageConfig::default(), ) .unwrap(); - let mut rng = thread_rng(); + let mut rng = SmallRng::seed_from_u64(42); let felt = Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(); for _ in 0..1000 { @@ -89,7 +124,7 @@ fn one_update(c: &mut Criterion) { BonsaiStorageConfig::default(), ) .unwrap(); - let mut rng = thread_rng(); + let mut rng = SmallRng::seed_from_u64(42); let felt = Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(); for _ in 0..1000 { @@ -126,7 +161,7 @@ fn five_updates(c: &mut Criterion) { BonsaiStorageConfig::default(), ) .unwrap(); - let mut rng = thread_rng(); + let mut rng = SmallRng::seed_from_u64(42); let felt = Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(); for _ in 0..1000 { @@ -226,6 +261,6 @@ fn hash(c: &mut Criterion) { criterion_group! { name = benches; config = Criterion::default(); // .with_profiler(flamegraph::FlamegraphProfiler::new(100)); - targets = storage, one_update, five_updates, hash, storage_with_insert, multiple_contracts + targets = storage, one_update, five_updates, hash, drop_storage, storage_with_insert, multiple_contracts } criterion_main!(benches); diff --git a/proptest-regressions/trie/merge.txt b/proptest-regressions/trie/merge.txt new file mode 100644 index 0000000..52e3057 --- /dev/null +++ b/proptest-regressions/trie/merge.txt @@ -0,0 +1,10 @@ +# 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 70ec50070d2ab1e02fbfc9e697b4db0c9c3584a4db39218dccca5b22e480378d # shrinks to pb = MerkleTreeMergeProblem { inserts_a: [], inserts_b: [(BitVec { addr: 0x778e9c046c30, head: 000, bits: 251, capacity: 256 } [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], Felt(FieldElement { value: UnsignedInteger { limbs: [576460752303388688, 18446744073709551615, 18446744073709551615, 18446744073709549569] } })), (BitVec { addr: 0x778e9c000b70, head: 000, bits: 251, capacity: 256 } [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0], Felt(FieldElement { value: UnsignedInteger { limbs: [576460752303388688, 18446744073709551615, 18446744073709551615, 18446744073709549569] } }))] } +cc c3a0408a5f08bf1f5d8e31050cc4d95e1f377d8a56ba6fc50d5ebf3682dcc02b # shrinks to pb = MerkleTreeInsertProblem([Insert([01111], 0x20), Insert([01110], 0x20), Commit, Insert([01110], 0x40)]) +cc 185dc53af11731e46e174743ebbf050242f75bf405d301b2fd8b506f382cc4f4 # shrinks to pb = MerkleTreeInsertProblem([Insert([10011], 0x20), Remove([10011]), Insert([00000], 0x0), Insert([00000], 0x20), Commit, Insert([00000], 0x0)]) +cc c6239e16e78d8c8ec7441ef221279eef9c6541bea5b43f2ae58d0848e0960271 # shrinks to pb = MerkleTreeInsertProblem([Insert([10000], 0x20), Insert([11000], 0x40), Insert([11010], 0x40), Remove([10000]), Remove([10000]), Commit]) diff --git a/proptest-regressions/trie/merkle_tree.txt b/proptest-regressions/trie/merkle_tree.txt new file mode 100644 index 0000000..aebcb4c --- /dev/null +++ b/proptest-regressions/trie/merkle_tree.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 8ed2164817fe54de5eb23b134024a200ceb9c6fe2fc0835740731e9cd4326e45 # shrinks to pb = MerkleTreeInsertProblem([Insert(InsertStep(BitVec { addr: 0x7254a0004ca0, head: 000, bits: 3, capacity: 64 } [1, 0, 1], Felt(FieldElement { value: UnsignedInteger { limbs: [576460752303388688, 18446744073709551615, 18446744073709551615, 18446744073709549569] } }))), Insert(InsertStep(BitVec { addr: 0x7254a00037b0, head: 000, bits: 3, capacity: 64 } [1, 0, 0], Felt(FieldElement { value: UnsignedInteger { limbs: [576460752303406096, 18446744073709551615, 18446744073709551615, 18446744073709550593] } }))), Insert(InsertStep(BitVec { addr: 0x7254a0001b10, head: 000, bits: 3, capacity: 64 } [1, 0, 0], Felt(FieldElement { value: UnsignedInteger { limbs: [576460752303284240, 18446744073709551615, 18446744073709551615, 18446744073709543425] } })))]) diff --git a/src/bonsai_database.rs b/src/bonsai_database.rs index 6c0b3ae..7118774 100644 --- a/src/bonsai_database.rs +++ b/src/bonsai_database.rs @@ -1,4 +1,4 @@ -use crate::{id::Id, Vec}; +use crate::{id::Id, ByteVec, Vec}; #[cfg(feature = "std")] use std::error::Error; @@ -38,14 +38,14 @@ pub trait BonsaiDatabase { fn create_batch(&self) -> Self::Batch; /// Returns the value of the key if it exists - fn get(&self, key: &DatabaseKey) -> Result>, Self::DatabaseError>; + fn get(&self, key: &DatabaseKey) -> Result, Self::DatabaseError>; #[allow(clippy::type_complexity)] /// Returns all values with keys that start with the given prefix fn get_by_prefix( &self, prefix: &DatabaseKey, - ) -> Result, Vec)>, Self::DatabaseError>; + ) -> Result, Self::DatabaseError>; /// Returns true if the key exists fn contains(&self, key: &DatabaseKey) -> Result; @@ -57,7 +57,7 @@ pub trait BonsaiDatabase { key: &DatabaseKey, value: &[u8], batch: Option<&mut Self::Batch>, - ) -> Result>, Self::DatabaseError>; + ) -> Result, Self::DatabaseError>; /// Remove a key-value pair, returns the old value if it existed. /// If a batch is provided, the change will be written in the batch instead of the database. @@ -65,7 +65,7 @@ pub trait BonsaiDatabase { &mut self, key: &DatabaseKey, batch: Option<&mut Self::Batch>, - ) -> Result>, Self::DatabaseError>; + ) -> Result, Self::DatabaseError>; /// Remove all keys that start with the given prefix fn remove_by_prefix(&mut self, prefix: &DatabaseKey) -> Result<(), Self::DatabaseError>; diff --git a/src/changes.rs b/src/changes.rs index 6413fc7..f0ce5f6 100644 --- a/src/changes.rs +++ b/src/changes.rs @@ -1,10 +1,11 @@ -use crate::{hash_map::Entry, id::Id, trie::TrieKey, HashMap, Vec, VecDeque}; +use crate::{hash_map::Entry, id::Id, trie::TrieKey, ByteVec, HashMap, Vec, VecDeque}; +use core::iter; use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Serialize, Deserialize, Default)] pub struct Change { - pub old_value: Option>, - pub new_value: Option>, + pub old_value: Option, + pub new_value: Option, } #[derive(Debug, Default)] @@ -31,7 +32,7 @@ impl ChangeBatch { } } - pub fn serialize(&self, id: &ID) -> Vec<(Vec, &[u8])> { + pub fn serialize(&self, id: &ID) -> Vec<(ByteVec, &[u8])> { self.0 .iter() .flat_map(|(change_key, change)| { @@ -56,7 +57,7 @@ impl ChangeBatch { .collect() } - pub fn deserialize(id: &ID, changes: Vec<(Vec, Vec)>) -> Self { + pub fn deserialize(id: &ID, changes: Vec<(ByteVec, ByteVec)>) -> Self { let id = id.to_bytes(); let mut change_batch = ChangeBatch(HashMap::new()); let mut current_change = Change::default(); @@ -69,8 +70,7 @@ impl ChangeBatch { let mut key = key.to_vec(); let change_type = key.pop().unwrap(); let key_type = key.pop().unwrap(); - let change_key = - TrieKey::from_variant_and_bytes(key_type, key[id.len() + 1..].to_vec()); + let change_key = TrieKey::from_variant_and_bytes(key_type, key[id.len() + 1..].into()); if let Some(last_key) = last_key { if last_key != change_key { change_batch.insert_in_place(last_key, current_change); @@ -93,29 +93,28 @@ impl ChangeBatch { } } -pub fn key_old_value(id: &ID, key: &TrieKey) -> Vec { - [ - id.to_bytes().as_slice(), - &[KEY_SEPARATOR], - key.as_slice(), - &[key.into()], - &[OLD_VALUE], - ] - .concat() +pub fn key_old_value(id: &ID, key: &TrieKey) -> ByteVec { + id.to_bytes() + .into_iter() + .chain(iter::once(KEY_SEPARATOR)) + .chain(key.as_slice().iter().copied()) + .chain(iter::once(key.into())) + .chain(iter::once(OLD_VALUE)) + .collect() } -pub fn key_new_value(id: &ID, key: &TrieKey) -> Vec { - [ - id.to_bytes().as_slice(), - &[KEY_SEPARATOR], - key.as_slice(), - &[key.into()], - &[NEW_VALUE], - ] - .concat() +pub fn key_new_value(id: &ID, key: &TrieKey) -> ByteVec { + id.to_bytes() + .into_iter() + .chain(iter::once(KEY_SEPARATOR)) + .chain(key.as_slice().iter().copied()) + .chain(iter::once(key.into())) + .chain(iter::once(NEW_VALUE)) + .collect() } #[cfg_attr(feature = "bench", derive(Clone))] +#[derive(Debug)] pub struct ChangeStore where ID: Id, diff --git a/src/databases/hashmap_db.rs b/src/databases/hashmap_db.rs index 65249fe..1b2ce3f 100644 --- a/src/databases/hashmap_db.rs +++ b/src/databases/hashmap_db.rs @@ -3,6 +3,7 @@ use crate::{ id::Id, BTreeMap, BonsaiDatabase, HashMap, Vec, }; +use crate::{ByteVec, DatabaseKey}; use core::{fmt, fmt::Display}; #[derive(Debug)] @@ -19,47 +20,69 @@ impl Display for HashMapDbError { impl DBError for HashMapDbError {} -#[derive(Clone, Default)] +#[derive(Clone, Default, Debug)] pub struct HashMapDb { - db: HashMap, Vec>, + trie_db: HashMap, + flat_db: HashMap, + trie_log_db: HashMap, snapshots: BTreeMap>, } +impl HashMapDb { + fn get_map(&self, key: &DatabaseKey) -> &HashMap { + match key { + DatabaseKey::Trie(_) => &self.trie_db, + DatabaseKey::Flat(_) => &self.flat_db, + DatabaseKey::TrieLog(_) => &self.trie_log_db, + } + } + fn get_map_mut(&mut self, key: &DatabaseKey) -> &mut HashMap { + match key { + DatabaseKey::Trie(_) => &mut self.trie_db, + DatabaseKey::Flat(_) => &mut self.flat_db, + DatabaseKey::TrieLog(_) => &mut self.trie_log_db, + } + } + + #[cfg(test)] + pub(crate) fn assert_empty(&self) { + assert_eq!(self.trie_db, [].into()); + assert_eq!(self.flat_db, [].into()); + } +} + impl BonsaiDatabase for HashMapDb { type Batch = (); type DatabaseError = HashMapDbError; fn create_batch(&self) -> Self::Batch {} - fn remove_by_prefix( - &mut self, - prefix: &crate::bonsai_database::DatabaseKey, - ) -> Result<(), Self::DatabaseError> { + fn remove_by_prefix(&mut self, prefix: &DatabaseKey) -> Result<(), Self::DatabaseError> { let mut keys_to_remove = Vec::new(); - for key in self.db.keys() { + let db = self.get_map_mut(prefix); + for key in db.keys() { if key.starts_with(prefix.as_slice()) { keys_to_remove.push(key.clone()); } } for key in keys_to_remove { - self.db.remove(&key); + db.remove(&key); } Ok(()) } - fn get( - &self, - key: &crate::bonsai_database::DatabaseKey, - ) -> Result>, Self::DatabaseError> { - Ok(self.db.get(key.as_slice()).cloned()) + fn get(&self, key: &DatabaseKey) -> Result, Self::DatabaseError> { + let db = &self.get_map(key); + Ok(db.get(key.as_slice()).cloned()) } fn get_by_prefix( &self, - prefix: &crate::bonsai_database::DatabaseKey, - ) -> Result, Vec)>, Self::DatabaseError> { + prefix: &DatabaseKey, + ) -> Result, Self::DatabaseError> { let mut result = Vec::new(); - for (key, value) in self.db.iter() { + let db = self.get_map(prefix); + for (key, value) in db.iter() { if key.starts_with(prefix.as_slice()) { result.push((key.clone(), value.clone())); } @@ -69,26 +92,26 @@ impl BonsaiDatabase for HashMapDb { fn insert( &mut self, - key: &crate::bonsai_database::DatabaseKey, + key: &DatabaseKey, value: &[u8], _batch: Option<&mut Self::Batch>, - ) -> Result>, Self::DatabaseError> { - Ok(self.db.insert(key.as_slice().to_vec(), value.to_vec())) + ) -> Result, Self::DatabaseError> { + let db = self.get_map_mut(key); + Ok(db.insert(key.as_slice().into(), value.into())) } fn remove( &mut self, - key: &crate::bonsai_database::DatabaseKey, + key: &DatabaseKey, _batch: Option<&mut Self::Batch>, - ) -> Result>, Self::DatabaseError> { - Ok(self.db.remove(key.as_slice())) + ) -> Result, Self::DatabaseError> { + let db = self.get_map_mut(key); + Ok(db.remove(key.as_slice())) } - fn contains( - &self, - key: &crate::bonsai_database::DatabaseKey, - ) -> Result { - Ok(self.db.contains_key(key.as_slice())) + fn contains(&self, key: &DatabaseKey) -> Result { + let db = self.get_map(key); + Ok(db.contains_key(key.as_slice())) } fn write_batch(&mut self, _batch: Self::Batch) -> Result<(), Self::DatabaseError> { @@ -97,7 +120,7 @@ impl BonsaiDatabase for HashMapDb { #[cfg(test)] fn dump_database(&self) { - log::debug!("{:?}", self.db); + log::debug!("{:?}", self); } } @@ -113,7 +136,9 @@ impl BonsaiPersistentDatabase for HashMapDb { } fn merge(&mut self, transaction: Self::Transaction) -> Result<(), Self::DatabaseError> { - self.db = transaction.db; + self.trie_db = transaction.trie_db; + self.flat_db = transaction.flat_db; + self.trie_log_db = transaction.trie_log_db; Ok(()) } } diff --git a/src/databases/rocks_db.rs b/src/databases/rocks_db.rs index dce95d3..1e2f176 100644 --- a/src/databases/rocks_db.rs +++ b/src/databases/rocks_db.rs @@ -14,6 +14,7 @@ use rocksdb::{ use crate::{ bonsai_database::{BonsaiDatabase, BonsaiPersistentDatabase, DBError, DatabaseKey}, id::Id, + ByteVec, }; use log::trace; @@ -53,6 +54,41 @@ pub struct RocksDB<'db, ID: Id> { snapshots: BTreeMap>, } +impl<'db, ID: Id> fmt::Debug for RocksDB<'db, ID> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + writeln!(f, "ROCKSDB_DATABASE_DUMP {{")?; + let handle_trie = self.db.cf_handle(TRIE_CF).expect(CF_ERROR); + let handle_flat = self.db.cf_handle(FLAT_CF).expect(CF_ERROR); + let handle_trie_log = self.db.cf_handle(TRIE_LOG_CF).expect(CF_ERROR); + let mut iter = self.db.raw_iterator_cf(&handle_trie); + iter.seek_to_first(); + while iter.valid() { + let key = iter.key().unwrap(); + let value = iter.value().unwrap(); + writeln!(f, "{:?} => {:?},", key, value)?; + iter.next(); + } + let mut iter = self.db.raw_iterator_cf(&handle_flat); + iter.seek_to_first(); + while iter.valid() { + let key = iter.key().unwrap(); + let value = iter.value().unwrap(); + writeln!(f, "{:?} => {:?},", key, value)?; + iter.next(); + } + let mut iter = self.db.raw_iterator_cf(&handle_trie_log); + iter.seek_to_first(); + while iter.valid() { + let key = iter.key().unwrap(); + let value = iter.value().unwrap(); + writeln!(f, "{:?} => {:?},", key, value)?; + iter.next(); + } + write!(f, "}}")?; + Ok(()) + } +} + /// Configuration for RocksDB database pub struct RocksDBConfig { /// Maximum number of snapshots kept in database @@ -150,33 +186,7 @@ where #[cfg(test)] fn dump_database(&self) { - let handle_trie = self.db.cf_handle(TRIE_CF).expect(CF_ERROR); - let handle_flat = self.db.cf_handle(FLAT_CF).expect(CF_ERROR); - let handle_trie_log = self.db.cf_handle(TRIE_LOG_CF).expect(CF_ERROR); - let mut iter = self.db.raw_iterator_cf(&handle_trie); - iter.seek_to_first(); - while iter.valid() { - let key = iter.key().unwrap(); - let value = iter.value().unwrap(); - println!("{:?} {:?}", key, value); - iter.next(); - } - let mut iter = self.db.raw_iterator_cf(&handle_flat); - iter.seek_to_first(); - while iter.valid() { - let key = iter.key().unwrap(); - let value = iter.value().unwrap(); - println!("{:?} {:?}", key, value); - iter.next(); - } - let mut iter = self.db.raw_iterator_cf(&handle_trie_log); - iter.seek_to_first(); - while iter.valid() { - let key = iter.key().unwrap(); - let value = iter.value().unwrap(); - println!("{:?} {:?}", key, value); - iter.next(); - } + println!("{:?}", self) } fn insert( @@ -184,7 +194,7 @@ where key: &DatabaseKey, value: &[u8], batch: Option<&mut Self::Batch>, - ) -> Result>, Self::DatabaseError> { + ) -> Result, Self::DatabaseError> { trace!("Inserting into RocksDB: {:?} {:?}", key, value); let handle_cf = self.db.cf_handle(key.get_cf()).expect(CF_ERROR); let old_value = self.db.get_cf(&handle_cf, key.as_slice())?; @@ -193,19 +203,19 @@ where } else { self.db.put_cf(&handle_cf, key.as_slice(), value)?; } - Ok(old_value) + Ok(old_value.map(Into::into)) } - fn get(&self, key: &DatabaseKey) -> Result>, Self::DatabaseError> { + fn get(&self, key: &DatabaseKey) -> Result, Self::DatabaseError> { trace!("Getting from RocksDB: {:?}", key); let handle = self.db.cf_handle(key.get_cf()).expect(CF_ERROR); - Ok(self.db.get_cf(&handle, key.as_slice())?) + Ok(self.db.get_cf(&handle, key.as_slice())?.map(Into::into)) } fn get_by_prefix( &self, prefix: &DatabaseKey, - ) -> Result, Vec)>, Self::DatabaseError> { + ) -> Result, Self::DatabaseError> { trace!("Getting from RocksDB: {:?}", prefix); let handle = self.db.cf_handle(prefix.get_cf()).expect(CF_ERROR); let iter = self.db.iterator_cf( @@ -216,7 +226,7 @@ where .map_while(|kv| { if let Ok((key, value)) = kv { if key.starts_with(prefix.as_slice()) { - Some((key.to_vec(), value.to_vec())) + Some(((*key).into(), (*value).into())) } else { None } @@ -240,7 +250,7 @@ where &mut self, key: &DatabaseKey, batch: Option<&mut Self::Batch>, - ) -> Result>, Self::DatabaseError> { + ) -> Result, Self::DatabaseError> { trace!("Removing from RocksDB: {:?}", key); let handle = self.db.cf_handle(key.get_cf()).expect(CF_ERROR); let old_value = self.db.get_cf(&handle, key.as_slice())?; @@ -249,7 +259,7 @@ where } else { self.db.delete_cf(&handle, key.as_slice())?; } - Ok(old_value) + Ok(old_value.map(Into::into)) } fn remove_by_prefix(&mut self, prefix: &DatabaseKey) -> Result<(), Self::DatabaseError> { @@ -326,7 +336,7 @@ impl<'db> BonsaiDatabase for RocksDBTransaction<'db> { key: &DatabaseKey, value: &[u8], batch: Option<&mut Self::Batch>, - ) -> Result>, Self::DatabaseError> { + ) -> Result, Self::DatabaseError> { trace!("Inserting into RocksDB: {:?} {:?}", key, value); let handle_cf = self.column_families.get(key.get_cf()).expect(CF_ERROR); let old_value = self @@ -337,21 +347,22 @@ impl<'db> BonsaiDatabase for RocksDBTransaction<'db> { } else { self.txn.put_cf(handle_cf, key.as_slice(), value)?; } - Ok(old_value) + Ok(old_value.map(Into::into)) } - fn get(&self, key: &DatabaseKey) -> Result>, Self::DatabaseError> { + fn get(&self, key: &DatabaseKey) -> Result, Self::DatabaseError> { trace!("Getting from RocksDB: {:?}", key); let handle = self.column_families.get(key.get_cf()).expect(CF_ERROR); Ok(self .txn - .get_cf_opt(handle, key.as_slice(), &self.read_options)?) + .get_cf_opt(handle, key.as_slice(), &self.read_options)? + .map(Into::into)) } fn get_by_prefix( &self, prefix: &DatabaseKey, - ) -> Result, Vec)>, Self::DatabaseError> { + ) -> Result, Self::DatabaseError> { trace!("Getting from RocksDB: {:?}", prefix); let handle = self.column_families.get(prefix.get_cf()).expect(CF_ERROR); let iter = self.txn.iterator_cf( @@ -362,7 +373,7 @@ impl<'db> BonsaiDatabase for RocksDBTransaction<'db> { .map_while(|kv| { if let Ok((key, value)) = kv { if key.starts_with(prefix.as_slice()) { - Some((key.to_vec(), value.to_vec())) + Some(((*key).into(), (*value).into())) } else { None } @@ -386,7 +397,7 @@ impl<'db> BonsaiDatabase for RocksDBTransaction<'db> { &mut self, key: &DatabaseKey, batch: Option<&mut Self::Batch>, - ) -> Result>, Self::DatabaseError> { + ) -> Result, Self::DatabaseError> { trace!("Removing from RocksDB: {:?}", key); let handle = self.column_families.get(key.get_cf()).expect(CF_ERROR); let old_value = self @@ -397,7 +408,7 @@ impl<'db> BonsaiDatabase for RocksDBTransaction<'db> { } else { self.txn.delete_cf(handle, key.as_slice())?; } - Ok(old_value) + Ok(old_value.map(Into::into)) } fn remove_by_prefix(&mut self, prefix: &DatabaseKey) -> Result<(), Self::DatabaseError> { diff --git a/src/id.rs b/src/id.rs index aa54311..54a366d 100644 --- a/src/id.rs +++ b/src/id.rs @@ -1,9 +1,9 @@ -use crate::Vec; +use crate::ByteVec; use core::{fmt::Debug, hash}; /// Trait to be implemented on any type that can be used as an ID. pub trait Id: hash::Hash + PartialEq + Eq + PartialOrd + Ord + Debug + Copy + Default { - fn to_bytes(&self) -> Vec; + fn to_bytes(&self) -> ByteVec; } /// A basic ID type that can be used for testing. @@ -17,8 +17,8 @@ impl BasicId { } impl Id for BasicId { - fn to_bytes(&self) -> Vec { - self.0.to_be_bytes().to_vec() + fn to_bytes(&self) -> ByteVec { + ByteVec::from(&self.0.to_be_bytes() as &[_]) } } diff --git a/src/key_value_db.rs b/src/key_value_db.rs index f9316c8..bde5402 100644 --- a/src/key_value_db.rs +++ b/src/key_value_db.rs @@ -1,6 +1,6 @@ use crate::{ - changes::key_new_value, format, trie::merkle_tree::bytes_to_bitvec, BTreeSet, - Change as ExternChange, ToString, Vec, + changes::key_new_value, format, trie::merkle_tree::bytes_to_bitvec, BTreeSet, ByteVec, + Change as ExternChange, ToString, }; use bitvec::{order::Msb0, vec::BitVec}; use hashbrown::HashMap; @@ -18,11 +18,8 @@ use crate::{ /// Crate Trie <= KeyValueDB => BonsaiDatabase #[cfg_attr(feature = "bench", derive(Clone))] -pub struct KeyValueDB -where - DB: BonsaiDatabase, - ID: Id, -{ +#[derive(Debug)] +pub struct KeyValueDB { pub(crate) db: DB, pub(crate) changes_store: ChangeStore, pub(crate) snap_holder: BTreeSet, @@ -31,7 +28,7 @@ where pub(crate) created_at: Option, } -#[derive(Clone)] +#[derive(Clone, Debug)] pub struct KeyValueDBConfig { /// Maximum number of trie logs to keep in the database (None = unlimited). pub max_saved_trie_logs: Option, @@ -165,7 +162,7 @@ where pub(crate) fn get( &self, key: &TrieKey, - ) -> Result>, BonsaiStorageError> { + ) -> Result, BonsaiStorageError> { trace!("Getting from KeyValueDB: {:?}", key); Ok(self.db.get(&key.into())?) } @@ -174,7 +171,7 @@ where &self, key: &TrieKey, id: ID, - ) -> Result>, BonsaiStorageError> { + ) -> Result, BonsaiStorageError> { trace!("Getting from KeyValueDB: {:?} at ID: {:?}", key, id); // makes sure given id exists @@ -226,7 +223,7 @@ where key.clone(), Change { old_value, - new_value: Some(value.to_vec()), + new_value: Some(value.into()), }, ); Ok(()) diff --git a/src/lib.rs b/src/lib.rs index 7595f01..2d7b4c5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -87,6 +87,7 @@ // hashbrown uses ahash by default instead of siphash pub(crate) type HashMap = hashbrown::HashMap; +pub(crate) type HashSet = hashbrown::HashSet; pub(crate) use hashbrown::hash_map; #[cfg(not(feature = "std"))] @@ -99,6 +100,8 @@ pub(crate) use alloc::{ vec, vec::Vec, }; +use core::{fmt, ops::Deref}; +use id::Id; #[cfg(feature = "std")] pub(crate) use std::{ collections::{BTreeMap, BTreeSet, VecDeque}, @@ -109,6 +112,26 @@ pub(crate) use std::{ }; use crate::trie::merkle_tree::MerkleTree; + +pub type ByteVec = smallvec::SmallVec<[u8; 32]>; + +pub(crate) trait EncodeExt: parity_scale_codec::Encode { + fn encode_sbytevec(&self) -> ByteVec { + struct Out(ByteVec); + impl parity_scale_codec::Output for Out { + #[inline] + fn write(&mut self, bytes: &[u8]) { + self.0.extend(bytes.iter().copied()) + } + } + + let mut v = Out(ByteVec::with_capacity(self.size_hint())); + self.encode_to(&mut v); + v.0 + } +} +impl EncodeExt for T {} + use bitvec::{order::Msb0, slice::BitSlice, vec::BitVec}; use changes::ChangeBatch; use key_value_db::KeyValueDB; @@ -177,15 +200,20 @@ pub struct Change { /// Structure that hold the trie and all the necessary information to work with it. /// /// This structure is the main entry point to work with this crate. -pub struct BonsaiStorage -where - DB: BonsaiDatabase, - ChangeID: id::Id, - H: StarkHash + Send + Sync, -{ +pub struct BonsaiStorage { tries: MerkleTrees, } +impl fmt::Debug + for BonsaiStorage +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BonsaiStorage") + .field("tries", &self.tries) + .finish() + } +} + #[cfg(feature = "bench")] impl Clone for BonsaiStorage where @@ -224,27 +252,13 @@ where db: DB, config: BonsaiStorageConfig, created_at: ChangeID, - identifiers: Vec>, + _identifiers: impl IntoIterator>, ) -> Result> { let key_value_db = KeyValueDB::new(db, config.into(), Some(created_at)); - let mut tries = MerkleTrees::::new(key_value_db); - for identifier in identifiers { - tries.init_tree(&identifier)?; - } + let tries = MerkleTrees::::new(key_value_db); Ok(Self { tries }) } - /// Initialize a new trie with the given identifier. - /// This function is useful when you want to create a new trie in the database without inserting any value. - /// If the trie already exists, it will do nothing. - /// When you insert a value in a trie, it will automatically create the trie if it doesn't exist. - pub fn init_tree( - &mut self, - identifier: &[u8], - ) -> Result<(), BonsaiStorageError> { - self.tries.init_tree(identifier) - } - /// Insert a new key/value in the trie, overwriting the previous value if it exists. /// If the value already exists it will overwrite it. pub fn insert( @@ -306,6 +320,8 @@ where &mut self, requested_id: ChangeID, ) -> Result<(), BonsaiStorageError> { + self.tries.reset_to_last_commit()?; + let kv = self.tries.db_mut(); // Clear current changes @@ -378,7 +394,6 @@ where // Write revert changes and trie logs truncation kv.db.write_batch(batch)?; - self.tries.reset_to_last_commit()?; Ok(()) } diff --git a/src/tests/merge.rs b/src/tests/merge.rs index 83181a2..976d212 100644 --- a/src/tests/merge.rs +++ b/src/tests/merge.rs @@ -1,15 +1,16 @@ -#![cfg(feature = "std")] +#![allow(clippy::type_complexity)] +#![cfg(all(feature = "std", feature = "rocksdb"))] + use crate::{ databases::{create_rocks_db, RocksDB, RocksDBConfig, RocksDBTransaction}, id::{BasicId, BasicIdBuilder}, BonsaiStorage, BonsaiStorageConfig, }; use bitvec::vec::BitVec; +use once_cell::sync::Lazy; use rocksdb::OptimisticTransactionDB; use starknet_types_core::{felt::Felt, hash::Pedersen}; -use once_cell::sync::Lazy; - static PAIR1: Lazy<(BitVec, Felt)> = Lazy::new(|| { ( BitVec::from_vec(vec![1, 2, 2]), @@ -48,21 +49,20 @@ static PAIR3: Lazy<(BitVec, Felt)> = Lazy::new(|| { /// * `id_builder` - An instance of `BasicIdBuilder`. /// * `start_id` - A `BasicId` representing the commit ID of the changes made in /// `bonsai_storage`. -fn init_test<'db>( - db: &'db OptimisticTransactionDB, +fn init_test( + db: &OptimisticTransactionDB, ) -> ( Vec, - BonsaiStorage, Pedersen>, - BonsaiStorage, Pedersen>, + BonsaiStorage, Pedersen>, + BonsaiStorage, Pedersen>, BasicIdBuilder, BasicId, ) { 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) + .expect("Failed to create BonsaiStorage"); let mut id_builder = BasicIdBuilder::new(); @@ -116,10 +116,7 @@ fn merge_before_simple_remove() { bonsai_storage.merge(bonsai_at_txn).unwrap(); bonsai_storage.commit(id_builder.new_id()).unwrap(); - assert_eq!( - bonsai_storage.contains(&identifier, &PAIR1.0).unwrap(), - false - ); + assert!(!bonsai_storage.contains(&identifier, &PAIR1.0).unwrap()); bonsai_storage.revert_to(start_id).unwrap(); @@ -142,10 +139,7 @@ fn merge_tx_commit_simple_remove() { bonsai_storage.merge(bonsai_at_txn).unwrap(); - assert_eq!( - bonsai_storage.contains(&identifier, &PAIR1.0).unwrap(), - false - ); + assert!(!bonsai_storage.contains(&identifier, &PAIR1.0).unwrap()); bonsai_storage.revert_to(start_id).unwrap(); @@ -168,10 +162,7 @@ fn merge_before_simple_revert_to() { bonsai_storage.commit(id_builder.new_id()).unwrap(); bonsai_storage.revert_to(start_id).unwrap(); - assert_eq!( - bonsai_storage.get(&identifier, &PAIR2.0).unwrap().is_none(), - true - ); + assert!(bonsai_storage.get(&identifier, &PAIR2.0).unwrap().is_none()); } #[test] @@ -206,19 +197,10 @@ fn merge_transactional_commit_in_txn_before() { bonsai_storage.get(&identifier, &PAIR2.0).unwrap(), Some(PAIR2.1) ); - assert_eq!( - bonsai_storage.get(&identifier, &PAIR3.0).unwrap().is_none(), - true - ); + assert!(bonsai_storage.get(&identifier, &PAIR3.0).unwrap().is_none()); bonsai_storage.revert_to(start_id).unwrap(); - assert_eq!( - bonsai_storage.get(&identifier, &PAIR2.0).unwrap().is_none(), - true - ); - assert_eq!( - bonsai_storage.get(&identifier, &PAIR3.0).unwrap().is_none(), - true - ); + assert!(bonsai_storage.get(&identifier, &PAIR2.0).unwrap().is_none()); + assert!(bonsai_storage.get(&identifier, &PAIR3.0).unwrap().is_none()); } #[test] @@ -234,10 +216,7 @@ fn merge_transactional_commit_in_txn_before_existing_key() { bonsai_storage.merge(bonsai_at_txn).unwrap(); bonsai_storage.revert_to(id2).unwrap(); - assert_eq!( - bonsai_storage.get(&identifier, &PAIR1.0).unwrap().is_none(), - true - ); + assert!(bonsai_storage.get(&identifier, &PAIR1.0).unwrap().is_none()); bonsai_storage.revert_to(start_id).unwrap(); diff --git a/src/tests/mod.rs b/src/tests/mod.rs index 81223fc..2c11aa9 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -1,6 +1,7 @@ mod madara_comparison; mod merge; mod proof; +mod proptest; mod simple; mod transactional_state; mod trie_log; diff --git a/src/tests/proptest.rs b/src/tests/proptest.rs new file mode 100644 index 0000000..476c7f3 --- /dev/null +++ b/src/tests/proptest.rs @@ -0,0 +1,282 @@ +#![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 bitvec::bitvec; +use bitvec::order::Msb0; +use bitvec::vec::BitVec; +use proptest::prelude::*; +use proptest_derive::Arbitrary; +use smallvec::smallvec; +use starknet_types_core::felt::Felt; +use starknet_types_core::hash::Pedersen; + +#[derive(PartialEq, Eq, Hash)] +struct Key(BitVec); +impl fmt::Debug for Key { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:b}", self.0) + } +} +impl Arbitrary for Key { + type Parameters = (); + type Strategy = BoxedStrategy; + + fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy { + <[bool; 5]>::arbitrary() + .prop_map(|arr| arr.into_iter().collect::>()) + .prop_map(Self) + .boxed() + } +} + +#[derive(PartialEq, Eq, Hash)] +struct Value(Felt); +impl fmt::Debug for Value { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:#x}", self.0) + } +} +impl Arbitrary for Value { + type Parameters = (); + type Strategy = BoxedStrategy; + + fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy { + <[bool; 251]>::arbitrary() + .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() + } +} + +#[derive(Debug, Arbitrary)] +enum Step { + Insert(Key, Value), + Remove(Key), + Commit, +} + +#[derive(Debug, Arbitrary)] +struct MerkleTreeInsertProblem(Vec); + +impl MerkleTreeInsertProblem { + fn check(&self) { + let mut hashmap_db = KeyValueDB::<_, BasicId>::new( + HashMapDb::::default(), + Default::default(), + None, + ); + + let mut ckv = HashMap::new(); + + // apply steps + let mut tree = MerkleTree::::new(smallvec![]); + for step in &self.0 { + match step { + Step::Insert(k, v) => { + log::trace!("== STEP == setting {k:?} => {v:?}"); + ckv.insert(k.0.clone(), v.0); + tree.set(&hashmap_db, &k.0, v.0).unwrap(); + } + Step::Remove(k) => { + log::trace!("== STEP == removing {k:?}"); + ckv.insert(k.0.clone(), Felt::ZERO); + tree.set(&hashmap_db, &k.0, Felt::ZERO).unwrap(); + } + Step::Commit => { + log::trace!("== STEP == commit"); + tree.commit(&mut hashmap_db).unwrap(); + } + } + log::trace!("TREE"); + tree.display(); + } + + // check + for (k, v) in &ckv { + log::trace!("checking {k:b}....."); + let v2 = tree.get(&hashmap_db, k).unwrap().unwrap_or_default(); + log::trace!("checking that {k:b} => {v:#x}, (tree returned {v2:#x})"); + assert_eq!(Value(*v), Value(v2)) + } + + // check for leaks + for (k, _v) in ckv { + log::trace!("removing {k:b}..... (check for leaks)"); + tree.set(&hashmap_db, &k, Felt::ZERO).unwrap(); + tree.commit(&mut hashmap_db).unwrap(); + } + + hashmap_db.db.assert_empty(); + tree.assert_empty(); + + log::trace!("okay!"); + log::trace!(""); + log::trace!(""); + log::trace!(""); + } +} + +proptest::proptest! { + #![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(); + log::set_max_level(log::LevelFilter::Trace); + pb.check(); + } +} + +#[test] +fn test_merkle_pb_1() { + 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; 1,0,0,1,1]), + Value(Felt::from_hex("0x20").unwrap()), + ), + Remove(Key(bitvec![u8, Msb0; 1,0,0,1,1])), + Remove(Key(bitvec![u8, Msb0; 0,0,0,0,0])), + Insert( + Key(bitvec![u8, Msb0; 0,0,0,0,0]), + Value(Felt::from_hex("0x20").unwrap()), + ), + Commit, + Remove(Key(bitvec![u8, Msb0; 0,0,0,0,0])), + ]); + + pb.check(); +} + +#[test] +fn test_merkle_pb_2() { + 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,1,0,0,0]), + Value(Felt::from_hex("0x20").unwrap()), + ), + // Remove( + // Key(bitvec![u8, Msb0; 0,0,0,0,0]), + // ), + Insert( + Key(bitvec![u8, Msb0; 0,0,0,0,0]), + Value(Felt::from_hex("0x80").unwrap()), + ), + Remove(Key(bitvec![u8, Msb0; 0,0,0,0,0])), + Commit, + ]); + + pb.check(); +} + +#[test] +fn test_merkle_pb_3() { + 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; 1,0,0,0,0]), + Value(Felt::from_hex("0x21").unwrap()), + ), + Insert( + Key(bitvec![u8, Msb0; 1,1,0,0,0]), + Value(Felt::from_hex("0x22").unwrap()), + ), + Insert( + Key(bitvec![u8, Msb0; 1,1,0,1,0]), + Value(Felt::from_hex("0x23").unwrap()), + ), + Remove(Key(bitvec![u8, Msb0; 1,0,0,0,0])), + Remove(Key(bitvec![u8, Msb0; 1,0,0,0,0])), + Commit, + ]); + + pb.check(); +} + +#[test] +fn test_merkle_pb_4() { + use Step::*; + let _ = env_logger::builder().is_test(true).try_init(); + log::set_max_level(log::LevelFilter::Trace); + let pb = MerkleTreeInsertProblem(vec![ + // Remove( + // Key(bitvec![u8, Msb0; 0,0,0,0,0]), + // ), + Insert( + Key(bitvec![u8, Msb0; 0,0,1,0,0]), + Value(Felt::from_hex("0x20").unwrap()), + ), + Commit, + Insert( + Key(bitvec![u8, Msb0; 0,0,0,0,0]), + Value(Felt::from_hex("0x21").unwrap()), + ), + Commit, + ]); + + pb.check(); +} + +#[test] +fn test_merkle_pb_5() { + 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,1]), + Value(Felt::from_hex("0x20").unwrap()), + ), + Insert( + Key(bitvec![u8, Msb0; 0,0,1,0,0]), + Value(Felt::from_hex("0x20").unwrap()), + ), + ]); + + pb.check(); +} + +#[test] +fn test_merkle_pb_6() { + 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; 1,0,0,0,0]), + Value(Felt::from_hex("0x20").unwrap()), + ), + Insert( + Key(bitvec![u8, Msb0; 1,1,0,0,0]), + Value(Felt::from_hex("0x20").unwrap()), + ), + Commit, + Insert( + Key(bitvec![u8, Msb0; 1,1,0,1,0]), + Value(Felt::from_hex("0x20").unwrap()), + ), + Insert( + Key(bitvec![u8, Msb0; 1,0,0,0,0]), + Value(Felt::from_hex("0x20").unwrap()), + ), + Remove(Key(bitvec![u8, Msb0; 1,0,0,0,0])), + Remove(Key(bitvec![u8, Msb0; 1,0,0,0,0])), + ]); + + pb.check(); +} diff --git a/src/tests/transactional_state.rs b/src/tests/transactional_state.rs index afb8dc6..a1a1a25 100644 --- a/src/tests/transactional_state.rs +++ b/src/tests/transactional_state.rs @@ -249,6 +249,8 @@ fn merge_with_uncommitted_insert() { ) .unwrap(); + println!("{:?}", bonsai_storage); + // revert to commit bonsai_storage.revert_to(revert_id).unwrap(); @@ -299,10 +301,7 @@ fn merge_with_uncommitted_remove() { // remove pair2 but don't commit in transational state bonsai_at_txn.remove(&identifier, &pair2.0).unwrap(); - assert_eq!( - bonsai_at_txn.contains(&identifier, &pair2.0).unwrap(), - false - ); + assert!(!bonsai_at_txn.contains(&identifier, &pair2.0).unwrap()); let merge = bonsai_storage.merge(bonsai_at_txn); match merge { @@ -316,10 +315,7 @@ fn merge_with_uncommitted_remove() { // commit after merge bonsai_storage.commit(id_builder.new_id()).unwrap(); - assert_eq!( - bonsai_storage.contains(&identifier, &pair2.0).unwrap(), - false - ); + assert!(!bonsai_storage.contains(&identifier, &pair2.0).unwrap()); } #[test] @@ -364,7 +360,7 @@ fn transactional_state_after_uncommitted() { // uncommitted changes, done after the transactional state was created, // are not included in the transactional state let contains = bonsai_at_txn.contains(&identifier, &pair2.0).unwrap(); - assert_eq!(contains, false); + assert!(!contains); } #[test] diff --git a/src/trie/merkle_node.rs b/src/trie/merkle_node.rs index fca4293..1fa1289 100644 --- a/src/trie/merkle_node.rs +++ b/src/trie/merkle_node.rs @@ -4,6 +4,8 @@ //! For more information about how these Starknet trees are structured, see //! [`MerkleTree`](super::merkle_tree::MerkleTree). +use core::fmt; + use bitvec::order::Msb0; use bitvec::slice::BitSlice; use parity_scale_codec::{Decode, Encode}; @@ -30,21 +32,25 @@ impl NodeId { /// A node in a Binary Merkle-Patricia Tree graph. #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Encode, Decode)] pub enum Node { - /// A node that has not been fetched from storage yet. - /// - /// As such, all we know is its hash. - Unresolved(Felt), /// A branch node with exactly two children. Binary(BinaryNode), /// Describes a path connecting two other nodes. Edge(EdgeNode), } -#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Encode, Decode)] +#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Encode, Decode)] pub enum NodeHandle { Hash(Felt), InMemory(NodeId), } +impl fmt::Debug for NodeHandle { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + NodeHandle::Hash(felt) => write!(f, "Hash({:#x})", felt), + NodeHandle::InMemory(node_id) => write!(f, "InMemory({:?})", node_id), + } + } +} /// Describes the [Node::Binary] variant. #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Encode, Decode)] @@ -176,22 +182,6 @@ impl BinaryNode { } impl Node { - /// Returns true if the node represents an empty node -- this is defined as a node - /// with the [Felt::ZERO]. - /// - /// This can occur for the root node in an empty graph. - pub fn is_empty(&self) -> bool { - match self { - Node::Unresolved(hash) => hash == &Felt::ZERO, - _ => false, - } - } - - /// Is the node a binary node. - pub fn is_binary(&self) -> bool { - matches!(self, Node::Binary(..)) - } - /// 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 { @@ -202,7 +192,6 @@ impl Node { pub fn hash(&self) -> Option { match self { - Node::Unresolved(hash) => Some(*hash), Node::Binary(binary) => binary.hash, Node::Edge(edge) => edge.hash, } diff --git a/src/trie/merkle_tree.rs b/src/trie/merkle_tree.rs index 7e729ea..6776afe 100644 --- a/src/trie/merkle_tree.rs +++ b/src/trie/merkle_tree.rs @@ -2,18 +2,17 @@ use bitvec::{ prelude::{BitSlice, BitVec, Msb0}, view::BitView, }; -use core::iter::once; -use core::marker::PhantomData; -use core::mem; +use core::{fmt, marker::PhantomData}; +use core::{iter, mem}; use derive_more::Constructor; -use parity_scale_codec::{Decode, Encode}; +use parity_scale_codec::Decode; #[cfg(feature = "std")] use rayon::prelude::*; use starknet_types_core::{felt::Felt, hash::StarkHash}; use crate::{ - error::BonsaiStorageError, format, id::Id, vec, BonsaiDatabase, HashMap, KeyValueDB, ToString, - Vec, + error::BonsaiStorageError, format, hash_map, id::Id, vec, BonsaiDatabase, ByteVec, EncodeExt, + HashMap, HashSet, KeyValueDB, ToString, Vec, }; use super::{ @@ -35,7 +34,7 @@ pub enum Membership { /// 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(HashMap); +pub struct NodesMapping(pub(crate) HashMap); /// A node used in proof generated by the trie. /// @@ -63,9 +62,26 @@ impl ProofNode { } } +#[derive(Debug, Clone)] +pub(crate) enum RootHandle { + Empty, + Loaded(NodeId), +} + pub(crate) struct MerkleTrees { pub db: KeyValueDB, - pub trees: HashMap, MerkleTree>, + 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")] @@ -88,30 +104,18 @@ impl MerkleTrees Result<(), BonsaiStorageError> { - let tree = MerkleTree::new(&mut self.db, identifier.to_vec())?; - self.trees.insert(identifier.to_vec(), tree); - Ok(()) - } - pub(crate) fn set( &mut self, identifier: &[u8], key: &BitSlice, value: Felt, ) -> Result<(), BonsaiStorageError> { - let tree = self.trees.get_mut(identifier); - if let Some(tree) = tree { - tree.set(&mut self.db, key, value) - } else { - let mut tree = MerkleTree::new(&mut self.db, identifier.to_vec())?; - tree.set(&mut self.db, key, value)?; - self.trees.insert(identifier.to_vec(), tree); - Ok(()) - } + let tree = self + .trees + .entry_ref(identifier) + .or_insert_with(|| MerkleTree::new(identifier.into())); + + tree.set(&self.db, key, value) } pub(crate) fn get( @@ -119,11 +123,10 @@ impl MerkleTrees, ) -> Result, BonsaiStorageError> { - let tree = self.trees.get(identifier); - if let Some(tree) = tree { + if let Some(tree) = self.trees.get(identifier) { tree.get(&self.db, key) } else { - Ok(None) + MerkleTree::::new(identifier.into()).get(&self.db, key) } } @@ -133,11 +136,10 @@ impl MerkleTrees, id: CommitID, ) -> Result, BonsaiStorageError> { - let tree = self.trees.get(identifier); - if let Some(tree) = tree { + if let Some(tree) = self.trees.get(identifier) { tree.get_at(&self.db, key, id) } else { - Ok(None) + MerkleTree::::new(identifier.into()).get_at(&self.db, key, id) } } @@ -146,11 +148,10 @@ impl MerkleTrees, ) -> Result> { - let tree = self.trees.get(identifier); - if let Some(tree) = tree { + if let Some(tree) = self.trees.get(identifier) { tree.contains(&self.db, key) } else { - Ok(false) + MerkleTree::::new(identifier.into()).contains(&self.db, key) } } @@ -161,9 +162,7 @@ impl MerkleTrees Result<(), BonsaiStorageError> { - for tree in self.trees.values_mut() { - tree.reset_to_last_commit(&mut self.db)?; - } + self.trees.clear(); // just clear the map Ok(()) } @@ -175,11 +174,10 @@ impl MerkleTrees Result> { - let tree = self.trees.get(identifier); - if let Some(tree) = tree { - Ok(tree.root_hash()) + if let Some(tree) = self.trees.get(identifier) { + Ok(tree.root_hash(&self.db)?) } else { - Err(BonsaiStorageError::Trie("Tree not found".to_string())) + MerkleTree::::new(identifier.into()).root_hash(&self.db) } } @@ -198,7 +196,7 @@ impl MerkleTrees identifier.len() { - Some(key[identifier.len() + 1..].to_vec()) + Some(key[identifier.len() + 1..].into()) } else { None } @@ -223,7 +221,7 @@ impl MerkleTrees identifier.len() { - Some((key[identifier.len() + 1..].to_vec(), value)) + Some((key[identifier.len() + 1..].into(), value.into_vec())) } else { None } @@ -270,16 +268,15 @@ impl MerkleTrees, ) -> Result, BonsaiStorageError> { - let tree = self.trees.get(identifier); - if let Some(tree) = tree { + if let Some(tree) = self.trees.get(identifier) { tree.get_proof(&self.db, key) } else { - Err(BonsaiStorageError::Trie("Tree not found".to_string())) + MerkleTree::::new(identifier.into()).get_proof(&self.db, key) } } pub(crate) fn get_identifiers(&self) -> Vec> { - self.trees.keys().cloned().collect() + self.trees.keys().cloned().map(ByteVec::into_vec).collect() } } @@ -290,35 +287,44 @@ impl MerkleTrees { - /// The handle to the current root node could be hash if no modifications has been done - /// since the last commit or in memory if there are some modifications. - root_handle: NodeHandle, - /// The last known root hash. Updated only each commit. (possibly outdated between two commits) - root_hash: Felt, + /// The root node. None means the node has not been loaded yet. + pub(crate) root_node: Option, /// Identifier of the tree in the database. - identifier: Vec, + pub(crate) identifier: ByteVec, /// This storage is used to avoid modifying the underlying database each time during a commit. - storage_nodes: NodesMapping, + pub(crate) storage_nodes: NodesMapping, /// The id of the last node that has been added to the temporary storage. - latest_node_id: NodeId, + pub(crate) latest_node_id: NodeId, /// The list of nodes that should be removed from the underlying database during the next commit. - death_row: Vec, + pub(crate) death_row: HashSet, /// The list of leaves that have been modified during the current commit. - cache_leaf_modified: HashMap, InsertOrRemove>, + 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_handle: self.root_handle.clone(), - root_hash: self.root_hash.clone(), + root_node: self.root_node.clone(), identifier: self.identifier.clone(), storage_nodes: self.storage_nodes.clone(), - latest_node_id: self.latest_node_id.clone(), + latest_node_id: self.latest_node_id, death_row: self.death_row.clone(), cache_leaf_modified: self.cache_leaf_modified.clone(), _hasher: PhantomData, @@ -337,105 +343,216 @@ enum NodeOrFelt<'a> { } impl MerkleTree { - /// Less visible initialization for `MerkleTree` as the main entry points should be - /// [`MerkleTree::::load`] for persistent trees and [`MerkleTree::empty`] for - /// transient ones. - - pub fn new( - db: &mut KeyValueDB, - identifier: Vec, - ) -> Result> { + pub fn new(identifier: ByteVec) -> Self { let nodes_mapping: HashMap = HashMap::new(); - let root_node = db.get(&TrieKey::new(&identifier, TrieKeyType::Trie, &[]))?; - let node = if let Some(root_node) = root_node { - Node::decode(&mut root_node.as_slice())? - } else { - db.insert( - &TrieKey::new(&identifier, TrieKeyType::Trie, &[]), - &Node::Unresolved(Felt::ZERO).encode(), - None, - )?; - Node::Unresolved(Felt::ZERO) - }; - // SAFETY: The root node has been created just above - let root = node.hash().unwrap(); - Ok(Self { - root_handle: NodeHandle::Hash(root), - root_hash: root, + 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: Vec::new(), + death_row: HashSet::new(), cache_leaf_modified: HashMap::new(), _hasher: PhantomData, - }) + } } - pub fn root_hash(&self) -> Felt { - self.root_hash + // 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)))), + } } - pub fn cache_leaf_modified(&self) -> &HashMap, InsertOrRemove> { - &self.cache_leaf_modified + /// 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) + } + } + } + } } - /// Remove all the modifications that have been done since the last commit. - pub fn reset_to_last_commit( - &mut self, - db: &mut KeyValueDB, - ) -> Result<(), BonsaiStorageError> { - let node = Self::get_trie_branch_in_db_from_path( - &self.identifier, - db, - &Path(BitVec::::new()), - )? - .ok_or(BonsaiStorageError::Trie( - "root node doesn't exist in the storage".to_string(), - ))?; - let node_hash = node.hash().ok_or(BonsaiStorageError::Trie( - "Root doesn't exist in the storage".to_string(), - ))?; - self.latest_node_id.reset(); - self.storage_nodes.0.clear(); - self.cache_leaf_modified.clear(); - self.root_handle = NodeHandle::Hash(node_hash); - self.root_hash = node_hash; - Ok(()) + /// # 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>)>, BonsaiStorageError> - { - let mut updates = vec![]; + ) -> Result< + impl Iterator)>, + BonsaiStorageError, + > { + let mut updates = HashMap::new(); for node_key in mem::take(&mut self.death_row) { - updates.push((node_key, InsertOrRemove::Remove)); + updates.insert(node_key, InsertOrRemove::Remove); } - let mut hashes = vec![]; - self.compute_root_hash::(&mut hashes)?; - let root_hash = self.commit_subtree::( - &mut updates, - self.root_handle, - Path(BitVec::new()), - &mut hashes.drain(..), - )?; + 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.push(( + updates.insert( TrieKey::new(&self.identifier, TrieKeyType::Flat, &key), match value { - InsertOrRemove::Insert(value) => InsertOrRemove::Insert(value.encode()), + InsertOrRemove::Insert(value) => { + InsertOrRemove::Insert(value.encode_sbytevec()) + } InsertOrRemove::Remove => InsertOrRemove::Remove, }, - )); + ); } self.latest_node_id.reset(); - self.root_hash = root_hash; - self.root_handle = NodeHandle::Hash(root_hash); - Ok(updates) + + #[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( @@ -460,10 +577,21 @@ impl MerkleTree { &self, hashes: &mut Vec, ) -> Result> { - match self.get_node_or_felt::(&self.root_handle)? { - NodeOrFelt::Felt(felt) => Ok(felt), - NodeOrFelt::Node(node) => self.compute_hashes::(node, Path(BitVec::new()), hashes), - } + 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 @@ -478,7 +606,6 @@ impl MerkleTree { use Node::*; match node { - Unresolved(hash) => Ok(*hash), Binary(binary) => { // we check if we have one or two changed children @@ -574,16 +701,11 @@ impl MerkleTree { /// Panics if the precomputed `hashes` do not match the length of the modified subtree. fn commit_subtree( &mut self, - updates: &mut Vec<(TrieKey, InsertOrRemove>)>, - node_handle: NodeHandle, + updates: &mut HashMap>, + node_id: NodeId, path: Path, hashes: &mut impl Iterator, ) -> Result> { - let node_id = match node_handle { - NodeHandle::Hash(hash) => return Ok(hash), - NodeHandle::InMemory(root_id) => root_id, - }; - match self .storage_nodes .0 @@ -591,48 +713,51 @@ impl MerkleTree { .ok_or(BonsaiStorageError::Trie( "Couldn't fetch node in the temporary storage".to_string(), ))? { - Node::Unresolved(hash) => { - if path.0.is_empty() { - updates.push(( - TrieKey::new(&self.identifier, TrieKeyType::Trie, &[]), - InsertOrRemove::Insert(Node::Unresolved(hash).encode()), - )); - Ok(hash) - } else { - Ok(hash) - } - } Node::Binary(mut binary) => { let left_path = path.new_with_direction(Direction::Left); - let left_hash = - self.commit_subtree::(updates, binary.left, left_path, hashes)?; + 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 = - self.commit_subtree::(updates, binary.right, right_path, hashes)?; + 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: Vec = path.into(); - updates.push(( + let key_bytes: ByteVec = path.into(); + updates.insert( TrieKey::new(&self.identifier, TrieKeyType::Trie, &key_bytes), - InsertOrRemove::Insert(Node::Binary(binary).encode()), - )); + 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 = - self.commit_subtree::(updates, edge.child, child_path, hashes)?; + 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: Vec = path.into(); - updates.push(( + let key_bytes: ByteVec = path.into(); + updates.insert( TrieKey::new(&self.identifier, TrieKeyType::Trie, &key_bytes), - InsertOrRemove::Insert(Node::Edge(edge).encode()), - )); + InsertOrRemove::Insert(Node::Edge(edge).encode_sbytevec()), + ); Ok(hash) } } @@ -646,7 +771,7 @@ impl MerkleTree { /// * `value` - The value to set. pub fn set( &mut self, - db: &mut KeyValueDB, + db: &KeyValueDB, key: &BitSlice, value: Felt, ) -> Result<(), BonsaiStorageError> { @@ -654,11 +779,18 @@ impl MerkleTree { return self.delete_leaf(db, key); } let key_bytes = bitslice_to_bytes(key); - if let Some(InsertOrRemove::Insert(value_db)) = self.cache_leaf_modified.get(&key_bytes) { - if &value == value_db { + 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, @@ -668,7 +800,7 @@ impl MerkleTree { return Ok(()); } } - let path = self.preload_nodes(db, key)?; + 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. @@ -685,8 +817,10 @@ impl MerkleTree { // 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.last() { + match path_nodes.last() { Some(node_id) => { let mut nodes_to_add = Vec::new(); self.storage_nodes.0.entry(*node_id).and_modify(|node| { @@ -698,6 +832,7 @@ impl MerkleTree { 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; @@ -712,6 +847,11 @@ impl MerkleTree { // 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)); @@ -771,14 +911,15 @@ impl MerkleTree { child: NodeHandle::InMemory(branch_id), }) }; - let path = key[..edge.height as usize].to_bitvec(); - let key_bytes = - [&[path.len() as u8], path.into_vec().as_slice()].concat(); - self.death_row.push(TrieKey::Trie(key_bytes)); + 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) => { - if (binary.height + 1) as usize == key.len() { + 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), @@ -788,7 +929,6 @@ impl MerkleTree { .insert(key_bytes, InsertOrRemove::Insert(value)); } } - _ => {} } }); self.storage_nodes.0.extend(nodes_to_add); @@ -809,7 +949,7 @@ impl MerkleTree { .0 .insert(self.latest_node_id.next_id(), edge); - self.root_handle = NodeHandle::InMemory(self.latest_node_id); + self.root_node = Some(RootHandle::Loaded(self.latest_node_id)); let key_bytes = bitslice_to_bytes(key); self.cache_leaf_modified @@ -829,7 +969,7 @@ impl MerkleTree { /// * `key` - The key to delete. fn delete_leaf( &mut self, - db: &mut KeyValueDB, + db: &KeyValueDB, key: &BitSlice, ) -> Result<(), BonsaiStorageError> { // Algorithm explanation: @@ -846,31 +986,38 @@ impl MerkleTree { // // Then we are done. let key_bytes = bitslice_to_bytes(key); - if db - .get(&TrieKey::new( + 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_none() - && !self.cache_leaf_modified.contains_key(&key_bytes) - { + .is_some() + }; + + if !tree_has_value { return Ok(()); } - self.cache_leaf_modified - .insert(key_bytes.clone(), InsertOrRemove::Remove); + leaf_entry.insert(InsertOrRemove::Remove); - let path = self.preload_nodes(db, key)?; + 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.into_iter().rev().skip_while(|node| { - // SAFETY: Has been populate by preload_nodes just above - let node = self.storage_nodes.0.get(node).unwrap(); - match node { - Node::Unresolved(_) => {} - Node::Binary(_) => {} + 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(); @@ -879,23 +1026,36 @@ impl MerkleTree { for i in last_binary_path.0.iter() { new_path.0.push(*i); } - last_binary_path = new_path; - let path: Vec = (&last_binary_path).into(); + 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 - .push(TrieKey::new(&self.identifier, TrieKeyType::Trie, &path)); + .insert(TrieKey::new(&self.identifier, TrieKeyType::Trie, &path)); + node.remove(); + + true } } - !node.is_binary() }); 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 = + 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) = @@ -905,7 +1065,8 @@ impl MerkleTree { // 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(once(bool::from(direction)).collect::>()); + let path = + Path(iter::once(bool::from(direction)).collect::>()); let mut edge = EdgeNode { hash: None, height, @@ -918,7 +1079,9 @@ impl MerkleTree { // Merge the remaining child if it's an edge. self.merge_edges::(&mut edge, db, &last_binary_path)?; - edge + 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 { @@ -929,6 +1092,16 @@ impl MerkleTree { 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)); } @@ -939,12 +1112,14 @@ impl MerkleTree { None => { // We reached the root without a hitting binary node. The new tree // must therefore be empty. - self.latest_node_id.next_id(); - self.storage_nodes - .0 - .insert(self.latest_node_id, Node::Unresolved(Felt::ZERO)); - self.root_handle = NodeHandle::InMemory(self.latest_node_id); - self.root_hash = Felt::ZERO; + + 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(()); } }; @@ -965,13 +1140,20 @@ impl MerkleTree { 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())) } @@ -1026,29 +1208,27 @@ impl MerkleTree { key: &BitSlice, ) -> Result, BonsaiStorageError> { let mut nodes = Vec::with_capacity(251); - let mut node = match self.root_handle { - NodeHandle::Hash(_) => { - let node = Self::get_trie_branch_in_db_from_path( - &self.identifier, - db, - &Path(BitVec::::new()), - )? - .ok_or(BonsaiStorageError::Trie( - "Couldn't fetch root node in db".to_string(), - ))?; - if node.is_empty() { - return Ok(Vec::new()); - } - node + let mut node = match self.root_node { + Some(RootHandle::Empty) => { + return Ok(Vec::new()); } - NodeHandle::InMemory(root_id) => self + Some(RootHandle::Loaded(node_id)) => self .storage_nodes .0 - .get(&root_id) + .get(&node_id) .ok_or(BonsaiStorageError::Trie( - "Couldn't fetch root node in the temporary storage".to_string(), + "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 { @@ -1057,6 +1237,7 @@ impl MerkleTree { 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), @@ -1104,6 +1285,7 @@ impl MerkleTree { let next_path = key[..binary.height as usize + 1].to_bitvec(); let next_node = match next { NodeHandle::Hash(_) => Self::get_trie_branch_in_db_from_path( + &self.death_row, &self.identifier, db, &Path(next_path), @@ -1157,14 +1339,13 @@ impl MerkleTree { } } node = next_node; - } - Node::Unresolved(hash) => { - nodes.push(ProofNode::Edge { - child: hash, - path: Path(BitVec::::new()), - }); - return Ok(nodes); - } + } // Node::Unresolved(hash) => { + // nodes.push(ProofNode::Edge { + // child: hash, + // path: Path(BitVec::::new()), + // }); + // return Ok(nodes); + // } } } } @@ -1189,63 +1370,78 @@ impl MerkleTree { /// The list of nodes along the path. fn preload_nodes( &mut self, - db: &mut KeyValueDB, + db: &KeyValueDB, dst: &BitSlice, - ) -> Result, BonsaiStorageError> { + ) -> Result<(Vec, Path), BonsaiStorageError> { let mut nodes = Vec::with_capacity(251); - - let mut prev_handle = &mut self.root_handle; 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 { - NodeHandle::Hash(_) => { - // load from db - let node = Self::get_trie_branch_in_db_from_path(&self.identifier, db, &path)? - .ok_or(BonsaiStorageError::Trie( - "Couldn't fetch node from db".to_string(), - ))?; - - if node.is_empty() { - // empty tree - break; + // 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)); + } } - - // 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; + // 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; + }; - let node = - self.storage_nodes - .0 - .get_mut(&node_id) - .ok_or(BonsaiStorageError::Trie( + // 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) - } + ), + )?; + + (node_id, node) + } + }, }; nodes.push(node_id); // visit the child match node { - // We are in a case where the trie is empty and so there is nothing to preload. - Node::Unresolved(_felt) => break, - Node::Binary(binary_node) => { let next_direction = binary_node.direction(dst); path.0.push(bool::from(next_direction)); - if path.0 == dst { - break; // found it :) - } - prev_handle = binary_node.get_child_mut(next_direction); + prev_handle = Some(binary_node.get_child_mut(next_direction)); } Node::Edge(edge_node) if edge_node.path_matches(dst) => { @@ -1254,7 +1450,7 @@ impl MerkleTree { break; // found it :) } - prev_handle = &mut edge_node.child; + 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. @@ -1262,18 +1458,28 @@ impl MerkleTree { } } - Ok(nodes) + 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> { - let path: Vec = path.into(); - db.get(&TrieKey::new(identifier, TrieKeyType::Trie, &path))? + 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)) }) @@ -1293,33 +1499,52 @@ impl MerkleTree { /// /// * `parent` - The parent node to merge the child with. fn merge_edges( - &self, + &mut self, parent: &mut EdgeNode, db: &KeyValueDB, path: &Path, ) -> Result<(), BonsaiStorageError> { - let child_node = match parent.child { + match parent.child { NodeHandle::Hash(_) => { - let node = Self::get_trie_branch_in_db_from_path(&self.identifier, db, path)?; - if let Some(node) = node { - node - } else { - return Ok(()); + 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)); } } - NodeHandle::InMemory(child_id) => self - .storage_nodes - .0 - .get(&child_id) - .ok_or(BonsaiStorageError::Trie( - "Couldn't fetch node in memory".to_string(), - ))? - .clone(), }; - if let Node::Edge(child_edge) = child_node { - parent.path.0.extend_from_bitslice(&child_edge.path.0); - parent.child = child_edge.child; - } Ok(()) } @@ -1405,15 +1630,16 @@ impl MerkleTree { #[cfg(test)] #[allow(dead_code)] - fn display(&self) { - match self.root_handle { - NodeHandle::Hash(hash) => { - trace!("root is hash: {:?}", hash); + pub(crate) fn display(&self) { + match self.root_node { + Some(RootHandle::Empty) => { + trace!("tree is empty") } - NodeHandle::InMemory(root_id) => { - trace!("root is node: {:?}", root_id); - self.print(&root_id); + Some(RootHandle::Loaded(node)) => { + trace!("root is node {:?}", node); + self.print(&node); } + None => trace!("root is not loaded"), } } @@ -1426,13 +1652,10 @@ impl MerkleTree { trace!("bonsai_node {:?} = {:?}", head, current_tmp); match current_tmp { - Unresolved(hash) => { - trace!("Unresolved: {:?}", hash); - } Binary(binary) => { match &binary.get_child(Direction::Left) { NodeHandle::Hash(hash) => { - trace!("left is hash {:?}", hash); + trace!("left is hash {:#x}", hash); } NodeHandle::InMemory(left_id) => { self.print(left_id); @@ -1440,7 +1663,7 @@ impl MerkleTree { } match &binary.get_child(Direction::Right) { NodeHandle::Hash(hash) => { - trace!("right is hash {:?}", hash); + trace!("right is hash {:#x}", hash); } NodeHandle::InMemory(right_id) => { self.print(right_id); @@ -1449,7 +1672,7 @@ impl MerkleTree { } Edge(edge) => match &edge.child { NodeHandle::Hash(hash) => { - trace!("child is hash {:?}", hash); + trace!("child is hash {:#x}", hash); } NodeHandle::InMemory(child_id) => { self.print(child_id); @@ -1459,8 +1682,14 @@ impl MerkleTree { } } -pub(crate) fn bitslice_to_bytes(bitslice: &BitSlice) -> Vec { - [&[bitslice.len() as u8], bitslice.to_bitvec().as_raw_slice()].concat() +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 { @@ -1477,7 +1706,7 @@ mod tests { use crate::{ databases::{create_rocks_db, RocksDB, RocksDBConfig}, id::BasicId, - BonsaiStorage, BonsaiStorageConfig, + BonsaiStorage, BonsaiStorageConfig, ByteVec, }; #[test_log::test] @@ -1688,7 +1917,7 @@ mod tests { ), ]; - let block_1 = vec![ + let block_1 = [ ( str_to_felt_bytes( "0x06538fdd3aa353af8a87f5fe77d1f533ea82815076e30a86d65b72d3eb4f0b80", @@ -1848,7 +2077,6 @@ mod tests { "contract address (write): {:#064x}", Felt::from_bytes_be_slice(contract_address) ); - assert!(bonsai.init_tree(contract_address).is_ok()); for (k, v) in storage { // truncate only keeps the first 251 bits in a key @@ -1868,11 +2096,9 @@ mod tests { // aggreates all storage changes to their latest state // (replacements are takent into account) - let mut storage_map = IndexMap::, IndexMap>::new(); + let mut storage_map = IndexMap::>::new(); for (contract_address, storage) in blocks.clone() { - let map = storage_map - .entry(contract_address.to_vec()) - .or_insert(IndexMap::new()); + let map = storage_map.entry((*contract_address).into()).or_default(); for (k, v) in storage { let k = Felt::from_bytes_be_slice(k); @@ -1927,6 +2153,7 @@ mod tests { fn truncate(key: &[u8]) -> BitVec { key.view_bits()[5..].to_owned() } + // use crate::{ // databases::{create_rocks_db, RocksDB, RocksDBConfig}, // id::BasicId, diff --git a/src/trie/path.rs b/src/trie/path.rs index de8716d..62e7c4e 100644 --- a/src/trie/path.rs +++ b/src/trie/path.rs @@ -1,21 +1,28 @@ 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::Vec; +use crate::{ByteVec, EncodeExt}; #[cfg(all(feature = "std", test))] use rstest::rstest; -#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Default)] pub struct Path(pub BitVec); +impl fmt::Debug for Path { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "Path({:b})", &self.0) + } +} + impl Encode for Path { fn encode_to(&self, dest: &mut T) { // Copied from scale_bits crate (https://github.com/paritytech/scale-bits/blob/820a3e8e0c9db18ef6acfa2a9a19f738400b0637/src/scale/encode_iter.rs#L28) // but don't use it directly to avoid copy and u32 length encoding - // How it works ? + // How it works: // 1. We encode the number of bits in the bitvec as a u8 // 2. We build elements of a size of u8 using bit shifting // 3. A last element, not full, is created if there is a remainder of bits @@ -103,26 +110,16 @@ impl Path { } } -/// Convert Path to Vec can be used, for example, to create keys for the database -impl From for Vec { +/// Convert Path to SByteVec can be used, for example, to create keys for the database +impl From for ByteVec { fn from(path: Path) -> Self { - let key = if path.0.is_empty() { - Vec::new() - } else { - [&[path.0.len() as u8], path.0.as_raw_slice()].concat() - }; - key + path.encode_sbytevec() } } -impl From<&Path> for Vec { +impl From<&Path> for ByteVec { fn from(path: &Path) -> Self { - let key = if path.0.is_empty() { - Vec::new() - } else { - [&[path.0.len() as u8], path.0.as_raw_slice()].concat() - }; - key + path.encode_sbytevec() } } diff --git a/src/trie/trie_db.rs b/src/trie/trie_db.rs index 857548c..60a64b7 100644 --- a/src/trie/trie_db.rs +++ b/src/trie/trie_db.rs @@ -1,12 +1,11 @@ -use crate::bonsai_database::DatabaseKey; -use crate::Vec; +use crate::{bonsai_database::DatabaseKey, ByteVec}; /// Key in the database of the different elements that are used in the storage of the trie data. /// Use `new` function to create a new key. #[derive(Debug, Clone, Hash, PartialEq, Eq)] pub(crate) enum TrieKey { - Trie(Vec), - Flat(Vec), + Trie(ByteVec), + Flat(ByteVec), } pub(crate) enum TrieKeyType { @@ -34,7 +33,7 @@ impl From<&TrieKey> for u8 { impl TrieKey { pub fn new(identifier: &[u8], key_type: TrieKeyType, key: &[u8]) -> Self { - let mut final_key = identifier.to_vec(); + let mut final_key = ByteVec::from(identifier); final_key.extend_from_slice(key); match key_type { TrieKeyType::Trie => TrieKey::Trie(final_key), @@ -42,7 +41,7 @@ impl TrieKey { } } - pub fn from_variant_and_bytes(variant: u8, bytes: Vec) -> Self { + pub fn from_variant_and_bytes(variant: u8, bytes: ByteVec) -> Self { match variant { x if x == TrieKeyType::Trie as u8 => TrieKey::Trie(bytes), x if x == TrieKeyType::Flat as u8 => TrieKey::Flat(bytes),