diff --git a/src/changes.rs b/src/changes.rs index 3018736..5de582b 100644 --- a/src/changes.rs +++ b/src/changes.rs @@ -39,11 +39,9 @@ impl ChangeBatch { } pub fn serialize(&self, id: &ID) -> Vec<(Vec, &[u8])> { - 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 { @@ -52,26 +50,12 @@ impl ChangeBatch { return changes; } } - let key = [ - id.as_slice(), - &[KEY_SEPARATOR], - key_slice, - &[change_key.into()], - &[OLD_VALUE], - ] - .concat(); + let key = key_old_value(id, change_key); changes.push((key, old_value.as_slice())); } if let Some(new_value) = &change.new_value { - let key = [ - id.as_slice(), - &[KEY_SEPARATOR], - key_slice, - &[change_key.into()], - &[NEW_VALUE], - ] - .concat(); + let key = key_new_value(id, change_key); changes.push((key, new_value.as_slice())); } changes @@ -116,6 +100,28 @@ impl ChangeBatch { } } +pub fn key_old_value(id: &ID, key: &TrieKey) -> Vec { + [ + id.to_bytes().as_slice(), + &[KEY_SEPARATOR], + key.as_slice(), + &[key.into()], + &[OLD_VALUE], + ] + .concat() +} + +pub fn key_new_value(id: &ID, key: &TrieKey) -> Vec { + [ + id.to_bytes().as_slice(), + &[KEY_SEPARATOR], + key.as_slice(), + &[key.into()], + &[NEW_VALUE], + ] + .concat() +} + #[cfg_attr(feature = "bench", derive(Clone))] pub struct ChangeStore where diff --git a/src/key_value_db.rs b/src/key_value_db.rs index 2b4d099..c5e5f71 100644 --- a/src/key_value_db.rs +++ b/src/key_value_db.rs @@ -1,4 +1,4 @@ -use crate::{trie::merkle_tree::bytes_to_bitvec, Change as ExternChange}; +use crate::{changes::key_new_value, 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}; @@ -171,6 +171,38 @@ where Ok(self.db.get(&key.into())?) } + pub(crate) fn get_at( + &self, + key: &TrieKey, + id: ID, + ) -> Result>, BonsaiStorageError> { + trace!("Getting from KeyValueDB: {:?} at ID: {:?}", key, id); + + // makes sure given id exists + let Ok(id_position) = self.changes_store.id_queue.binary_search(&id) else { + return Err(BonsaiStorageError::Transaction(format!( + "invalid id {:?}", + id + ))); + }; + + // looking for the first storage insertion with given key + let iter = self + .changes_store + .id_queue + .iter() + .take(id_position + 1) + .rev(); + for id in iter { + let key = key_new_value(id, key); + if let Some(value) = self.db.get(&DatabaseKey::TrieLog(&key))? { + return Ok(Some(value)); + } + } + + Ok(None) + } + pub(crate) fn get_latest_id(&self) -> Option { self.changes_store.id_queue.back().cloned() } diff --git a/src/lib.rs b/src/lib.rs index 0272a72..b51edc8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -259,6 +259,19 @@ where self.tries.get(identifier, key) } + /// Gets a value in a trie at a given commit ID. + /// + /// Note that this is much faster that calling `revert_to1 + /// as it only reverts storage for a single key. + pub fn get_at( + &self, + identifier: &[u8], + key: &BitSlice, + id: ChangeID, + ) -> Result, BonsaiStorageError> { + self.tries.get_at(identifier, key, id) + } + /// Checks if the key exists in the trie. pub fn contains( &self, diff --git a/src/trie/merkle_tree.rs b/src/trie/merkle_tree.rs index 122d506..77452d0 100644 --- a/src/trie/merkle_tree.rs +++ b/src/trie/merkle_tree.rs @@ -130,6 +130,20 @@ impl MerkleTrees, + id: CommitID, + ) -> Result, BonsaiStorageError> { + let tree = self.trees.get(identifier); + if let Some(tree) = tree { + tree.get_at(&self.db, key, id) + } else { + Ok(None) + } + } + pub(crate) fn contains( &self, identifier: &[u8], @@ -961,6 +975,17 @@ impl MerkleTree { .map(|r| r.map(|opt| Felt::decode(&mut opt.as_slice()).unwrap())) } + pub fn get_at( + &self, + db: &KeyValueDB, + key: &BitSlice, + id: ID, + ) -> Result, BonsaiStorageError> { + let key = bitslice_to_bytes(key); + db.get_at(&TrieKey::new(&self.identifier, TrieKeyType::Flat, &key), id) + .map(|r| r.map(|opt| Felt::decode(&mut opt.as_slice()).unwrap())) + } + pub fn contains( &self, db: &KeyValueDB,