refacto: code structure, doc, tests (#5)

* modifications following starket review * Format + Clippy after merge * Remove useless prints, add more comments and tests --------- Co-authored-by: AurelienFT <[email protected]> Co-authored-by: AurelienFT <[email protected]>
madara-alliance · Jan 28, 2024 · abb735a · abb735a
1 parent 5b8b67a
commit abb735a
Show file tree

Hide file tree

Showing 18 changed files with 803 additions and 559 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -15,6 +15,8 @@ derive_more = { version = "0.99.17", default-features = false, features = [
 ] }
 hashbrown = "0.14.3"
 log = "0.4.20"
+smallvec = "1.11.2"
+
 parity-scale-codec = { version = "3.0.0", default-features = false, features = [
     "derive",
 ] }
@@ -44,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/ensure_no_std/Cargo.lock b/ensure_no_std/Cargo.lock
diff --git a/src/bonsai_database.rs b/src/bonsai_database.rs
@@ -1,62 +1,62 @@
-use crate::{changes::ChangeKeyType, error::BonsaiStorageError, id::Id};
+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)]
-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,
         }
     }
 }
 
+#[cfg(feature = "std")]
+pub trait DBError: Error + Send + Sync {}
+
+#[cfg(not(feature = "std"))]
+pub trait DBError: Send + Sync {}
+
 /// Trait to be implemented on any type that can be used as a database.
 pub trait BonsaiDatabase {
     type Batch: Default;
     #[cfg(feature = "std")]
-    type DatabaseError: std::error::Error + Into<BonsaiStorageError>;
+    type DatabaseError: Error + DBError;
     #[cfg(not(feature = "std"))]
-    type DatabaseError: Into<BonsaiStorageError>;
+    type DatabaseError: 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<Option<Vec<u8>>, Self::DatabaseError>;
+    fn get(&self, key: &DatabaseKey) -> Result<Option<Vec<u8>>, 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<(Vec<u8>, Vec<u8>)>, Self::DatabaseError>;
 
     /// Returns true if the key exists
-    fn contains(&self, key: &KeyType) -> Result<bool, Self::DatabaseError>;
+    fn contains(&self, key: &DatabaseKey) -> Result<bool, Self::DatabaseError>;
 
     /// 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<Option<Vec<u8>>, Self::DatabaseError>;
@@ -65,12 +65,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<Option<Vec<u8>>, 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>;
@@ -81,11 +81,11 @@ pub trait BonsaiDatabase {
 }
 
 pub trait BonsaiPersistentDatabase<ID: Id> {
+    type Transaction: BonsaiDatabase<DatabaseError = Self::DatabaseError>;
     #[cfg(feature = "std")]
-    type DatabaseError: std::error::Error + Into<BonsaiStorageError>;
+    type DatabaseError: Error + DBError;
     #[cfg(not(feature = "std"))]
-    type DatabaseError: Into<BonsaiStorageError>;
-    type Transaction: BonsaiDatabase;
+    type DatabaseError: DBError;
     /// 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
@@ -1,4 +1,4 @@
-use crate::id::Id;
+use crate::{id::Id, trie::TrieKey};
 use serde::{Deserialize, Serialize};
 #[cfg(feature = "std")]
 use std::collections::{hash_map::Entry, HashMap, VecDeque};
@@ -14,45 +14,15 @@ pub struct Change {
     pub new_value: Option<Vec<u8>>,
 }
 
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
-pub enum ChangeKeyType {
-    Trie(Vec<u8>),
-    Flat(Vec<u8>),
-}
-
-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<u8>) -> Self {
-        match id {
-            0 => ChangeKeyType::Trie(key),
-            1 => ChangeKeyType::Flat(key),
-            _ => panic!("Invalid id"),
-        }
-    }
-}
-
 #[derive(Debug, Default)]
-pub struct ChangeBatch(pub(crate) HashMap<ChangeKeyType, Change>);
+pub struct ChangeBatch(pub(crate) HashMap<TrieKey, Change>);
 
 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) {
             Entry::Occupied(mut entry) => {
                 let e = entry.get_mut();
@@ -68,19 +38,24 @@ impl ChangeBatch {
     }
 
     pub fn serialize<ID: Id>(&self, id: &ID) -> Vec<(Vec<u8>, &[u8])> {
-        let id = id.serialize();
+        let id = id.to_bytes();
         self.0
             .iter()
             .flat_map(|(change_key, change)| {
                 let key_slice = change_key.as_slice();
                 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();
@@ -92,7 +67,7 @@ impl ChangeBatch {
                         id.as_slice(),
                         &[KEY_SEPARATOR],
                         key_slice,
-                        &[change_key.get_id()],
+                        &[change_key.into()],
                         &[NEW_VALUE],
                     ]
                     .concat();
@@ -104,7 +79,7 @@ impl ChangeBatch {
     }
 
     pub fn deserialize<ID: Id>(id: &ID, changes: Vec<(Vec<u8>, Vec<u8>)>) -> 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;
@@ -116,7 +91,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
@@ -1,11 +1,10 @@
 use crate::{
-    bonsai_database::BonsaiPersistentDatabase, error::BonsaiStorageError, id::Id, BonsaiDatabase,
+    bonsai_database::{BonsaiPersistentDatabase, DBError},
+    id::Id,
+    BonsaiDatabase,
 };
 #[cfg(not(feature = "std"))]
-use alloc::{
-    vec::Vec,
-    {collections::BTreeMap, string::ToString},
-};
+use alloc::{collections::BTreeMap, vec::Vec};
 use core::{fmt, fmt::Display};
 #[cfg(not(feature = "std"))]
 use hashbrown::HashMap;
@@ -24,11 +23,7 @@ impl Display for HashMapDbError {
     }
 }
 
-impl From<HashMapDbError> for BonsaiStorageError {
-    fn from(err: HashMapDbError) -> Self {
-        Self::Database(err.to_string())
-    }
-}
+impl DBError for HashMapDbError {}
 
 #[derive(Clone, Default)]
 pub struct HashMapDb<ID: Id> {
@@ -44,7 +39,7 @@ impl<ID: Id> BonsaiDatabase for HashMapDb<ID> {
 
     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() {
@@ -60,14 +55,14 @@ impl<ID: Id> BonsaiDatabase for HashMapDb<ID> {
 
     fn get(
         &self,
-        key: &crate::bonsai_database::KeyType,
+        key: &crate::bonsai_database::DatabaseKey,
     ) -> Result<Option<Vec<u8>>, 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<(Vec<u8>, Vec<u8>)>, Self::DatabaseError> {
         let mut result = Vec::new();
         for (key, value) in self.db.iter() {
@@ -80,7 +75,7 @@ impl<ID: Id> BonsaiDatabase for HashMapDb<ID> {
 
     fn insert(
         &mut self,
-        key: &crate::bonsai_database::KeyType,
+        key: &crate::bonsai_database::DatabaseKey,
         value: &[u8],
         _batch: Option<&mut Self::Batch>,
     ) -> Result<Option<Vec<u8>>, Self::DatabaseError> {
@@ -89,13 +84,16 @@ impl<ID: Id> BonsaiDatabase for HashMapDb<ID> {
 
     fn remove(
         &mut self,
-        key: &crate::bonsai_database::KeyType,
+        key: &crate::bonsai_database::DatabaseKey,
         _batch: Option<&mut Self::Batch>,
     ) -> Result<Option<Vec<u8>>, Self::DatabaseError> {
         Ok(self.db.remove(key.as_slice()))
     }
 
-    fn contains(&self, key: &crate::bonsai_database::KeyType) -> Result<bool, Self::DatabaseError> {
+    fn contains(
+        &self,
+        key: &crate::bonsai_database::DatabaseKey,
+    ) -> Result<bool, Self::DatabaseError> {
         Ok(self.db.contains_key(key.as_slice()))
     }