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_addresses
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::models::{
200        protocol::{ProtocolComponent, ProtocolComponentState},
201        Chain, ChangeType,
202    };
203
204    use super::*;
205    use crate::evm::{
206        engine_db::{create_engine, engine_db_interface::EngineDatabaseInterface},
207        protocol::vm::constants::{BALANCER_V2, CURVE},
208        tycho_models::AccountUpdate,
209    };
210
211    #[test]
212    fn test_to_adapter_file_name() {
213        assert_eq!(get_adapter_file("balancer_v2").unwrap(), BALANCER_V2);
214        assert_eq!(get_adapter_file("curve").unwrap(), CURVE);
215    }
216
217    fn vm_component() -> ProtocolComponent {
218        let creation_time = DateTime::from_timestamp(1622526000, 0)
219            .unwrap()
220            .naive_utc(); //Sample timestamp
221
222        let mut static_attributes: HashMap<String, Bytes> = HashMap::new();
223        static_attributes.insert("manual_updates".to_string(), Bytes::from_str("0x01").unwrap());
224
225        let dai_addr = Bytes::from_str("0x6b175474e89094c44da98b954eedeac495271d0f").unwrap();
226        let bal_addr = Bytes::from_str("0xba100000625a3754423978a60c9317c58a424e3d").unwrap();
227        let tokens = vec![dai_addr, bal_addr];
228
229        ProtocolComponent {
230            id: "0x4626d81b3a1711beb79f4cecff2413886d461677000200000000000000000011".to_string(),
231            protocol_system: "vm:balancer_v2".to_string(),
232            protocol_type_name: "balancer_v2_pool".to_string(),
233            chain: Chain::Ethereum,
234            tokens,
235            contract_addresses: vec![
236                Bytes::from_str("0xBA12222222228d8Ba445958a75a0704d566BF2C8").unwrap()
237            ],
238            static_attributes,
239            change: ChangeType::Creation,
240            creation_tx: Bytes::from_str("0x0000").unwrap(),
241            created_at: creation_time,
242        }
243    }
244
245    fn load_balancer_account_data() -> Vec<AccountUpdate> {
246        let project_root = env!("CARGO_MANIFEST_DIR");
247        let asset_path =
248            Path::new(project_root).join("tests/assets/decoder/balancer_v2_snapshot.json");
249        let json_data = fs::read_to_string(asset_path).expect("Failed to read test asset");
250        let data: Value = serde_json::from_str(&json_data).expect("Failed to parse JSON");
251
252        let accounts: Vec<AccountUpdate> = serde_json::from_value(data["accounts"].clone())
253            .expect("Expected accounts to match AccountUpdate structure");
254        accounts
255    }
256
257    #[tokio::test]
258    async fn test_try_from_with_header() {
259        let attributes: HashMap<String, Bytes> = vec![
260            (
261                "balance_owner".to_string(),
262                Bytes::from_str("0xBA12222222228d8Ba445958a75a0704d566BF2C8").unwrap(),
263            ),
264            ("reserve1".to_string(), Bytes::from(200_u64.to_le_bytes().to_vec())),
265        ]
266        .into_iter()
267        .collect();
268        let tokens = [
269            Token::new(
270                &Bytes::from_str("0x6b175474e89094c44da98b954eedeac495271d0f").unwrap(),
271                "DAI",
272                18,
273                0,
274                &[Some(10_000)],
275                tycho_common::models::Chain::Ethereum,
276                100,
277            ),
278            Token::new(
279                &Bytes::from_str("0xba100000625a3754423978a60c9317c58a424e3d").unwrap(),
280                "BAL",
281                18,
282                0,
283                &[Some(10_000)],
284                tycho_common::models::Chain::Ethereum,
285                100,
286            ),
287        ]
288        .into_iter()
289        .map(|t| (t.address.clone(), t))
290        .collect::<HashMap<_, _>>();
291        let snapshot = ComponentWithState {
292            state: ProtocolComponentState {
293                component_id: "0x4626d81b3a1711beb79f4cecff2413886d461677000200000000000000000011"
294                    .to_owned(),
295                attributes,
296                balances: HashMap::new(),
297            },
298            component: vm_component(),
299            component_tvl: None,
300            entrypoints: Vec::new(),
301        };
302        // Initialize engine with balancer storage
303        let block = BlockHeader::default();
304        let accounts = load_balancer_account_data();
305        let db = SHARED_TYCHO_DB.clone();
306        let engine = create_engine(db.clone(), false).unwrap();
307        for account in accounts.clone() {
308            engine
309                .state
310                .init_account(
311                    account.address,
312                    AccountInfo {
313                        balance: account.balance.unwrap_or_default(),
314                        nonce: 0u64,
315                        code_hash: KECCAK_EMPTY,
316                        code: account
317                            .code
318                            .clone()
319                            .map(|arg0: Vec<u8>| Bytecode::new_raw(arg0.into())),
320                    },
321                    None,
322                    false,
323                )
324                .expect("Failed to init account");
325        }
326        db.update(accounts, Some(block.clone()))
327            .unwrap();
328        let account_balances = HashMap::from([(
329            Bytes::from("0xBA12222222228d8Ba445958a75a0704d566BF2C8"),
330            HashMap::from([
331                (
332                    Bytes::from("0x6b175474e89094c44da98b954eedeac495271d0f"),
333                    Bytes::from(100_u64.to_le_bytes().to_vec()),
334                ),
335                (
336                    Bytes::from("0xba100000625a3754423978a60c9317c58a424e3d"),
337                    Bytes::from(100_u64.to_le_bytes().to_vec()),
338                ),
339            ]),
340        )]);
341
342        let decoder_context = DecoderContext::new();
343        let res = EVMPoolState::try_from_with_header(
344            snapshot,
345            block,
346            &account_balances,
347            &tokens,
348            &decoder_context,
349        )
350        .await
351        .unwrap();
352
353        let res_pool = res;
354
355        assert_eq!(
356            res_pool.get_balance_owner(),
357            Some(Address::from_str("0xBA12222222228d8Ba445958a75a0704d566BF2C8").unwrap())
358        );
359        let mut exp_involved_contracts = HashSet::new();
360        exp_involved_contracts
361            .insert(Address::from_str("0xBA12222222228d8Ba445958a75a0704d566BF2C8").unwrap());
362        assert_eq!(res_pool.get_involved_contracts(), exp_involved_contracts);
363        assert!(res_pool.get_manual_updates());
364    }
365}
tycho_simulation/evm/protocol/vm/decoder.rs

tycho_simulation/evm/protocol/vm/
decoder.rs
