Merge pull request #7 from propeller-heads/ah/dynamic-token-storage-s…

…lot-detection

feat: Dynamically detect token storage slots when necessary

protosim_py/python/protosim_py/assets/ERC20.abi

@@ -0,0 +1,263 @@
+[
+  {
+    "constant": true,
+    "inputs": [],
+    "name": "name",
+    "outputs": [
+      {
+        "name": "",
+        "type": "string"
+      }
+    ],
+    "payable": false,
+    "stateMutability": "view",
+    "type": "function"
+  },
+  {
+    "constant": true,
+    "inputs": [],
+    "name": "symbol",
+    "outputs": [
+      {
+        "name": "",
+        "type": "string"
+      }
+    ],
+    "payable": false,
+    "stateMutability": "view",
+    "type": "function"
+  },
+  {
+    "constant": true,
+    "inputs": [],
+    "name": "decimals",
+    "outputs": [
+      {
+        "name": "",
+        "type": "uint8"
+      }
+    ],
+    "payable": false,
+    "stateMutability": "view",
+    "type": "function"
+  },
+  {
+    "constant": false,
+    "inputs": [
+      {
+        "name": "spender",
+        "type": "address"
+      },
+      {
+        "name": "value",
+        "type": "uint256"
+      }
+    ],
+    "name": "approve",
+    "outputs": [
+      {
+        "name": "",
+        "type": "bool"
+      }
+    ],
+    "payable": false,
+    "stateMutability": "nonpayable",
+    "type": "function"
+  },
+  {
+    "constant": true,
+    "inputs": [],
+    "name": "totalSupply",
+    "outputs": [
+      {
+        "name": "",
+        "type": "uint256"
+      }
+    ],
+    "payable": false,
+    "stateMutability": "view",
+    "type": "function"
+  },
+  {
+    "constant": false,
+    "inputs": [
+      {
+        "name": "from",
+        "type": "address"
+      },
+      {
+        "name": "to",
+        "type": "address"
+      },
+      {
+        "name": "value",
+        "type": "uint256"
+      }
+    ],
+    "name": "transferFrom",
+    "outputs": [
+      {
+        "name": "",
+        "type": "bool"
+      }
+    ],
+    "payable": false,
+    "stateMutability": "nonpayable",
+    "type": "function"
+  },
+  {
+    "constant": false,
+    "inputs": [
+      {
+        "name": "spender",
+        "type": "address"
+      },
+      {
+        "name": "addedValue",
+        "type": "uint256"
+      }
+    ],
+    "name": "increaseAllowance",
+    "outputs": [
+      {
+        "name": "",
+        "type": "bool"
+      }
+    ],
+    "payable": false,
+    "stateMutability": "nonpayable",
+    "type": "function"
+  },
+  {
+    "constant": true,
+    "inputs": [
+      {
+        "name": "owner",
+        "type": "address"
+      }
+    ],
+    "name": "balanceOf",
+    "outputs": [
+      {
+        "name": "",
+        "type": "uint256"
+      }
+    ],
+    "payable": false,
+    "stateMutability": "view",
+    "type": "function"
+  },
+  {
+    "constant": false,
+    "inputs": [
+      {
+        "name": "spender",
+        "type": "address"
+      },
+      {
+        "name": "subtractedValue",
+        "type": "uint256"
+      }
+    ],
+    "name": "decreaseAllowance",
+    "outputs": [
+      {
+        "name": "",
+        "type": "bool"
+      }
+    ],
+    "payable": false,
+    "stateMutability": "nonpayable",
+    "type": "function"
+  },
+  {
+    "constant": false,
+    "inputs": [
+      {
+        "name": "to",
+        "type": "address"
+      },
+      {
+        "name": "value",
+        "type": "uint256"
+      }
+    ],
+    "name": "transfer",
+    "outputs": [
+      {
+        "name": "",
+        "type": "bool"
+      }
+    ],
+    "payable": false,
+    "stateMutability": "nonpayable",
+    "type": "function"
+  },
+  {
+    "constant": true,
+    "inputs": [
+      {
+        "name": "owner",
+        "type": "address"
+      },
+      {
+        "name": "spender",
+        "type": "address"
+      }
+    ],
+    "name": "allowance",
+    "outputs": [
+      {
+        "name": "",
+        "type": "uint256"
+      }
+    ],
+    "payable": false,
+    "stateMutability": "view",
+    "type": "function"
+  },
+  {
+    "anonymous": false,
+    "inputs": [
+      {
+        "indexed": true,
+        "name": "from",
+        "type": "address"
+      },
+      {
+        "indexed": true,
+        "name": "to",
+        "type": "address"
+      },
+      {
+        "indexed": false,
+        "name": "value",
+        "type": "uint256"
+      }
+    ],
+    "name": "Transfer",
+    "type": "event"
+  },
+  {
+    "anonymous": false,
+    "inputs": [
+      {
+        "indexed": true,
+        "name": "owner",
+        "type": "address"
+      },
+      {
+        "indexed": true,
+        "name": "spender",
+        "type": "address"
+      },
+      {
+        "indexed": false,
+        "name": "value",
+        "type": "uint256"
+      }
+    ],
+    "name": "Approval",
+    "type": "event"
+  }
+]

protosim_py/python/protosim_py/evm/decoders.py

@@ -165,6 +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),
             **optional_attributes,
         )
 

protosim_py/python/protosim_py/evm/pool_state.py

@@ -12,6 +12,7 @@
 from eth_utils import keccak
 from eth_typing import HexStr
 
+from . import token
 from . import SimulationEngine, AccountInfo, SimulationParameters
 from .adapter_contract import AdapterContract
 from .constants import MAX_BALANCE, EXTERNAL_ACCOUNT
@@ -46,6 +47,8 @@ def __init__(
         block_lasting_overwrites: defaultdict[Address, dict[int, int]] = None,
         manual_updates: bool = False,
         trace: bool = False,
+        involved_contracts=None,
+        token_storage_slots=None,
     ):
         self.id_ = id_
         """The pools identifier."""
@@ -91,10 +94,22 @@ def __init__(
         self.trace: bool = trace
         """If set, vm will emit detailed traces about the execution."""
 
+        self.involved_contracts: set[Address] = involved_contracts or set()
+        """A set of all contract addresses involved in the simulation of this pool."""
+
+        self.token_storage_slots: dict[Address, tuple[int, int]] = (
+            token_storage_slots or {}
+        )
+        """Allows the specification of custom storage slots for token allowances and
+        balances. This is particularly useful for token contracts involved in protocol
+        logic that extends beyond simple transfer functionality.
+        """
+
         self._engine: Optional[SimulationEngine] = None
         self._set_engine()
         self._adapter_contract = AdapterContract(ADAPTER_ADDRESS, self._engine)
         self._set_capabilities()
+        self._init_token_storage_slots()
         if len(self.marginal_prices) == 0:
             self._set_marginal_prices()
 
@@ -105,11 +120,6 @@ def _set_engine(self):
 
         The engine will have the specified adapter contract mocked, as well as the
         tokens used by the pool.
-
-        Parameters
-        ----------
-        engine
-            Optional simulation engine instance.
         """
         if self._engine is not None:
             return
@@ -194,6 +204,16 @@ def _set_capabilities(self):
                 f"Pool {self.id_} hash different capabilities depending on the token pair!"
             )
 
+    def _init_token_storage_slots(self):
+        for t in self.tokens:
+            if (
+                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(
+                    t, self.block, self._engine
+                )
+
     def get_amount_out(
         self: TPoolState,
         sell_token: EthereumToken,
@@ -291,8 +311,11 @@ 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:
-            overwrites = ERC20OverwriteFactory(t)
+            if t.address in self.involved_contracts:
+                slots = self.token_storage_slots.get(t.address)
+            overwrites = ERC20OverwriteFactory(t, token_slots=slots)
             overwrites.set_balance(
                 t.to_onchain_amount(self.balances[t.address]), address
             )