Skip to content

Commit 8a552d5

Browse files
ArielElpArielElp
committed
index layout forest creation tests
1 parent 97e36dc commit 8a552d5

File tree

2 files changed

+123
-14
lines changed

2 files changed

+123
-14
lines changed

‎crates/starknet_committer/src/forest/skeleton_forest_test.rs‎

Lines changed: 98 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use std::collections::HashMap;
22

33
use pretty_assertions::assert_eq;
44
use rstest::rstest;
5+
use rstest_reuse::{apply, template};
56
use starknet_api::core::{ClassHash, ContractAddress, Nonce};
67
use starknet_api::hash::{HashOutput, StateRoots};
78
use starknet_patricia::patricia_merkle_tree::external_test_utils::{
@@ -21,11 +22,13 @@ use starknet_patricia_storage::db_object::{DBObject, EmptyKeyContext};
2122
use starknet_patricia_storage::map_storage::MapStorage;
2223
use starknet_patricia_storage::storage_trait::{DbHashMap, DbKey, DbValue};
2324
use starknet_types_core::felt::Felt;
25+
use starknet_types_core::hash::StarkHash;
2426

2527
use crate::block_committer::commit::get_all_modified_indices;
2628
use crate::block_committer::input::{
2729
contract_address_into_node_index,
2830
Input,
31+
InputContext,
2932
ReaderConfig,
3033
StarknetStorageKey,
3134
StarknetStorageValue,
@@ -34,6 +37,8 @@ use crate::block_committer::input::{
3437
use crate::db::facts_db::db::FactsDb;
3538
use crate::db::facts_db::types::FactsDbInitialRead;
3639
use crate::db::forest_trait::ForestReader;
40+
use crate::db::index_db::db::IndexDb;
41+
use crate::db::index_db::test_utils::convert_facts_forest_db_to_index_db;
3742
use crate::forest::original_skeleton_forest::{ForestSortedIndices, OriginalSkeletonForest};
3843
use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;
3944
use crate::patricia_merkle_tree::types::CompiledClassHash;
@@ -61,14 +66,24 @@ pub(crate) fn create_compiled_class_leaf_entry(val: u128) -> (DbKey, DbValue) {
6166
(leaf.get_db_key(&EmptyKeyContext, &leaf.0.to_bytes_be()), leaf.serialize().unwrap())
6267
}
6368

64-
pub(crate) fn create_contract_state_leaf_entry(val: u128) -> (DbKey, DbValue) {
65-
let felt = Felt::from(val);
69+
pub(crate) fn create_contract_state_leaf_entry(
70+
nonce: u128,
71+
storage_root: u128,
72+
class_hash: u128,
73+
hash: Option<u128>,
74+
) -> (DbKey, DbValue) {
75+
let (nonce, storage_root, class_hash) =
76+
(Felt::from(nonce), Felt::from(storage_root), Felt::from(class_hash));
6677
let leaf = ContractState {
67-
nonce: Nonce(felt),
68-
storage_root_hash: HashOutput(felt),
69-
class_hash: ClassHash(felt),
78+
nonce: Nonce(nonce),
79+
storage_root_hash: HashOutput(storage_root),
80+
class_hash: ClassHash(class_hash),
7081
};
71-
(leaf.get_db_key(&EmptyKeyContext, &felt.to_bytes_be()), leaf.serialize().unwrap())
82+
let mut hash = hash.map(Felt::from);
83+
if hash.is_none() {
84+
hash = Some(AdditionHash::hash_array(&[class_hash, storage_root, nonce]));
85+
}
86+
(leaf.get_db_key(&EmptyKeyContext, &hash.unwrap().to_bytes_be()), leaf.serialize().unwrap())
7287
}
7388

7489
// This test uses addition hash for simplicity (i.e hash(a,b) = a + b).
@@ -133,7 +148,7 @@ pub(crate) fn create_contract_state_leaf_entry(val: u128) -> (DbKey, DbValue) {
133148
/// / \ \ \ / \ \ \
134149
/// NZ 2 NZ NZ NZ 9 16 15
135150
136-
#[tokio::test]
151+
#[template]
137152
#[rstest]
138153
#[case(
139154
Input {
@@ -165,9 +180,11 @@ pub(crate) fn create_contract_state_leaf_entry(val: u128) -> (DbKey, DbValue) {
165180
create_edge_entry_from_u128::<AdditionHash>(554, 1, 1),
166181
create_binary_entry_from_u128::<AdditionHash>(305, 556),
167182
// Contracts trie leaves.
168-
create_contract_state_leaf_entry(277),
169-
create_contract_state_leaf_entry(303),
170-
create_contract_state_leaf_entry(1),
183+
// TODO(Ariel): Why both layouts pass when we lie about the hash?
184+
// probably because the original skeleton doesn't care about these hashes (there were modifications).
185+
create_contract_state_leaf_entry(277, 277, 277, Some(277)),
186+
create_contract_state_leaf_entry(303, 303, 303, Some(303)),
187+
create_contract_state_leaf_entry(0, 0, 1, None),
171188
// Classes trie inner nodes.
172189
create_binary_entry_from_u128::<AdditionHash>(33, 47),
173190
create_edge_entry_from_u128::<AdditionHash>(72, 1, 1),
@@ -298,17 +315,85 @@ pub(crate) fn create_contract_state_leaf_entry(val: u128) -> (DbKey, DbValue) {
298315
vec![6, 7, 0],
299316
vec![7, 6, 0],
300317
)]
301-
async fn test_create_original_skeleton_forest(
318+
fn create_original_skeleton_forest_cases(
319+
#[case] input: Input<FactsDbInitialRead>,
320+
#[case] storage: MapStorage,
321+
#[case] expected_forest: OriginalSkeletonForest<'_>,
322+
#[case] expected_original_contracts_trie_leaves: HashMap<ContractAddress, ContractState>,
323+
#[case] expected_storage_tries_sorted_indices: HashMap<u128, Vec<u128>>,
324+
#[case] expected_contracts_trie_sorted_indices: Vec<u128>,
325+
#[case] expected_classes_trie_sorted_indices: Vec<u128>,
326+
) {
327+
}
328+
329+
#[apply(create_original_skeleton_forest_cases)]
330+
#[rstest]
331+
#[tokio::test]
332+
async fn test_create_original_skeleton_forest_facts_layout(
302333
#[case] input: Input<FactsDbInitialRead>,
303334
#[case] storage: MapStorage,
304335
#[case] expected_forest: OriginalSkeletonForest<'_>,
305336
#[case] expected_original_contracts_trie_leaves: HashMap<ContractAddress, ContractState>,
306337
#[case] expected_storage_tries_sorted_indices: HashMap<u128, Vec<u128>>,
307338
#[case] expected_contracts_trie_sorted_indices: Vec<u128>,
308339
#[case] expected_classes_trie_sorted_indices: Vec<u128>,
340+
) {
341+
test_create_original_skeleton_forest::<FactsDbInitialRead, FactsDb<MapStorage>>(
342+
input,
343+
storage,
344+
expected_forest,
345+
expected_original_contracts_trie_leaves,
346+
expected_storage_tries_sorted_indices,
347+
expected_contracts_trie_sorted_indices,
348+
expected_classes_trie_sorted_indices,
349+
|storage| FactsDb::new(storage),
350+
)
351+
.await;
352+
}
353+
354+
#[apply(create_original_skeleton_forest_cases)]
355+
#[rstest]
356+
#[tokio::test]
357+
async fn test_create_original_skeleton_forest_index_layout(
358+
#[case] input: Input<FactsDbInitialRead>,
359+
#[case] mut storage: MapStorage,
360+
#[case] expected_forest: OriginalSkeletonForest<'_>,
361+
#[case] expected_original_contracts_trie_leaves: HashMap<ContractAddress, ContractState>,
362+
#[case] expected_storage_tries_sorted_indices: HashMap<u128, Vec<u128>>,
363+
#[case] expected_contracts_trie_sorted_indices: Vec<u128>,
364+
#[case] expected_classes_trie_sorted_indices: Vec<u128>,
365+
) {
366+
let index_storage =
367+
convert_facts_forest_db_to_index_db(&mut storage, input.initial_read_context.0).await;
368+
test_create_original_skeleton_forest::<FactsDbInitialRead, IndexDb<MapStorage>>(
369+
input,
370+
index_storage,
371+
expected_forest,
372+
expected_original_contracts_trie_leaves,
373+
expected_storage_tries_sorted_indices,
374+
expected_contracts_trie_sorted_indices,
375+
expected_classes_trie_sorted_indices,
376+
|storage| IndexDb::new(storage),
377+
)
378+
.await;
379+
}
380+
381+
async fn test_create_original_skeleton_forest<
382+
ReaderInputContext: InputContext,
383+
Reader: ForestReader<ReaderInputContext>,
384+
>(
385+
input: Input<ReaderInputContext>,
386+
storage: MapStorage,
387+
expected_forest: OriginalSkeletonForest<'_>,
388+
expected_original_contracts_trie_leaves: HashMap<ContractAddress, ContractState>,
389+
expected_storage_tries_sorted_indices: HashMap<u128, Vec<u128>>,
390+
expected_contracts_trie_sorted_indices: Vec<u128>,
391+
expected_classes_trie_sorted_indices: Vec<u128>,
392+
reader_generator: impl FnOnce(MapStorage) -> Reader,
309393
) {
310394
let (mut storage_tries_indices, mut contracts_trie_indices, mut classes_trie_indices) =
311395
get_all_modified_indices(&input.state_diff);
396+
312397
let forest_sorted_indices = ForestSortedIndices {
313398
storage_tries_sorted_indices: storage_tries_indices
314399
.iter_mut()
@@ -320,7 +405,8 @@ async fn test_create_original_skeleton_forest(
320405

321406
let actual_storage_updates = input.state_diff.actual_storage_updates();
322407
let actual_classes_updates = input.state_diff.actual_classes_updates();
323-
let (actual_forest, original_contracts_trie_leaves) = FactsDb::new(storage)
408+
let mut reader = reader_generator(storage);
409+
let (actual_forest, original_contracts_trie_leaves) = reader
324410
.read(
325411
input.initial_read_context,
326412
&actual_storage_updates,

‎crates/starknet_committer/src/hash_function/hash.rs‎

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,22 @@ impl TreeHashFunction<ContractState> for TreeHashFunctionImpl {
5656

5757
impl TreeHashFunction<IndexLayoutContractState> for TreeHashFunctionImpl {
5858
fn compute_leaf_hash(contract_state: &IndexLayoutContractState) -> HashOutput {
59-
compute_contract_state_leaf_hash(contract_state)
59+
#[cfg(not(test))]
60+
{
61+
compute_contract_state_leaf_hash(contract_state)
62+
}
63+
#[cfg(test)]
64+
{
65+
use starknet_patricia::patricia_merkle_tree::external_test_utils::AdditionHash;
66+
HashOutput(AdditionHash::hash_array(
67+
vec![
68+
contract_state.0.class_hash.0,
69+
contract_state.0.storage_root_hash.0,
70+
contract_state.0.nonce.0,
71+
]
72+
.as_slice(),
73+
))
74+
}
6075
}
6176
fn compute_node_hash(node_data: &NodeData<IndexLayoutContractState, HashOutput>) -> HashOutput {
6277
Self::compute_node_hash_with_inner_hash_function::<PedersenHashFunction>(node_data)
@@ -90,7 +105,15 @@ impl TreeHashFunction<CompiledClassHash> for TreeHashFunctionImpl {
90105

91106
impl TreeHashFunction<IndexLayoutCompiledClassHash> for TreeHashFunctionImpl {
92107
fn compute_leaf_hash(compiled_class_hash: &IndexLayoutCompiledClassHash) -> HashOutput {
93-
compute_compiled_class_leaf_hash(compiled_class_hash)
108+
#[cfg(not(test))]
109+
{
110+
compute_compiled_class_leaf_hash(compiled_class_hash)
111+
}
112+
#[cfg(test)]
113+
{
114+
use starknet_patricia::patricia_merkle_tree::external_test_utils::AdditionHash;
115+
HashOutput(AdditionHash::hash_array(vec![compiled_class_hash.0.0].as_slice()))
116+
}
94117
}
95118
fn compute_node_hash(
96119
node_data: &NodeData<IndexLayoutCompiledClassHash, HashOutput>,

0 commit comments

Comments
 (0)