kabu_types_entities/
required_state.rs

1use alloy_network::Network;
2use alloy_primitives::{Address, BlockNumber, Bytes, U256};
3use alloy_provider::Provider;
4use alloy_rpc_types::{BlockId, BlockNumberOrTag};
5use alloy_rpc_types_trace::geth::AccountState;
6use eyre::{eyre, Result};
7use std::collections::BTreeMap;
8use std::fmt::Debug;
9use std::marker::PhantomData;
10use tracing::{error, trace};
11
12use kabu_node_debug_provider::DebugProviderExt;
13use kabu_types_blockchain::{debug_trace_call_pre_state, GethStateUpdate, GethStateUpdateVec, KabuDataTypesEVM};
14use kabu_types_blockchain::{KabuDataTypes, KabuDataTypesEthereum, KabuTransactionRequest};
15
16#[derive(Clone, Debug, Default)]
17pub struct RequiredState {
18    calls: Vec<(Address, Bytes)>,
19    slots: Vec<(Address, U256)>,
20    empty_slots: Vec<(Address, U256)>,
21}
22
23impl RequiredState {
24    pub fn new() -> Self {
25        Self::default()
26    }
27
28    pub fn add_call<T: Into<Bytes>, A: Into<Address>>(&mut self, to: A, call_data: T) -> &mut Self {
29        self.calls.push((to.into(), call_data.into()));
30        self
31    }
32    pub fn add_slot<A: Into<Address>>(&mut self, address: A, slot: U256) -> &mut Self {
33        self.slots.push((address.into(), slot));
34        self
35    }
36
37    pub fn add_empty_slot<A: Into<Address>>(&mut self, address: A, slot: U256) -> &mut Self {
38        self.empty_slots.push((address.into(), slot));
39        self
40    }
41
42    pub fn add_empty_slot_range<A: Into<Address> + Clone>(&mut self, address: A, start_slot: U256, size: usize) -> &mut Self {
43        let mut cur_slot = start_slot;
44        for _ in 0..size {
45            self.add_empty_slot(address.clone(), cur_slot);
46            cur_slot += U256::from(1);
47        }
48        self
49    }
50
51    pub fn add_slot_range<A: Into<Address> + Clone>(&mut self, address: A, start_slot: U256, size: usize) -> &mut Self {
52        let mut cur_slot = start_slot;
53        for _ in 0..size {
54            self.add_slot(address.clone(), cur_slot);
55            cur_slot += U256::from(1);
56        }
57        self
58    }
59}
60
61pub struct RequiredStateReader<LDT: KabuDataTypes = KabuDataTypesEthereum> {
62    _ldt: PhantomData<LDT>,
63}
64
65impl<LDT> RequiredStateReader<LDT>
66where
67    LDT: KabuDataTypesEVM,
68{
69    pub async fn fetch_calls_and_slots<
70        N: Network<TransactionRequest = LDT::TransactionRequest>,
71        C: DebugProviderExt<N> + Provider<N> + Clone + 'static,
72    >(
73        client: C,
74        required_state: RequiredState,
75        block_number: Option<BlockNumber>,
76    ) -> Result<LDT::StateUpdate> {
77        let block_id = if block_number.is_none() {
78            BlockId::Number(BlockNumberOrTag::Latest)
79        } else {
80            BlockId::Number(BlockNumberOrTag::Number(block_number.unwrap_or_default()))
81        };
82
83        let mut ret: GethStateUpdate = GethStateUpdate::new();
84        for req in required_state.calls.into_iter() {
85            let to: Address = req.0;
86            let req: LDT::TransactionRequest = LDT::TransactionRequest::build_call(to, req.1);
87
88            let call_result = debug_trace_call_pre_state(client.clone(), req, block_id, None).await;
89            trace!("trace_call_result: {:?}", call_result);
90            match call_result {
91                Ok(update) => {
92                    for (address, account_state) in update.into_iter() {
93                        let entry = ret.entry(address).or_insert(account_state.clone());
94                        for (slot, value) in account_state.storage.clone().into_iter() {
95                            entry.storage.insert(slot, value);
96                            trace!(%address, %slot, %value, "Inserting storage");
97                        }
98                    }
99                }
100                Err(e) => {
101                    error!("Contract call failed {} {}", to, e);
102                    return Err(eyre!("CONTRACT_CALL_FAILED"));
103                }
104            }
105        }
106        for (address, slot) in required_state.slots.into_iter() {
107            let value_result = client.get_storage_at(address, slot).block_id(block_id).await;
108            trace!("get_storage_at_result {} slot {} :  {:?}", address, slot, value_result);
109            match value_result {
110                Ok(value) => {
111                    let entry = ret.entry(address).or_default();
112                    entry.storage.insert(slot.into(), value.into());
113                }
114                Err(e) => {
115                    error!("{}", e)
116                }
117            }
118        }
119
120        for (address, slot) in required_state.empty_slots.into_iter() {
121            let value = U256::ZERO;
122
123            let entry = ret.entry(address).or_default();
124            entry.storage.insert(slot.into(), value.into());
125        }
126
127        Ok(ret)
128    }
129}
130
131pub fn accounts_len(state: &BTreeMap<Address, AccountState>) -> (usize, usize) {
132    let accounts = state.len();
133    let storage = state.values().map(|item| item.storage.clone().len()).sum();
134    (accounts, storage)
135}
136
137pub fn accounts_vec_len(state: &GethStateUpdateVec) -> usize {
138    state.iter().map(|item| accounts_len(item).0).sum()
139}
140
141pub fn storage_vec_len(state: &GethStateUpdateVec) -> usize {
142    state.iter().map(|item| accounts_len(item).1).sum()
143}
kabu_types_entities/required_state.rs

kabu_types_entities/
required_state.rs
