tycho_simulation/evm/protocol/vm/
decoder.rs

1use std::{
2    collections::{HashMap, HashSet},
3    str::FromStr,
4};
5
6use alloy::primitives::{Address, U256};
7use revm::state::Bytecode;
8use tycho_client::feed::{synchronizer::ComponentWithState, BlockHeader};
9use tycho_common::{models::token::Token, simulation::errors::SimulationError, Bytes};
10
11use super::{state::EVMPoolState, state_builder::EVMPoolStateBuilder};
12use crate::{
13    evm::{
14        engine_db::{tycho_db::PreCachedDB, SHARED_TYCHO_DB},
15        protocol::vm::{constants::get_adapter_file, utils::json_deserialize_address_list},
16    },
17    protocol::{
18        errors::InvalidSnapshotError,
19        models::{DecoderContext, TryFromWithBlock},
20    },
21};
22
23impl TryFromWithBlock<ComponentWithState, BlockHeader> for EVMPoolState<PreCachedDB> {
24    type Error = InvalidSnapshotError;
25
26    /// Decodes a `ComponentWithState`, block `BlockHeader` and HashMap of all available tokens into
27    /// an `EVMPoolState`.
28    ///
29    /// Errors with a `InvalidSnapshotError`.
30    #[allow(deprecated)]
31    async fn try_from_with_header(
32        snapshot: ComponentWithState,
33        _block: BlockHeader,
34        account_balances: &HashMap<Bytes, HashMap<Bytes, Bytes>>,
35        all_tokens: &HashMap<Bytes, Token>,
36        decoder_context: &DecoderContext,
37    ) -> Result<Self, Self::Error> {
38        let id = snapshot.component.id.clone();
39        let tokens = snapshot.component.tokens.clone();
40
41        // Decode involved contracts
42        let mut stateless_contracts = HashMap::new();
43        let mut index = 0;
44
45        loop {
46            let address_key = format!("stateless_contract_addr_{index}");
47            if let Some(encoded_address_bytes) = snapshot
48                .state
49                .attributes
50                .get(&address_key)
51            {
52                let encoded_address = hex::encode(encoded_address_bytes);
53                // Stateless contracts address are UTF-8 encoded
54                let address_hex = encoded_address
55                    .strip_prefix("0x")
56                    .unwrap_or(&encoded_address);
57
58                let decoded = match hex::decode(address_hex) {
59                    Ok(decoded_bytes) => match String::from_utf8(decoded_bytes) {
60                        Ok(decoded_string) => decoded_string,
61                        Err(_) => continue,
62                    },
63                    Err(_) => continue,
64                };
65
66                let code_key = format!("stateless_contract_code_{index}");
67                let code = snapshot
68                    .state
69                    .attributes
70                    .get(&code_key)
71                    .map(|value| value.to_vec());
72
73                stateless_contracts.insert(decoded, code);
74                index += 1;
75            } else {
76                break;
77            }
78        }
79        let involved_contracts = snapshot
80            .component
81            .contract_ids
82            .iter()
83            .map(|bytes: &Bytes| Address::from_slice(bytes.as_ref()))
84            .collect::<HashSet<Address>>();
85
86        let potential_rebase_tokens: HashSet<Address> = if let Some(bytes) = snapshot
87            .component
88            .static_attributes
89            .get("rebase_tokens")
90        {
91            if let Ok(vecs) = json_deserialize_address_list(bytes) {
92                vecs.into_iter()
93                    .map(|addr| Address::from_slice(&addr))
94                    .collect()
95            } else {
96                HashSet::new()
97            }
98        } else {
99            HashSet::new()
100        };
101
102        // Decode balances
103        let balance_owner = snapshot
104            .state
105            .attributes
106            .get("balance_owner")
107            .map(|owner| Address::from_slice(owner.as_ref()));
108        let component_balances = snapshot
109            .state
110            .balances
111            .iter()
112            .map(|(k, v)| (Address::from_slice(k), U256::from_be_slice(v)))
113            .collect::<HashMap<_, _>>();
114        let account_balances = account_balances
115            .iter()
116            .filter(|(k, _)| involved_contracts.contains(&Address::from_slice(k)))
117            .map(|(k, v)| {
118                let addr = Address::from_slice(k);
119                let balances = v
120                    .iter()
121                    .map(|(k, v)| (Address::from_slice(k), U256::from_be_slice(v)))
122                    .collect();
123                (addr, balances)
124            })
125            .collect::<HashMap<_, _>>();
126
127        let manual_updates = snapshot
128            .component
129            .static_attributes
130            .contains_key("manual_updates");
131
132        let protocol_name = snapshot
133            .component
134            .protocol_system
135            .strip_prefix("vm:")
136            .unwrap_or({
137                snapshot
138                    .component
139                    .protocol_system
140                    .as_str()
141            });
142        let adapter_bytecode;
143        if let Some(adapter_bytecode_path) = &decoder_context.adapter_path {
144            let bytecode_bytes = std::fs::read(adapter_bytecode_path).map_err(|e| {
145                SimulationError::FatalError(format!(
146                    "Failed to read adapter bytecode from {adapter_bytecode_path}: {e}"
147                ))
148            })?;
149            adapter_bytecode = Bytecode::new_raw(bytecode_bytes.into());
150        } else {
151            adapter_bytecode = Bytecode::new_raw(get_adapter_file(protocol_name)?.into());
152        }
153        let adapter_contract_address = Address::from_str(&format!(
154            "{hex_protocol_name:0>40}",
155            hex_protocol_name = hex::encode(protocol_name)
156        ))
157        .map_err(|_| {
158            InvalidSnapshotError::ValueError(
159                "Error converting protocol name to address".to_string(),
160            )
161        })?;
162        let mut vm_traces = false;
163        if let Some(trace) = &decoder_context.vm_traces {
164            vm_traces = *trace;
165        }
166        let mut pool_state_builder =
167            EVMPoolStateBuilder::new(id.clone(), tokens.clone(), adapter_contract_address)
168                .balances(component_balances)
169                .disable_overwrite_tokens(potential_rebase_tokens)
170                .account_balances(account_balances)
171                .adapter_contract_bytecode(adapter_bytecode)
172                .involved_contracts(involved_contracts)
173                .stateless_contracts(stateless_contracts)
174                .manual_updates(manual_updates)
175                .trace(vm_traces);
176
177        if let Some(balance_owner) = balance_owner {
178            pool_state_builder = pool_state_builder.balance_owner(balance_owner)
179        };
180
181        let mut pool_state = pool_state_builder
182            .build(SHARED_TYCHO_DB.clone())
183            .await
184            .map_err(InvalidSnapshotError::VMError)?;
185
186        pool_state.set_spot_prices(all_tokens)?;
187
188        Ok(pool_state)
189    }
190}
191
192#[cfg(test)]
193mod tests {
194    use std::{collections::HashSet, fs, path::Path};
195
196    use chrono::DateTime;
197    use revm::{primitives::KECCAK_EMPTY, state::AccountInfo};
198    use serde_json::Value;
199    use tycho_common::dto::{Chain, ChangeType, ProtocolComponent, ResponseProtocolState};
200
201    use super::*;
202    use crate::evm::{
203        engine_db::{create_engine, engine_db_interface::EngineDatabaseInterface},
204        protocol::vm::constants::{BALANCER_V2, CURVE},
205        tycho_models::AccountUpdate,
206    };
207
208    #[test]
209    fn test_to_adapter_file_name() {
210        assert_eq!(get_adapter_file("balancer_v2").unwrap(), BALANCER_V2);
211        assert_eq!(get_adapter_file("curve").unwrap(), CURVE);
212    }
213
214    fn vm_component() -> ProtocolComponent {
215        let creation_time = DateTime::from_timestamp(1622526000, 0)
216            .unwrap()
217            .naive_utc(); //Sample timestamp
218
219        let mut static_attributes: HashMap<String, Bytes> = HashMap::new();
220        static_attributes.insert("manual_updates".to_string(), Bytes::from_str("0x01").unwrap());
221
222        let dai_addr = Bytes::from_str("0x6b175474e89094c44da98b954eedeac495271d0f").unwrap();
223        let bal_addr = Bytes::from_str("0xba100000625a3754423978a60c9317c58a424e3d").unwrap();
224        let tokens = vec![dai_addr, bal_addr];
225
226        ProtocolComponent {
227            id: "0x4626d81b3a1711beb79f4cecff2413886d461677000200000000000000000011".to_string(),
228            protocol_system: "vm:balancer_v2".to_string(),
229            protocol_type_name: "balancer_v2_pool".to_string(),
230            chain: Chain::Ethereum,
231            tokens,
232            contract_ids: vec![
233                Bytes::from_str("0xBA12222222228d8Ba445958a75a0704d566BF2C8").unwrap()
234            ],
235            static_attributes,
236            change: ChangeType::Creation,
237            creation_tx: Bytes::from_str("0x0000").unwrap(),
238            created_at: creation_time,
239        }
240    }
241
242    fn header() -> BlockHeader {
243        BlockHeader {
244            number: 1,
245            hash: Bytes::from(vec![0; 32]),
246            parent_hash: Bytes::from(vec![0; 32]),
247            revert: false,
248            timestamp: 1,
249        }
250    }
251
252    fn load_balancer_account_data() -> Vec<AccountUpdate> {
253        let project_root = env!("CARGO_MANIFEST_DIR");
254        let asset_path =
255            Path::new(project_root).join("tests/assets/decoder/balancer_v2_snapshot.json");
256        let json_data = fs::read_to_string(asset_path).expect("Failed to read test asset");
257        let data: Value = serde_json::from_str(&json_data).expect("Failed to parse JSON");
258
259        let accounts: Vec<AccountUpdate> = serde_json::from_value(data["accounts"].clone())
260            .expect("Expected accounts to match AccountUpdate structure");
261        accounts
262    }
263
264    #[allow(deprecated)]
265    #[tokio::test]
266    async fn test_try_from_with_header() {
267        let attributes: HashMap<String, Bytes> = vec![
268            (
269                "balance_owner".to_string(),
270                Bytes::from_str("0xBA12222222228d8Ba445958a75a0704d566BF2C8").unwrap(),
271            ),
272            ("reserve1".to_string(), Bytes::from(200_u64.to_le_bytes().to_vec())),
273        ]
274        .into_iter()
275        .collect();
276        let tokens = [
277            Token::new(
278                &Bytes::from_str("0x6b175474e89094c44da98b954eedeac495271d0f").unwrap(),
279                "DAI",
280                18,
281                0,
282                &[Some(10_000)],
283                tycho_common::models::Chain::Ethereum,
284                100,
285            ),
286            Token::new(
287                &Bytes::from_str("0xba100000625a3754423978a60c9317c58a424e3d").unwrap(),
288                "BAL",
289                18,
290                0,
291                &[Some(10_000)],
292                tycho_common::models::Chain::Ethereum,
293                100,
294            ),
295        ]
296        .into_iter()
297        .map(|t| (t.address.clone(), t))
298        .collect::<HashMap<_, _>>();
299        let snapshot = ComponentWithState {
300            state: ResponseProtocolState {
301                component_id: "0x4626d81b3a1711beb79f4cecff2413886d461677000200000000000000000011"
302                    .to_owned(),
303                attributes,
304                balances: HashMap::new(),
305            },
306            component: vm_component(),
307            component_tvl: None,
308            entrypoints: Vec::new(),
309        };
310        // Initialize engine with balancer storage
311        let block = header();
312        let accounts = load_balancer_account_data();
313        let db = SHARED_TYCHO_DB.clone();
314        let engine = create_engine(db.clone(), false).unwrap();
315        for account in accounts.clone() {
316            engine
317                .state
318                .init_account(
319                    account.address,
320                    AccountInfo {
321                        balance: account.balance.unwrap_or_default(),
322                        nonce: 0u64,
323                        code_hash: KECCAK_EMPTY,
324                        code: account
325                            .code
326                            .clone()
327                            .map(|arg0: Vec<u8>| Bytecode::new_raw(arg0.into())),
328                    },
329                    None,
330                    false,
331                )
332                .expect("Failed to init account");
333        }
334        db.update(accounts, Some(block))
335            .unwrap();
336        let account_balances = HashMap::from([(
337            Bytes::from("0xBA12222222228d8Ba445958a75a0704d566BF2C8"),
338            HashMap::from([
339                (
340                    Bytes::from("0x6b175474e89094c44da98b954eedeac495271d0f"),
341                    Bytes::from(100_u64.to_le_bytes().to_vec()),
342                ),
343                (
344                    Bytes::from("0xba100000625a3754423978a60c9317c58a424e3d"),
345                    Bytes::from(100_u64.to_le_bytes().to_vec()),
346                ),
347            ]),
348        )]);
349
350        let decoder_context = DecoderContext::new();
351        let res = EVMPoolState::try_from_with_header(
352            snapshot,
353            header(),
354            &account_balances,
355            &tokens,
356            &decoder_context,
357        )
358        .await
359        .unwrap();
360
361        let res_pool = res;
362
363        assert_eq!(
364            res_pool.get_balance_owner(),
365            Some(Address::from_str("0xBA12222222228d8Ba445958a75a0704d566BF2C8").unwrap())
366        );
367        let mut exp_involved_contracts = HashSet::new();
368        exp_involved_contracts
369            .insert(Address::from_str("0xBA12222222228d8Ba445958a75a0704d566BF2C8").unwrap());
370        assert_eq!(res_pool.get_involved_contracts(), exp_involved_contracts);
371        assert!(res_pool.get_manual_updates());
372    }
373}