Skip to main content

revive_strategy/tracing/
expect_create.rs

1use alloy_primitives::{B256, Bytes, U256 as RU256, hex};
2use foundry_cheatcodes::ExpectedCreate;
3use foundry_compilers::resolc::dual_compiled_contracts::DualCompiledContracts;
4use itertools::Itertools;
5use polkadot_sdk::{
6    pallet_revive::{AccountInfo, Code, Pallet, tracing::Tracing},
7    sp_core::{H160, U256},
8};
9use revive_env::Runtime;
10use revm::context::CreateScheme;
11
12#[derive(Debug, Clone)]
13#[allow(dead_code)]
14pub(crate) struct Create {
15    pub(crate) addr: H160,
16    pub(crate) from: H160,
17    scheme: CreateScheme,
18}
19
20#[derive(Debug)]
21pub(crate) struct CreateTracer {
22    pub calls: Vec<polkadot_sdk::sp_core::H160>,
23    pub is_create: Option<(Code, Option<[u8; 32]>)>,
24    pub call_types: Vec<Type>,
25    pub data: Vec<ExpectedCreate>,
26    creates: Vec<Create>,
27}
28
29#[derive(Debug)]
30pub enum Type {
31    Create { salt: Option<[u8; 32]> },
32    Rest,
33}
34
35impl CreateTracer {
36    fn current_addr(&self) -> H160 {
37        self.calls.last().copied().unwrap_or_default()
38    }
39    pub fn new(data: Vec<ExpectedCreate>) -> Self {
40        Self { data, creates: vec![], calls: vec![], is_create: None, call_types: vec![] }
41    }
42}
43
44impl Tracing for CreateTracer {
45    fn instantiate_code(&mut self, code: &Code, salt: Option<&[u8; 32]>) {
46        self.is_create = Some((code.to_owned(), salt.to_owned().copied()));
47    }
48
49    fn enter_child_span(
50        &mut self,
51        _from: H160,
52        to: H160,
53        _is_delegate_call: Option<H160>,
54        _is_read_only: bool,
55        _value: U256,
56        _input: &[u8],
57        _gas: U256,
58    ) {
59        self.call_types.push(if let Some((_, salt)) = self.is_create.take() {
60            Type::Create { salt }
61        } else {
62            Type::Rest
63        });
64        if self.calls.is_empty() {
65            self.calls.push(_from);
66        }
67        self.calls.push(if _is_delegate_call.is_some() { self.current_addr() } else { to });
68    }
69
70    fn exit_child_span(
71        &mut self,
72        _output: &polkadot_sdk::pallet_revive::ExecReturnValue,
73        _gas_left: U256,
74    ) {
75        let addr = self.calls.pop().unwrap_or_default();
76
77        let typ = self.call_types.pop();
78        if typ.as_ref().is_some_and(|x| matches!(x, Type::Create { .. })) {
79            self.is_create = None;
80        }
81        if let Some(Type::Create { salt, .. }) = typ {
82            let mut create =
83                Create { addr, from: self.current_addr(), scheme: CreateScheme::Create };
84            if let Some(salt) = salt {
85                let salt = RU256::from_be_bytes(
86                    B256::from_slice(&alloy_primitives::keccak256::<&[u8]>(&salt)[..]).0,
87                );
88                create.scheme = CreateScheme::Create2 { salt };
89            };
90            self.creates.push(create);
91        }
92    }
93}
94
95impl CreateTracer {
96    pub fn finalize(
97        &mut self,
98        dual_compiled_contracts: &DualCompiledContracts,
99    ) -> Vec<ExpectedCreate> {
100        let mut created = vec![];
101        for c in self.creates.iter().cloned() {
102            let Some(info) = AccountInfo::<Runtime>::load_contract(&c.addr) else {
103                continue;
104            };
105            let hash = hex::encode(info.code_hash);
106            let code = Pallet::<Runtime>::code(&c.addr);
107            let cb = Bytes::from(code.clone()).0;
108            let bytecode_result = dual_compiled_contracts
109                .find_by_evm_bytecode_hash(hash.clone())
110                .and_then(|(_, contract)| contract.evm_deployed_bytecode.as_bytes())
111                .or_else(|| {
112                    dual_compiled_contracts
113                        .find_bytecode(&cb)
114                        .and_then(|f| f.contract().evm_deployed_bytecode.as_bytes())
115                })
116                .cloned()
117                .unwrap_or_else(|| Bytes::from(code.clone()));
118            created.push(ExpectedCreate {
119                deployer: c.from.0.into(),
120                bytecode: bytecode_result,
121                create_scheme: c.scheme.into(),
122            })
123        }
124        for c in created {
125            if let Some((index, _)) = self.data.iter().find_position(|expected_create| {
126                expected_create.deployer == c.deployer
127                    && expected_create.create_scheme.eq(c.create_scheme.clone())
128                    && expected_create.bytecode == c.bytecode
129            }) {
130                self.data.swap_remove(index);
131            }
132        }
133        self.data.clone()
134    }
135}