Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(committer): use FilledTree::create() in filled forest #495

Merged
merged 1 commit into from
Aug 20, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 22 additions & 0 deletions crates/starknet_committer/src/block_committer/input.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,10 +10,32 @@ use tracing::level_filters::LevelFilter;

use crate::patricia_merkle_tree::types::{ClassHash, CompiledClassHash, Nonce};

#[cfg(test)]
#[path = "input_test.rs"]
pub mod input_test;

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
// TODO(Nimrod, 1/6/2025): Use the ContractAddress defined in starknet-types-core when available.
pub struct ContractAddress(pub Felt);

impl TryFrom<&NodeIndex> for ContractAddress {
type Error = String;

fn try_from(node_index: &NodeIndex) -> Result<ContractAddress, Self::Error> {
if !node_index.is_leaf() {
return Err("NodeIndex is not a leaf.".to_string());
}
let result = Felt::try_from(*node_index - NodeIndex::FIRST_LEAF);
match result {
Ok(felt) => Ok(ContractAddress(felt)),
Err(error) => Err(format!(
"Tried to convert node index to felt and got the following error: {:?}",
error.to_string()
)),
}
}
}

impl From<&ContractAddress> for NodeIndex {
fn from(address: &ContractAddress) -> NodeIndex {
NodeIndex::from_leaf_felt(&address.0)
Expand Down
25 changes: 25 additions & 0 deletions crates/starknet_committer/src/block_committer/input_test.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
use rstest::rstest;
use starknet_patricia::felt::Felt;
use starknet_patricia::patricia_merkle_tree::types::NodeIndex;

use crate::block_committer::input::ContractAddress;

#[rstest]
fn test_node_index_to_contract_address_conversion() {
// Positive flow.
assert_eq!(ContractAddress::try_from(&NodeIndex::FIRST_LEAF), Ok(ContractAddress(Felt::ZERO)));
assert_eq!(
ContractAddress::try_from(&(NodeIndex::FIRST_LEAF + NodeIndex(1_u32.into()))),
Ok(ContractAddress(Felt::ONE))
);
assert_eq!(
ContractAddress::try_from(&NodeIndex::MAX),
Ok(ContractAddress(Felt::try_from(NodeIndex::MAX - NodeIndex::FIRST_LEAF).unwrap()))
);

// Negative flow.
assert_eq!(
ContractAddress::try_from(&(NodeIndex::FIRST_LEAF - NodeIndex(1_u32.into()))),
Err("NodeIndex is not a leaf.".to_string())
);
}
144 changes: 86 additions & 58 deletions crates/starknet_committer/src/forest/filled_forest.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,11 @@ use std::collections::HashMap;

use starknet_patricia::hash::hash_trait::HashOutput;
use starknet_patricia::patricia_merkle_tree::filled_tree::tree::FilledTree;
use starknet_patricia::patricia_merkle_tree::node_data::leaf::LeafModifications;
use starknet_patricia::patricia_merkle_tree::node_data::leaf::{Leaf, LeafModifications};
use starknet_patricia::patricia_merkle_tree::types::NodeIndex;
use starknet_patricia::patricia_merkle_tree::updated_skeleton_tree::tree::UpdatedSkeletonTreeImpl;
use starknet_patricia::storage::storage_trait::Storage;
use tokio::task::JoinSet;
use tracing::info;

use crate::block_committer::input::{ContractAddress, StarknetStorageValue};
use crate::forest::forest_errors::{ForestError, ForestResult};
Expand All @@ -19,7 +19,6 @@ use crate::patricia_merkle_tree::types::{
CompiledClassHash,
ContractsTrie,
Nonce,
StorageTrie,
StorageTrieMap,
};

Expand Down Expand Up @@ -52,8 +51,10 @@ impl FilledForest {
self.classes_trie.get_root_hash()
}

/// Creates a filled forest. Assumes the storage updates and the updated skeletons of the
/// storage tries include all modified contracts, including those with unmodified storage.
pub(crate) async fn create<TH: ForestHashFunction + 'static>(
mut updated_forest: UpdatedSkeletonForest,
updated_forest: UpdatedSkeletonForest,
storage_updates: HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
classes_updates: LeafModifications<CompiledClassHash>,
original_contracts_trie_leaves: &HashMap<NodeIndex, ContractState>,
Expand All @@ -64,67 +65,94 @@ impl FilledForest {
updated_forest.classes_trie,
classes_updates,
));
let mut contracts_trie_modifications = HashMap::new();
let mut filled_storage_tries = HashMap::new();
let mut contracts_state_tasks = JoinSet::new();

for (address, inner_updates) in storage_updates {
let updated_storage_trie = updated_forest
.storage_tries
.remove(&address)
.ok_or(ForestError::MissingUpdatedSkeleton(address))?;

let original_contract_state = original_contracts_trie_leaves
.get(&((&address).into()))
.ok_or(ForestError::MissingContractCurrentState(address))?;
contracts_state_tasks.spawn(Self::new_contract_state::<TH>(
address,
*(address_to_nonce.get(&address).unwrap_or(&original_contract_state.nonce)),
*(address_to_class_hash
.get(&address)
.unwrap_or(&original_contract_state.class_hash)),
updated_storage_trie,
inner_updates,
));
}

while let Some(result) = contracts_state_tasks.join_next().await {
let (address, new_contract_state, filled_storage_trie) = result??;
contracts_trie_modifications.insert((&address).into(), new_contract_state);
filled_storage_tries.insert(address, filled_storage_trie);
}

let contracts_trie_task = tokio::spawn(ContractsTrie::create_with_existing_leaves::<TH>(
let contracts_trie_task = tokio::task::spawn(ContractsTrie::create::<TH>(
updated_forest.contracts_trie,
contracts_trie_modifications,
FilledForest::get_contracts_trie_leaf_input(
original_contracts_trie_leaves,
storage_updates,
updated_forest.storage_tries,
address_to_class_hash,
address_to_nonce,
)?,
));

let contracts_trie = contracts_trie_task.await?.map_err(ForestError::ContractsTrie)?;
let classes_trie = classes_trie_task.await?.map_err(ForestError::ClassesTrie)?;
info!(
"Classes trie update complete; {:?} new facts computed.",
classes_trie.tree_map.len()
);
let (contracts_trie, storage_tries) =
contracts_trie_task.await?.map_err(ForestError::ContractsTrie)?;
info!(
"Contracts trie update complete; {:?} new facts computed.",
contracts_trie.tree_map.len()
);

Ok(Self { storage_tries: filled_storage_tries, contracts_trie, classes_trie })
Ok(Self {
storage_tries: storage_tries
.into_iter()
.map(|(node_index, storage_trie)| {
(
ContractAddress::try_from(&node_index).unwrap_or_else(|error| {
panic!(
"Got the following error when trying to convert node index \
{node_index:?} to a contract address: {error:?}",
)
}),
storage_trie,
)
})
.collect(),
contracts_trie,
classes_trie,
})
}

async fn new_contract_state<TH: ForestHashFunction + 'static>(
contract_address: ContractAddress,
new_nonce: Nonce,
new_class_hash: ClassHash,
updated_storage_trie: UpdatedSkeletonTreeImpl,
inner_updates: LeafModifications<StarknetStorageValue>,
) -> ForestResult<(ContractAddress, ContractState, StorageTrie)> {
let filled_storage_trie =
StorageTrie::create_with_existing_leaves::<TH>(updated_storage_trie, inner_updates)
.await
.map_err(ForestError::StorageTrie)?;
let new_root_hash = filled_storage_trie.get_root_hash();
Ok((
contract_address,
ContractState {
nonce: new_nonce,
storage_root_hash: new_root_hash,
class_hash: new_class_hash,
},
filled_storage_trie,
))
fn get_contracts_trie_leaf_input(
original_contracts_trie_leaves: &HashMap<NodeIndex, ContractState>,
contract_address_to_storage_updates: HashMap<
ContractAddress,
LeafModifications<StarknetStorageValue>,
>,
mut contract_address_to_storage_skeleton: HashMap<ContractAddress, UpdatedSkeletonTreeImpl>,
address_to_class_hash: &HashMap<ContractAddress, ClassHash>,
address_to_nonce: &HashMap<ContractAddress, Nonce>,
) -> ForestResult<HashMap<NodeIndex, <ContractState as Leaf>::Input>> {
let mut leaf_index_to_leaf_input = HashMap::new();
assert_eq!(
contract_address_to_storage_updates.len(),
contract_address_to_storage_skeleton.len(),
"Mismatch between number of updated storage trie skeletons and number of storage \
leaf-modification maps. Number of updated storage trie skeletons: {0:?}, number of \
storage leaf-modification maps: {1:?}",
contract_address_to_storage_skeleton.len(),
contract_address_to_storage_updates.len()
);
// `contract_address_to_storage_updates` includes all modified contracts, even those with
// unmodified storage, see StateDiff::actual_storage_updates().
for (contract_address, storage_updates) in contract_address_to_storage_updates {
let node_index = NodeIndex::from(&contract_address);
let original_contract_state = original_contracts_trie_leaves
.get(&node_index)
.ok_or(ForestError::MissingContractCurrentState(contract_address))?;
leaf_index_to_leaf_input.insert(
node_index,
(
node_index,
*(address_to_nonce
.get(&contract_address)
.unwrap_or(&original_contract_state.nonce)),
*(address_to_class_hash
.get(&contract_address)
.unwrap_or(&original_contract_state.class_hash)),
contract_address_to_storage_skeleton
.remove(&contract_address)
.ok_or(ForestError::MissingUpdatedSkeleton(contract_address))?,
storage_updates,
),
);
}
Ok(leaf_index_to_leaf_input)
}
}
5 changes: 4 additions & 1 deletion crates/starknet_committer/src/forest/forest_errors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,10 @@ pub enum ForestError {
address {0:?}"
)]
MissingOriginalSkeleton(ContractAddress),
#[error("Can't fill storage trie, because there is no updated skeleton at address {0:?}")]
#[error(
"Can't create Contracts trie, because there is no updated skeleton for storage trie at \
address {0:?}"
)]
MissingUpdatedSkeleton(ContractAddress),
#[error(
"Can't build storage trie, because there are no sorted leaf indices of the contract at \
Expand Down
2 changes: 1 addition & 1 deletion crates/starknet_patricia/src/patricia_merkle_tree/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ impl NodeIndex {
Self(index)
}

pub(crate) fn is_leaf(&self) -> bool {
pub fn is_leaf(&self) -> bool {
Self::FIRST_LEAF <= *self && *self <= Self::MAX
}

Expand Down
Loading