Trie log fixes for madara

Cargo.toml

@@ -10,6 +10,7 @@ std = [
   "parity-scale-codec/std",
   "bitvec/std",
   "starknet-types-core/std",
+  "thiserror/std",
   "rayon",
   "hashbrown/rayon",
 ]
@@ -26,6 +27,7 @@ log = "0.4.20"
 rayon = { version = "1.9.0", optional = true }
 smallvec = { version = "1.11.2", features = ["serde"] }
 slotmap = "1.0.7"
+thiserror = { version = "2.0", default-features = false }
 
 parity-scale-codec = { version = "3.0.0", default-features = false, features = [
   "derive",
@@ -34,7 +36,7 @@ serde = { version = "1.0.195", default-features = false, features = [
   "derive",
   "alloc",
 ] }
-starknet-types-core = { version = "0.1.5", default-features = false, features = [
+starknet-types-core = { version = "0.1.7", default-features = false, features = [
   "hash",
   "parity-scale-codec",
   "alloc",

src/bonsai_database.rs

@@ -79,7 +79,9 @@ pub trait BonsaiDatabase: core::fmt::Debug {
 }
 
 pub trait BonsaiPersistentDatabase<ID: Id> {
-    type Transaction: BonsaiDatabase<DatabaseError = Self::DatabaseError>;
+    type Transaction<'a>: BonsaiDatabase<DatabaseError = Self::DatabaseError>
+    where
+        Self: 'a;
     #[cfg(feature = "std")]
     type DatabaseError: Error + DBError;
     #[cfg(not(feature = "std"))]
@@ -89,8 +91,10 @@ pub trait BonsaiPersistentDatabase<ID: Id> {
     fn snapshot(&mut self, id: ID);
 
     /// Create a transaction based on the given snapshot id
-    fn transaction(&self, id: ID) -> Option<Self::Transaction>;
+    fn transaction(&self, id: ID) -> Option<(ID, Self::Transaction<'_>)>;
 
     /// Merge a transaction in the current persistent database
-    fn merge(&mut self, transaction: Self::Transaction) -> Result<(), Self::DatabaseError>;
+    fn merge<'a>(&mut self, transaction: Self::Transaction<'a>) -> Result<(), Self::DatabaseError>
+    where
+        Self: 'a;
 }

src/changes.rs

@@ -1,4 +1,4 @@
-use crate::{hash_map::Entry, id::Id, trie::TrieKey, ByteVec, HashMap, Vec, VecDeque};
+use crate::{hash_map::Entry, id::Id, trie::TrieKey, ByteVec, HashMap, Vec};
 use core::iter;
 use serde::{Deserialize, Serialize};
 
@@ -115,22 +115,13 @@ pub fn key_new_value<ID: Id>(id: &ID, key: &TrieKey) -> ByteVec {
 
 #[cfg_attr(feature = "bench", derive(Clone))]
 #[derive(Debug)]
-pub struct ChangeStore<ID>
-where
-    ID: Id,
-{
-    // Newest are inserted at the back
-    pub id_queue: VecDeque<ID>,
+pub struct ChangeStore {
     pub current_changes: ChangeBatch,
 }
 
-impl<ID> ChangeStore<ID>
-where
-    ID: Id,
-{
+impl ChangeStore {
     pub fn new() -> Self {
         Self {
-            id_queue: VecDeque::new(),
             current_changes: ChangeBatch(HashMap::new()),
         }
     }

src/databases/hashmap_db.rs

@@ -126,16 +126,22 @@ impl<ID: Id> BonsaiDatabase for HashMapDb<ID> {
 
 impl<ID: Id> BonsaiPersistentDatabase<ID> for HashMapDb<ID> {
     type DatabaseError = HashMapDbError;
-    type Transaction = HashMapDb<ID>;
+    type Transaction<'a> = HashMapDb<ID> where ID: 'a;
     fn snapshot(&mut self, id: ID) {
         self.snapshots.insert(id, self.clone());
     }
 
-    fn transaction(&self, id: ID) -> Option<Self::Transaction> {
-        self.snapshots.get(&id).cloned()
+    fn transaction(&self, id: ID) -> Option<(ID, Self::Transaction<'_>)> {
+        self.snapshots
+            .range(..&id)
+            .next()
+            .map(|(id, snapshot)| (*id, snapshot.clone()))
     }
 
-    fn merge(&mut self, transaction: Self::Transaction) -> Result<(), Self::DatabaseError> {
+    fn merge<'a>(&mut self, transaction: Self::Transaction<'a>) -> Result<(), Self::DatabaseError>
+    where
+        ID: 'a,
+    {
         self.trie_db = transaction.trie_db;
         self.flat_db = transaction.flat_db;
         self.trie_log_db = transaction.trie_log_db;

src/databases/rocks_db.rs

@@ -451,7 +451,7 @@ impl<'db, ID> BonsaiPersistentDatabase<ID> for RocksDB<'db, ID>
 where
     ID: Id,
 {
-    type Transaction = RocksDBTransaction<'db>;
+    type Transaction<'a> = RocksDBTransaction<'a> where Self: 'a;
     type DatabaseError = RocksDBError;
 
     fn snapshot(&mut self, id: ID) {
@@ -465,9 +465,9 @@ where
         }
     }
 
-    fn transaction(&self, id: ID) -> Option<Self::Transaction> {
+    fn transaction(&self, id: ID) -> Option<(ID, Self::Transaction<'_>)> {
         trace!("Generating RocksDB transaction");
-        if let Some(snapshot) = self.snapshots.get(&id) {
+        if let Some((id, snapshot)) = self.snapshots.range(..&id).next() {
             let write_opts = WriteOptions::default();
             let mut txn_opts = OptimisticTransactionOptions::default();
             txn_opts.set_snapshot(true);
@@ -494,13 +494,16 @@ where
                 column_families,
                 read_options,
             };
-            Some(boxed_txn)
+            Some((*id, boxed_txn))
         } else {
             None
         }
     }
 
-    fn merge(&mut self, transaction: Self::Transaction) -> Result<(), Self::DatabaseError> {
+    fn merge<'a>(&mut self, transaction: Self::Transaction<'a>) -> Result<(), Self::DatabaseError>
+    where
+        Self: 'a,
+    {
         transaction.txn.commit()?;
         Ok(())
     }

src/error.rs

@@ -1,7 +1,7 @@
 #[cfg(feature = "std")]
 use std::{error::Error, fmt::Display};
 
-use crate::{bonsai_database::DBError, BitVec, String};
+use crate::{bonsai_database::DBError, String};
 
 /// All errors that can be returned by BonsaiStorage.
 #[derive(Debug)]
@@ -21,8 +21,6 @@ where
     Database(DatabaseError),
     /// Error when decoding a node
     NodeDecodeError(parity_scale_codec::Error),
-    /// Error when creating a storage proof.
-    CreateProofKeyNotInTree { key: BitVec },
     /// Malformated trie key.
     KeyLength { expected: usize, got: usize },
 }
@@ -56,9 +54,6 @@ where
             BonsaiStorageError::Merge(e) => write!(f, "Merge error: {}", e),
             BonsaiStorageError::Database(e) => write!(f, "Database error: {}", e),
             BonsaiStorageError::NodeDecodeError(e) => write!(f, "Node decode error: {}", e),
-            BonsaiStorageError::CreateProofKeyNotInTree { key } => {
-                write!(f, "Key not in tree: {key:b}")
-            }
             BonsaiStorageError::KeyLength { expected, got } => {
                 write!(f, "Malformated key length: expected {expected}, got {got}")
             }

src/id.rs

@@ -4,6 +4,8 @@ use core::{fmt::Debug, hash};
 /// Trait to be implemented on any type that can be used as an ID.
 pub trait Id: hash::Hash + PartialEq + Eq + PartialOrd + Ord + Debug + Copy + Default {
     fn to_bytes(&self) -> ByteVec;
+    fn as_u64(self) -> u64;
+    fn from_u64(v: u64) -> Self;
 }
 
 /// A basic ID type that can be used for testing.
@@ -20,6 +22,12 @@ impl Id for BasicId {
     fn to_bytes(&self) -> ByteVec {
         ByteVec::from(&self.0.to_be_bytes() as &[_])
     }
+    fn as_u64(self) -> u64 {
+        self.0
+    }
+    fn from_u64(v: u64) -> Self {
+        Self(v)
+    }
 }
 
 /// A builder for basic IDs.