diff --git a/crates/orchestrator/src/jobs/da_job/mod.rs b/crates/orchestrator/src/jobs/da_job/mod.rs index 2f87b6bc..30517538 100644 --- a/crates/orchestrator/src/jobs/da_job/mod.rs +++ b/crates/orchestrator/src/jobs/da_job/mod.rs @@ -396,33 +396,33 @@ fn da_word(class_flag: bool, nonce_change: Option, num_changes: u64) -> Fe } fn refactor_state_update(state_update: &mut StateDiff) { - let addresses_in_nonces: Vec = state_update.nonces.clone().iter().map(|item| item.contract_address).collect(); - let addresses_in_storage_diffs: Vec = - state_update.storage_diffs.clone().iter().map(|item| item.address).collect(); + // Create set of all addresses that need storage diffs + let mut required_addresses: HashSet = HashSet::new(); - let address_to_insert = find_unique_addresses(addresses_in_nonces, addresses_in_storage_diffs.clone()); + // First add all existing storage addresses + required_addresses.extend(state_update.storage_diffs.iter().map(|item| item.address)); - for address in address_to_insert { - state_update.storage_diffs.push(ContractStorageDiffItem { address, storage_entries: vec![] }) - } - - let addresses_in_deployed_contracts = - state_update.deployed_contracts.clone().iter().map(|item| item.address).collect(); - let addresses_to_insert = find_unique_addresses(addresses_in_deployed_contracts, addresses_in_storage_diffs); + // Add addresses from nonces + required_addresses.extend(state_update.nonces.iter().map(|item| item.contract_address)); - for address in addresses_to_insert { - state_update.storage_diffs.push(ContractStorageDiffItem { address, storage_entries: vec![] }) - } -} + // Add addresses from deployed contracts + required_addresses.extend(state_update.deployed_contracts.iter().map(|item| item.address)); -fn find_unique_addresses(nonce_addresses: Vec, storage_diff_addresses: Vec) -> Vec { - let storage_set: HashSet<_> = storage_diff_addresses.into_iter().collect(); + // Create a set of existing storage addresses for quick lookup + let existing_storage_addresses: HashSet = + state_update.storage_diffs.iter().map(|item| item.address).collect(); - nonce_addresses.into_iter().filter(|addr| !storage_set.contains(addr)).collect() + // Add storage diffs only for addresses that don't already have them + for address in required_addresses { + if !existing_storage_addresses.contains(&address) { + state_update.storage_diffs.push(ContractStorageDiffItem { address, storage_entries: Vec::new() }); + } + } } #[cfg(test)] pub mod test { + use std::collections::HashSet; use std::fs; use std::fs::File; use std::io::Read; @@ -435,12 +435,14 @@ pub mod test { use majin_blob_types::serde; use rstest::rstest; use serde_json::json; - use starknet::core::types::{Felt, StateUpdate}; + use starknet::core::types::{ + ContractStorageDiffItem, DeployedContractItem, Felt, NonceUpdate, StateDiff, StateUpdate, StorageEntry, + }; use starknet::providers::jsonrpc::HttpTransport; use starknet::providers::JsonRpcClient; use url::Url; - use crate::jobs::da_job::da_word; + use crate::jobs::da_job::{da_word, refactor_state_update}; /// Tests `da_word` function with various inputs for class flag, new nonce, and number of /// changes. Verifies that `da_word` produces the correct Felt based on the provided @@ -494,8 +496,12 @@ pub mod test { "src/tests/jobs/da_job/test_data/test_blob/671070.txt", "src/tests/jobs/da_job/test_data/nonces/671070.txt" )] + // Block from pragma madara and orch test run. Here we faced an issue where our + // blob building logic was not able to take the contract addresses from + // `deployed_contracts` field in state diff from state update. This test case + // was added after the fix #[case( - 178, // Block from pragma madara and orch test run + 178, "src/tests/jobs/da_job/test_data/state_update/178.txt", "src/tests/jobs/da_job/test_data/test_blob/178.txt", "src/tests/jobs/da_job/test_data/nonces/178.txt" @@ -584,6 +590,71 @@ pub mod test { assert_eq!(data, deserialize_data); } + #[rstest] + #[case(vec![], vec![], vec![], 0)] // Empty case + #[case( + vec![(Felt::from(1), Felt::from(10)), (Felt::from(2), Felt::from(20))], + vec![], + vec![], + 2 + )] // Only nonces + #[case( + vec![], + vec![], + vec![(Felt::from(1), vec![1]), (Felt::from(2), vec![2])], + 2 + )] // Only deployed + #[case( + vec![(Felt::from(1), Felt::from(10))], + vec![(Felt::from(1), vec![(Felt::from(1), Felt::from(100))])], + vec![(Felt::from(1), vec![1])], + 1 + )] // Overlapping addresses + #[case( + vec![(Felt::from(1), Felt::from(10)), (Felt::from(1), Felt::from(20))], + vec![], + vec![(Felt::from(1), vec![1]), (Felt::from(1), vec![2])], + 1 + )] // Duplicate addresses + fn test_refactor_state_update( + #[case] nonces: Vec<(Felt, Felt)>, + #[case] storage_diffs: Vec<(Felt, Vec<(Felt, Felt)>)>, + #[case] deployed_contracts: Vec<(Felt, Vec)>, + #[case] expected_storage_count: usize, + ) { + // Create initial state + let mut state_diff = create_state_diff(nonces, storage_diffs.clone(), deployed_contracts); + let initial_storage = state_diff.storage_diffs.clone(); + + // Run the function + refactor_state_update(&mut state_diff); + + // Verify results + assert!(verify_addresses_have_storage_diffs(&state_diff)); + + // Verify no duplicate addresses in storage_diffs + let unique_addresses: HashSet<_> = state_diff.storage_diffs.iter().map(|item| &item.address).collect(); + assert_eq!( + unique_addresses.len(), + state_diff.storage_diffs.len(), + "Number of unique addresses should match number of storage diffs" + ); + assert_eq!( + unique_addresses.len(), + expected_storage_count, + "Number of storage diffs should match expected count" + ); + + // Verify storage preservation for existing entries + for orig_storage in initial_storage { + if let Some(current_storage) = + state_diff.storage_diffs.iter().find(|item| item.address == orig_storage.address) + { + assert_eq!(orig_storage.storage_entries, current_storage.storage_entries); + } + } + } + pub(crate) fn read_state_update_from_file(file_path: &str) -> Result { // let file_path = format!("state_update_block_no_{}.txt", block_no); let mut file = File::open(file_path)?; @@ -630,4 +701,35 @@ pub mod test { new_hex_chars = new_hex_chars.trim_start_matches('0').to_string(); if new_hex_chars.is_empty() { "0x0".to_string() } else { format!("0x{}", new_hex_chars) } } + + fn create_state_diff( + nonces: Vec<(Felt, Felt)>, + storage_diffs: Vec<(Felt, Vec<(Felt, Felt)>)>, + deployed_contracts: Vec<(Felt, Vec)>, + ) -> StateDiff { + StateDiff { + nonces: nonces.into_iter().map(|(addr, nonce)| NonceUpdate { contract_address: addr, nonce }).collect(), + storage_diffs: storage_diffs + .into_iter() + .map(|(addr, entries)| ContractStorageDiffItem { + address: addr, + storage_entries: entries.into_iter().map(|(key, value)| StorageEntry { key, value }).collect(), + }) + .collect(), + deprecated_declared_classes: vec![], + declared_classes: vec![], + deployed_contracts: deployed_contracts + .into_iter() + .map(|(addr, _code)| DeployedContractItem { address: addr, class_hash: Default::default() }) + .collect(), + replaced_classes: vec![], + } + } + + fn verify_addresses_have_storage_diffs(state_diff: &StateDiff) -> bool { + let storage_addresses: HashSet<_> = state_diff.storage_diffs.iter().map(|item| &item.address).collect(); + + state_diff.nonces.iter().all(|item| storage_addresses.contains(&item.contract_address)) + && state_diff.deployed_contracts.iter().all(|item| storage_addresses.contains(&item.address)) + } }