Merge pull request #29 from Trantorian1/storage_access

Added new storage acessors

src/changes.rs

@@ -39,11 +39,9 @@ impl ChangeBatch {
     }
 
     pub fn serialize<ID: Id>(&self, id: &ID) -> Vec<(Vec<u8>, &[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>(id: &ID, key: &TrieKey) -> Vec<u8> {
+    [
+        id.to_bytes().as_slice(),
+        &[KEY_SEPARATOR],
+        key.as_slice(),
+        &[key.into()],
+        &[OLD_VALUE],
+    ]
+    .concat()
+}
+
+pub fn key_new_value<ID: Id>(id: &ID, key: &TrieKey) -> Vec<u8> {
+    [
+        id.to_bytes().as_slice(),
+        &[KEY_SEPARATOR],
+        key.as_slice(),
+        &[key.into()],
+        &[NEW_VALUE],
+    ]
+    .concat()
+}
+
 #[cfg_attr(feature = "bench", derive(Clone))]
 pub struct ChangeStore<ID>
 where

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<Option<Vec<u8>>, BonsaiStorageError<DB::DatabaseError>> {
+        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<ID> {
         self.changes_store.id_queue.back().cloned()
     }

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<u8, Msb0>,
+        id: ChangeID,
+    ) -> Result<Option<Felt>, BonsaiStorageError<DB::DatabaseError>> {
+        self.tries.get_at(identifier, key, id)
+    }
+
     /// Checks if the key exists in the trie.
     pub fn contains(
         &self,

src/trie/merkle_tree.rs

@@ -130,6 +130,20 @@ impl<H: StarkHash + Send + Sync, DB: BonsaiDatabase, CommitID: Id> MerkleTrees<H
         }
     }
 
+    pub(crate) fn get_at(
+        &self,
+        identifier: &[u8],
+        key: &BitSlice<u8, Msb0>,
+        id: CommitID,
+    ) -> Result<Option<Felt>, BonsaiStorageError<DB::DatabaseError>> {
+        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<H: StarkHash + Send + Sync> MerkleTree<H> {
             .map(|r| r.map(|opt| Felt::decode(&mut opt.as_slice()).unwrap()))
     }
 
+    pub fn get_at<DB: BonsaiDatabase, ID: Id>(
+        &self,
+        db: &KeyValueDB<DB, ID>,
+        key: &BitSlice<u8, Msb0>,
+        id: ID,
+    ) -> Result<Option<Felt>, BonsaiStorageError<DB::DatabaseError>> {
+        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<DB: BonsaiDatabase, ID: Id>(
         &self,
         db: &KeyValueDB<DB, ID>,