kabu_types_blockchain/
state_update.rs

1use alloy_primitives::{Address, TxHash};
2use alloy_provider::ext::DebugApi;
3use alloy_provider::{Network, Provider};
4use alloy_rpc_types::{BlockId, TransactionRequest};
5use alloy_rpc_types_trace::common::TraceResult;
6use alloy_rpc_types_trace::geth::GethDebugBuiltInTracerType::PreStateTracer;
7use alloy_rpc_types_trace::geth::GethDebugTracerType::BuiltInTracer;
8use alloy_rpc_types_trace::geth::{
9    AccountState, GethDebugBuiltInTracerType, GethDebugTracerConfig, GethDebugTracerType, GethDebugTracingCallOptions,
10    GethDebugTracingOptions, GethDefaultTracingOptions, GethTrace, PreStateConfig, PreStateFrame,
11};
12use eyre::Result;
13use lazy_static::lazy_static;
14use std::collections::BTreeMap;
15use tracing::{debug, trace};
16
17use kabu_node_debug_provider::DebugProviderExt;
18
19pub type GethStateUpdate = BTreeMap<Address, AccountState>;
20
21pub type GethStateUpdateVec = Vec<BTreeMap<Address, AccountState>>;
22
23lazy_static! {
24    pub static ref TRACING_OPTS: GethDebugTracingOptions = GethDebugTracingOptions {
25        tracer: Some(GethDebugTracerType::BuiltInTracer(GethDebugBuiltInTracerType::PreStateTracer,)),
26        tracer_config: GethDebugTracerConfig::default(),
27        config: GethDefaultTracingOptions::default().disable_storage().disable_stack().disable_memory().disable_return_data(),
28        timeout: None,
29    };
30    pub static ref TRACING_CALL_OPTS: GethDebugTracingCallOptions =
31        GethDebugTracingCallOptions { tracing_options: TRACING_OPTS.clone(), state_overrides: None, block_overrides: None, tx_index: None };
32}
33
34pub fn get_touched_addresses(state_update: &GethStateUpdate) -> Vec<Address> {
35    let mut ret: Vec<Address> = Vec::new();
36
37    for (address, state) in state_update.iter() {
38        if !state.storage.is_empty() {
39            ret.push(*address)
40        }
41    }
42
43    ret
44}
45
46pub fn debug_log_geth_state_update(state_update: &GethStateUpdate) {
47    for (address, state) in state_update {
48        debug!("{} nonce {:?} balance {:?} is_code {}", address, state.nonce, state.balance, state.code.is_some())
49    }
50}
51
52pub async fn debug_trace_block<N: Network, P: Provider<N> + DebugProviderExt<N>>(
53    client: P,
54    block_id: BlockId,
55    diff_mode: bool,
56) -> eyre::Result<(GethStateUpdateVec, GethStateUpdateVec)> {
57    let tracer_opts = GethDebugTracingOptions { config: GethDefaultTracingOptions::default(), ..GethDebugTracingOptions::default() }
58        .with_tracer(BuiltInTracer(PreStateTracer))
59        .with_prestate_config(PreStateConfig { diff_mode: Some(diff_mode), disable_code: Some(false), disable_storage: Some(false) });
60
61    let trace_result_vec = match block_id {
62        BlockId::Number(block_number) => client.geth_debug_trace_block_by_number(block_number, tracer_opts).await?,
63        BlockId::Hash(rpc_block_hash) => {
64            //client.debug_trace_block_by_number(BlockNumber::from(19776525u32), tracer_opts).await?
65            client.geth_debug_trace_block_by_hash(rpc_block_hash.block_hash, tracer_opts).await?
66        }
67    };
68
69    trace!("block trace {}", trace_result_vec.len());
70
71    let mut pre: GethStateUpdateVec = Default::default();
72    let mut post: GethStateUpdateVec = Default::default();
73
74    for trace_result in trace_result_vec.into_iter() {
75        if let TraceResult::Success { result, .. } = trace_result {
76            match result {
77                GethTrace::PreStateTracer(geth_trace_frame) => match geth_trace_frame {
78                    PreStateFrame::Diff(diff_frame) => {
79                        pre.push(diff_frame.pre);
80                        post.push(diff_frame.post);
81                    }
82                    PreStateFrame::Default(diff_frame) => {
83                        pre.push(diff_frame.0.into_iter().collect());
84                    }
85                },
86                _ => {
87                    return Err(eyre::eyre!("TRACE_RESULT_FAILED"));
88                }
89            }
90        }
91    }
92    Ok((pre, post))
93}
94
95async fn debug_trace_call<N: Network, C: DebugProviderExt<N>, TR: Into<N::TransactionRequest> + Send + Sync>(
96    client: C,
97    req: TR,
98    block: BlockId,
99    opts: Option<GethDebugTracingCallOptions>,
100    diff_mode: bool,
101) -> Result<(GethStateUpdate, GethStateUpdate)> {
102    let tracer_opts = GethDebugTracingOptions { config: GethDefaultTracingOptions::default(), ..GethDebugTracingOptions::default() }
103        .with_tracer(BuiltInTracer(PreStateTracer))
104        .with_prestate_config(PreStateConfig { diff_mode: Some(diff_mode), disable_code: Some(false), disable_storage: Some(false) });
105
106    let tracer_call_opts = GethDebugTracingCallOptions {
107        tracing_options: tracer_opts.clone(),
108        state_overrides: opts.as_ref().and_then(|x| x.state_overrides.clone()),
109        block_overrides: opts.as_ref().and_then(|x| x.block_overrides.clone()),
110        tx_index: opts.and_then(|x| x.tx_index),
111    };
112
113    let trace_result = client.geth_debug_trace_call(req.into(), block, tracer_call_opts.clone()).await?;
114    trace!(
115        "{} {} {:?} {:?}",
116        tracer_opts.config.is_stack_enabled(),
117        tracer_opts.config.is_storage_enabled(),
118        tracer_call_opts.clone(),
119        trace_result
120    );
121
122    match trace_result {
123        GethTrace::PreStateTracer(geth_trace_frame) => match geth_trace_frame {
124            PreStateFrame::Diff(diff_frame) => Ok((diff_frame.pre, diff_frame.post)),
125            PreStateFrame::Default(diff_frame) => Ok((diff_frame.0, Default::default())),
126        },
127        _ => Err(eyre::eyre!("TRACE_RESULT_FAILED")),
128    }
129}
130
131pub async fn debug_trace_call_pre_state<N: Network, C: DebugProviderExt<N>, TR: Into<N::TransactionRequest> + Send + Sync>(
132    client: C,
133    req: TR,
134    block: BlockId,
135    opts: Option<GethDebugTracingCallOptions>,
136) -> eyre::Result<GethStateUpdate> {
137    Ok(debug_trace_call(client, req, block, opts, false).await?.0)
138}
139
140pub async fn debug_trace_call_post_state<
141    N: Network<TransactionRequest = TransactionRequest>,
142    C: DebugProviderExt<N>,
143    TR: Into<TransactionRequest> + Send + Sync,
144>(
145    client: C,
146    req: TR,
147    block: BlockId,
148    opts: Option<GethDebugTracingCallOptions>,
149) -> eyre::Result<GethStateUpdate>
150where
151    <N as Network>::TransactionRequest: From<TR>,
152{
153    Ok(debug_trace_call(client, req, block, opts, true).await?.1)
154}
155
156pub async fn debug_trace_call_diff<N: Network, C: DebugProviderExt<N>, TR: Into<N::TransactionRequest> + Send + Sync>(
157    client: C,
158    req: TR,
159    block: BlockId,
160    call_opts: Option<GethDebugTracingCallOptions>,
161) -> eyre::Result<(GethStateUpdate, GethStateUpdate)> {
162    debug_trace_call(client, req, block, call_opts, true).await
163}
164
165pub async fn debug_trace_transaction<N: Network, P: Provider<N> + DebugApi<N>>(
166    client: P,
167    req: TxHash,
168    diff_mode: bool,
169) -> Result<(GethStateUpdate, GethStateUpdate)> {
170    let tracer_opts = GethDebugTracingOptions { config: GethDefaultTracingOptions::default(), ..GethDebugTracingOptions::default() }
171        .with_tracer(BuiltInTracer(PreStateTracer))
172        .with_prestate_config(PreStateConfig { diff_mode: Some(diff_mode), disable_code: Some(false), disable_storage: Some(false) });
173
174    let trace_result = client.debug_trace_transaction(req, tracer_opts).await?;
175    trace!("{:?}", trace_result);
176
177    match trace_result {
178        GethTrace::PreStateTracer(geth_trace_frame) => match geth_trace_frame {
179            PreStateFrame::Diff(diff_frame) => Ok((diff_frame.pre.into_iter().collect(), diff_frame.post.into_iter().collect())),
180            PreStateFrame::Default(diff_frame) => Ok((diff_frame.0.into_iter().collect(), Default::default())),
181        },
182        _ => Err(eyre::eyre!("TRACE_RESULT_FAILED")),
183    }
184}
185
186#[cfg(test)]
187mod test {
188    use super::*;
189    use alloy_primitives::map::B256HashMap;
190    use alloy_primitives::{B256, U256};
191
192    use alloy_provider::ProviderBuilder;
193    use alloy_rpc_client::{ClientBuilder, WsConnect};
194    use alloy_rpc_types::state::{AccountOverride, StateOverride};
195    use env_logger::Env as EnvLog;
196    use tracing::{debug, error};
197
198    #[tokio::test]
199    async fn test_debug_block() -> Result<()> {
200        dotenvy::from_filename(".env.test").ok();
201        let node_url = std::env::var("MAINNET_WS")?;
202
203        let _ = env_logger::try_init_from_env(EnvLog::default().default_filter_or("info,tokio_tungstenite=off,tungstenite=off"));
204        let node_url = url::Url::parse(node_url.as_str())?;
205
206        let ws_connect = WsConnect::new(node_url);
207        let client = ClientBuilder::default().ws(ws_connect).await?;
208
209        let client = ProviderBuilder::new().disable_recommended_fillers().connect_client(client);
210
211        let blocknumber = client.get_block_number().await?;
212        let _block = client.get_block_by_number(blocknumber.into()).await?.unwrap();
213
214        let _ret = debug_trace_block(client, BlockId::Number(blocknumber.into()), true).await?;
215
216        Ok(())
217    }
218
219    #[test]
220    fn test_encode_override() {
221        let mut state_override: StateOverride = StateOverride::default();
222        let address = Address::default();
223        let mut account_override: AccountOverride = AccountOverride::default();
224        let mut state_update_hashmap: B256HashMap<B256> = B256HashMap::default();
225        state_update_hashmap.insert(B256::from(U256::from(1)), B256::from(U256::from(3)));
226        account_override.state_diff = Some(state_update_hashmap);
227
228        state_override.insert(address, account_override);
229
230        match serde_json::to_string_pretty(&state_override) {
231            Ok(data) => {
232                debug!("{}", data);
233            }
234            Err(e) => {
235                error!("{}", e);
236                panic!("DESERIALIZATION_ERROR");
237            }
238        }
239    }
240}
kabu_types_blockchain/state_update.rs

kabu_types_blockchain/
state_update.rs
