diff --git a/protocol_decoder/src/decoding.rs b/protocol_decoder/src/decoding.rs index ba98ab2..e070c33 100644 --- a/protocol_decoder/src/decoding.rs +++ b/protocol_decoder/src/decoding.rs @@ -1,7 +1,7 @@ use std::{ collections::HashMap, fmt::{self, Display, Formatter}, - iter::once, + iter::{self, empty, once}, }; use eth_trie_utils::{ @@ -342,6 +342,10 @@ impl ProcessedBlockTrace { withdrawals: Vec<(Address, U256)>, dummies_already_added: bool, ) -> TraceParsingResult<()> { + let withdrawals_with_hashed_addrs_iter = withdrawals + .iter() + .map(|(addr, v)| (*addr, hash(addr.as_bytes()), *v)); + match dummies_already_added { // If we have no actual dummy proofs, then we create one and append it to the // end of the block. @@ -350,12 +354,26 @@ impl ProcessedBlockTrace { let txn_idx_of_dummy_entry = txn_ir.last().unwrap().txn_number_before.low_u64() as usize + 1; - // Dummy state will be the state after the final txn. - let mut withdrawal_dummy = - create_dummy_gen_input(other_data, extra_data, final_trie_state); + // To avoid double hashing the addrs, but I don't know if the extra `Vec` + // allocation is worth it. + let withdrawals_with_hashed_addrs: Vec<_> = + withdrawals_with_hashed_addrs_iter.collect(); + + // Dummy state will be the state after the final txn. Also need to include the + // account nodes that were accessed by the withdrawals. + let withdrawal_addrs = withdrawals_with_hashed_addrs + .iter() + .cloned() + .map(|(_, h_addr, _)| h_addr); + let mut withdrawal_dummy = create_dummy_gen_input_with_state_addrs_accessed( + other_data, + extra_data, + final_trie_state, + withdrawal_addrs, + )?; Self::update_trie_state_from_withdrawals( - &withdrawals, + withdrawals_with_hashed_addrs, &mut final_trie_state.state, )?; @@ -368,7 +386,7 @@ impl ProcessedBlockTrace { } true => { Self::update_trie_state_from_withdrawals( - &withdrawals, + withdrawals_with_hashed_addrs_iter, &mut final_trie_state.state, )?; @@ -385,22 +403,21 @@ impl ProcessedBlockTrace { /// Withdrawals update balances in the account trie, so we need to update /// our local trie state. fn update_trie_state_from_withdrawals<'a>( - withdrawals: impl IntoIterator + 'a, + withdrawals: impl IntoIterator + 'a, state: &mut HashedPartialTrie, ) -> TraceParsingResult<()> { - for (addr, amt) in withdrawals { - let h_addr = hash(addr.as_bytes()); + for (addr, h_addr, amt) in withdrawals { let h_addr_nibs = Nibbles::from_h256_be(h_addr); let acc_bytes = state .get(h_addr_nibs) .ok_or(TraceParsingError::MissingWithdrawalAccount( - *addr, h_addr, *amt, + addr, h_addr, amt, ))?; let mut acc_data = account_from_rlped_bytes(acc_bytes)?; - acc_data.balance += *amt; + acc_data.balance += amt; state.insert(h_addr_nibs, rlp::encode(&acc_data).to_vec()); } @@ -467,12 +484,37 @@ fn create_dummy_gen_input( extra_data: &ExtraBlockData, final_tries: &PartialTrieState, ) -> TxnProofGenIR { - let tries = create_dummy_proof_trie_inputs(final_tries); + let sub_tries = create_dummy_proof_trie_inputs( + final_tries, + create_fully_hashed_out_sub_partial_trie(&final_tries.state), + ); + create_dummy_gen_input_common(other_data, extra_data, sub_tries) +} +fn create_dummy_gen_input_with_state_addrs_accessed( + other_data: &OtherBlockData, + extra_data: &ExtraBlockData, + final_tries: &PartialTrieState, + account_addrs_accessed: impl Iterator, +) -> TraceParsingResult { + let sub_tries = create_dummy_proof_trie_inputs( + final_tries, + create_minimal_state_partial_trie(&final_tries.state, account_addrs_accessed)?, + ); + Ok(create_dummy_gen_input_common( + other_data, extra_data, sub_tries, + )) +} + +fn create_dummy_gen_input_common( + other_data: &OtherBlockData, + extra_data: &ExtraBlockData, + sub_tries: TrieInputs, +) -> GenerationInputs { let trie_roots_after = TrieRoots { - state_root: tries.state_trie.hash(), - transactions_root: tries.transactions_trie.hash(), - receipts_root: tries.receipts_trie.hash(), + state_root: sub_tries.state_trie.hash(), + transactions_root: sub_tries.transactions_trie.hash(), + receipts_root: sub_tries.receipts_trie.hash(), }; // Sanity checks @@ -487,7 +529,7 @@ fn create_dummy_gen_input( GenerationInputs { signed_txn: None, - tries, + tries: sub_tries, trie_roots_after, checkpoint_state_trie_root: extra_data.checkpoint_state_trie_root, block_metadata: other_data.b_data.b_meta.clone(), @@ -500,16 +542,10 @@ fn create_dummy_gen_input( } } -impl TxnMetaState { - fn txn_bytes(&self) -> Vec { - match self.txn_bytes.as_ref() { - Some(v) => v.clone(), - None => Vec::default(), - } - } -} - -fn create_dummy_proof_trie_inputs(final_trie_state: &PartialTrieState) -> TrieInputs { +fn create_dummy_proof_trie_inputs( + final_trie_state: &PartialTrieState, + state_trie: HashedPartialTrie, +) -> TrieInputs { let partial_sub_storage_tries: Vec<_> = final_trie_state .storage .iter() @@ -522,7 +558,7 @@ fn create_dummy_proof_trie_inputs(final_trie_state: &PartialTrieState) -> TrieIn .collect(); TrieInputs { - state_trie: create_fully_hashed_out_sub_partial_trie(&final_trie_state.state), + state_trie, transactions_trie: create_fully_hashed_out_sub_partial_trie(&final_trie_state.txn), receipts_trie: create_fully_hashed_out_sub_partial_trie(&final_trie_state.receipt), storage_tries: partial_sub_storage_tries, @@ -586,3 +622,12 @@ fn account_from_rlped_bytes(bytes: &[u8]) -> TraceParsingResult { rlp::decode(bytes) .map_err(|err| TraceParsingError::AccountDecode(hex::encode(bytes), err.to_string())) } + +impl TxnMetaState { + fn txn_bytes(&self) -> Vec { + match self.txn_bytes.as_ref() { + Some(v) => v.clone(), + None => Vec::default(), + } + } +}