Skip to content

Commit 13e8f7b

Browse files
heemankvheemankv
authored
Merge pull request #48 from madara-alliance/feat/root-hash-staged
feat: add root_hash_staged() for non-destructive root computation
2 parents 43829ef + 5068696 commit 13e8f7b

5 files changed

Lines changed: 267 additions & 0 deletions

File tree

‎src/lib.rs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -409,6 +409,15 @@ where
409409
self.tries.root_hash(identifier)
410410
}
411411

412+
/// Compute trie root hash from staged (uncommitted) changes without persisting.
413+
/// Falls back to the committed root when there are no pending modifications.
414+
pub fn root_hash_staged(
415+
&self,
416+
identifier: &[u8],
417+
) -> Result<BonsaiTrieHash, BonsaiStorageError<DB::DatabaseError>> {
418+
self.tries.root_hash_staged(identifier)
419+
}
420+
412421
/// This function must be used with transactional state only.
413422
/// Similar to `commit` but without optimizations.
414423
pub fn transactional_commit(

‎src/tests/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,5 +3,6 @@ mod madara_comparison;
33
mod merkle_tree;
44
mod proptest;
55
mod simple;
6+
mod staged;
67
// mod transactional_state;
78
mod trie_log;

‎src/tests/staged.rs‎

Lines changed: 227 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,227 @@
1+
#![cfg(all(feature = "std", feature = "rocksdb"))]
2+
use crate::{
3+
databases::{create_rocks_db, RocksDB, RocksDBConfig},
4+
id::{BasicId, BasicIdBuilder},
5+
BitVec, BonsaiStorage, BonsaiStorageConfig,
6+
};
7+
use starknet_types_core::{felt::Felt, hash::Pedersen};
8+
9+
#[test]
10+
fn staged_root_matches_committed_root() {
11+
let identifier = vec![];
12+
let tempdir = tempfile::tempdir().unwrap();
13+
let db = create_rocks_db(tempdir.path()).unwrap();
14+
let config = BonsaiStorageConfig::default();
15+
let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
16+
BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24);
17+
let mut id_builder = BasicIdBuilder::new();
18+
19+
let pair1 = (
20+
vec![1, 2, 1],
21+
Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(),
22+
);
23+
let pair2 = (
24+
vec![1, 2, 2],
25+
Felt::from_hex("0x66342762FD54D033c195fec3ce2568b62052e").unwrap(),
26+
);
27+
28+
bonsai_storage
29+
.insert(&identifier, &BitVec::from_vec(pair1.0.clone()), &pair1.1)
30+
.unwrap();
31+
bonsai_storage
32+
.insert(&identifier, &BitVec::from_vec(pair2.0.clone()), &pair2.1)
33+
.unwrap();
34+
35+
// Staged root should be computable before commit
36+
let staged_root = bonsai_storage.root_hash_staged(&identifier).unwrap();
37+
assert_ne!(staged_root, Felt::ZERO);
38+
39+
// Commit and get the committed root
40+
bonsai_storage.commit(id_builder.new_id()).unwrap();
41+
let committed_root = bonsai_storage.root_hash(&identifier).unwrap();
42+
43+
assert_eq!(staged_root, committed_root);
44+
}
45+
46+
#[test]
47+
fn staged_root_after_incremental_inserts() {
48+
let identifier = vec![];
49+
let tempdir = tempfile::tempdir().unwrap();
50+
let db = create_rocks_db(tempdir.path()).unwrap();
51+
let config = BonsaiStorageConfig::default();
52+
let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
53+
BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24);
54+
let mut id_builder = BasicIdBuilder::new();
55+
56+
// Commit an initial state
57+
let pair1 = (
58+
vec![1, 2, 1],
59+
Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(),
60+
);
61+
bonsai_storage
62+
.insert(&identifier, &BitVec::from_vec(pair1.0.clone()), &pair1.1)
63+
.unwrap();
64+
bonsai_storage.commit(id_builder.new_id()).unwrap();
65+
66+
// Insert more data without committing
67+
let pair2 = (
68+
vec![1, 2, 2],
69+
Felt::from_hex("0x66342762FD54D033c195fec3ce2568b62052e").unwrap(),
70+
);
71+
bonsai_storage
72+
.insert(&identifier, &BitVec::from_vec(pair2.0.clone()), &pair2.1)
73+
.unwrap();
74+
75+
// Staged root reflects both pair1 (committed) and pair2 (staged)
76+
let staged_root = bonsai_storage.root_hash_staged(&identifier).unwrap();
77+
assert_ne!(staged_root, Felt::ZERO);
78+
79+
// Commit and verify match
80+
bonsai_storage.commit(id_builder.new_id()).unwrap();
81+
let committed_root = bonsai_storage.root_hash(&identifier).unwrap();
82+
83+
assert_eq!(staged_root, committed_root);
84+
}
85+
86+
#[test]
87+
fn staged_root_does_not_mutate_tree() {
88+
let identifier = vec![];
89+
let tempdir = tempfile::tempdir().unwrap();
90+
let db = create_rocks_db(tempdir.path()).unwrap();
91+
let config = BonsaiStorageConfig::default();
92+
let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
93+
BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24);
94+
let mut id_builder = BasicIdBuilder::new();
95+
96+
let pair1 = (
97+
vec![1, 2, 1],
98+
Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(),
99+
);
100+
bonsai_storage
101+
.insert(&identifier, &BitVec::from_vec(pair1.0.clone()), &pair1.1)
102+
.unwrap();
103+
104+
// Call staged root multiple times — should be idempotent
105+
let root1 = bonsai_storage.root_hash_staged(&identifier).unwrap();
106+
let root2 = bonsai_storage.root_hash_staged(&identifier).unwrap();
107+
assert_eq!(root1, root2);
108+
109+
// Commit should still work after staged computations
110+
bonsai_storage.commit(id_builder.new_id()).unwrap();
111+
let committed_root = bonsai_storage.root_hash(&identifier).unwrap();
112+
assert_eq!(root1, committed_root);
113+
}
114+
115+
#[test]
116+
fn staged_root_on_empty_trie() {
117+
let identifier = vec![];
118+
let tempdir = tempfile::tempdir().unwrap();
119+
let db = create_rocks_db(tempdir.path()).unwrap();
120+
let config = BonsaiStorageConfig::default();
121+
let bonsai_storage: BonsaiStorage<BasicId, RocksDB<'_, BasicId>, Pedersen> =
122+
BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24);
123+
124+
let root = bonsai_storage.root_hash_staged(&identifier).unwrap();
125+
assert_eq!(root, Felt::ZERO);
126+
}
127+
128+
#[test]
129+
fn staged_root_falls_back_to_committed_when_no_changes() {
130+
let identifier = vec![];
131+
let tempdir = tempfile::tempdir().unwrap();
132+
let db = create_rocks_db(tempdir.path()).unwrap();
133+
let config = BonsaiStorageConfig::default();
134+
let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
135+
BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24);
136+
let mut id_builder = BasicIdBuilder::new();
137+
138+
let pair1 = (
139+
vec![1, 2, 1],
140+
Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(),
141+
);
142+
bonsai_storage
143+
.insert(&identifier, &BitVec::from_vec(pair1.0.clone()), &pair1.1)
144+
.unwrap();
145+
bonsai_storage.commit(id_builder.new_id()).unwrap();
146+
147+
let committed_root = bonsai_storage.root_hash(&identifier).unwrap();
148+
149+
// No new inserts — staged should return the committed root
150+
let staged_root = bonsai_storage.root_hash_staged(&identifier).unwrap();
151+
assert_eq!(staged_root, committed_root);
152+
}
153+
154+
#[test]
155+
fn staged_root_with_removal() {
156+
let identifier = vec![];
157+
let tempdir = tempfile::tempdir().unwrap();
158+
let db = create_rocks_db(tempdir.path()).unwrap();
159+
let config = BonsaiStorageConfig::default();
160+
let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
161+
BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24);
162+
let mut id_builder = BasicIdBuilder::new();
163+
164+
let pair1 = (
165+
vec![1, 2, 1],
166+
Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(),
167+
);
168+
let pair2 = (
169+
vec![1, 2, 2],
170+
Felt::from_hex("0x66342762FD54D033c195fec3ce2568b62052e").unwrap(),
171+
);
172+
bonsai_storage
173+
.insert(&identifier, &BitVec::from_vec(pair1.0.clone()), &pair1.1)
174+
.unwrap();
175+
bonsai_storage
176+
.insert(&identifier, &BitVec::from_vec(pair2.0.clone()), &pair2.1)
177+
.unwrap();
178+
bonsai_storage.commit(id_builder.new_id()).unwrap();
179+
180+
// Remove one key without committing
181+
bonsai_storage
182+
.remove(&identifier, &BitVec::from_vec(pair1.0.clone()))
183+
.unwrap();
184+
185+
let staged_root = bonsai_storage.root_hash_staged(&identifier).unwrap();
186+
assert_ne!(staged_root, Felt::ZERO);
187+
188+
bonsai_storage.commit(id_builder.new_id()).unwrap();
189+
let committed_root = bonsai_storage.root_hash(&identifier).unwrap();
190+
191+
assert_eq!(staged_root, committed_root);
192+
}
193+
194+
#[test]
195+
fn staged_root_with_multiple_identifiers() {
196+
let id_a = vec![0x01];
197+
let id_b = vec![0x02];
198+
let tempdir = tempfile::tempdir().unwrap();
199+
let db = create_rocks_db(tempdir.path()).unwrap();
200+
let config = BonsaiStorageConfig::default();
201+
let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
202+
BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 24);
203+
let mut id_builder = BasicIdBuilder::new();
204+
205+
let val_a = Felt::from_hex("0xaa").unwrap();
206+
let val_b = Felt::from_hex("0xbb").unwrap();
207+
let key = vec![1, 2, 1];
208+
209+
bonsai_storage
210+
.insert(&id_a, &BitVec::from_vec(key.clone()), &val_a)
211+
.unwrap();
212+
bonsai_storage
213+
.insert(&id_b, &BitVec::from_vec(key.clone()), &val_b)
214+
.unwrap();
215+
216+
let staged_a = bonsai_storage.root_hash_staged(&id_a).unwrap();
217+
let staged_b = bonsai_storage.root_hash_staged(&id_b).unwrap();
218+
assert_ne!(staged_a, staged_b);
219+
220+
bonsai_storage.commit(id_builder.new_id()).unwrap();
221+
222+
let committed_a = bonsai_storage.root_hash(&id_a).unwrap();
223+
let committed_b = bonsai_storage.root_hash(&id_b).unwrap();
224+
225+
assert_eq!(staged_a, committed_a);
226+
assert_eq!(staged_b, committed_b);
227+
}

‎src/trie/tree.rs‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -379,6 +379,23 @@ impl<H: StarkHash + Send + Sync> MerkleTree<H> {
379379
Ok(NodeOrFelt::Node(node))
380380
}
381381

382+
/// Compute the root hash from staged (uncommitted) in-memory changes without
383+
/// persisting anything. If no staged changes exist, falls back to reading the
384+
/// committed root from the database.
385+
pub(crate) fn root_hash_staged<DB: BonsaiDatabase, ID: Id>(
386+
&self,
387+
db: &KeyValueDB<DB, ID>,
388+
) -> Result<Felt, BonsaiStorageError<DB::DatabaseError>> {
389+
match &self.root_node {
390+
Some(RootHandle::Loaded(_)) => {
391+
let mut hashes = vec![];
392+
self.compute_root_hash::<DB>(&mut hashes)
393+
}
394+
Some(RootHandle::Empty) => Ok(Felt::ZERO),
395+
None => self.root_hash(db),
396+
}
397+
}
398+
382399
fn compute_root_hash<DB: BonsaiDatabase>(
383400
&self,
384401
hashes: &mut Vec<Felt>,

‎src/trie/trees.rs‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,19 @@ impl<H: StarkHash + Send + Sync, DB: BonsaiDatabase, CommitID: Id> MerkleTrees<H
131131
}
132132
}
133133

134+
/// Compute root hash from staged (uncommitted) changes. Falls back to the
135+
/// committed root when the identified tree has no pending modifications.
136+
pub(crate) fn root_hash_staged(
137+
&self,
138+
identifier: &[u8],
139+
) -> Result<Felt, BonsaiStorageError<DB::DatabaseError>> {
140+
if let Some(tree) = self.trees.get(identifier) {
141+
tree.root_hash_staged(&self.db)
142+
} else {
143+
MerkleTree::<H>::new(identifier.into(), self.max_height).root_hash(&self.db)
144+
}
145+
}
146+
134147
pub(crate) fn get_keys(
135148
&self,
136149
identifier: &[u8],

0 commit comments

Comments
 (0)