From 571c2e572d9c27ad68eac7b47f4f347376b88dc5 Mon Sep 17 00:00:00 2001 From: Thomas Plisson Date: Mon, 8 Jan 2024 18:09:16 +0100 Subject: [PATCH 1/3] modifications following starket review --- Cargo.lock | 1 + Cargo.toml | 1 + src/bonsai_database.rs | 42 ++- src/changes.rs | 53 +--- src/databases/hashmap_db.rs | 25 +- src/databases/rocks_db.rs | 43 ++- src/error.rs | 36 ++- src/id.rs | 9 +- src/key_value_db.rs | 58 ++-- src/lib.rs | 67 +++-- src/tests/trie_log.rs | 1 + src/trie/merkle_node.rs | 155 +++++++++- src/trie/merkle_tree.rs | 571 ++++++++++++++---------------------- src/trie/mod.rs | 3 +- src/trie/path.rs | 126 ++++++++ src/trie/trie_db.rs | 58 ++-- 16 files changed, 725 insertions(+), 524 deletions(-) create mode 100644 src/trie/path.rs diff --git a/Cargo.lock b/Cargo.lock index 1a1ef0a..6e84299 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -603,6 +603,7 @@ dependencies = [ "serde", "smallvec", "tempfile", + "thiserror", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 7e46b15..b9ce410 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,6 +18,7 @@ rustc-hex = "2.1.0" env_logger = "0.10.1" lru = "0.12.1" parking_lot = "0.12.1" +thiserror = "1.0" bitvec = { version = "1", default-features = false } derive_more = { version = "0.99.17", features = ["constructor"] } diff --git a/src/bonsai_database.rs b/src/bonsai_database.rs index 6317c78..dcdaa9f 100644 --- a/src/bonsai_database.rs +++ b/src/bonsai_database.rs @@ -1,59 +1,53 @@ use std::error::Error; -use crate::{changes::ChangeKeyType, error::BonsaiStorageError, id::Id}; +use crate::id::Id; +/// Key in the database of the different elements that can be stored in the database. #[derive(Debug, Hash, PartialEq, Eq)] -pub enum KeyType<'a> { +pub enum DatabaseKey<'a> { Trie(&'a [u8]), Flat(&'a [u8]), TrieLog(&'a [u8]), } -impl<'a> From<&'a ChangeKeyType> for KeyType<'a> { - fn from(change_key: &'a ChangeKeyType) -> Self { - match change_key { - ChangeKeyType::Trie(key) => KeyType::Trie(key.as_slice()), - ChangeKeyType::Flat(key) => KeyType::Flat(key.as_slice()), - } - } -} - -impl KeyType<'_> { +impl DatabaseKey<'_> { pub fn as_slice(&self) -> &[u8] { match self { - KeyType::Trie(slice) => slice, - KeyType::Flat(slice) => slice, - KeyType::TrieLog(slice) => slice, + DatabaseKey::Trie(slice) => slice, + DatabaseKey::Flat(slice) => slice, + DatabaseKey::TrieLog(slice) => slice, } } } +pub trait DBError: Error + Send + Sync {} + /// Trait to be implemented on any type that can be used as a database. pub trait BonsaiDatabase { type Batch: Default; - type DatabaseError: Error + Into; + type DatabaseError: Error + DBError; /// Create a new empty batch of changes to be used in `insert`, `remove` and applied in database using `write_batch`. fn create_batch(&self) -> Self::Batch; /// Returns the value of the key if it exists - fn get(&self, key: &KeyType) -> 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: &KeyType, + prefix: &DatabaseKey, ) -> Result, Vec)>, Self::DatabaseError>; /// Returns true if the key exists - fn contains(&self, key: &KeyType) -> Result; + fn contains(&self, key: &DatabaseKey) -> Result; /// Insert a new 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. fn insert( &mut self, - key: &KeyType, + key: &DatabaseKey, value: &[u8], batch: Option<&mut Self::Batch>, ) -> Result>, Self::DatabaseError>; @@ -62,12 +56,12 @@ pub trait BonsaiDatabase { /// If a batch is provided, the change will be written in the batch instead of the database. fn remove( &mut self, - key: &KeyType, + key: &DatabaseKey, batch: Option<&mut Self::Batch>, ) -> Result>, Self::DatabaseError>; /// Remove all keys that start with the given prefix - fn remove_by_prefix(&mut self, prefix: &KeyType) -> Result<(), Self::DatabaseError>; + fn remove_by_prefix(&mut self, prefix: &DatabaseKey) -> Result<(), Self::DatabaseError>; /// Write batch of changes directly in the database fn write_batch(&mut self, batch: Self::Batch) -> Result<(), Self::DatabaseError>; @@ -78,8 +72,8 @@ pub trait BonsaiDatabase { } pub trait BonsaiPersistentDatabase { - type DatabaseError: Error + Into; - type Transaction: BonsaiDatabase; + type DatabaseError: Error + DBError; + type Transaction: BonsaiDatabase; /// Save a snapshot of the current database state /// This function returns a snapshot id that can be used to create a transaction fn snapshot(&mut self, id: ID); diff --git a/src/changes.rs b/src/changes.rs index 1195e48..adc0c73 100644 --- a/src/changes.rs +++ b/src/changes.rs @@ -1,4 +1,4 @@ -use crate::id::Id; +use crate::{id::Id, trie::TrieKey}; use serde::{Deserialize, Serialize}; use std::collections::{HashMap, VecDeque}; @@ -8,45 +8,15 @@ pub struct Change { pub new_value: Option>, } -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub enum ChangeKeyType { - Trie(Vec), - Flat(Vec), -} - -impl ChangeKeyType { - pub fn get_id(&self) -> u8 { - match self { - ChangeKeyType::Trie(_) => 0, - ChangeKeyType::Flat(_) => 1, - } - } - - pub fn as_slice(&self) -> &[u8] { - match self { - ChangeKeyType::Trie(key) => key.as_slice(), - ChangeKeyType::Flat(key) => key.as_slice(), - } - } - - pub fn from_id(id: u8, key: Vec) -> Self { - match id { - 0 => ChangeKeyType::Trie(key), - 1 => ChangeKeyType::Flat(key), - _ => panic!("Invalid id"), - } - } -} - #[derive(Debug, Default)] -pub struct ChangeBatch(pub(crate) HashMap); +pub struct ChangeBatch(pub(crate) HashMap); const KEY_SEPARATOR: u8 = 0x00; const NEW_VALUE: u8 = 0x00; const OLD_VALUE: u8 = 0x01; impl ChangeBatch { - pub fn insert_in_place(&mut self, key: ChangeKeyType, change: Change) { + pub fn insert_in_place(&mut self, key: TrieKey, change: Change) { match self.0.entry(key) { std::collections::hash_map::Entry::Occupied(mut entry) => { let e = entry.get_mut(); @@ -61,8 +31,9 @@ impl ChangeBatch { } } + //TODO: Use serde pub fn serialize(&self, id: &ID) -> Vec<(Vec, &[u8])> { - let id = id.serialize(); + let id = id.to_bytes(); self.0 .iter() .flat_map(|(change_key, change)| { @@ -70,11 +41,16 @@ impl ChangeBatch { let mut changes = Vec::new(); if let Some(old_value) = &change.old_value { + if let Some(new_value) = &change.new_value { + if old_value == new_value { + return changes; + } + } let key = [ id.as_slice(), &[KEY_SEPARATOR], key_slice, - &[change_key.get_id()], + &[change_key.into()], &[OLD_VALUE], ] .concat(); @@ -86,7 +62,7 @@ impl ChangeBatch { id.as_slice(), &[KEY_SEPARATOR], key_slice, - &[change_key.get_id()], + &[change_key.into()], &[NEW_VALUE], ] .concat(); @@ -98,7 +74,7 @@ impl ChangeBatch { } pub fn deserialize(id: &ID, changes: Vec<(Vec, Vec)>) -> Self { - let id = id.serialize(); + let id = id.to_bytes(); let mut change_batch = ChangeBatch(HashMap::new()); let mut current_change = Change::default(); let mut last_key = None; @@ -110,7 +86,8 @@ impl ChangeBatch { let mut key = key.to_vec(); let change_type = key.pop().unwrap(); let key_type = key.pop().unwrap(); - let change_key = ChangeKeyType::from_id(key_type, key[id.len() + 1..].to_vec()); + let change_key = + TrieKey::from_variant_and_bytes(key_type, key[id.len() + 1..].to_vec()); if let Some(last_key) = last_key { if last_key != change_key { change_batch.insert_in_place(last_key, current_change); diff --git a/src/databases/hashmap_db.rs b/src/databases/hashmap_db.rs index df37e1f..672ddee 100644 --- a/src/databases/hashmap_db.rs +++ b/src/databases/hashmap_db.rs @@ -5,7 +5,9 @@ use std::{ }; use crate::{ - bonsai_database::BonsaiPersistentDatabase, error::BonsaiStorageError, id::Id, BonsaiDatabase, + bonsai_database::{BonsaiPersistentDatabase, DBError}, + id::Id, + BonsaiDatabase, }; #[derive(Debug)] @@ -19,11 +21,7 @@ impl Display for HashMapDbError { impl Error for HashMapDbError {} -impl From for BonsaiStorageError { - fn from(err: HashMapDbError) -> Self { - Self::Database(err.to_string()) - } -} +impl DBError for HashMapDbError {} #[derive(Clone)] pub struct HashMapDbConfig {} @@ -53,7 +51,7 @@ impl BonsaiDatabase for HashMapDb { fn remove_by_prefix( &mut self, - prefix: &crate::bonsai_database::KeyType, + prefix: &crate::bonsai_database::DatabaseKey, ) -> Result<(), Self::DatabaseError> { let mut keys_to_remove = Vec::new(); for key in self.db.keys() { @@ -69,14 +67,14 @@ impl BonsaiDatabase for HashMapDb { fn get( &self, - key: &crate::bonsai_database::KeyType, + key: &crate::bonsai_database::DatabaseKey, ) -> Result>, Self::DatabaseError> { Ok(self.db.get(key.as_slice()).cloned()) } fn get_by_prefix( &self, - prefix: &crate::bonsai_database::KeyType, + prefix: &crate::bonsai_database::DatabaseKey, ) -> Result, Vec)>, Self::DatabaseError> { let mut result = Vec::new(); for (key, value) in self.db.iter() { @@ -89,7 +87,7 @@ impl BonsaiDatabase for HashMapDb { fn insert( &mut self, - key: &crate::bonsai_database::KeyType, + key: &crate::bonsai_database::DatabaseKey, value: &[u8], _batch: Option<&mut Self::Batch>, ) -> Result>, Self::DatabaseError> { @@ -98,13 +96,16 @@ impl BonsaiDatabase for HashMapDb { fn remove( &mut self, - key: &crate::bonsai_database::KeyType, + key: &crate::bonsai_database::DatabaseKey, _batch: Option<&mut Self::Batch>, ) -> Result>, Self::DatabaseError> { Ok(self.db.remove(key.as_slice())) } - fn contains(&self, key: &crate::bonsai_database::KeyType) -> Result { + fn contains( + &self, + key: &crate::bonsai_database::DatabaseKey, + ) -> Result { Ok(self.db.contains_key(key.as_slice())) } diff --git a/src/databases/rocks_db.rs b/src/databases/rocks_db.rs index 1739501..1b58c91 100644 --- a/src/databases/rocks_db.rs +++ b/src/databases/rocks_db.rs @@ -13,9 +13,8 @@ use rocksdb::{ }; use crate::{ - bonsai_database::{BonsaiDatabase, BonsaiPersistentDatabase, KeyType}, + bonsai_database::{BonsaiDatabase, BonsaiPersistentDatabase, DBError, DatabaseKey}, id::Id, - BonsaiStorageError, }; const TRIE_LOG_CF: &str = "trie_log"; @@ -89,12 +88,6 @@ pub enum RocksDBError { Custom(String), } -impl From for BonsaiStorageError { - fn from(err: RocksDBError) -> Self { - Self::Database(err.to_string()) - } -} - impl From for RocksDBError { fn from(err: Error) -> Self { Self::RocksDB(err) @@ -110,6 +103,8 @@ impl fmt::Display for RocksDBError { } } +impl DBError for RocksDBError {} + impl StdError for RocksDBError { fn cause(&self) -> Option<&dyn StdError> { match self { @@ -126,12 +121,12 @@ impl StdError for RocksDBError { } } -impl KeyType<'_> { +impl DatabaseKey<'_> { fn get_cf(&self) -> &'static str { match self { - KeyType::Trie(_) => TRIE_CF, - KeyType::Flat(_) => FLAT_CF, - KeyType::TrieLog(_) => TRIE_LOG_CF, + DatabaseKey::Trie(_) => TRIE_CF, + DatabaseKey::Flat(_) => FLAT_CF, + DatabaseKey::TrieLog(_) => TRIE_LOG_CF, } } } @@ -186,7 +181,7 @@ where fn insert( &mut self, - key: &KeyType, + key: &DatabaseKey, value: &[u8], batch: Option<&mut Self::Batch>, ) -> Result>, Self::DatabaseError> { @@ -201,7 +196,7 @@ where Ok(old_value) } - fn get(&self, key: &KeyType) -> 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())?) @@ -209,7 +204,7 @@ where fn get_by_prefix( &self, - prefix: &KeyType, + prefix: &DatabaseKey, ) -> Result, Vec)>, Self::DatabaseError> { trace!("Getting from RocksDB: {:?}", prefix); let handle = self.db.cf_handle(prefix.get_cf()).expect(CF_ERROR); @@ -232,7 +227,7 @@ where .collect()) } - fn contains(&self, key: &KeyType) -> Result { + fn contains(&self, key: &DatabaseKey) -> Result { trace!("Checking if RocksDB contains: {:?}", key); let handle = self.db.cf_handle(key.get_cf()).expect(CF_ERROR); Ok(self @@ -243,7 +238,7 @@ where fn remove( &mut self, - key: &KeyType, + key: &DatabaseKey, batch: Option<&mut Self::Batch>, ) -> Result>, Self::DatabaseError> { trace!("Removing from RocksDB: {:?}", key); @@ -257,7 +252,7 @@ where Ok(old_value) } - fn remove_by_prefix(&mut self, prefix: &KeyType) -> Result<(), Self::DatabaseError> { + fn remove_by_prefix(&mut self, prefix: &DatabaseKey) -> 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( @@ -328,7 +323,7 @@ impl<'db> BonsaiDatabase for RocksDBTransaction<'db> { fn insert( &mut self, - key: &KeyType, + key: &DatabaseKey, value: &[u8], batch: Option<&mut Self::Batch>, ) -> Result>, Self::DatabaseError> { @@ -345,7 +340,7 @@ impl<'db> BonsaiDatabase for RocksDBTransaction<'db> { Ok(old_value) } - fn get(&self, key: &KeyType) -> 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 @@ -355,7 +350,7 @@ impl<'db> BonsaiDatabase for RocksDBTransaction<'db> { fn get_by_prefix( &self, - prefix: &KeyType, + prefix: &DatabaseKey, ) -> Result, Vec)>, Self::DatabaseError> { trace!("Getting from RocksDB: {:?}", prefix); let handle = self.column_families.get(prefix.get_cf()).expect(CF_ERROR); @@ -378,7 +373,7 @@ impl<'db> BonsaiDatabase for RocksDBTransaction<'db> { .collect()) } - fn contains(&self, key: &KeyType) -> Result { + fn contains(&self, key: &DatabaseKey) -> Result { trace!("Checking if RocksDB contains: {:?}", key); let handle = self.column_families.get(key.get_cf()).expect(CF_ERROR); Ok(self @@ -389,7 +384,7 @@ impl<'db> BonsaiDatabase for RocksDBTransaction<'db> { fn remove( &mut self, - key: &KeyType, + key: &DatabaseKey, batch: Option<&mut Self::Batch>, ) -> Result>, Self::DatabaseError> { trace!("Removing from RocksDB: {:?}", key); @@ -405,7 +400,7 @@ impl<'db> BonsaiDatabase for RocksDBTransaction<'db> { Ok(old_value) } - fn remove_by_prefix(&mut self, prefix: &KeyType) -> Result<(), Self::DatabaseError> { + fn remove_by_prefix(&mut self, prefix: &DatabaseKey) -> Result<(), Self::DatabaseError> { trace!("Getting from RocksDB: {:?}", prefix); let mut batch = self.create_batch(); { diff --git a/src/error.rs b/src/error.rs index 8055cda..2dba0bc 100644 --- a/src/error.rs +++ b/src/error.rs @@ -1,6 +1,15 @@ +use std::{error::Error, fmt::Display}; + +use mp_felt::Felt252WrapperError; + +use crate::bonsai_database::DBError; + /// All errors that can be returned by BonsaiStorage. -#[derive(Debug)] -pub enum BonsaiStorageError { +#[derive(Debug, thiserror::Error)] +pub enum BonsaiStorageError +where + DatabaseError: Error + DBError, +{ /// Error from the underlying trie. Trie(String), /// Error when trying to go to a specific commit ID. @@ -10,5 +19,26 @@ pub enum BonsaiStorageError { /// Error when trying to merge a transactional state. Merge(String), /// Error from the underlying database. - Database(String), + Database(#[from] DatabaseError), + /// Error from Felt conversion + Felt252WrapperError(#[from] Felt252WrapperError), + /// Error when decoding a node + NodeDecodeError(#[from] parity_scale_codec::Error), +} + +impl Display for BonsaiStorageError +where + DatabaseError: Error + DBError, +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + BonsaiStorageError::Trie(e) => write!(f, "Trie error: {}", e), + BonsaiStorageError::GoTo(e) => write!(f, "GoTo error: {}", e), + BonsaiStorageError::Transaction(e) => write!(f, "Transaction error: {}", e), + BonsaiStorageError::Merge(e) => write!(f, "Merge error: {}", e), + BonsaiStorageError::Database(e) => write!(f, "Database error: {}", e), + BonsaiStorageError::Felt252WrapperError(e) => write!(f, "Felt error: {}", e), + BonsaiStorageError::NodeDecodeError(e) => write!(f, "Node decode error: {}", e), + } + } } diff --git a/src/id.rs b/src/id.rs index 8945a01..e6ca76f 100644 --- a/src/id.rs +++ b/src/id.rs @@ -2,7 +2,7 @@ use std::{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 { - fn serialize(&self) -> Vec; + fn to_bytes(&self) -> Vec; } /// A basic ID type that can be used for testing. @@ -10,7 +10,7 @@ pub trait Id: hash::Hash + PartialEq + Eq + PartialOrd + Ord + Debug + Copy { pub struct BasicId(u64); impl Id for BasicId { - fn serialize(&self) -> Vec { + fn to_bytes(&self) -> Vec { self.0.to_be_bytes().to_vec() } } @@ -34,7 +34,8 @@ impl BasicIdBuilder { /// Create a new ID (unique). pub fn new_id(&mut self) -> BasicId { - self.last_id += 1; - BasicId(self.last_id) + let id = BasicId(self.last_id); + self.last_id = self.last_id.checked_add(1).expect("Id overflow"); + id } } diff --git a/src/key_value_db.rs b/src/key_value_db.rs index 9bd831a..3dfaed7 100644 --- a/src/key_value_db.rs +++ b/src/key_value_db.rs @@ -2,10 +2,10 @@ use log::trace; use std::collections::BTreeSet; use crate::{ - bonsai_database::{BonsaiDatabase, BonsaiPersistentDatabase, KeyType}, + bonsai_database::{BonsaiDatabase, BonsaiPersistentDatabase, DatabaseKey}, changes::{Change, ChangeBatch, ChangeStore}, id::Id, - trie::TrieKeyType, + trie::TrieKey, BonsaiStorageConfig, BonsaiStorageError, }; @@ -67,7 +67,6 @@ impl KeyValueDB where DB: BonsaiDatabase, ID: Id, - BonsaiStorageError: std::convert::From<::DatabaseError>, { pub(crate) fn new(underline_db: DB, config: KeyValueDBConfig, created_at: Option) -> Self { let mut changes_store = ChangeStore::new(); @@ -84,7 +83,7 @@ where } } - pub(crate) fn commit(&mut self, id: ID) -> Result<(), BonsaiStorageError> { + pub(crate) fn commit(&mut self, id: ID) -> Result<(), BonsaiStorageError> { if Some(&id) > self.changes_store.id_queue.back() { self.changes_store.id_queue.push_back(id); } else { @@ -99,15 +98,16 @@ where let current_changes = std::mem::take(&mut self.changes_store.current_changes); for (key, change) in current_changes.serialize(&id).iter() { self.db - .insert(&KeyType::TrieLog(key), change, Some(&mut batch))?; + .insert(&DatabaseKey::TrieLog(key), change, Some(&mut batch))?; } self.db.write_batch(batch)?; if let Some(max_saved_trie_logs) = self.config.max_saved_trie_logs { while self.changes_store.id_queue.len() > max_saved_trie_logs { // verified by previous conditional statement - let id = self.changes_store.id_queue.pop_front().unwrap().serialize(); - self.db.remove_by_prefix(&KeyType::TrieLog(&id))?; + let id = self.changes_store.id_queue.pop_front().unwrap(); + self.db + .remove_by_prefix(&DatabaseKey::TrieLog(&id.to_bytes()))?; } } Ok(()) @@ -121,26 +121,32 @@ where self.config.clone() } - pub(crate) fn get(&self, key: &TrieKeyType) -> Result>, BonsaiStorageError> { + pub(crate) fn get( + &self, + key: &TrieKey, + ) -> Result>, BonsaiStorageError> { trace!("Getting from KeyValueDB: {:?}", key); Ok(self.db.get(&key.into())?) } - pub(crate) fn contains(&self, key: &TrieKeyType) -> Result { + pub(crate) fn contains( + &self, + key: &TrieKey, + ) -> Result> { trace!("Contains from KeyValueDB: {:?}", key); Ok(self.db.contains(&key.into())?) } pub(crate) fn insert( &mut self, - key: &TrieKeyType, + key: &TrieKey, value: &[u8], batch: Option<&mut DB::Batch>, - ) -> Result<(), BonsaiStorageError> { + ) -> Result<(), BonsaiStorageError> { trace!("Inserting into KeyValueDB: {:?} {:?}", key, value); let old_value = self.db.insert(&key.into(), value, batch)?; self.changes_store.current_changes.insert_in_place( - key.into(), + key.clone(), Change { old_value, new_value: Some(value.to_vec()), @@ -151,13 +157,13 @@ where pub(crate) fn remove( &mut self, - key: &TrieKeyType, + key: &TrieKey, batch: Option<&mut DB::Batch>, - ) -> Result<(), BonsaiStorageError> { + ) -> Result<(), BonsaiStorageError> { trace!("Removing from KeyValueDB: {:?}", key); let old_value = self.db.remove(&key.into(), batch)?; self.changes_store.current_changes.insert_in_place( - key.into(), + key.clone(), Change { old_value, new_value: None, @@ -166,7 +172,10 @@ where Ok(()) } - pub(crate) fn write_batch(&mut self, batch: DB::Batch) -> Result<(), BonsaiStorageError> { + pub(crate) fn write_batch( + &mut self, + batch: DB::Batch, + ) -> Result<(), BonsaiStorageError> { trace!("Writing batch into KeyValueDB"); Ok(self.db.write_batch(batch)?) } @@ -193,10 +202,10 @@ where pub(crate) fn get_transaction( &self, id: ID, - ) -> Result, BonsaiStorageError> - where - BonsaiStorageError: std::convert::From<::DatabaseError>, - { + ) -> Result< + Option, + BonsaiStorageError<::DatabaseError>, + > { let Some(change_id) = self.snap_holder.range(..=id).last() else { return Ok(None); }; @@ -221,7 +230,7 @@ where let changes = ChangeBatch::deserialize( id, self.db - .get_by_prefix(&KeyType::TrieLog(id.serialize().as_ref())) + .get_by_prefix(&DatabaseKey::TrieLog(&id.to_bytes())) .map_err(|_| { BonsaiStorageError::Transaction(format!( "database is missing trie logs for {:?}", @@ -230,7 +239,7 @@ where })?, ); for (key, change) in changes.0 { - let key = KeyType::from(&key); + let key = DatabaseKey::from(&key); match (&change.old_value, &change.new_value) { (Some(_), Some(new_value)) => { txn.insert(&key, new_value, Some(&mut batch))?; @@ -252,10 +261,7 @@ where pub(crate) fn merge( &mut self, transaction: KeyValueDB, - ) -> Result<(), BonsaiStorageError> - where - BonsaiStorageError: std::convert::From<>::DatabaseError>, - { + ) -> Result<(), BonsaiStorageError<>::DatabaseError>> { let Some(created_at) = transaction.created_at else { return Err(BonsaiStorageError::Merge( "Transaction has no created_at".to_string(), diff --git a/src/lib.rs b/src/lib.rs index a89f353..1a64952 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -89,7 +89,7 @@ use key_value_db::KeyValueDB; use mp_felt::Felt252Wrapper; use mp_hashers::pedersen::PedersenHasher; -use bonsai_database::KeyType; +use bonsai_database::DatabaseKey; use trie::merkle_tree::MerkleTree; mod changes; @@ -156,10 +156,12 @@ impl BonsaiStorage where DB: BonsaiDatabase, ChangeID: id::Id, - BonsaiStorageError: std::convert::From<::DatabaseError>, { /// Create a new bonsai storage instance - pub fn new(db: DB, config: BonsaiStorageConfig) -> Result { + pub fn new( + db: DB, + config: BonsaiStorageConfig, + ) -> Result> { let key_value_db = KeyValueDB::new(db, config.into(), None); Ok(Self { trie: MerkleTree::new(key_value_db)?, @@ -170,7 +172,7 @@ where db: DB, config: BonsaiStorageConfig, created_at: ChangeID, - ) -> Result { + ) -> Result> { let key_value_db = KeyValueDB::new(db, config.into(), Some(created_at)); Ok(Self { trie: MerkleTree::new(key_value_db)?, @@ -183,14 +185,17 @@ where &mut self, key: &BitSlice, value: &Felt252Wrapper, - ) -> Result<(), BonsaiStorageError> { + ) -> Result<(), BonsaiStorageError> { self.trie.set(key, *value)?; Ok(()) } /// Remove a key/value in the trie /// If the value doesn't exist it will do nothing - pub fn remove(&mut self, key: &BitSlice) -> Result<(), BonsaiStorageError> { + pub fn remove( + &mut self, + key: &BitSlice, + ) -> Result<(), BonsaiStorageError> { self.trie.set(key, Felt252Wrapper::ZERO)?; Ok(()) } @@ -199,19 +204,25 @@ where pub fn get( &self, key: &BitSlice, - ) -> Result, BonsaiStorageError> { + ) -> Result, BonsaiStorageError> { self.trie.get(key) } /// Checks if the key exists in the trie. - pub fn contains(&self, key: &BitSlice) -> Result { + pub fn contains( + &self, + key: &BitSlice, + ) -> Result> { self.trie.contains(key) } /// Go to a specific commit ID. /// If insert/remove is called between the last `commit()` and a call to this function, /// the in-memory changes will be discarded. - pub fn revert_to(&mut self, requested_id: ChangeID) -> Result<(), BonsaiStorageError> { + pub fn revert_to( + &mut self, + requested_id: ChangeID, + ) -> Result<(), BonsaiStorageError> { let kv = self.trie.db_mut(); // Clear current changes @@ -241,7 +252,7 @@ where full.extend( ChangeBatch::deserialize( id, - kv.db.get_by_prefix(&KeyType::TrieLog(&id.serialize()))?, + kv.db.get_by_prefix(&DatabaseKey::TrieLog(&id.to_bytes()))?, ) .0, ); @@ -250,7 +261,7 @@ where // Revert changes let mut batch = kv.db.create_batch(); for (key, change) in full.iter().rev() { - let key = KeyType::from(key); + let key = DatabaseKey::from(key); match (&change.old_value, &change.new_value) { (Some(old_value), Some(_)) => { kv.db.insert(&key, old_value, Some(&mut batch))?; @@ -271,12 +282,13 @@ where kv.changes_store.id_queue.push_back(current); } for id in truncated.iter() { - kv.db.remove_by_prefix(&KeyType::TrieLog(&id.serialize()))?; + kv.db + .remove_by_prefix(&DatabaseKey::TrieLog(&id.to_bytes()))?; } // Write revert changes and trie logs truncation kv.db.write_batch(batch)?; - self.trie.reset_root_from_db()?; + self.trie.reset_to_last_commit()?; Ok(()) } @@ -286,13 +298,16 @@ where } /// Get trie root hash at the latest commit - pub fn root_hash(&self) -> Result { + pub fn root_hash(&self) -> Result> { Ok(self.trie.root_hash()) } /// This function must be used with transactional state only. /// Similar to `commit` but without optimizations. - pub fn transactional_commit(&mut self, id: ChangeID) -> Result<(), BonsaiStorageError> { + pub fn transactional_commit( + &mut self, + id: ChangeID, + ) -> Result<(), BonsaiStorageError> { self.trie.commit()?; self.trie.db_mut().commit(id)?; Ok(()) @@ -312,7 +327,7 @@ where pub fn get_proof( &self, key: &BitSlice, - ) -> Result, BonsaiStorageError> { + ) -> Result, BonsaiStorageError> { self.trie.get_proof(key) } @@ -331,16 +346,19 @@ impl BonsaiStorage where DB: BonsaiDatabase + BonsaiPersistentDatabase, ChangeID: id::Id, - BonsaiStorageError: std::convert::From<::DatabaseError>, { /// Update trie and database using all changes since the last commit. - pub fn commit(&mut self, id: ChangeID) -> Result<(), BonsaiStorageError> { + pub fn commit( + &mut self, + id: ChangeID, + ) -> Result<(), BonsaiStorageError<::DatabaseError>> { self.trie.commit()?; self.trie.db_mut().commit(id)?; self.trie.db_mut().create_snapshot(id); Ok(()) } + #[allow(clippy::type_complexity)] /// Get a transactional state of the trie at a specific commit ID. /// /// Transactional state allow you to fetch a point-in-time state of the trie. You can @@ -349,10 +367,10 @@ where &self, change_id: ChangeID, config: BonsaiStorageConfig, - ) -> Result>, BonsaiStorageError> - where - BonsaiStorageError: std::convert::From<::DatabaseError>, - { + ) -> Result< + Option>, + BonsaiStorageError<::DatabaseError>, + > { if let Some(transaction) = self.trie.db_ref().get_transaction(change_id)? { Ok(Some(BonsaiStorage::new_from_transactional_state( transaction, @@ -373,10 +391,7 @@ where pub fn merge( &mut self, transactional_bonsai_storage: BonsaiStorage, - ) -> Result<(), BonsaiStorageError> - where - BonsaiStorageError: - std::convert::From<>::DatabaseError>, + ) -> Result<(), BonsaiStorageError<>::DatabaseError>> { self.trie .db_mut() diff --git a/src/tests/trie_log.rs b/src/tests/trie_log.rs index 7d50eb4..630b76a 100644 --- a/src/tests/trie_log.rs +++ b/src/tests/trie_log.rs @@ -183,6 +183,7 @@ fn remove_and_reinsert() { bonsai_storage.remove(&bitvec).unwrap(); bonsai_storage.commit(id1).unwrap(); let root_hash1 = bonsai_storage.root_hash().unwrap(); + println!("root_hash1: {:?}", root_hash1); let id2 = id_builder.new_id(); println!("before second insert"); diff --git a/src/trie/merkle_node.rs b/src/trie/merkle_node.rs index 9ec8e1e..bf81032 100644 --- a/src/trie/merkle_node.rs +++ b/src/trie/merkle_node.rs @@ -9,7 +9,7 @@ use bitvec::slice::BitSlice; use mp_felt::Felt252Wrapper; use parity_scale_codec::{Decode, Encode}; -use super::merkle_tree::Path; +use super::path::Path; /// Id of a Node within the tree #[derive(Copy, Clone, Debug, PartialEq, Eq, Default, PartialOrd, Ord, Hash, Encode, Decode)] @@ -18,7 +18,7 @@ pub struct NodeId(pub u64); impl NodeId { /// Mutates the given NodeId to be the next one and returns it. pub fn next_id(&mut self) -> NodeId { - self.0 += 1; + self.0 = self.0.checked_add(1).expect("Node id overflow"); NodeId(self.0) } @@ -49,7 +49,8 @@ pub enum NodeHandle { /// Describes the [Node::Binary] variant. #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Encode, Decode)] pub struct BinaryNode { - /// The hash of this node. + /// The hash of this node. Is [None] if the node + /// has not yet been committed. pub hash: Option, /// The height of this node in the tree. pub height: u64, @@ -217,3 +218,151 @@ impl EdgeNode { &self.path.0[..common_length] } } + +#[test] +fn test_path_matches_basic() { + let path = Path( + BitSlice::::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]) + .to_bitvec(), + ); + let edge = EdgeNode { + hash: None, + height: 0, + path, + child: NodeHandle::Hash(Felt252Wrapper::ZERO), + }; + + let key = BitSlice::::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]); + assert!(edge.path_matches(key)); +} + +#[test] +fn test_path_matches_with_height() { + let path = Path( + BitSlice::::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]) + .to_bitvec(), + ); + let edge = EdgeNode { + hash: None, + height: 8, + path, + child: NodeHandle::Hash(Felt252Wrapper::ZERO), + }; + + let key = BitSlice::::from_slice(&[ + 0b10101010, 0b10101010, 0b01010101, 0b10101010, 0b01010101, + ]); + assert!(edge.path_matches(key)); +} + +#[test] +fn test_path_matches_only_part_with_height() { + let path = Path( + BitSlice::::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]) + .to_bitvec(), + ); + let edge = EdgeNode { + hash: None, + height: 8, + path, + child: NodeHandle::Hash(Felt252Wrapper::ZERO), + }; + + let key = BitSlice::::from_slice(&[ + 0b10101010, 0b10101010, 0b01010101, 0b10101010, 0b01010101, 0b10101010, + ]); + assert!(edge.path_matches(key)); +} + +#[test] +fn test_path_dont_match() { + let path = Path( + BitSlice::::from_slice(&[0b10111010, 0b01010101, 0b10101010, 0b01010101]) + .to_bitvec(), + ); + let edge = EdgeNode { + hash: None, + height: 0, + path, + child: NodeHandle::Hash(Felt252Wrapper::ZERO), + }; + + let key = BitSlice::::from_slice(&[ + 0b10101010, 0b01010101, 0b10101010, 0b01010101, 0b10101010, + ]); + assert!(!edge.path_matches(key)); +} + +#[test] +fn test_common_path_basic() { + let path = Path( + BitSlice::::from_slice(&[0b10101010, 0b01010101, 0b10101010, 0b01010101]) + .to_bitvec(), + ); + let edge = EdgeNode { + hash: None, + height: 0, + path: path.clone(), + child: NodeHandle::Hash(Felt252Wrapper::ZERO), + }; + + 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 edge = EdgeNode { + hash: None, + height: 0, + path, + child: NodeHandle::Hash(Felt252Wrapper::ZERO), + }; + + let key = BitSlice::::from_slice(&[0b10101010, 0b01010101]); + assert_eq!( + edge.common_path(key), + 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 edge = EdgeNode { + hash: None, + height: 8, + path, + child: NodeHandle::Hash(Felt252Wrapper::ZERO), + }; + + 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 edge = EdgeNode { + hash: None, + height: 0, + path, + child: NodeHandle::Hash(Felt252Wrapper::ZERO), + }; + + 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 index b990120..64ea8f8 100644 --- a/src/trie/merkle_tree.rs +++ b/src/trie/merkle_tree.rs @@ -7,15 +7,16 @@ use bitvec::{ }; use derive_more::Constructor; use mp_felt::Felt252Wrapper; -use mp_hashers::HasherT; -use parity_scale_codec::{Decode, Encode, Error, Input, Output}; +use mp_hashers::{pedersen::PedersenHasher, HasherT}; +use parity_scale_codec::{Decode, Encode}; use std::{collections::HashMap, mem}; use crate::{error::BonsaiStorageError, id::Id, BonsaiDatabase, KeyValueDB}; use super::{ merkle_node::{BinaryNode, Direction, EdgeNode, Node, NodeHandle, NodeId}, - TrieKeyType, + path::Path, + TrieKey, }; #[cfg(test)] @@ -30,97 +31,11 @@ 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(pub HashMap); - -#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct Path(pub BitVec); - -impl Encode for Path { - fn encode_to(&self, dest: &mut T) { - // Inspired from scale_bits crate but don't use it to avoid copy and u32 length encoding - let iter = self.0.iter(); - let len = iter.len(); - // SAFETY: len is <= 251 - dest.push_byte(len as u8); - let mut next_store: u8 = 0; - let mut pos_in_next_store: u8 = 7; - for b in iter { - let bit = match *b { - true => 1, - false => 0, - }; - next_store |= bit << pos_in_next_store; - - if pos_in_next_store == 0 { - pos_in_next_store = 8; - dest.push_byte(next_store); - next_store = 0; - } - pos_in_next_store -= 1; - } - - if pos_in_next_store < 7 { - dest.push_byte(next_store); - } - } -} +pub struct NodesMapping(HashMap); -impl Decode for Path { - fn decode(input: &mut I) -> Result { - // Inspired from scale_bits crate but don't use it to avoid copy and u32 length encoding - // SAFETY: len is <= 251 - let len: u8 = input.read_byte()?; - let mut remaining_bits = len as usize; - let mut current_byte = None; - let mut bit = 7; - 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: - let store = match current_byte { - Some(store) => store, - None => { - let store = match input.read_byte() { - Ok(s) => s, - Err(e) => return Err(e), - }; - current_byte = Some(store); - store - } - }; - - // Extract a bit: - let res = match (store >> bit) & 1 { - 0 => false, - 1 => true, - _ => unreachable!("Can only be 0 or 1 owing to &1"), - }; - bits.push(res); - - // Update records for next bit: - remaining_bits -= 1; - if bit == 0 { - current_byte = None; - bit = 8; - } - bit -= 1; - } - Ok(Self(bits)) - } -} - -#[test] -fn test_shared_path_encode_decode() { - let path = Path(BitVec::::from_slice(&[0b10101010, 0b10101010])); - let mut encoded = Vec::new(); - path.encode_to(&mut encoded); - - let decoded = Path::decode(&mut &encoded[..]).unwrap(); - assert_eq!(path, decoded); -} /// A node used in proof generated by the trie. /// -/// Each node hold only the minimum of data that need to be known for the proof: the child hashes (and path for edge node) +/// See pathfinders merkle-tree crate for more information. #[derive(Debug, Clone, PartialEq)] pub enum ProofNode { Binary { @@ -140,9 +55,8 @@ impl ProofNode { ProofNode::Edge { child, path } => { let mut bytes = [0u8; 32]; bytes.view_bits_mut::()[256 - path.0.len()..].copy_from_bitslice(&path.0); - // SAFETY: byte array is 32 bytes long let path_hash = Felt252Wrapper::try_from(&bytes).unwrap(); - + // Safe as len() is guaranteed to be <= 251 let length = Felt252Wrapper::from(path.0.len() as u8); Felt252Wrapper(H::hash_elements(child.0, path_hash.0) + length.0) } @@ -157,17 +71,27 @@ impl ProofNode { /// /// For more information on how this functions internally, see [here](super::merkle_node). pub struct MerkleTree { + /// The handle to the current root node could be hash if no modifications has been done + /// since the last commit or in memory if there are some modifications. root_handle: NodeHandle, + /// The last known root hash. Updated only each commit. (possibly outdated between two commits) root_hash: Felt252Wrapper, + /// Temporary storage used to store the nodes that are modified during a commit. + /// This storage is used to avoid modifying the underlying database each time during a commit. storage_nodes: NodesMapping, + /// The underlying database used to store the nodes. db: KeyValueDB, + /// The id of the last node that has been added to the temporary storage. latest_node_id: NodeId, - death_row: Vec, + /// The list of nodes that should be removed from the underlying database during the next commit. + death_row: Vec, + /// The list of leaves that have been modified during the current commit. cache_leaf_modified: HashMap, InsertOrRemove>, + /// The hasher used to hash the nodes. _hasher: PhantomData, } -#[derive(Debug, PartialEq)] +#[derive(Debug, PartialEq, Eq)] enum InsertOrRemove { Insert(T), Remove, @@ -177,27 +101,21 @@ impl MerkleTree { /// Less visible initialization for `MerkleTree` as the main entry points should be /// [`MerkleTree::::load`] for persistent trees and [`MerkleTree::empty`] for /// transient ones. - pub fn new(mut db: KeyValueDB) -> Result - where - BonsaiStorageError: std::convert::From<::DatabaseError>, - { + pub fn new(mut db: KeyValueDB) -> Result> { let nodes_mapping: HashMap = HashMap::new(); - let root_node = db.get(&TrieKeyType::Trie(vec![]))?; + let root_node = db.get(&TrieKey::Trie(vec![]))?; let node = if let Some(root_node) = root_node { - Node::decode(&mut root_node.as_slice()).map_err(|err| { - BonsaiStorageError::Trie(format!("Couldn't decode root node: {}", err)) - })? + Node::decode(&mut root_node.as_slice())? } else { db.insert( - &TrieKeyType::Trie(vec![]), + &TrieKey::Trie(vec![]), &Node::Unresolved(Felt252Wrapper::ZERO).encode(), None, )?; Node::Unresolved(Felt252Wrapper::ZERO) }; - let root = node.hash().ok_or(BonsaiStorageError::Trie( - "Root doesn't exist in the storage".to_string(), - ))?; + // SAFETY: The root node has been created just above + let root = node.hash().unwrap(); Ok(Self { root_handle: NodeHandle::Hash(root), root_hash: root, @@ -214,12 +132,10 @@ impl MerkleTree { self.root_hash } - pub fn reset_root_from_db(&mut self) -> Result<(), BonsaiStorageError> - where - BonsaiStorageError: std::convert::From<::DatabaseError>, - { + /// Remove all the modifications that have been done since the last commit. + pub fn reset_to_last_commit(&mut self) -> Result<(), BonsaiStorageError> { let node = self - .get_tree_branch_in_db_from_path(&BitVec::::new())? + .get_trie_branch_in_db_from_path(&Path(BitVec::::new()))? .ok_or(BonsaiStorageError::Trie( "root node doesn't exist in the storage".to_string(), ))?; @@ -235,23 +151,20 @@ impl MerkleTree { } /// Persists all changes to storage and returns the new root hash. - pub fn commit(&mut self) -> Result - where - BonsaiStorageError: std::convert::From<::DatabaseError>, - { + pub fn commit(&mut self) -> Result> { let mut batch = self.db.create_batch(); for node_key in mem::take(&mut self.death_row) { self.db.remove(&node_key, Some(&mut batch))?; } - let root_hash = self.commit_subtree(self.root_handle, BitVec::new(), &mut batch)?; + let root_hash = self.commit_subtree(self.root_handle, Path(BitVec::new()), &mut batch)?; for (key, value) in mem::take(&mut self.cache_leaf_modified) { match value { InsertOrRemove::Insert(value) => { self.db - .insert(&TrieKeyType::Flat(key), &value.encode(), Some(&mut batch))?; + .insert(&TrieKey::Flat(key), &value.encode(), Some(&mut batch))?; } InsertOrRemove::Remove => { - self.db.remove(&TrieKeyType::Flat(key), Some(&mut batch))?; + self.db.remove(&TrieKey::Flat(key), Some(&mut batch))?; } } } @@ -276,12 +189,9 @@ impl MerkleTree { fn commit_subtree( &mut self, node_handle: NodeHandle, - path: BitVec, + path: Path, batch: &mut DB::Batch, - ) -> Result - where - BonsaiStorageError: std::convert::From<::DatabaseError>, - { + ) -> Result> { use Node::*; let node_id = match node_handle { NodeHandle::Hash(hash) => return Ok(hash), @@ -296,9 +206,9 @@ impl MerkleTree { "Couldn't fetch node in the temporary storage".to_string(), ))? { Unresolved(hash) => { - if path.is_empty() { + if path.0.is_empty() { self.db.insert( - &TrieKeyType::Trie(vec![]), + &TrieKey::Trie(vec![]), &Node::Unresolved(hash).encode(), Some(batch), )?; @@ -308,19 +218,17 @@ impl MerkleTree { } } Binary(mut binary) => { - let mut left_path = path.clone(); - left_path.push(false); + let left_path = path.new_with_direction(Direction::Left); let left_hash = self.commit_subtree(binary.left, left_path, batch)?; - let mut right_path = path.clone(); - right_path.push(true); + let right_path = path.new_with_direction(Direction::Right); let right_hash = self.commit_subtree(binary.right, right_path, batch)?; let hash = Felt252Wrapper(H::hash_elements(left_hash.0, right_hash.0)); binary.hash = Some(hash); binary.left = NodeHandle::Hash(left_hash); binary.right = NodeHandle::Hash(right_hash); - let key_bytes = [&[path.len() as u8], path.as_raw_slice()].concat(); + let key_bytes = [&[path.0.len() as u8], path.0.as_raw_slice()].concat(); self.db.insert( - &TrieKeyType::Trie(key_bytes), + &TrieKey::Trie(key_bytes), &Node::Binary(binary).encode(), Some(batch), )?; @@ -329,32 +237,28 @@ impl MerkleTree { Edge(mut edge) => { let mut child_path = path.clone(); - child_path.extend(&edge.path.0); + child_path.0.extend(&edge.path.0); let child_hash = self.commit_subtree(edge.child, child_path, batch)?; let mut bytes = [0u8; 32]; bytes.view_bits_mut::()[256 - edge.path.0.len()..] .copy_from_bitslice(&edge.path.0); - let felt_path = Felt252Wrapper::try_from(&bytes).map_err(|err| { - BonsaiStorageError::Trie(format!("Couldn't convert path to felt: {}", err)) - })?; + let felt_path = Felt252Wrapper::try_from(&bytes)?; let mut length = [0; 32]; // Safe as len() is guaranteed to be <= 251 length[31] = edge.path.0.len() as u8; - let length = Felt252Wrapper::try_from(&length).map_err(|err| { - BonsaiStorageError::Trie(format!("Couldn't convert length to felt: {}", err)) - })?; + let length = Felt252Wrapper::try_from(&length)?; let hash = Felt252Wrapper(H::hash_elements(child_hash.0, felt_path.0) + length.0); edge.hash = Some(hash); edge.child = NodeHandle::Hash(child_hash); - let key_bytes = if path.is_empty() { + let key_bytes = if path.0.is_empty() { vec![] } else { - [&[path.len() as u8], path.as_raw_slice()].concat() + [&[path.0.len() as u8], path.0.as_raw_slice()].concat() }; self.db.insert( - &TrieKeyType::Trie(key_bytes), + &TrieKey::Trie(key_bytes), &Node::Edge(edge).encode(), Some(batch), )?; @@ -373,10 +277,7 @@ impl MerkleTree { &mut self, key: &BitSlice, value: Felt252Wrapper, - ) -> Result<(), BonsaiStorageError> - where - BonsaiStorageError: std::convert::From<::DatabaseError>, - { + ) -> Result<(), BonsaiStorageError> { if value == Felt252Wrapper::ZERO { return self.delete_leaf(key); } @@ -400,102 +301,94 @@ impl MerkleTree { use Node::*; match path.last() { Some(node_id) => { - let mut nodes_to_add = Vec::with_capacity(4); + 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. - let key_bytes = - [&[key.len() as u8], key.to_bitvec().as_raw_slice()].concat(); - self.cache_leaf_modified - .insert(key_bytes, InsertOrRemove::Insert(value)); - return; - } - // Height of the binary node's children - let child_height = branch_height + 1; - - // Path from binary node to new leaf - let new_path = key[child_height..].to_bitvec(); - // Path from binary node to existing child - let old_path = edge.path.0[common.len() + 1..].to_bitvec(); - - // The new leaf branch of the binary node. - // (this may be edge -> leaf, or just leaf depending). - let key_bytes = - [&[key.len() as u8], key.to_bitvec().as_raw_slice()].concat(); + if let Edge(edge) = node { + let common = edge.common_path(key); + // Height of the binary node + let branch_height = edge.height as usize + common.len(); + if branch_height == key.len() { + edge.child = NodeHandle::Hash(value); + // The leaf already exists, we simply change its value. + let key_bytes = bitslice_to_bytes(key); self.cache_leaf_modified .insert(key_bytes, InsertOrRemove::Insert(value)); + return; + } + // Height of the binary node's children + let child_height = branch_height + 1; + + // Path from binary node to new leaf + let new_path = key[child_height..].to_bitvec(); + // Path from binary node to existing child + let old_path = edge.path.0[common.len() + 1..].to_bitvec(); + + // The new leaf branch of the binary node. + // (this may be edge -> leaf, or just leaf depending). + let key_bytes = bitslice_to_bytes(key); + self.cache_leaf_modified + .insert(key_bytes, InsertOrRemove::Insert(value)); + + let new = if new_path.is_empty() { + NodeHandle::Hash(value) + } else { + let new_edge = Node::Edge(EdgeNode { + hash: None, + height: child_height as u64, + path: Path(new_path), + child: NodeHandle::Hash(value), + }); + let edge_id = self.latest_node_id.next_id(); + nodes_to_add.push((edge_id, new_edge)); + NodeHandle::InMemory(edge_id) + }; - 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 { + // 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: branch_height as u64, - left, - right, + 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) + }; - // We may require an edge leading to the binary node. - let new_node = if common.is_empty() { - branch - } else { - let branch_id = self.latest_node_id.next_id(); - nodes_to_add.push((branch_id, branch)); - - Node::Edge(EdgeNode { - hash: None, - height: edge.height, - path: Path(common.to_bitvec()), - child: NodeHandle::InMemory(branch_id), - }) - }; - let path = key[..edge.height as usize].to_bitvec(); - let key_bytes = - [&[path.len() as u8], path.into_vec().as_slice()].concat(); - self.death_row.push(TrieKeyType::Trie(key_bytes)); - *node = new_node; - } - Unresolved(_) | Binary(_) => { - unreachable!("The end of a traversion cannot be unresolved or binary") - } + 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 path = key[..edge.height as usize].to_bitvec(); + let key_bytes = [&[path.len() as u8], path.into_vec().as_slice()].concat(); + self.death_row.push(TrieKey::Trie(key_bytes)); + *node = new_node; }; }); for (id, node) in nodes_to_add { @@ -520,7 +413,7 @@ impl MerkleTree { self.root_handle = NodeHandle::InMemory(self.latest_node_id); - let key_bytes = [&[key.len() as u8], key.to_bitvec().as_raw_slice()].concat(); + let key_bytes = bitslice_to_bytes(key); self.cache_leaf_modified .insert(key_bytes, InsertOrRemove::Insert(value)); Ok(()) @@ -548,10 +441,10 @@ impl MerkleTree { /// # Arguments /// /// * `key` - The key to delete. - fn delete_leaf(&mut self, key: &BitSlice) -> Result<(), BonsaiStorageError> - where - BonsaiStorageError: std::convert::From<::DatabaseError>, - { + fn delete_leaf( + &mut self, + key: &BitSlice, + ) -> Result<(), BonsaiStorageError> { // Algorithm explanation: // // The leaf's parent node is either an edge, or a binary node. @@ -566,16 +459,13 @@ impl MerkleTree { // // Then we are done. - let key_bytes = [&[key.len() as u8], key.to_bitvec().as_raw_slice()].concat(); + let key_bytes = bitslice_to_bytes(key); self.cache_leaf_modified .insert(key_bytes.clone(), InsertOrRemove::Remove); - if !self.db.contains(&TrieKeyType::Flat(key_bytes))? { - return Ok(()); - } let path = self.preload_nodes(key)?; - let mut last_binary_path = key.to_bitvec(); + 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| { @@ -586,18 +476,9 @@ impl MerkleTree { Node::Binary(_) => {} Node::Edge(edge) => { for _ in 0..edge.path.0.len() { - last_binary_path.pop(); - } - let key_bytes = [ - &[last_binary_path.len() as u8], - last_binary_path.as_raw_slice(), - ] - .concat(); - if last_binary_path.is_empty() { - self.death_row.push(TrieKeyType::Trie(vec![])); - } else { - self.death_row.push(TrieKeyType::Trie(key_bytes)); + last_binary_path.0.pop(); } + self.death_row.push((&last_binary_path).into()); } } !node.is_binary() @@ -619,7 +500,7 @@ 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. - last_binary_path.push(direction.into()); + last_binary_path.0.push(direction.into()); let path = Path(once(bool::from(direction)).collect::>()); let mut edge = EdgeNode { hash: None, @@ -650,47 +531,45 @@ impl MerkleTree { }; // Check the parent of the new edge. If it is also an edge, then they must merge. - if let Some(node) = parent_branch_node { - let child = if let Node::Edge(edge) = - self.storage_nodes - .0 - .get(&node) - .ok_or(BonsaiStorageError::Trie( - "Node not found in memory".to_string(), - ))? { - let child_node = match edge.child { - NodeHandle::Hash(_) => return Ok(()), - NodeHandle::InMemory(child_id) => { - self.storage_nodes - .0 - .get(&child_id) - .ok_or(BonsaiStorageError::Trie( - "Node not found in memory".to_string(), - ))? + if let Some(node_id) = parent_branch_node { + let node = self + .storage_nodes + .0 + .get(&node_id) + .ok_or(BonsaiStorageError::Trie( + "Node not found in memory".to_string(), + ))?; + // If it's an edge node and the child is in memory and it's an edge too we + // return the child otherwise we leave + let child = + if let Node::Edge(edge) = node { + match edge.child { + NodeHandle::Hash(_) => return Ok(()), + NodeHandle::InMemory(child_id) => { + let child_node = self.storage_nodes.0.get(&child_id).ok_or( + BonsaiStorageError::Trie("Node not found in memory".to_string()), + )?; + if let Node::Edge(child_edge) = child_node { + child_edge.clone() + } else { + return Ok(()); + } + } } + } else { + return Ok(()); }; - match child_node { - Node::Edge(child_edge) => child_edge.clone(), - _ => { - return Ok(()); - } - } - } else { - return Ok(()); - }; + // Get a mutable reference to the parent node to merge them let edge = self .storage_nodes .0 - .get_mut(&node) + .get_mut(&node_id) .ok_or(BonsaiStorageError::Trie( "Node not found in memory".to_string(), ))?; - match edge { - Node::Edge(edge) => { - edge.path.0.extend_from_bitslice(&child.path.0); - edge.child = child.child; - } - _ => unreachable!(), + if let Node::Edge(edge) = edge { + edge.path.0.extend_from_bitslice(&child.path.0); + edge.child = child.child; } } Ok(()) @@ -708,34 +587,31 @@ impl MerkleTree { pub fn get( &self, key: &BitSlice, - ) -> Result, BonsaiStorageError> - where - BonsaiStorageError: std::convert::From<::DatabaseError>, - { - let key = &[&[key.len() as u8], key.to_bitvec().as_raw_slice()].concat(); - if let Some(value) = self.cache_leaf_modified.get(key) { + ) -> Result, BonsaiStorageError> { + let key = bitslice_to_bytes(key); + if let Some(value) = self.cache_leaf_modified.get(&key) { match value { InsertOrRemove::Remove => return Ok(None), InsertOrRemove::Insert(value) => return Ok(Some(*value)), } } self.db - .get(&TrieKeyType::Flat(key.to_vec())) + .get(&TrieKey::Flat(key.to_vec())) .map(|r| r.map(|opt| Felt252Wrapper::decode(&mut opt.as_slice()).unwrap())) } - pub fn contains(&self, key: &BitSlice) -> Result - where - BonsaiStorageError: std::convert::From<::DatabaseError>, - { - let key = &[&[key.len() as u8], key.to_bitvec().as_raw_slice()].concat(); - if let Some(value) = self.cache_leaf_modified.get(key) { + pub fn contains( + &self, + 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), } } - self.db.contains(&TrieKeyType::Flat(key.to_vec())) + self.db.contains(&TrieKey::Flat(key.to_vec())) } /// Generates a merkle-proof for a given `key`. @@ -756,15 +632,15 @@ impl MerkleTree { /// # Returns /// /// The merkle proof and all the child nodes hashes. - pub fn get_proof(&self, key: &BitSlice) -> Result, BonsaiStorageError> - where - BonsaiStorageError: std::convert::From<::DatabaseError>, - { + pub fn get_proof( + &self, + key: &BitSlice, + ) -> Result, BonsaiStorageError> { let mut nodes = Vec::with_capacity(251); let mut node = match self.root_handle { NodeHandle::Hash(_) => { let node = self - .get_tree_branch_in_db_from_path(&BitVec::::new())? + .get_trie_branch_in_db_from_path(&Path(BitVec::::new()))? .ok_or(BonsaiStorageError::Trie( "Couldn't fetch root node in db".to_string(), ))?; @@ -788,7 +664,7 @@ impl MerkleTree { let child_path = key[..edge.height as usize + edge.path.0.len()].to_bitvec(); let child_node = match edge.child { NodeHandle::Hash(hash) => { - let node = self.get_tree_branch_in_db_from_path(&child_path)?; + let node = self.get_trie_branch_in_db_from_path(&Path(child_path))?; if let Some(node) = node { node } else { @@ -832,7 +708,7 @@ impl MerkleTree { let next_path = key[..binary.height as usize + 1].to_bitvec(); let next_node = match next { NodeHandle::Hash(_) => self - .get_tree_branch_in_db_from_path(&next_path)? + .get_trie_branch_in_db_from_path(&Path(next_path))? .ok_or(BonsaiStorageError::Trie( "Couldn't fetch next node in db".to_string(), ))?, @@ -911,15 +787,15 @@ impl MerkleTree { /// # Returns /// /// The list of nodes along the path. - fn preload_nodes(&mut self, dst: &BitSlice) -> Result, BonsaiStorageError> - where - BonsaiStorageError: std::convert::From<::DatabaseError>, - { + fn preload_nodes( + &mut self, + dst: &BitSlice, + ) -> Result, BonsaiStorageError> { let mut nodes = Vec::with_capacity(251); let node_id = match self.root_handle { NodeHandle::Hash(_) => { let node = self - .get_tree_branch_in_db_from_path(&BitVec::::new())? + .get_trie_branch_in_db_from_path(&Path(BitVec::::new()))? .ok_or(BonsaiStorageError::Trie( "Couldn't fetch root node in db".to_string(), ))?; @@ -937,7 +813,7 @@ impl MerkleTree { root_id } }; - self.preload_nodes_subtree(dst, node_id, BitVec::::new(), &mut nodes)?; + self.preload_nodes_subtree(dst, node_id, Path(BitVec::::new()), &mut nodes)?; Ok(nodes) } @@ -945,12 +821,9 @@ impl MerkleTree { &mut self, dst: &BitSlice, root_id: NodeId, - mut path: BitVec, + mut path: Path, nodes: &mut Vec, - ) -> Result<(), BonsaiStorageError> - where - BonsaiStorageError: std::convert::From<::DatabaseError>, - { + ) -> Result<(), BonsaiStorageError> { let node = self .storage_nodes .0 @@ -960,14 +833,21 @@ impl MerkleTree { ))? .clone(); match node { + // We are in a case where the trie is empty and so there is nothing to preload. Node::Unresolved(_hash) => Ok(()), + // We are checking which side of the binary we should load in memory (if we don't have it already) + // We load this "child-side" node in the memory and refer his memory handle in the binary node. + // We also add the "child-side" node in the list that accumulate all the nodes we want to preload. + // We override the binary node in the memory with this new version that has the "child-side" memory handle + // instead of the hash. + // We call recursively the function with the "child-side" node. Node::Binary(mut binary_node) => { let next_direction = binary_node.direction(dst); - path.push(bool::from(next_direction)); + path.0.push(bool::from(next_direction)); let next = binary_node.get_child(next_direction); match next { NodeHandle::Hash(_) => { - let node = self.get_tree_branch_in_db_from_path(&path)?.ok_or( + let node = self.get_trie_branch_in_db_from_path(&path)?.ok_or( BonsaiStorageError::Trie("Couldn't fetch node in db".to_string()), )?; self.latest_node_id.next_id(); @@ -992,15 +872,19 @@ impl MerkleTree { } } } + // If the edge node match the path we want to preload then we load the child node in memory (if we don't have it already) + // and we override the edge node in the memory with this new version that has the child memory handle instead of the hash. + // We also add the child node in the list that accumulate all the nodes we want to preload. + // We call recursively the function with the child node. Node::Edge(mut edge_node) if edge_node.path_matches(dst) => { - path.extend_from_bitslice(&edge_node.path.0); - if path == dst { + path.0.extend_from_bitslice(&edge_node.path.0); + if path.0 == dst { return Ok(()); } let next = edge_node.child; match next { NodeHandle::Hash(_) => { - let node = self.get_tree_branch_in_db_from_path(&path)?; + let node = self.get_trie_branch_in_db_from_path(&path)?; if let Some(node) = node { self.latest_node_id.next_id(); self.storage_nodes.0.insert(self.latest_node_id, node); @@ -1018,24 +902,18 @@ impl MerkleTree { } } } + // We are in a case where the edge node doesn't match the path we want to preload so we return nothing. Node::Edge(_) => Ok(()), } } - fn get_tree_branch_in_db_from_path( + /// Get the node of the trie that corresponds to the path. + fn get_trie_branch_in_db_from_path( &self, - path: &BitVec, - ) -> Result, BonsaiStorageError> - where - BonsaiStorageError: std::convert::From<::DatabaseError>, - { - let key = if path.is_empty() { - vec![] - } else { - [&[path.len() as u8], path.as_raw_slice()].concat() - }; + path: &Path, + ) -> Result, BonsaiStorageError> { self.db - .get(&TrieKeyType::Trie(key))? + .get(&path.into())? .map(|node| { Node::decode(&mut node.as_slice()).map_err(|err| { BonsaiStorageError::Trie(format!("Couldn't decode node: {}", err)) @@ -1055,11 +933,10 @@ impl MerkleTree { /// # Arguments /// /// * `parent` - The parent node to merge the child with. - fn merge_edges(&self, parent: &mut EdgeNode) -> Result<(), BonsaiStorageError> - where - BonsaiStorageError: std::convert::From<::DatabaseError>, - { - //TODO: Add deletion of unused nodes + fn merge_edges( + &self, + parent: &mut EdgeNode, + ) -> Result<(), BonsaiStorageError> { let child_node = match parent.child { NodeHandle::Hash(_) => return Ok(()), NodeHandle::InMemory(child_id) => { @@ -1112,7 +989,7 @@ impl MerkleTree { for proof_node in proofs.iter() { // Hash mismatch? Return None. - if proof_node.hash::() != expected_hash { + if proof_node.hash::() != expected_hash { return None; } match proof_node { @@ -1212,6 +1089,10 @@ impl MerkleTree { } } +fn bitslice_to_bytes(bitslice: &BitSlice) -> Vec { + [&[bitslice.len() as u8], bitslice.to_bitvec().as_raw_slice()].concat() +} + #[cfg(test)] mod tests { use bitvec::vec::BitVec; diff --git a/src/trie/mod.rs b/src/trie/mod.rs index bb341ea..0d06bb8 100644 --- a/src/trie/mod.rs +++ b/src/trie/mod.rs @@ -1,5 +1,6 @@ mod merkle_node; pub mod merkle_tree; +mod path; mod trie_db; -pub use trie_db::TrieKeyType; +pub use trie_db::TrieKey; diff --git a/src/trie/path.rs b/src/trie/path.rs new file mode 100644 index 0000000..9a6f2cd --- /dev/null +++ b/src/trie/path.rs @@ -0,0 +1,126 @@ +use bitvec::{order::Msb0, vec::BitVec}; +use parity_scale_codec::{Decode, Encode, Error, Input, Output}; + +use super::{merkle_node::Direction, TrieKey}; + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct Path(pub BitVec); + +impl Encode for Path { + fn encode_to(&self, dest: &mut T) { + // Inspired from scale_bits crate but don't use it to avoid copy and u32 length encoding + let iter = self.0.iter(); + let len = iter.len(); + // SAFETY: len is <= 251 + dest.push_byte(len as u8); + let mut next_store: u8 = 0; + let mut pos_in_next_store: u8 = 7; + for b in iter { + let bit = match *b { + true => 1, + false => 0, + }; + next_store |= bit << pos_in_next_store; + + if pos_in_next_store == 0 { + pos_in_next_store = 8; + dest.push_byte(next_store); + next_store = 0; + } + pos_in_next_store -= 1; + } + + if pos_in_next_store < 7 { + dest.push_byte(next_store); + } + } + + fn size_hint(&self) -> usize { + // Inspired from scale_bits crate but don't use it to avoid copy and u32 length encoding + 1 + (self.0.len() + 7) / 8 + } +} + +impl Decode for Path { + fn decode(input: &mut I) -> Result { + // Inspired from scale_bits crate but don't use it to avoid copy and u32 length encoding + // SAFETY: len is <= 251 + let len: u8 = input.read_byte()?; + let mut remaining_bits = len as usize; + let mut current_byte = None; + let mut bit = 7; + 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: + let store = match current_byte { + Some(store) => store, + None => { + let store = match input.read_byte() { + Ok(s) => s, + Err(e) => return Err(e), + }; + current_byte = Some(store); + store + } + }; + + // Extract a bit: + let res = match (store >> bit) & 1 { + 0 => false, + 1 => true, + _ => unreachable!("Can only be 0 or 1 owing to &1"), + }; + bits.push(res); + + // Update records for next bit: + remaining_bits -= 1; + if bit == 0 { + current_byte = None; + bit = 8; + } + bit -= 1; + } + Ok(Self(bits)) + } +} + +impl Path { + pub(crate) fn new_with_direction(&self, direction: Direction) -> Path { + let mut path = self.0.clone(); + path.push(direction.into()); + Path(path) + } +} + +impl From for TrieKey { + fn from(path: Path) -> Self { + let key = if path.0.is_empty() { + vec![] + } else { + [&[path.0.len() as u8], path.0.as_raw_slice()].concat() + }; + TrieKey::Trie(key) + } +} + +impl From<&Path> for TrieKey { + fn from(path: &Path) -> Self { + let key = if path.0.is_empty() { + vec![] + } else { + [&[path.0.len() as u8], path.0.as_raw_slice()].concat() + }; + TrieKey::Trie(key) + } +} + +#[test] +fn test_shared_path_encode_decode() { + let path = Path(BitVec::::from_slice(&[0b10101010, 0b10101010])); + let mut encoded = Vec::new(); + path.encode_to(&mut encoded); + + let decoded = Path::decode(&mut &encoded[..]).unwrap(); + assert_eq!(path, decoded); +} diff --git a/src/trie/trie_db.rs b/src/trie/trie_db.rs index 84bc876..c2d5ccb 100644 --- a/src/trie/trie_db.rs +++ b/src/trie/trie_db.rs @@ -1,36 +1,58 @@ -use crate::{bonsai_database::KeyType, changes::ChangeKeyType}; +use crate::bonsai_database::DatabaseKey; -#[derive(Debug, Hash, PartialEq, Eq)] -pub enum TrieKeyType { +/// Key in the database of the different elements that are used in the storage of the trie data. +#[derive(Debug, Clone, Hash, PartialEq, Eq)] +pub enum TrieKey { Trie(Vec), Flat(Vec), } -impl TrieKeyType { - pub fn as_slice(&self) -> &[u8] { - match self { - TrieKeyType::Trie(slice) => slice, - TrieKeyType::Flat(slice) => slice, +enum TrieKeyType { + Trie = 0, + Flat = 1, +} + +impl From for u8 { + fn from(value: TrieKey) -> Self { + match value { + TrieKey::Trie(_) => TrieKeyType::Trie as u8, + TrieKey::Flat(_) => TrieKeyType::Flat as u8, } } } -impl<'a> From<&'a TrieKeyType> for KeyType<'a> { - fn from(key: &'a TrieKeyType) -> Self { - let key_slice = key.as_slice(); - match key { - TrieKeyType::Trie(_) => KeyType::Trie(key_slice), - TrieKeyType::Flat(_) => KeyType::Flat(key_slice), +impl From<&TrieKey> for u8 { + fn from(value: &TrieKey) -> Self { + match value { + TrieKey::Trie(_) => TrieKeyType::Trie as u8, + TrieKey::Flat(_) => TrieKeyType::Flat as u8, + } + } +} + +impl TrieKey { + pub fn from_variant_and_bytes(variant: u8, bytes: Vec) -> Self { + match variant { + x if x == TrieKeyType::Trie as u8 => TrieKey::Trie(bytes), + x if x == TrieKeyType::Flat as u8 => TrieKey::Flat(bytes), + _ => panic!("Invalid trie key type"), + } + } + + pub fn as_slice(&self) -> &[u8] { + match self { + TrieKey::Trie(slice) => slice, + TrieKey::Flat(slice) => slice, } } } -impl<'a> From<&'a TrieKeyType> for ChangeKeyType { - fn from(key: &'a TrieKeyType) -> Self { +impl<'a> From<&'a TrieKey> for DatabaseKey<'a> { + fn from(key: &'a TrieKey) -> Self { let key_slice = key.as_slice(); match key { - TrieKeyType::Trie(_) => ChangeKeyType::Trie(key_slice.to_vec()), - TrieKeyType::Flat(_) => ChangeKeyType::Flat(key_slice.to_vec()), + TrieKey::Trie(_) => DatabaseKey::Trie(key_slice), + TrieKey::Flat(_) => DatabaseKey::Flat(key_slice), } } } From cdc7dff7f7a0519c553f9ccaf8177c56f87587ae Mon Sep 17 00:00:00 2001 From: AurelienFT Date: Fri, 26 Jan 2024 15:45:53 +0100 Subject: [PATCH 2/3] Format + Clippy after merge --- src/bonsai_database.rs | 4 ++-- src/databases/hashmap_db.rs | 5 +---- src/error.rs | 5 +++-- src/lib.rs | 10 ++++++++-- 4 files changed, 14 insertions(+), 10 deletions(-) diff --git a/src/bonsai_database.rs b/src/bonsai_database.rs index 4509929..39b07ae 100644 --- a/src/bonsai_database.rs +++ b/src/bonsai_database.rs @@ -1,8 +1,8 @@ -#[cfg(feature = "std")] -use std::error::Error; use crate::id::Id; #[cfg(not(feature = "std"))] use alloc::vec::Vec; +#[cfg(feature = "std")] +use std::error::Error; /// Key in the database of the different elements that can be stored in the database. #[derive(Debug, Hash, PartialEq, Eq)] diff --git a/src/databases/hashmap_db.rs b/src/databases/hashmap_db.rs index 21b7cee..f11914e 100644 --- a/src/databases/hashmap_db.rs +++ b/src/databases/hashmap_db.rs @@ -4,10 +4,7 @@ use crate::{ BonsaiDatabase, }; #[cfg(not(feature = "std"))] -use alloc::{ - vec::Vec, - collections::BTreeMap, -}; +use alloc::{collections::BTreeMap, vec::Vec}; use core::{fmt, fmt::Display}; #[cfg(not(feature = "std"))] use hashbrown::HashMap; diff --git a/src/error.rs b/src/error.rs index 8ecb8a7..92e8ff4 100644 --- a/src/error.rs +++ b/src/error.rs @@ -25,7 +25,9 @@ where NodeDecodeError(parity_scale_codec::Error), } -impl core::convert::From for BonsaiStorageError { +impl core::convert::From + for BonsaiStorageError +{ fn from(value: DatabaseError) -> Self { Self::Database(value) } @@ -39,7 +41,6 @@ impl core::convert::From } } - #[cfg(feature = "std")] impl Display for BonsaiStorageError where diff --git a/src/lib.rs b/src/lib.rs index 3d2a2e6..3a8b09c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -200,13 +200,19 @@ where /// Remove a key/value in the trie /// If the value doesn't exist it will do nothing - pub fn remove(&mut self, key: &BitSlice) -> Result<(), BonsaiStorageError> { + pub fn remove( + &mut self, + key: &BitSlice, + ) -> Result<(), BonsaiStorageError> { self.trie.set(key, Felt::ZERO)?; Ok(()) } /// Get a value in the trie. - pub fn get(&self, key: &BitSlice) -> Result, BonsaiStorageError> { + pub fn get( + &self, + key: &BitSlice, + ) -> Result, BonsaiStorageError> { self.trie.get(key) } From fe355fe8c33cb64859f786dae87420e38cae18ae Mon Sep 17 00:00:00 2001 From: AurelienFT Date: Fri, 26 Jan 2024 16:29:31 +0100 Subject: [PATCH 3/3] Remove useless prints, add more comments and tests --- Cargo.toml | 1 + src/changes.rs | 1 - src/tests/madara_comparison.rs | 3 +-- src/tests/simple.rs | 12 ------------ src/tests/trie_log.rs | 3 --- src/trie/path.rs | 23 +++++++++++++++++++---- 6 files changed, 21 insertions(+), 22 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index dd4f93a..85e7393 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -46,3 +46,4 @@ pathfinder-merkle-tree = { git = "https://github.com/massalabs/pathfinder.git", pathfinder-storage = { git = "https://github.com/massalabs/pathfinder.git", package = "pathfinder-storage", rev = "b7b6d76a76ab0e10f92e5f84ce099b5f727cb4db" } rand = "0.8.5" tempfile = "3.8.0" +rstest = "0.18.2" diff --git a/src/changes.rs b/src/changes.rs index 3af5c3e..3d9a19c 100644 --- a/src/changes.rs +++ b/src/changes.rs @@ -37,7 +37,6 @@ impl ChangeBatch { } } - //TODO: Use serde pub fn serialize(&self, id: &ID) -> Vec<(Vec, &[u8])> { let id = id.to_bytes(); self.0 diff --git a/src/tests/madara_comparison.rs b/src/tests/madara_comparison.rs index 6a525f5..48b1a82 100644 --- a/src/tests/madara_comparison.rs +++ b/src/tests/madara_comparison.rs @@ -24,8 +24,7 @@ fn trie_height_251() { let mut id_builder = BasicIdBuilder::new(); let id = id_builder.new_id(); bonsai_storage.commit(id).unwrap(); - let root_hash = bonsai_storage.root_hash().unwrap(); - println!("root_hash: {:?}", root_hash); + bonsai_storage.root_hash().unwrap(); } // Test to add on Madara side to check with a tree of height 251 and see that we have same hash // #[test]// fn test_height_251() { diff --git a/src/tests/simple.rs b/src/tests/simple.rs index a67902f..9ae3eb6 100644 --- a/src/tests/simple.rs +++ b/src/tests/simple.rs @@ -35,19 +35,7 @@ fn basics() { ); let bitvec = BitVec::from_vec(pair3.0.clone()); bonsai_storage.insert(&bitvec, &pair3.1).unwrap(); - println!( - "get: {:?}", - bonsai_storage.get(&BitVec::from_vec(vec![1, 2, 1])) - ); bonsai_storage.commit(id_builder.new_id()).unwrap(); - println!( - "get: {:?}", - bonsai_storage.get(&BitVec::from_vec(vec![1, 2, 2])) - ); - println!( - "get: {:?}", - bonsai_storage.get(&BitVec::from_vec(vec![1, 2, 3])) - ); let bitvec = BitVec::from_vec(vec![1, 2, 1]); bonsai_storage.remove(&bitvec).unwrap(); assert_eq!( diff --git a/src/tests/trie_log.rs b/src/tests/trie_log.rs index 06477ee..56c6ecc 100644 --- a/src/tests/trie_log.rs +++ b/src/tests/trie_log.rs @@ -184,10 +184,7 @@ fn remove_and_reinsert() { bonsai_storage.remove(&bitvec).unwrap(); bonsai_storage.commit(id1).unwrap(); let root_hash1 = bonsai_storage.root_hash().unwrap(); - println!("root_hash1: {:?}", root_hash1); - let id2 = id_builder.new_id(); - println!("before second insert"); bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); bonsai_storage.commit(id2).unwrap(); diff --git a/src/trie/path.rs b/src/trie/path.rs index 2604244..78906e4 100644 --- a/src/trie/path.rs +++ b/src/trie/path.rs @@ -6,12 +6,20 @@ use super::{merkle_node::Direction, TrieKey}; #[cfg(not(feature = "std"))] use alloc::vec::Vec; +#[cfg(all(feature = "std", test))] +use rstest::rstest; + #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct Path(pub BitVec); impl Encode for Path { fn encode_to(&self, dest: &mut T) { - // Inspired from scale_bits crate but don't use it to avoid copy and u32 length encoding + // 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 ? + // 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 let iter = self.0.iter(); let len = iter.len(); // SAFETY: len is <= 251 @@ -118,9 +126,16 @@ impl From<&Path> for TrieKey { } } -#[test] -fn test_shared_path_encode_decode() { - let path = Path(BitVec::::from_slice(&[0b10101010, 0b10101010])); +#[cfg(all(feature = "std", test))] +#[rstest] +#[case(&[0b10101010, 0b10101010])] +#[case(&[])] +#[case(&[0b10101010])] +#[case(&[0b00000000])] +#[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 mut encoded = Vec::new(); path.encode_to(&mut encoded);