Revert "Simplify withdrawals logic (#168)"

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Changed
+- Revert "Simplify withdrawals logic (#175)" ([#168](https://github.com/0xPolygonZero/zk_evm/pull/175))
 
 ## [0.3.0] - 2024-04-19
 

trace_decoder/src/decoding.rs

@@ -283,8 +283,11 @@ impl ProcessedBlockTrace {
         if !self.withdrawals.is_empty() {
             Self::add_withdrawals_to_txns(
                 &mut txn_gen_inputs,
+                &other_data,
+                &extra_data,
                 &mut curr_block_tries,
                 self.withdrawals,
+                dummies_added,
             )?;
         }
 
@@ -575,27 +578,56 @@ impl ProcessedBlockTrace {
     /// - If no dummy proofs are already present, then a dummy proof that just
     ///   contains the withdrawals is appended to the end of the IR vec.
     fn add_withdrawals_to_txns(
-        txn_ir: &mut [GenerationInputs],
+        txn_ir: &mut Vec<GenerationInputs>,
+        other_data: &OtherBlockData,
+        extra_data: &ExtraBlockData,
         final_trie_state: &mut PartialTrieState,
         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))
         };
 
-        Self::update_trie_state_from_withdrawals(
-            withdrawals_with_hashed_addrs_iter(),
-            &mut final_trie_state.state,
-        )?;
+        match dummies_already_added {
+            // If we have no actual dummy proofs, then we create one and append it to the
+            // end of the block.
+            false => {
+                let withdrawal_addrs =
+                    withdrawals_with_hashed_addrs_iter().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_with_hashed_addrs_iter(),
+                    &mut final_trie_state.state,
+                )?;
+
+                withdrawal_dummy.withdrawals = withdrawals;
 
-        let last_inputs = txn_ir
-            .last_mut()
-            .expect("We cannot have an empty list of payloads.");
+                // Only the state root hash needs to be updated from the withdrawals.
+                withdrawal_dummy.trie_roots_after.state_root = final_trie_state.state.hash();
 
-        last_inputs.withdrawals = withdrawals;
-        last_inputs.trie_roots_after.state_root = final_trie_state.state.hash();
+                txn_ir.push(withdrawal_dummy);
+            }
+            true => {
+                Self::update_trie_state_from_withdrawals(
+                    withdrawals_with_hashed_addrs_iter(),
+                    &mut final_trie_state.state,
+                )?;
+
+                // If we have dummy proofs (note: `txn_ir[1]` is always a dummy txn in this
+                // case), then this dummy will get the withdrawals.
+                txn_ir[1].withdrawals = withdrawals;
+                txn_ir[1].trie_roots_after.state_root = final_trie_state.state.hash();
+            }
+        }
 
         Ok(())
     }