Skip to main content

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 load_balancer_account_data() -> Vec<AccountUpdate> {
243        let project_root = env!("CARGO_MANIFEST_DIR");
244        let asset_path =
245            Path::new(project_root).join("tests/assets/decoder/balancer_v2_snapshot.json");
246        let json_data = fs::read_to_string(asset_path).expect("Failed to read test asset");
247        let data: Value = serde_json::from_str(&json_data).expect("Failed to parse JSON");
248
249        let accounts: Vec<AccountUpdate> = serde_json::from_value(data["accounts"].clone())
250            .expect("Expected accounts to match AccountUpdate structure");
251        accounts
252    }
253
254    #[allow(deprecated)]
255    #[tokio::test]
256    async fn test_try_from_with_header() {
257        let attributes: HashMap<String, Bytes> = vec![
258            (
259                "balance_owner".to_string(),
260                Bytes::from_str("0xBA12222222228d8Ba445958a75a0704d566BF2C8").unwrap(),
261            ),
262            ("reserve1".to_string(), Bytes::from(200_u64.to_le_bytes().to_vec())),
263        ]
264        .into_iter()
265        .collect();
266        let tokens = [
267            Token::new(
268                &Bytes::from_str("0x6b175474e89094c44da98b954eedeac495271d0f").unwrap(),
269                "DAI",
270                18,
271                0,
272                &[Some(10_000)],
273                tycho_common::models::Chain::Ethereum,
274                100,
275            ),
276            Token::new(
277                &Bytes::from_str("0xba100000625a3754423978a60c9317c58a424e3d").unwrap(),
278                "BAL",
279                18,
280                0,
281                &[Some(10_000)],
282                tycho_common::models::Chain::Ethereum,
283                100,
284            ),
285        ]
286        .into_iter()
287        .map(|t| (t.address.clone(), t))
288        .collect::<HashMap<_, _>>();
289        let snapshot = ComponentWithState {
290            state: ResponseProtocolState {
291                component_id: "0x4626d81b3a1711beb79f4cecff2413886d461677000200000000000000000011"
292                    .to_owned(),
293                attributes,
294                balances: HashMap::new(),
295            },
296            component: vm_component(),
297            component_tvl: None,
298            entrypoints: Vec::new(),
299        };
300        // Initialize engine with balancer storage
301        let block = BlockHeader::default();
302        let accounts = load_balancer_account_data();
303        let db = SHARED_TYCHO_DB.clone();
304        let engine = create_engine(db.clone(), false).unwrap();
305        for account in accounts.clone() {
306            engine
307                .state
308                .init_account(
309                    account.address,
310                    AccountInfo {
311                        balance: account.balance.unwrap_or_default(),
312                        nonce: 0u64,
313                        code_hash: KECCAK_EMPTY,
314                        code: account
315                            .code
316                            .clone()
317                            .map(|arg0: Vec<u8>| Bytecode::new_raw(arg0.into())),
318                    },
319                    None,
320                    false,
321                )
322                .expect("Failed to init account");
323        }
324        db.update(accounts, Some(block.clone()))
325            .unwrap();
326        let account_balances = HashMap::from([(
327            Bytes::from("0xBA12222222228d8Ba445958a75a0704d566BF2C8"),
328            HashMap::from([
329                (
330                    Bytes::from("0x6b175474e89094c44da98b954eedeac495271d0f"),
331                    Bytes::from(100_u64.to_le_bytes().to_vec()),
332                ),
333                (
334                    Bytes::from("0xba100000625a3754423978a60c9317c58a424e3d"),
335                    Bytes::from(100_u64.to_le_bytes().to_vec()),
336                ),
337            ]),
338        )]);
339
340        let decoder_context = DecoderContext::new();
341        let res = EVMPoolState::try_from_with_header(
342            snapshot,
343            block,
344            &account_balances,
345            &tokens,
346            &decoder_context,
347        )
348        .await
349        .unwrap();
350
351        let res_pool = res;
352
353        assert_eq!(
354            res_pool.get_balance_owner(),
355            Some(Address::from_str("0xBA12222222228d8Ba445958a75a0704d566BF2C8").unwrap())
356        );
357        let mut exp_involved_contracts = HashSet::new();
358        exp_involved_contracts
359            .insert(Address::from_str("0xBA12222222228d8Ba445958a75a0704d566BF2C8").unwrap());
360        assert_eq!(res_pool.get_involved_contracts(), exp_involved_contracts);
361        assert!(res_pool.get_manual_updates());
362    }
363}