|
| 1 | +use super::run::fetch_contracts_bytecode_from_trace; |
1 | 2 | use crate::{ |
2 | 3 | Cast, |
3 | 4 | traces::TraceKind, |
4 | 5 | tx::{CastTxBuilder, SenderKind}, |
5 | 6 | }; |
6 | 7 | use alloy_ens::NameOrAddress; |
7 | | -use alloy_primitives::{Address, Bytes, TxKind, U256}; |
| 8 | +use alloy_primitives::{Address, B256, Bytes, TxKind, U256, map::HashMap}; |
8 | 9 | use alloy_provider::Provider; |
9 | 10 | use alloy_rpc_types::{ |
10 | 11 | BlockId, BlockNumberOrTag, BlockOverrides, |
@@ -34,8 +35,6 @@ use regex::Regex; |
34 | 35 | use revm::context::TransactionType; |
35 | 36 | use std::{str::FromStr, sync::LazyLock}; |
36 | 37 |
|
37 | | -use super::run::fetch_contracts_bytecode_from_trace; |
38 | | - |
39 | 38 | // matches override pattern <address>:<slot>:<value> |
40 | 39 | // e.g. 0x123:0x1:0x1234 |
41 | 40 | static OVERRIDE_PATTERN: LazyLock<Regex> = |
@@ -131,29 +130,29 @@ pub struct CallArgs { |
131 | 130 | #[arg(long, visible_alias = "la")] |
132 | 131 | pub with_local_artifacts: bool, |
133 | 132 |
|
134 | | - /// Override the balance of an account. |
135 | | - /// Format: address:balance |
136 | | - #[arg(long = "override-balance", value_name = "ADDRESS:BALANCE")] |
| 133 | + /// Override the accounts balance. |
| 134 | + /// Format: "address:balance,address:balance" |
| 135 | + #[arg(long = "override-balance", value_name = "ADDRESS:BALANCE", value_delimiter = ',')] |
137 | 136 | pub balance_overrides: Option<Vec<String>>, |
138 | 137 |
|
139 | | - /// Override the nonce of an account. |
140 | | - /// Format: address:nonce |
141 | | - #[arg(long = "override-nonce", value_name = "ADDRESS:NONCE")] |
| 138 | + /// Override the accounts nonce. |
| 139 | + /// Format: "address:nonce,address:nonce" |
| 140 | + #[arg(long = "override-nonce", value_name = "ADDRESS:NONCE", value_delimiter = ',')] |
142 | 141 | pub nonce_overrides: Option<Vec<String>>, |
143 | 142 |
|
144 | | - /// Override the code of an account. |
145 | | - /// Format: address:code |
146 | | - #[arg(long = "override-code", value_name = "ADDRESS:CODE")] |
| 143 | + /// Override the accounts code. |
| 144 | + /// Format: "address:code,address:code" |
| 145 | + #[arg(long = "override-code", value_name = "ADDRESS:CODE", value_delimiter = ',')] |
147 | 146 | pub code_overrides: Option<Vec<String>>, |
148 | 147 |
|
149 | | - /// Override the state of an account. |
150 | | - /// Format: address:slot:value |
151 | | - #[arg(long = "override-state", value_name = "ADDRESS:SLOT:VALUE")] |
| 148 | + /// Override the accounts state and replace the current state entirely with the new one. |
| 149 | + /// Format: "address:slot:value,address:slot:value" |
| 150 | + #[arg(long = "override-state", value_name = "ADDRESS:SLOT:VALUE", value_delimiter = ',')] |
152 | 151 | pub state_overrides: Option<Vec<String>>, |
153 | 152 |
|
154 | | - /// Override the state diff of an account. |
155 | | - /// Format: address:slot:value |
156 | | - #[arg(long = "override-state-diff", value_name = "ADDRESS:SLOT:VALUE")] |
| 153 | + /// Override the accounts state specific slots and preserve the rest of the state. |
| 154 | + /// Format: "address:slot:value,address:slot:value" |
| 155 | + #[arg(long = "override-state-diff", value_name = "ADDRESS:SLOT:VALUE", value_delimiter = ',')] |
157 | 156 | pub state_diff_overrides: Option<Vec<String>>, |
158 | 157 |
|
159 | 158 | /// Override the block timestamp. |
@@ -393,18 +392,29 @@ impl CallArgs { |
393 | 392 | state_overrides_builder.with_code(addr.parse()?, Bytes::from_str(code_str)?); |
394 | 393 | } |
395 | 394 |
|
396 | | - // Parse state overrides |
397 | | - for override_str in self.state_overrides.iter().flatten() { |
398 | | - let (addr, slot, value) = address_slot_value_override(override_str)?; |
399 | | - state_overrides_builder = |
400 | | - state_overrides_builder.with_state(addr, [(slot.into(), value.into())]); |
| 395 | + type StateOverrides = HashMap<Address, HashMap<B256, B256>>; |
| 396 | + let parse_state_overrides = |
| 397 | + |overrides: &Option<Vec<String>>| -> Result<StateOverrides, eyre::Report> { |
| 398 | + let mut state_overrides: StateOverrides = StateOverrides::default(); |
| 399 | + |
| 400 | + overrides.iter().flatten().try_for_each(|s| -> Result<(), eyre::Report> { |
| 401 | + let (addr, slot, value) = address_slot_value_override(s)?; |
| 402 | + state_overrides.entry(addr).or_default().insert(slot.into(), value.into()); |
| 403 | + Ok(()) |
| 404 | + })?; |
| 405 | + |
| 406 | + Ok(state_overrides) |
| 407 | + }; |
| 408 | + |
| 409 | + // Parse and apply state overrides |
| 410 | + for (addr, entries) in parse_state_overrides(&self.state_overrides)? { |
| 411 | + state_overrides_builder = state_overrides_builder.with_state(addr, entries.into_iter()); |
401 | 412 | } |
402 | 413 |
|
403 | | - // Parse state diff overrides |
404 | | - for override_str in self.state_diff_overrides.iter().flatten() { |
405 | | - let (addr, slot, value) = address_slot_value_override(override_str)?; |
| 414 | + // Parse and apply state diff overrides |
| 415 | + for (addr, entries) in parse_state_overrides(&self.state_diff_overrides)? { |
406 | 416 | state_overrides_builder = |
407 | | - state_overrides_builder.with_state_diff(addr, [(slot.into(), value.into())]); |
| 417 | + state_overrides_builder.with_state_diff(addr, entries.into_iter()) |
408 | 418 | } |
409 | 419 |
|
410 | 420 | Ok(Some(state_overrides_builder.build())) |
|
0 commit comments