diff --git a/src/key_value_db.rs b/src/key_value_db.rs index 7e786f4..84ce41f 100644 --- a/src/key_value_db.rs +++ b/src/key_value_db.rs @@ -1,6 +1,11 @@ +use crate::{trie::merkle_tree::bytes_to_bitvec, Change as ExternChange}; #[cfg(not(feature = "std"))] use alloc::{collections::BTreeSet, format, string::ToString, vec::Vec}; +use bitvec::{order::Msb0, vec::BitVec}; +use hashbrown::HashMap; use log::trace; +use parity_scale_codec::Decode; +use starknet_types_core::felt::Felt; #[cfg(feature = "std")] use std::collections::BTreeSet; @@ -86,6 +91,39 @@ where } } + #[allow(clippy::type_complexity)] + pub(crate) fn get_changes( + &self, + id: ID, + ) -> Result, ExternChange>, BonsaiStorageError> + { + if self.changes_store.id_queue.contains(&id) { + let mut leaf_changes = HashMap::new(); + let changes = ChangeBatch::deserialize( + &id, + self.db + .get_by_prefix(&DatabaseKey::TrieLog(&id.to_bytes()))?, + ); + for (k, v) in changes.0 { + if let TrieKey::Flat(k) = k { + leaf_changes.insert( + bytes_to_bitvec(&k), + ExternChange { + // SAFETY: We are sure that the values are valid Felt because they can be saved only by our crate + old_value: v.old_value.map(|x| Felt::decode(&mut x.as_ref()).unwrap()), + new_value: v.new_value.map(|x| Felt::decode(&mut x.as_ref()).unwrap()), + }, + ); + } + } + Ok(leaf_changes) + } else { + Err(BonsaiStorageError::GoTo( + "ID asked isn't in our ID records".to_string(), + )) + } + } + 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); diff --git a/src/lib.rs b/src/lib.rs index 8932bfd..f157172 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -89,9 +89,10 @@ extern crate alloc; use crate::trie::merkle_tree::MerkleTree; #[cfg(not(feature = "std"))] use alloc::{format, vec::Vec}; -use bitvec::{order::Msb0, slice::BitSlice}; +use bitvec::{order::Msb0, slice::BitSlice, vec::BitVec}; use bonsai_database::{BonsaiPersistentDatabase, DatabaseKey}; use changes::ChangeBatch; +use hashbrown::HashMap; use key_value_db::KeyValueDB; use starknet_types_core::{ felt::Felt, @@ -144,6 +145,16 @@ impl Default for BonsaiStorageConfig { } } +/// Structure used to represent a change in the trie for a specific value. +/// It contains the old value and the new value. +/// If the `old_value` is None, it means that the key was not present in the trie before the change. +/// If the `new_value` is None, it means that the key was removed from the trie. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Change { + pub old_value: Option, + pub new_value: Option, +} + /// Structure that hold the trie and all the necessary information to work with it. /// /// This structure is the main entry point to work with this crate. @@ -300,6 +311,15 @@ where Ok(()) } + /// Get all changes applied at a certain commit ID. + #[allow(clippy::type_complexity)] + pub fn get_changes( + &self, + id: ChangeID, + ) -> Result, Change>, BonsaiStorageError> { + self.trie.db_ref().get_changes(id) + } + #[cfg(test)] pub fn dump_database(&self) { self.trie.db_ref().db.dump_database(); diff --git a/src/tests/simple.rs b/src/tests/simple.rs index 9ae3eb6..e8c238a 100644 --- a/src/tests/simple.rs +++ b/src/tests/simple.rs @@ -2,7 +2,7 @@ use crate::{ databases::{create_rocks_db, RocksDB, RocksDBConfig}, id::BasicIdBuilder, - BonsaiStorage, BonsaiStorageConfig, + BonsaiStorage, BonsaiStorageConfig, Change, }; use bitvec::vec::BitVec; use starknet_types_core::{felt::Felt, hash::Pedersen}; @@ -52,3 +52,44 @@ fn basics() { None ); } + +#[test] +fn get_changes() { + let tempdir = tempfile::tempdir().unwrap(); + 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(); + 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()); + bonsai_storage.insert(&bitvec, &pair1.1).unwrap(); + bonsai_storage.commit(id_builder.new_id()).unwrap(); + let pair2 = (vec![1, 2, 2], Felt::from_hex("0x01").unwrap()); + let bitvec = BitVec::from_vec(pair2.0.clone()); + bonsai_storage.insert(&bitvec, &pair2.1).unwrap(); + let pair1_edited_1 = (vec![1, 2, 1], Felt::from_hex("0x02").unwrap()); + let bitvec = BitVec::from_vec(pair1_edited_1.0.clone()); + bonsai_storage.insert(&bitvec, &pair1_edited_1.1).unwrap(); + let pair1_edited_2 = (vec![1, 2, 1], Felt::from_hex("0x03").unwrap()); + let bitvec = BitVec::from_vec(pair1_edited_2.0.clone()); + bonsai_storage.insert(&bitvec, &pair1_edited_2.1).unwrap(); + let id = id_builder.new_id(); + bonsai_storage.commit(id).unwrap(); + let changes = bonsai_storage.get_changes(id).unwrap(); + assert_eq!(changes.len(), 2); + assert_eq!( + changes.get(&BitVec::from_vec(pair1.0)).unwrap(), + &Change { + old_value: Some(pair1.1), + new_value: Some(pair1_edited_2.1), + } + ); + assert_eq!( + changes.get(&BitVec::from_vec(pair2.0)).unwrap(), + &Change { + old_value: None, + new_value: Some(pair2.1), + } + ); +} diff --git a/src/trie/merkle_tree.rs b/src/trie/merkle_tree.rs index 6219cd7..0cd861d 100644 --- a/src/trie/merkle_tree.rs +++ b/src/trie/merkle_tree.rs @@ -1091,10 +1091,14 @@ impl MerkleTree { } } -fn bitslice_to_bytes(bitslice: &BitSlice) -> Vec { +pub(crate) fn bitslice_to_bytes(bitslice: &BitSlice) -> Vec { [&[bitslice.len() as u8], bitslice.to_bitvec().as_raw_slice()].concat() } +pub(crate) fn bytes_to_bitvec(bytes: &[u8]) -> BitVec { + BitSlice::from_slice(&bytes[1..]).to_bitvec() +} + #[cfg(test)] #[cfg(all(test, feature = "std"))] mod tests {