Fix clippy and batch db writes

madara-alliance · Mar 14, 2024 · 2754970 · 2754970
1 parent 5b2a080
commit 2754970
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 21 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -320,7 +320,6 @@ where
     }
 
     /// Get all changes applied at a certain commit ID.
-    /// TODO: Split changes by identifier
     #[allow(clippy::type_complexity)]
     pub fn get_changes(
         &self,

diff --git a/src/tests/trie_log.rs b/src/tests/trie_log.rs
@@ -23,7 +23,9 @@ fn basics() {
     );
     let id1 = id_builder.new_id();
     let bitvec = BitVec::from_vec(pair1.0.clone());
-    bonsai_storage.insert(&identifier, &bitvec, pair1.1).unwrap();
+    bonsai_storage
+        .insert(&identifier, &bitvec, pair1.1)
+        .unwrap();
     bonsai_storage.commit(id1).unwrap();
     let root_hash1 = bonsai_storage.root_hash(&identifier).unwrap();
 
@@ -33,7 +35,9 @@ fn basics() {
         &Felt::from_hex("0x66342762FDD54D3c195fec3ce2568b62052e").unwrap(),
     );
     let bitvec = BitVec::from_vec(pair2.0.clone());
-    bonsai_storage.insert(&identifier, &bitvec, pair2.1).unwrap();
+    bonsai_storage
+        .insert(&identifier, &bitvec, pair2.1)
+        .unwrap();
     bonsai_storage.commit(id2).unwrap();
     let root_hash2 = bonsai_storage.root_hash(&identifier).unwrap();
 
@@ -68,7 +72,9 @@ fn unrecorded_revert() {
     );
     let id1 = id_builder.new_id();
     let bitvec = BitVec::from_vec(pair1.0.clone());
-    bonsai_storage.insert(&identifier, &bitvec, &pair1.1).unwrap();
+    bonsai_storage
+        .insert(&identifier, &bitvec, &pair1.1)
+        .unwrap();
     bonsai_storage.commit(id1).unwrap();
 
     let uncommited_id = id_builder.new_id();
@@ -88,7 +94,9 @@ fn in_place_revert() {
     let pair1 = (vec![1, 2, 3], &BonsaiTrieHash::default());
     let id1 = id_builder.new_id();
     let bitvec = BitVec::from_vec(pair1.0.clone());
-    bonsai_storage.insert(&identifier, &bitvec, pair1.1).unwrap();
+    bonsai_storage
+        .insert(&identifier, &bitvec, pair1.1)
+        .unwrap();
     bonsai_storage.commit(id1).unwrap();
     let root_hash1 = bonsai_storage.root_hash(&identifier).unwrap();
 
@@ -112,7 +120,9 @@ fn truncated_revert() {
     );
     let id1 = id_builder.new_id();
     let bitvec = BitVec::from_vec(pair1.0.clone());
-    bonsai_storage.insert(&identifier, &bitvec, pair1.1).unwrap();
+    bonsai_storage
+        .insert(&identifier, &bitvec, pair1.1)
+        .unwrap();
     bonsai_storage.commit(id1).unwrap();
     let root_hash1 = bonsai_storage.root_hash(&identifier).unwrap();
 
@@ -122,7 +132,9 @@ fn truncated_revert() {
         &Felt::from_hex("0x66342762FDD54D3c195fec3ce2568b62052e").unwrap(),
     );
     let bitvec = BitVec::from_vec(pair2.0.clone());
-    bonsai_storage.insert(&identifier, &bitvec, pair2.1).unwrap();
+    bonsai_storage
+        .insert(&identifier, &bitvec, pair2.1)
+        .unwrap();
     bonsai_storage.commit(id2).unwrap();
 
     bonsai_storage.revert_to(id1).unwrap();
@@ -148,7 +160,9 @@ fn double_revert() {
     );
     let id1 = id_builder.new_id();
     let bitvec = BitVec::from_vec(pair1.0.clone());
-    bonsai_storage.insert(&identifier, &bitvec, pair1.1).unwrap();
+    bonsai_storage
+        .insert(&identifier, &bitvec, pair1.1)
+        .unwrap();
     bonsai_storage.commit(id1).unwrap();
     let root_hash1 = bonsai_storage.root_hash(&identifier).unwrap();
 
@@ -158,7 +172,9 @@ fn double_revert() {
         &Felt::from_hex("0x66342762FDD54D3c195fec3ce2568b62052e").unwrap(),
     );
     let bitvec = BitVec::from_vec(pair2.0.clone());
-    bonsai_storage.insert(&identifier, &bitvec, pair2.1).unwrap();
+    bonsai_storage
+        .insert(&identifier, &bitvec, pair2.1)
+        .unwrap();
     bonsai_storage.commit(id2).unwrap();
 
     bonsai_storage.revert_to(id1).unwrap();
@@ -186,12 +202,16 @@ fn remove_and_reinsert() {
     );
     let id1 = id_builder.new_id();
     let bitvec = BitVec::from_vec(pair1.0.clone());
-    bonsai_storage.insert(&identifier, &bitvec, &pair1.1).unwrap();
+    bonsai_storage
+        .insert(&identifier, &bitvec, &pair1.1)
+        .unwrap();
     bonsai_storage.remove(&identifier, &bitvec).unwrap();
     bonsai_storage.commit(id1).unwrap();
     let root_hash1 = bonsai_storage.root_hash(&identifier).unwrap();
     let id2 = id_builder.new_id();
-    bonsai_storage.insert(&identifier, &bitvec, &pair1.1).unwrap();
+    bonsai_storage
+        .insert(&identifier, &bitvec, &pair1.1)
+        .unwrap();
     bonsai_storage.commit(id2).unwrap();
 
     bonsai_storage.revert_to(id1).unwrap();

diff --git a/src/trie/merkle_tree.rs b/src/trie/merkle_tree.rs
@@ -163,11 +163,12 @@ impl<H: StarkHash, DB: BonsaiDatabase, CommitID: Id> MerkleTrees<H, DB, CommitID
         }
     }
 
-    //TODO: Use a common batch for all
     pub(crate) fn commit(&mut self) -> Result<(), BonsaiStorageError<DB::DatabaseError>> {
+        let mut batch = self.db.create_batch();
         for tree in self.trees.values_mut() {
-            tree.commit(&mut self.db)?;
+            tree.commit(&mut self.db, &mut batch)?;
         }
+        self.db.write_batch(batch)?;
         Ok(())
     }
 
@@ -222,8 +223,10 @@ enum InsertOrRemove<T> {
 }
 
 /// Note to developers : Must be used to build a key for the database.
-pub fn build_db_key(identifier: &[u8], key: &[u8]) -> Vec<u8> {
-    let mut db_key = identifier.to_owned();
+// Allow because needed for no-std
+#[allow(clippy::ptr_arg)]
+pub fn build_db_key(identifier: &Vec<u8>, key: &[u8]) -> Vec<u8> {
+    let mut db_key = identifier.clone();
     db_key.extend_from_slice(key);
     db_key
 }
@@ -292,31 +295,30 @@ impl<H: StarkHash> MerkleTree<H> {
     pub fn commit<DB: BonsaiDatabase, ID: Id>(
         &mut self,
         db: &mut KeyValueDB<DB, ID>,
+        batch: &mut DB::Batch,
     ) -> Result<Felt, BonsaiStorageError<DB::DatabaseError>> {
-        let mut batch = db.create_batch();
         for node_key in mem::take(&mut self.death_row) {
-            db.remove(&node_key, Some(&mut batch))?;
+            db.remove(&node_key, Some(batch))?;
         }
         let root_hash =
-            self.commit_subtree(db, self.root_handle, Path(BitVec::new()), &mut batch)?;
+            self.commit_subtree(db, self.root_handle, Path(BitVec::new()), batch)?;
         for (key, value) in mem::take(&mut self.cache_leaf_modified) {
             match value {
                 InsertOrRemove::Insert(value) => {
                     db.insert(
                         &TrieKey::Flat(build_db_key(&self.identifier, &key)),
                         &value.encode(),
-                        Some(&mut batch),
+                        Some(batch),
                     )?;
                 }
                 InsertOrRemove::Remove => {
                     db.remove(
                         &TrieKey::Flat(build_db_key(&self.identifier, &key)),
-                        Some(&mut batch),
+                        Some(batch),
                     )?;
                 }
             }
         }
-        db.write_batch(batch)?;
         self.latest_node_id.reset();
         self.root_hash = root_hash;
         self.root_handle = NodeHandle::Hash(root_hash);