diff --git a/protosim_py/python/protosim_py/evm/decoders.py b/protosim_py/python/protosim_py/evm/decoders.py index 86fd04d6..1660ef9f 100644 --- a/protosim_py/python/protosim_py/evm/decoders.py +++ b/protosim_py/python/protosim_py/evm/decoders.py @@ -165,7 +165,7 @@ def decode_pool_state( adapter_contract_path=self.adapter_contract, trace=self.trace, manual_updates=manual_updates, - involved_contracts=set(component.contract_ids), + involved_contracts=set(to_checksum_address(b.hex()) for b in component.contract_ids), **optional_attributes, ) diff --git a/protosim_py/python/protosim_py/evm/pool_state.py b/protosim_py/python/protosim_py/evm/pool_state.py index 3155d05d..a255e98c 100644 --- a/protosim_py/python/protosim_py/evm/pool_state.py +++ b/protosim_py/python/protosim_py/evm/pool_state.py @@ -210,9 +210,10 @@ def _init_token_storage_slots(self): t.address in self.involved_contracts and t.address not in self.token_storage_slots ): - self.token_storage_slots[t.address] = token.brute_force_slots( + self.token_storage_slots[t.address] = slots = token.brute_force_slots( t, self.block, self._engine ) + log.debug(f"Using custom storage slots for {t.address}: {slots}") def get_amount_out( self: TPoolState, @@ -240,6 +241,7 @@ def _get_amount_out( sell_amount: Decimal, buy_token: EthereumToken, ) -> tuple[Decimal, int, TPoolState]: + overwrites = self._get_overwrites(sell_token, buy_token) trade, state_changes = self._adapter_contract.swap( cast(HexStr, self.id_), sell_token, @@ -247,7 +249,7 @@ def _get_amount_out( False, sell_token.to_onchain_amount(sell_amount), block=self.block, - overwrites=self._get_overwrites(sell_token, buy_token), + overwrites=overwrites, ) new_state = self._duplicate() for address, state_update in state_changes.items(): @@ -274,7 +276,7 @@ def _get_overwrites( level, and token-specific overwrites that depend on passed tokens. """ token_overwrites = self._get_token_overwrites(sell_token, buy_token, **kwargs) - return _merge(self.block_lasting_overwrites, token_overwrites) + return _merge(self.block_lasting_overwrites.copy(), token_overwrites) def _get_token_overwrites( self, sell_token: EthereumToken, buy_token: EthereumToken, max_amount=None @@ -296,7 +298,10 @@ def _get_token_overwrites( max_amount = sell_token.to_onchain_amount( self.get_sell_amount_limit(sell_token, buy_token) ) - overwrites = ERC20OverwriteFactory(sell_token) + overwrites = ERC20OverwriteFactory( + sell_token, + token_slots=self.token_storage_slots.get(sell_token.address, (0,1)) + ) overwrites.set_balance(max_amount, EXTERNAL_ACCOUNT) overwrites.set_allowance( allowance=max_amount, owner=EXTERNAL_ACCOUNT, spender=ADAPTER_ADDRESS @@ -311,8 +316,8 @@ def _get_token_overwrites( def _get_balance_overwrites(self) -> dict[Address, dict[int, int]]: balance_overwrites = {} address = self.balance_owner or self.id_ - slots = (0, 1) for t in self.tokens: + slots = (0, 1) if t.address in self.involved_contracts: slots = self.token_storage_slots.get(t.address) overwrites = ERC20OverwriteFactory(t, token_slots=slots) diff --git a/protosim_py/python/test/assets/sfrxEthToken.evm.runtime b/protosim_py/python/test/assets/sfrxEthToken.evm.runtime new file mode 100644 index 00000000..b67cb96d Binary files /dev/null and b/protosim_py/python/test/assets/sfrxEthToken.evm.runtime differ diff --git a/protosim_py/python/test/test_third_party_pool.py b/protosim_py/python/test/test_third_party_pool.py index 127624ea..71240fc3 100644 --- a/protosim_py/python/test/test_third_party_pool.py +++ b/protosim_py/python/test/test_third_party_pool.py @@ -2,18 +2,20 @@ from decimal import Decimal from pathlib import Path from unittest.mock import patch, call +from venv import create import pytest from hexbytes import HexBytes -from protosim_py._protosim_py import AccountInfo -from protosim_py.evm import BlockHeader +from tycho_indexer_client.dto import ChangeType + +from protosim_py.evm import AccountInfo, AccountUpdate, BlockHeader from protosim_py.evm.adapter_contract import AdapterContract from protosim_py.evm.pool_state import ThirdPartyPool from protosim_py.evm.storage import TychoDBSingleton from protosim_py.evm.utils import parse_account_info, create_engine from protosim_py.exceptions import RecoverableSimulationException -from protosim_py.models import EVMBlock, Capability, EthereumToken +from protosim_py.models import EVMBlock, Capability, EthereumToken, Blockchain ADDRESS_ZERO = "0x0000000000000000000000000000000000000000" @@ -31,6 +33,11 @@ def Token(name: str) -> EthereumToken: address="0xba100000625a3754423978a60c9317c58a424e3D", decimals=18, ), + "sfrxETH": EthereumToken( + symbol="sfrxETH", + address="0xac3e018457b222d93114458476f3e3416abbe38f", + decimals=18, + ), }[name] @@ -222,3 +229,86 @@ def test_stateless_contract_pool(asset_dir): ) assert mock_get_code.call_count == 2 + + +def test_overwrites_mock_erc20(pool_state): + overwrites = pool_state._get_overwrites(pool_state.tokens[0], pool_state.tokens[1]) + + assert overwrites == { + "0x6b175474e89094c44da98b954eedeac495271d0f": { + # pool balance, caller balance & allowance + 24060209162895628919861412957428278191632570471602070876674374646072182449944: 178754012737301800000, + 58546993237423525698686728856645416951692145960565761888391937184176623942864: 578960446186580977117854925043439539266349923328202820197287920039565648199, + 110136159478993350616340414857413728709904511599989695046923576775517543504731: 578960446186580977117854925043439539266349923328202820197287920039565648199, + }, + "0xba100000625a3754423978a60c9317c58a424e3d": { + # pool balance + 24060209162895628919861412957428278191632570471602070876674374646072182449944: 91082987763369890000 + }, + } + + +def test_overwrites_custom_erc20(asset_dir): + dai, sfrxEth = Token("DAI"), Token("sfrxETH") + with open(asset_dir / "sfrxEthToken.evm.runtime", "rb") as fp: + code = fp.read() + block = EVMBlock( + id=18485417, + hash_="0x28d41d40f2ac275a4f5f621a636b9016b527d11d37d610a45ac3a821346ebf8c", + ) + engine = create_engine([]) + engine.init_account( + address=sfrxEth.address, + account=AccountInfo( + nonce=0, + balance=0, + code=bytearray(code), + ), + mocked=True, + permanent_storage=None, + ) + + pool = ThirdPartyPool( + block=block, + id_="0x4626d81b3a1711beb79f4cecff2413886d461677000200000000000000000011", + tokens=(dai, sfrxEth), + marginal_prices={ + (sfrxEth, dai): Decimal("7.071503245428245871486924221"), + (dai, sfrxEth): Decimal("0.1377789143190479049114331557"), + }, + balances={ + dai.address: "178.7540127373018", + sfrxEth.address: "91.08298776336989", + }, + adapter_contract_path=str(asset_dir / "BalancerV2SwapAdapter.evm.runtime"), + balance_owner="0xBA12222222228d8Ba445958a75a0704d566BF2C8", + involved_contracts={sfrxEth.address}, + ) + overwrites_zero2one = dict(pool._get_overwrites(pool.tokens[0], pool.tokens[1])) + overwrites_one2zero = dict(pool._get_overwrites(pool.tokens[1], pool.tokens[0])) + + assert pool.token_storage_slots == { + "0xac3e018457b222d93114458476f3e3416abbe38f": (3, 4) + } + assert overwrites_zero2one == { + "0x6b175474e89094c44da98b954eedeac495271d0f": { + # same as test above + 24060209162895628919861412957428278191632570471602070876674374646072182449944: 178754012737301800000, + 58546993237423525698686728856645416951692145960565761888391937184176623942864: 100279494253364362835, + 110136159478993350616340414857413728709904511599989695046923576775517543504731: 100279494253364362835, + }, + "0xac3e018457b222d93114458476f3e3416abbe38f": { + # different key compared to our mocked contract in previous test! + 26260780229836133480733911416306220824002130525338603371401637394485347754320: 91082987763369890000 + }, + } + assert overwrites_one2zero == { + "0x6b175474e89094c44da98b954eedeac495271d0f": { + 24060209162895628919861412957428278191632570471602070876674374646072182449944: 178754012737301800000, + }, + "0xac3e018457b222d93114458476f3e3416abbe38f": { + 26260780229836133480733911416306220824002130525338603371401637394485347754320: 91082987763369890000, + 39404303412979837706729056057326168057508607640016353144739821671269873586962: 0, + 55680195869663131363475621218576616366970057593641329011542336081174417563938: 0, + }, + }