Liquidity.Party/tycho-protocol-sdk
2
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: Misc improvements to code (#277)

Browse Source 
* refactor: Misc improvements to code

- Decouple validating logic from TychoRunner
- Move all data fetching and decoding the tycho message into the same method
- Split validate_state into validate_state, validate_token_balances and simulate_and_execute
- Make rpc_provider and runtime attributes of TestRunner
- Add references where possible to avoid clones
- Remove unnecessary code
- Make clippy happy

#time 2h 36m

#time 0m


#time 3m

* chore: Use tycho deps and foundry from tycho_simulation

This is to try to decrease the risk of using conflicting versions in the different repositories

#time 32m


#time 0m

* chore: Read RPC_URL in main.rs

#time 10m

* fix: Support eth trades (skip balance and allowance overwrites) and set balance overwrite to amount in

For tokens like USDC setting the balance super high was making us getting blacklisted

#time 1h 12m

* fix: Fix curve tests and filter components_by_id with the expected_component_ids

#time 1h 30m


#time 0m

* fix: Don't use all the possible executor addresses. Hardcode just one for the test

Refactor overwrites logic:
- renamed functions
- moved logic around that fits together
- don't use StateOverrides and then convert to alloy overrides. Use alloy's directly

#time 1h 21m

* fix: Assume that the executors mapping starts at storage value 1

Move setup_router_overwrites away from the rpc and into the execution file
Delete unnecessary get_storage_at

#time 33m
This commit is contained in:
dianacarvalho1
2025-09-25 17:27:05 +01:00
committed by GitHub
parent 12369c3981
commit b577e7d6b2
13 changed files with 1161 additions and 1087 deletions
Download Patch File Download Diff File Expand all files Collapse all files

371
protocol-testing/src/execution.rs
View File

@@ -7,25 +7,29 @@
use std::{collections::HashMap, str::FromStr, sync::LazyLock};
use alloy::{
primitives::{Address, U256},
rpc::types::{Block, TransactionRequest},
primitives::{keccak256, map::AddressHashMap, Address, FixedBytes, U256},
rpc::types::{state::AccountOverride, Block, TransactionRequest},
};
use miette::{miette, IntoDiagnostic, WrapErr};
use num_bigint::BigUint;
use serde_json::Value;
use tracing::info;
use tycho_common::traits::{AllowanceSlotDetector, BalanceSlotDetector};
use tycho_ethereum::entrypoint_tracer::{
allowance_slot_detector::{AllowanceSlotDetectorConfig, EVMAllowanceSlotDetector},
balance_slot_detector::{BalanceSlotDetectorConfig, EVMBalanceSlotDetector},
};
use tycho_simulation::{
evm::protocol::u256_num::u256_to_biguint, tycho_execution::encoding::models::Solution,
evm::protocol::u256_num::{biguint_to_u256, u256_to_biguint},
tycho_common::{
traits::{AllowanceSlotDetector, BalanceSlotDetector},
Bytes,
},
tycho_ethereum::entrypoint_tracer::{
allowance_slot_detector::{AllowanceSlotDetectorConfig, EVMAllowanceSlotDetector},
balance_slot_detector::{BalanceSlotDetectorConfig, EVMBalanceSlotDetector},
},
tycho_execution::encoding::models::Solution,
};
use crate::rpc::RPCProvider;
const ROUTER_BYTECODE_JSON: &str = include_str!("../../evm/test/router/TychoRouter.runtime.json");
pub const EXECUTORS_JSON: &str = include_str!("../test_executor_addresses.json");
pub const ROUTER_BYTECODE_JSON: &str =
include_str!("../../evm/test/router/TychoRouter.runtime.json");
pub const EXECUTOR_ADDRESS: &str = "0xaE04CA7E9Ed79cBD988f6c536CE11C621166f41B";
// Include all executor bytecode files at compile time
const UNISWAP_V2_BYTECODE_JSON: &str =
@@ -60,62 +64,6 @@ static EXECUTOR_MAPPING: LazyLock<HashMap<&'static str, &'static str>> = LazyLoc
map
});
/// Executor addresses loaded from test_executor_addresses.json at startup
pub static EXECUTOR_ADDRESSES: LazyLock<HashMap<String, Address>> = LazyLock::new(|| {
let json_value: Value = serde_json::from_str(&EXECUTORS_JSON)
.expect("Failed to parse test_executor_addresses.json");
let ethereum_addresses = json_value["ethereum"]
.as_object()
.expect("Missing 'ethereum' key in test_executor_addresses.json");
let mut addresses = HashMap::new();
for (protocol_name, address_value) in ethereum_addresses {
let address_str = address_value
.as_str()
.unwrap_or_else(|| panic!("Invalid address format for protocol '{protocol_name}'"));
let address = Address::from_str(address_str).unwrap_or_else(|_| {
panic!("Invalid address '{address_str}' for protocol '{protocol_name}'")
});
addresses.insert(protocol_name.clone(), address);
}
addresses
});
#[derive(Debug, Clone)]
pub struct StateOverride {
pub code: Option<Vec<u8>>,
pub balance: Option<U256>,
pub state_diff: HashMap<alloy::primitives::Bytes, alloy::primitives::Bytes>,
}
impl StateOverride {
pub fn new() -> Self {
Self { code: None, balance: None, state_diff: HashMap::new() }
}
pub fn with_code(mut self, code: Vec<u8>) -> Self {
self.code = Some(code);
self
}
pub fn with_balance(mut self, balance: U256) -> Self {
self.balance = Some(balance);
self
}
pub fn with_state_diff(
mut self,
slot: alloy::primitives::Bytes,
value: alloy::primitives::Bytes,
) -> Self {
self.state_diff.insert(slot, value);
self
}
}
/// Get executor bytecode JSON based on component ID
fn get_executor_bytecode_json(component_id: &str) -> miette::Result<&'static str> {
for (pattern, executor_json) in EXECUTOR_MAPPING.iter() {
@@ -126,20 +74,9 @@ fn get_executor_bytecode_json(component_id: &str) -> miette::Result<&'static str
Err(miette!("Unknown component type '{}' - no matching executor found", component_id))
}
/// Get executor address for a given component ID
fn get_executor_address(component_id: &str) -> miette::Result<Address> {
if let Some(&address) = EXECUTOR_ADDRESSES.get(component_id) {
return Ok(address);
}
Err(miette!("No executor address found for component type '{}'", component_id))
}
/// Load executor bytecode from embedded constants based on solution component
fn load_executor_bytecode(solution: &Solution) -> miette::Result<Vec<u8>> {
let first_swap = solution.swaps.first().unwrap();
let component_id = &first_swap.component;
let executor_json = get_executor_bytecode_json(&component_id.protocol_system)?;
/// Load executor bytecode from embedded constants based on the protocol system
pub fn load_executor_bytecode(protocol_system: &str) -> miette::Result<Vec<u8>> {
let executor_json = get_executor_bytecode_json(protocol_system)?;
let json_value: serde_json::Value = serde_json::from_str(executor_json)
.into_diagnostic()
@@ -177,93 +114,206 @@ fn calculate_gas_fees(block_header: &Block) -> miette::Result<(U256, U256)> {
Ok((max_fee_per_gas, max_priority_fee_per_gas))
}
/// Set up all state overrides needed for simulation
async fn setup_state_overrides(
/// Calculate storage slot for Solidity mapping.
///
/// The solidity code:
/// keccak256(abi.encodePacked(bytes32(key), bytes32(slot)))
pub fn calculate_executor_storage_slot(key: Address) -> FixedBytes<32> {
// Convert key (20 bytes) to 32-byte left-padded array (uint256)
let mut key_bytes = [0u8; 32];
key_bytes[12..].copy_from_slice(key.as_slice());
// The base of the executor storage slot is 1, since there is only one
// variable that is initialized before it (which is _roles in AccessControl.sol).
// In this case, _roles gets slot 0.
// The slots are given in order to the parent contracts' variables first and foremost.
let slot = U256::from(1);
// Convert U256 slot to 32-byte big-endian array
let slot_bytes = slot.to_be_bytes::<32>();
// Concatenate key_bytes + slot_bytes, then keccak hash
let mut buf = [0u8; 64];
buf[..32].copy_from_slice(&key_bytes);
buf[32..].copy_from_slice(&slot_bytes);
keccak256(buf)
}
/// Sets up state overwrites for the Tycho router and its associated executor.
///
/// This method prepares the router for simulation by:
/// 1. Overriding the router's bytecode with the embedded runtime bytecode
/// 2. Copying executor approval storage from the current block to maintain permissions
/// 3. Overriding the executor's bytecode based on the protocol system
///
/// # Arguments
/// * `router_address` - The address of the Tycho router contract
/// * `protocol_system` - The protocol system identifier (e.g., "uniswap_v2", "vm:balancer_v2")
///
/// # Returns
/// A HashMap containing account overwrites for both the router and executor addresses.
/// The router override includes bytecode and executor approval storage.
/// The executor override includes the appropriate bytecode for the protocol.
///
/// # Errors
/// Returns an error if:
/// - Router bytecode JSON parsing fails
/// - Executor address parsing fails
/// - Storage slot fetching fails
/// - Executor bytecode loading fails
pub async fn setup_router_overwrites(
router_address: Address,
protocol_system: &str,
) -> miette::Result<AddressHashMap<AccountOverride>> {
let json_value: serde_json::Value = serde_json::from_str(ROUTER_BYTECODE_JSON)
.into_diagnostic()
.wrap_err("Failed to parse router JSON")?;
let bytecode_str = json_value["runtimeBytecode"]
.as_str()
.ok_or_else(|| miette::miette!("No runtimeBytecode field found in router JSON"))?;
// Remove 0x prefix if present
let bytecode_hex =
if let Some(stripped) = bytecode_str.strip_prefix("0x") { stripped } else { bytecode_str };
let router_bytecode = hex::decode(bytecode_hex)
.into_diagnostic()
.wrap_err("Failed to decode router bytecode from hex")?;
// Start with the router bytecode override
let mut state_overwrites = AddressHashMap::default();
let mut tycho_router_override = AccountOverride::default().with_code(router_bytecode);
// Find executor address approval storage slot
let executor_address = Address::from_str(EXECUTOR_ADDRESS).into_diagnostic()?;
let storage_slot = calculate_executor_storage_slot(executor_address);
// The executors mapping starts at storage value 1
let storage_value = FixedBytes::<32>::from(U256::ONE);
tycho_router_override =
tycho_router_override.with_state_diff(vec![(storage_slot, storage_value)]);
state_overwrites.insert(router_address, tycho_router_override);
// Add bytecode overwrite for the executor
let executor_bytecode = load_executor_bytecode(protocol_system)?;
state_overwrites
.insert(executor_address, AccountOverride::default().with_code(executor_bytecode.to_vec()));
Ok(state_overwrites)
}
/// Sets up state overwrites for user accounts and tokens required for swap simulation.
///
/// This method prepares the user environment for historical block simulation by:
/// 1. Providing the user with sufficient ETH balance (100 ETH) for gas payments
/// 2. For ETH swaps: Adding the swap amount to the user's ETH balance
/// 3. For ERC20 swaps: Overriding token balance and allowance storage slots to ensure:
/// - User has sufficient tokens for the swap
/// - Router has unlimited allowance to spend user's tokens
///
/// The function uses EVM storage slot detection to find the correct storage locations
/// for token balances and allowances, then applies state overwrites to simulate the
/// required pre-conditions without executing actual token transfers.
///
/// # Arguments
/// * `solution` - The encoded swap solution containing token and amount information
/// * `transaction` - The transaction details for determining router address
/// * `user_address` - The address of the user performing the swap
/// * `rpc_url` - RPC endpoint URL for storage slot detection
/// * `block` - The historical block context for storage queries
///
/// # Returns
/// A HashMap containing account overwrites for:
/// - User account: ETH balance override
/// - Token contract: Balance and allowance storage slot overwrites (for ERC20 swaps)
///
/// # Errors
/// Returns an error if:
/// - Storage slot detection fails for balance or allowance
/// - Token address parsing fails
/// - RPC queries for storage detection fail
async fn setup_user_overwrites(
solution: &Solution,
transaction: &tycho_simulation::tycho_execution::encoding::models::Transaction,
user_address: Address,
executor_bytecode: &[u8],
rpc_url: String,
block: &Block,
) -> miette::Result<HashMap<Address, StateOverride>> {
let mut state_overwrites = HashMap::new();
let token_address = Address::from_slice(&solution.given_token[..20]);
// Extract executor address from the encoded solution's swaps data.
// The solution should only have one swap for the test, so this should be safe.
let executor_address = if let Some(first_swap) = solution.swaps.first() {
get_executor_address(&first_swap.component.protocol_system)?
} else {
return Err(miette!("No swaps in solution - cannot determine executor address"));
};
// Add bytecode overwrite for the executor
state_overwrites
.insert(executor_address, StateOverride::new().with_code(executor_bytecode.to_vec()));
) -> miette::Result<AddressHashMap<AccountOverride>> {
let mut overwrites = AddressHashMap::default();
// Add ETH balance override for the user to ensure they have enough gas funds
state_overwrites.insert(
user_address,
StateOverride::new().with_balance(U256::from_str("100000000000000000000").unwrap()), // 100 ETH
);
let mut eth_balance = U256::from_str("100000000000000000000").unwrap(); // 100 ETH
let detector = EVMBalanceSlotDetector::new(BalanceSlotDetectorConfig {
rpc_url: rpc_url.clone(),
..Default::default()
})
.into_diagnostic()?;
let token_address = Address::from_slice(&solution.given_token[..20]);
// If given token is ETH, add the given amount to the balance
if solution.given_token == Bytes::zero(20) {
eth_balance += biguint_to_u256(&solution.given_amount);
// if the given token is not ETH, do balance and allowance slots overwrites
} else {
let detector = EVMBalanceSlotDetector::new(BalanceSlotDetectorConfig {
rpc_url: rpc_url.clone(),
..Default::default()
})
.into_diagnostic()?;
let results = detector
.detect_balance_slots(
&[solution.given_token.clone()],
(**user_address).into(),
(*block.header.hash).into(),
)
.await;
let results = detector
.detect_balance_slots(
std::slice::from_ref(&solution.given_token),
(**user_address).into(),
(*block.header.hash).into(),
)
.await;
let balance_slot =
if let Some(Ok((_storage_addr, slot))) = results.get(&solution.given_token.clone()) {
slot
} else {
return Err(miette!("Couldn't find balance storage slot for token {token_address}"));
};
let balance_slot =
if let Some(Ok((_storage_addr, slot))) = results.get(&solution.given_token.clone()) {
slot
} else {
return Err(miette!("Couldn't find balance storage slot for token {token_address}"));
};
let detector = EVMAllowanceSlotDetector::new(AllowanceSlotDetectorConfig {
rpc_url,
..Default::default()
})
.into_diagnostic()?;
let detector = EVMAllowanceSlotDetector::new(AllowanceSlotDetectorConfig {
rpc_url,
..Default::default()
})
.into_diagnostic()?;
let results = detector
.detect_allowance_slots(
&[solution.given_token.clone()],
(**user_address).into(),
transaction.to.clone(), // tycho router
(*block.header.hash).into(),
)
.await;
let results = detector
.detect_allowance_slots(
std::slice::from_ref(&solution.given_token),
(**user_address).into(),
transaction.to.clone(), // tycho router
(*block.header.hash).into(),
)
.await;
let allowance_slot =
if let Some(Ok((_storage_addr, slot))) = results.get(&solution.given_token.clone()) {
let allowance_slot = if let Some(Ok((_storage_addr, slot))) =
results.get(&solution.given_token.clone())
{
slot
} else {
return Err(miette!("Couldn't find allowance storage slot for token {token_address}"));
};
state_overwrites.insert(
token_address,
StateOverride::new()
.with_state_diff(
alloy::primitives::Bytes::from(allowance_slot.to_vec()),
alloy::primitives::Bytes::from(U256::MAX.to_be_bytes::<32>()),
)
.with_state_diff(
alloy::primitives::Bytes::from(balance_slot.to_vec()),
alloy::primitives::Bytes::from(U256::MAX.to_be_bytes::<32>()),
),
);
overwrites.insert(
token_address,
AccountOverride::default().with_state_diff(vec![
(
alloy::primitives::B256::from_slice(allowance_slot),
alloy::primitives::B256::from_slice(&U256::MAX.to_be_bytes::<32>()),
),
(
alloy::primitives::B256::from_slice(balance_slot),
alloy::primitives::B256::from_slice(
&biguint_to_u256(&solution.given_amount).to_be_bytes::<32>(),
),
),
]),
);
}
overwrites.insert(user_address, AccountOverride::default().with_balance(eth_balance));
Ok(state_overwrites)
Ok(overwrites)
}
/// Simulate a trade using eth_call for historical blocks
@@ -271,10 +321,11 @@ pub async fn simulate_trade_with_eth_call(
rpc_provider: &RPCProvider,
transaction: &tycho_simulation::tycho_execution::encoding::models::Transaction,
solution: &Solution,
block_number: u64,
block: &Block,
) -> miette::Result<BigUint> {
let executor_bytecode = load_executor_bytecode(solution)?;
let first_swap = solution.swaps.first().unwrap();
let protocol_system = &first_swap.component.protocol_system;
let user_address = Address::from_slice(&solution.sender[..20]);
let (max_fee_per_gas, max_priority_fee_per_gas) = calculate_gas_fees(block)?;
// Convert main transaction to alloy TransactionRequest
@@ -295,28 +346,24 @@ pub async fn simulate_trade_with_eth_call(
);
let tycho_router_address = Address::from_slice(&transaction.to[..20]);
// Copy router storage and code from current block to historical block
let router_override = rpc_provider
.copy_contract_storage_and_code(tycho_router_address, ROUTER_BYTECODE_JSON)
let router_overwrites = setup_router_overwrites(tycho_router_address, protocol_system)
.await
.wrap_err("Failed to create router override")?;
// Set up state overrides including router override
let mut state_overwrites = setup_state_overrides(
let mut user_overwrites = setup_user_overwrites(
solution,
transaction,
user_address,
&executor_bytecode,
rpc_provider.url.to_string(),
block,
)
.await?; // Include executor override for historical blocks
.await?;
// Add the router override
state_overwrites.insert(tycho_router_address, router_override);
// Merge router overwrites with user overwrites
user_overwrites.extend(router_overwrites);
let execution_amount_out = rpc_provider
.simulate_transactions_with_tracing(execution_tx, block_number, state_overwrites)
.simulate_transactions_with_tracing(execution_tx, block.number(), user_overwrites)
.await
.map_err(|e| {
info!("Execution transaction failed with error: {}", e);