feat: tree max_height

madara-alliance · Sep 30, 2024 · bcab08a · bcab08a
1 parent b0ff3da
commit bcab08a
Show file tree

Hide file tree

Showing 13 changed files with 97 additions and 57 deletions.
diff --git a/src/error.rs b/src/error.rs
@@ -23,6 +23,8 @@ where
     NodeDecodeError(parity_scale_codec::Error),
     /// Error when creating a storage proof.
     CreateProof(String),
+    /// Malformated trie key.
+    KeyLength { expected: usize, got: usize },
 }
 
 impl<DatabaseError: DBError> core::convert::From<DatabaseError>
@@ -55,6 +57,9 @@ where
             BonsaiStorageError::Database(e) => write!(f, "Database error: {}", e),
             BonsaiStorageError::NodeDecodeError(e) => write!(f, "Node decode error: {}", e),
             BonsaiStorageError::CreateProof(e) => write!(f, "Proof creation error: {}", e),
+            BonsaiStorageError::KeyLength { expected, got } => {
+                write!(f, "Malformated key length: expected {expected}, got {got}")
+            }
         }
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -237,21 +237,23 @@ where
     pub fn new(
         db: DB,
         config: BonsaiStorageConfig,
+        max_height: u8,
     ) -> Result<Self, BonsaiStorageError<DB::DatabaseError>> {
         let key_value_db = KeyValueDB::new(db, config.into(), None);
         Ok(Self {
-            tries: MerkleTrees::new(key_value_db),
+            tries: MerkleTrees::new(key_value_db, max_height),
         })
     }
 
     pub fn new_from_transactional_state(
         db: DB,
         config: BonsaiStorageConfig,
+        max_height: u8,
         created_at: ChangeID,
         _identifiers: impl IntoIterator<Item = impl Deref<Target = [u8]>>,
     ) -> Result<Self, BonsaiStorageError<DB::DatabaseError>> {
         let key_value_db = KeyValueDB::new(db, config.into(), Some(created_at));
-        let tries = MerkleTrees::<H, DB, ChangeID>::new(key_value_db);
+        let tries = MerkleTrees::<H, DB, ChangeID>::new(key_value_db, max_height);
         Ok(Self { tries })
     }
 
@@ -516,6 +518,7 @@ where
             Ok(Some(BonsaiStorage::new_from_transactional_state(
                 transaction,
                 config,
+                self.tries.max_height,
                 change_id,
                 self.tries.get_identifiers(),
             )?))

diff --git a/src/tests/madara_comparison.rs b/src/tests/madara_comparison.rs
@@ -15,7 +15,7 @@ fn trie_height_251() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     for i in 0..251 {
         let mut key: BitVec = bits![u8, Msb0; 0; 251].to_bitvec();
         key.set(i, true);

diff --git a/src/tests/merge.rs b/src/tests/merge.rs
@@ -60,8 +60,9 @@ fn init_test(
     let identifier = vec![];
 
     let config = BonsaiStorageConfig::default();
-    let mut bonsai_storage = BonsaiStorage::new(RocksDB::new(db, RocksDBConfig::default()), config)
-        .expect("Failed to create BonsaiStorage");
+    let mut bonsai_storage =
+        BonsaiStorage::new(RocksDB::new(db, RocksDBConfig::default()), config, 251)
+            .expect("Failed to create BonsaiStorage");
 
     let mut id_builder = BasicIdBuilder::new();
 

diff --git a/src/tests/merkle_tree.rs b/src/tests/merkle_tree.rs
@@ -17,7 +17,8 @@ fn test_key_retrieval() {
     let rocksdb = create_rocks_db(tempdir.path()).unwrap();
     let db = RocksDB::new(&rocksdb, RocksDBConfig::default());
     let mut bonsai =
-        BonsaiStorage::<BasicId, _, Pedersen>::new(db, BonsaiStorageConfig::default()).unwrap();
+        BonsaiStorage::<BasicId, _, Pedersen>::new(db, BonsaiStorageConfig::default(), 251)
+            .unwrap();
 
     let block_0 = vec![
         (

diff --git a/src/tests/proptest.rs b/src/tests/proptest.rs
@@ -73,7 +73,7 @@ impl MerkleTreeInsertProblem {
         let mut ckv = HashMap::new();
 
         // apply steps
-        let mut tree = MerkleTree::<Pedersen>::new(smallvec![]);
+        let mut tree = MerkleTree::<Pedersen>::new(smallvec![], 251);
         for step in &self.0 {
             match step {
                 Step::Insert(k, v) => {

diff --git a/src/tests/simple.rs b/src/tests/simple.rs
@@ -14,7 +14,7 @@ fn basics() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
     let pair1 = (
         vec![1, 2, 1],
@@ -69,7 +69,7 @@ fn root_hash_similar_rocks_db() {
         let db = create_rocks_db(tempdir.path()).unwrap();
         let config = BonsaiStorageConfig::default();
         let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
-            BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+            BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
         let mut id_builder = BasicIdBuilder::new();
         let pair1 = (
             vec![1, 2, 1],
@@ -116,7 +116,7 @@ fn root_hash_similar_rocks_db() {
         let db = create_rocks_db(tempdir.path()).unwrap();
         let config = BonsaiStorageConfig::default();
         let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
-            BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+            BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
         let mut id_builder = BasicIdBuilder::new();
         let pair1 = (
             vec![1, 2, 3],
@@ -156,13 +156,13 @@ fn starknet_specific() {
     let db1 = create_rocks_db(tempdir1.path()).unwrap();
     let config1 = BonsaiStorageConfig::default();
     let mut bonsai_storage1: BonsaiStorage<_, _, Pedersen> =
-        BonsaiStorage::new(RocksDB::new(&db1, RocksDBConfig::default()), config1).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db1, RocksDBConfig::default()), config1, 251).unwrap();
 
     let tempdir2 = tempfile::tempdir().unwrap();
     let db2 = create_rocks_db(tempdir2.path()).unwrap();
     let config2 = BonsaiStorageConfig::default();
     let mut bonsai_storage2: BonsaiStorage<_, _, Pedersen> =
-        BonsaiStorage::new(RocksDB::new(&db2, RocksDBConfig::default()), config2).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db2, RocksDBConfig::default()), config2, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
 
     let contract_states = vec![
@@ -237,7 +237,7 @@ fn root_hash_similar_hashmap_db() {
         let db = HashMapDb::<BasicId>::default();
         let config = BonsaiStorageConfig::default();
         let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
-            BonsaiStorage::new(db, config).unwrap();
+            BonsaiStorage::new(db, config, 251).unwrap();
         let mut id_builder = BasicIdBuilder::new();
         let pair1 = (
             vec![1, 2, 1],
@@ -283,7 +283,7 @@ fn root_hash_similar_hashmap_db() {
         let db = HashMapDb::<BasicId>::default();
         let config = BonsaiStorageConfig::default();
         let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
-            BonsaiStorage::new(db, config).unwrap();
+            BonsaiStorage::new(db, config, 251).unwrap();
         let mut id_builder = BasicIdBuilder::new();
         let pair1 = (
             vec![1, 2, 3],
@@ -322,7 +322,7 @@ fn double_insert() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
     let contract_states = vec![
         ContractState {
@@ -385,7 +385,7 @@ fn double_identifier() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
     let contract_states = vec![
         ContractState {
@@ -450,7 +450,7 @@ fn get_changes() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
     let pair1 = (vec![1, 2, 1], Felt::from_hex("0x01").unwrap());
     let bitvec = BitVec::from_vec(pair1.0.clone());
@@ -501,7 +501,7 @@ fn keyer(felt: Felt) -> BitVec {
 fn test_insert_zero() {
     let config = BonsaiStorageConfig::default();
     let bonsai_db = HashMapDb::<BasicId>::default();
-    let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config)
+    let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config, 251)
         .expect("Failed to create bonsai storage");
     let identifier =
         "0x056e4fed965fccd7fb01fcadd827470338f35ced62275328929d0d725b5707ba".as_bytes();
@@ -645,7 +645,7 @@ fn test_block_7_starknet() {
     let _ = env_logger::builder().is_test(true).try_init();
     let config = BonsaiStorageConfig::default();
     let bonsai_db = HashMapDb::<BasicId>::default();
-    let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config)
+    let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config, 251)
         .expect("Failed to create bonsai storage");
     let identifier =
         "0x056e4fed965fccd7fb01fcadd827470338f35ced62275328929d0d725b5707ba".as_bytes();
@@ -857,7 +857,7 @@ fn test_block_7_starknet() {
 fn test_block_7_starknet_2() {
     let config = BonsaiStorageConfig::default();
     let bonsai_db = HashMapDb::<BasicId>::default();
-    let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config)
+    let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config, 251)
         .expect("Failed to create bonsai storage");
     let identifier = "0x421203c58e1b4a6c3675be26cfaa18d2b6b42695ca206be1f08ce29f7f1bc7c".as_bytes();
 
@@ -988,7 +988,7 @@ fn test_block_7_starknet_2() {
 fn test_block_9() {
     let config = BonsaiStorageConfig::default();
     let bonsai_db = HashMapDb::<BasicId>::default();
-    let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config)
+    let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config, 251)
         .expect("Failed to create bonsai storage");
     let identifier =
         "0x06F3C934BA4EC49245CB9A42FC715E4D589AA502AF69BE13916127A538D525CE".as_bytes();

diff --git a/src/tests/transactional_state.rs b/src/tests/transactional_state.rs
@@ -14,7 +14,7 @@ fn basics() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
 
     let pair1 = (
@@ -58,8 +58,12 @@ fn test_thread() {
     let tempdir = tempfile::tempdir().unwrap();
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
-    let mut bonsai_storage =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config.clone()).unwrap();
+    let mut bonsai_storage = BonsaiStorage::new(
+        RocksDB::new(&db, RocksDBConfig::default()),
+        config.clone(),
+        251,
+    )
+    .unwrap();
     let mut id_builder = BasicIdBuilder::new();
 
     let pair1 = (
@@ -119,7 +123,7 @@ fn remove() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
 
     let pair1 = (
@@ -164,7 +168,7 @@ fn merge() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
 
     let pair1 = (
@@ -207,7 +211,7 @@ fn merge_with_uncommitted_insert() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
 
     let pair1 = (
@@ -269,7 +273,7 @@ fn merge_with_uncommitted_remove() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
 
     let pair1 = (
@@ -327,7 +331,7 @@ fn transactional_state_after_uncommitted() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
 
     let pair1 = (
@@ -369,7 +373,7 @@ fn merge_override() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
 
     let pair1 = (
@@ -412,7 +416,7 @@ fn merge_remove() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
 
     let pair1 = (
@@ -451,7 +455,7 @@ fn merge_txn_revert() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
 
     let pair1 = (
@@ -511,7 +515,7 @@ fn merge_invalid() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
 
     let pair1 = (
@@ -558,7 +562,7 @@ fn many_snapshots() {
         ..Default::default()
     };
     let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
 
     let pair1 = (

diff --git a/src/tests/trie_log.rs b/src/tests/trie_log.rs
@@ -13,7 +13,7 @@ fn basics() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
 
     let pair1 = (
@@ -62,7 +62,7 @@ fn unrecorded_revert() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
 
     let pair1 = (
@@ -87,7 +87,7 @@ fn in_place_revert() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
 
     let pair1 = (vec![1, 2, 3], &BonsaiTrieHash::default());
@@ -110,7 +110,7 @@ fn truncated_revert() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
 
     let pair1 = (
@@ -150,7 +150,7 @@ fn double_revert() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
 
     let pair1 = (
@@ -192,7 +192,7 @@ fn remove_and_reinsert() {
     let db = create_rocks_db(tempdir.path()).unwrap();
     let config = BonsaiStorageConfig::default();
     let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
-        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
+        BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap();
     let mut id_builder = BasicIdBuilder::new();
 
     let pair1 = (

diff --git a/src/trie/iterator.rs b/src/trie/iterator.rs
@@ -283,6 +283,7 @@ mod tests {
         let mut bonsai_storage: BonsaiStorage<BasicId, _, Pedersen> = BonsaiStorage::new(
             RocksDB::<BasicId>::new(&db, RocksDBConfig::default()),
             BonsaiStorageConfig::default(),
+            8,
         )
         .unwrap();