@paraswap/dex-lib 4.1.16 → 4.1.17-arena

package/tests/tenderly-simulation.ts

@@ -1,289 +1,698 @@
 /* eslint-disable no-console */
+import 'dotenv/config';
 import axios from 'axios';
-import { TxObject } from '../src/types';
-import { StateOverrides, StateSimulateApiOverride } from './smart-tokens';
-import { Provider, StaticJsonRpcProvider } from '@ethersproject/providers';
 import { ethers } from 'ethers';
-import { Network } from '../build/constants';
-import { Address } from '@paraswap/core';
-
-const TENDERLY_TOKEN = process.env.TENDERLY_TOKEN;
-const TENDERLY_ACCOUNT_ID = process.env.TENDERLY_ACCOUNT_ID;
-const TENDERLY_PROJECT = process.env.TENDERLY_PROJECT;
-const TENDERLY_VNET_ID = process.env.TENDERLY_VNET_ID;
-
-export type SimulationResult = {
-  success: boolean;
-  gasUsed?: string;
-  url?: string;
-  transaction?: any;
-};
-
-export interface TransactionSimulator {
-  vnetId: string;
-  rpcURL: string;
-  setup(): Promise<void>;
-
-  getChainNameByChainId(network: number): string;
-
-  simulate(
-    params: TxObject,
-    stateOverrides?: StateOverrides,
-  ): Promise<SimulationResult>;
+import * as fs from 'fs/promises';
+import * as path from 'path';
+import { ETHER_ADDRESS } from '../src/constants';
+
+const TENDERLY_TOKEN = process.env.TENDERLY_TOKEN!;
+const TENDERLY_ACCOUNT_ID = process.env.TENDERLY_ACCOUNT_ID!;
+const TENDERLY_PROJECT = process.env.TENDERLY_PROJECT!;
+
+interface StateObject {
+  // Overrides of storage slots.
+  // In this mapping, the key is variable storage slot,
+  // and the value contains the override.
+  storage?: Record<string, string>; // storage slot -> value
+  balance?: string;
+  nonce?: number;
+  // Bytecode that will override the code associated to the given account.
+  code?: string;
 }
+export type StateOverride = Record<string, StateObject>; // contract -> storage override
 
-export class EstimateGasSimulation implements TransactionSimulator {
-  vnetId: string = '0';
-  rpcURL: string = '';
+interface TokenStorageSlots {
+  balanceSlot: string;
+  allowanceSlot: string;
+  isVyper?: boolean;
+}
+
+interface SimulateTransactionRequest {
+  from: string | null;
+  to: string | null;
+  value?: string;
+  data: string;
+  chainId: number;
+  timestamp?: number;
+  blockNumber?: number;
+  stateOverride?: StateOverride;
+}
+
+// not fully complete
+interface SimulatedTransactionCall {
+  hash: string;
+  contract_name: string;
+  function_name: string;
+  function_pc: number;
+  function_op: string;
+  function_file_index: number;
+  function_code_start: number;
+  function_line_number: number;
+  function_code_length: number;
+  absolute_position: number;
+  caller_pc: number;
+  caller_op: string;
+  call_type: string;
+  address: string;
+  from: string;
+  from_balance: string;
+  to: string;
+  to_balance: string;
+  value: string | null;
+  caller: {
+    address: string;
+    balance: string;
+  };
+  block_timestamp: string;
+  gas: number;
+  gas_used: number;
+  intrinsic_gas: number;
+  storage_address: string;
+  input: string;
+  storage_slot: string[] | undefined;
+  calls: SimulatedTransactionCall[] | null;
+}
 
-  constructor(private provider: Provider) {}
+// not fully complete
+interface Simulation {
+  id: string;
+  project_id: string;
+  owner_id: string;
+  network_id: string;
+  block_number: number;
+  transaction_index: number;
+  from: string;
+  to: string;
+  input: string;
+  gas: number;
+  gas_price: string;
+  gas_used: number;
+  value: string;
+  method: string;
+  status: boolean;
+}
 
-  async setup() {}
+// not complete, all details include a lot more info
+interface SimulatedTransactionDetails {
+  transaction: {
+    hash: string;
+    block_hash: string;
+    block_number: number;
+    from: string;
+    gas: number;
+    gas_price: number;
+    gas_fee_cap: number;
+    gas_tip_cap: number;
+    cumulative_gas_used: number;
+    gas_used: number;
+    effective_gas_price: number;
+    input: string;
+    nonce: number;
+    to: string;
+    index: number;
+    value: string;
+    access_list: null;
+    status: boolean;
+    transaction_info: {
+      call_trace: SimulatedTransactionCall;
+    };
+  };
+}
 
-  getChainNameByChainId(network: number) {
-    return '';
+class TokenStorageSlotsCache {
+  private static TOKEN_STORAGE_SLOTS: Record<
+    number,
+    Record<string, TokenStorageSlots>
+  > | null = null;
+  private static readonly TOKEN_STORAGE_SLOTS_FILEPATH = path.join(
+    __dirname,
+    'token-storage-slots.json',
+  );
+
+  private static async loadTokenStorageSlots(): Promise<
+    Record<number, Record<string, TokenStorageSlots>>
+  > {
+    return JSON.parse(
+      await fs.readFile(TokenStorageSlotsCache.TOKEN_STORAGE_SLOTS_FILEPATH, {
+        encoding: 'utf-8',
+      }),
+    );
   }
 
-  async simulate(
-    params: TxObject,
-    _: StateOverrides,
-  ): Promise<SimulationResult> {
-    try {
-      const result = await this.provider.estimateGas(params);
-      return {
-        success: true,
-        gasUsed: result.toNumber().toString(),
-      };
-    } catch (e) {
-      console.error(`Estimate gas simulation failed:`, e);
-      return {
-        success: false,
-      };
+  private static async saveTokenStorageSlots(): Promise<void> {
+    await fs.writeFile(
+      TokenStorageSlotsCache.TOKEN_STORAGE_SLOTS_FILEPATH,
+      JSON.stringify(TokenStorageSlotsCache.TOKEN_STORAGE_SLOTS, null, 2),
+      { encoding: 'utf-8' },
+    );
+  }
+
+  static async getTokenStorageSlots(
+    chainId: number,
+    token: string,
+  ): Promise<TokenStorageSlots | null> {
+    if (!TokenStorageSlotsCache.TOKEN_STORAGE_SLOTS) {
+      TokenStorageSlotsCache.TOKEN_STORAGE_SLOTS =
+        await TokenStorageSlotsCache.loadTokenStorageSlots();
     }
+
+    return (
+      TokenStorageSlotsCache.TOKEN_STORAGE_SLOTS[chainId]?.[
+        token.toLowerCase()
+      ] ?? null
+    );
   }
-}
 
-export class TenderlySimulation implements TransactionSimulator {
-  vnetId: string = '';
-  rpcURL: string = '';
-  maxGasLimit = 80000000;
-
-  private readonly chainIdToChainNameMap: { [key: number]: string } = {
-    [Network.MAINNET]: 'mainnet',
-    [Network.BSC]: 'bnb',
-    [Network.POLYGON]: 'polygon',
-    [Network.AVALANCHE]: 'avalanche-mainnet',
-    [Network.FANTOM]: 'fantom',
-    [Network.ARBITRUM]: 'arbitrum',
-    [Network.OPTIMISM]: 'optimistic',
-    [Network.GNOSIS]: 'gnosis-chain',
-    [Network.BASE]: 'base',
-  };
+  static async setTokenStorageSlots(
+    chainId: number,
+    token: string,
+    slots: TokenStorageSlots,
+  ): Promise<void> {
+    if (!TokenStorageSlotsCache.TOKEN_STORAGE_SLOTS) {
+      TokenStorageSlotsCache.TOKEN_STORAGE_SLOTS =
+        await TokenStorageSlotsCache.loadTokenStorageSlots();
+    }
 
-  constructor(private network: number = 1) {}
+    TokenStorageSlotsCache.TOKEN_STORAGE_SLOTS ||= {};
+    TokenStorageSlotsCache.TOKEN_STORAGE_SLOTS[chainId] ||= {};
+    TokenStorageSlotsCache.TOKEN_STORAGE_SLOTS[chainId][token.toLowerCase()] =
+      slots;
 
-  getChainNameByChainId(network: number): string {
-    return this.chainIdToChainNameMap[network];
+    void TokenStorageSlotsCache.saveTokenStorageSlots();
   }
+}
 
-  async setup() {
-    if (!TENDERLY_TOKEN)
-      throw new Error(
-        `TenderlySimulation_setup: TENDERLY_TOKEN not found in the env`,
-      );
+export class TenderlySimulator {
+  // public constants
+  static readonly DEFAULT_OWNER = '0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa';
+  static readonly DEFAULT_SPENDER =
+    '0xbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb';
 
-    const findAdminRPC = (rpcs: { name: string; url: string }[]) => {
-      return rpcs.find(
-        rpc => rpc.name.toLowerCase() === 'Admin RPC'.toLowerCase(),
-      );
+  // singleton
+  private static instance: TenderlySimulator;
+
+  private constructor() {}
+
+  public static getInstance(): TenderlySimulator {
+    if (!TenderlySimulator.instance) {
+      TenderlySimulator.instance = new TenderlySimulator();
+    }
+
+    return TenderlySimulator.instance;
+  }
+
+  public async simulateTransaction(
+    request: SimulateTransactionRequest,
+  ): Promise<Simulation> {
+    const data = {
+      network_id: request.chainId,
+      from: request.from,
+      to: request.to,
+      input: request.data,
+      value: request.value,
+      save: true,
+      save_if_fails: true,
+      state_objects: request.stateOverride,
     };
 
-    if (TENDERLY_VNET_ID) {
-      this.vnetId = TENDERLY_VNET_ID;
-
-      try {
-        let res = await axios.get(
-          `https://api.tenderly.co/api/v1/account/${TENDERLY_ACCOUNT_ID}/project/${TENDERLY_PROJECT}/vnets/${this.vnetId}`,
-          {
-            timeout: 200000,
-            headers: {
-              'x-access-key': TENDERLY_TOKEN,
-            },
-          },
-        );
+    console.log('Sending transaction simulation with params:');
+    console.log(JSON.stringify(data, null, 2));
 
-        const rpc = findAdminRPC(res.data.rpcs);
+    const response = await axios.request({
+      method: 'POST',
+      url: `https://api.tenderly.co/api/v1/account/${TENDERLY_ACCOUNT_ID}/project/${TENDERLY_PROJECT}/simulate`,
+      headers: {
+        'X-Access-Key': TENDERLY_TOKEN,
+      },
+      data,
+    });
 
-        if (!rpc) {
-          throw new Error(`RPC url was not found for testnet: ${this.vnetId}`);
-        }
+    const simulation = response.data.simulation as Simulation;
+    const url = `https://dashboard.tenderly.co/${TENDERLY_ACCOUNT_ID}/${TENDERLY_PROJECT}/simulator/${simulation.id}`;
 
-        this.rpcURL = rpc.url;
-      } catch (e) {
-        console.error(`TenderlySimulation_setup:`, e);
-        throw e;
-      }
+    console.log('Successfully simulated a transaction:');
+    console.log(`Simulation URL - ${url}`);
 
-      return;
-    }
+    return simulation;
+  }
 
+  public async getSimulatedTransactionDetails(
+    id: string,
+  ): Promise<SimulatedTransactionDetails | null> {
     try {
-      await process.nextTick(() => {}); // https://stackoverflow.com/questions/69169492/async-external-function-leaves-open-handles-jest-supertest-express
-      let res = await axios.post(
-        `https://api.tenderly.co/api/v1/account/${TENDERLY_ACCOUNT_ID}/project/${TENDERLY_PROJECT}/vnets`,
-        {
-          slug: `e2e-tests-testnetwork-${this.network.toString()}-${Date.now()}`,
-          fork_config: {
-            network_id: this.network,
-          },
-          virtual_network_config: {
-            chain_config: {
-              chain_id: this.network,
-            },
-          },
-          sync_state_config: {
-            enabled: false,
-          },
-          explorer_page_config: {
-            enabled: true,
-            verification_visibility: 'bytecode',
-          },
-        },
-        {
-          timeout: 200000,
-          headers: {
-            'x-access-key': TENDERLY_TOKEN,
-          },
-        },
+      const { data } = await axios.post(
+        `https://api.tenderly.co/api/v1/account/${TENDERLY_ACCOUNT_ID}/project/${TENDERLY_PROJECT}/simulations/${id}`,
+        {},
+        { headers: { 'X-Access-Key': TENDERLY_TOKEN } },
       );
 
-      this.vnetId = res.data.id;
-      const rpc = findAdminRPC(res.data.rpcs);
-      if (!rpc) {
-        throw new Error(`RPC url was not found for testnet: ${this.vnetId}`);
-      }
-
-      this.rpcURL = rpc.url;
+      return data;
     } catch (e) {
-      console.error(`TenderlySimulation_setup:`, e);
-      throw e;
+      console.error(e);
+      return null;
     }
   }
 
-  async simulate(params: TxObject, stateOverrides?: StateOverrides) {
-    try {
-      let stateOverridesParams = {};
-
-      if (stateOverrides) {
-        await process.nextTick(() => {}); // https://stackoverflow.com/questions/69169492/async-external-function-leaves-open-handles-jest-supertest-express
-        const result = await axios.post(
-          `https://api.tenderly.co/api/v1/account/${TENDERLY_ACCOUNT_ID}/project/${TENDERLY_PROJECT}/contracts/encode-states`,
-          stateOverrides,
-          {
-            headers: {
-              'x-access-key': TENDERLY_TOKEN!,
-            },
-          },
-        );
+  getSLOADCalls = (
+    callTrace: SimulatedTransactionCall,
+  ): SimulatedTransactionCall[] => {
+    const results: SimulatedTransactionCall[] = [];
+
+    if (callTrace.call_type === 'SLOAD') {
+      results.push(callTrace);
+    }
+
+    if (callTrace.calls) {
+      for (const call of callTrace.calls) {
+        results.push(...this.getSLOADCalls(call));
+      }
+    }
+
+    return results;
+  };
 
-        stateOverridesParams = Object.keys(result.data.stateOverrides).reduce(
-          (acc, contract) => {
-            const _storage = result.data.stateOverrides[contract].value;
+  /**
+   *
+   * @param balanceOfSlot storage slot of `balanceOf` mapping
+   * @param owner account's address
+   * @param isVyper `true` if contract is written in Vyper
+   */
+  calculateAddressBalanceSlot(
+    balanceOfSlot: string,
+    owner: string,
+    isVyper = false,
+  ) {
+    return isVyper
+      ? this.calculateVyperAddressBalanceSlot(balanceOfSlot, owner)
+      : this.calculateSolidityAddressBalanceSlot(balanceOfSlot, owner);
+  }
+
+  /**
+   *
+   * @param balanceOfSlot storage slot of `balanceOf` mapping
+   * @param owner account's address
+   */
+  calculateSolidityAddressBalanceSlot(balanceOfSlot: string, owner: string) {
+    return ethers.utils.keccak256(
+      ethers.utils.concat([ethers.utils.hexZeroPad(owner, 32), balanceOfSlot]),
+    );
+  }
+
+  /**
+   *
+   * @param balanceOfSlot storage slot of `balanceOf` mapping
+   * @param owner account's address
+   */
+  calculateVyperAddressBalanceSlot(balanceOfSlot: string, owner: string) {
+    return ethers.utils.keccak256(
+      ethers.utils.concat([balanceOfSlot, ethers.utils.hexZeroPad(owner, 32)]),
+    );
+  }
 
-            acc[contract] = {
-              storage: _storage,
-            };
-            return acc;
-          },
-          {} as Record<Address, StateSimulateApiOverride>,
+  /**
+   *
+   * @param allowanceSlot storage slot of `allowance` mapping
+   * @param owner account's address
+   * @param spender spender's address
+   * @param isVyper `true` if contract is written in Vyper
+   */
+  calculateAddressAllowanceSlot(
+    allowanceSlot: string,
+    owner: string,
+    spender: string,
+    isVyper = false,
+  ) {
+    return isVyper
+      ? this.calculateVyperAddressAllowanceSlot(allowanceSlot, owner, spender)
+      : this.calculateSolidityAddressAllowanceSlot(
+          allowanceSlot,
+          owner,
+          spender,
         );
+  }
 
-        await this.executeStateOverrides(stateOverridesParams);
-      }
+  /**
+   *
+   * @param allowanceSlot storage slot of `allowance` mapping
+   * @param owner account's address
+   * @param spender spender's address
+   */
+  calculateSolidityAddressAllowanceSlot(
+    allowanceSlot: string,
+    owner: string,
+    spender: string,
+  ) {
+    const slotHash = ethers.utils.keccak256(
+      ethers.utils.concat([
+        ethers.utils.hexZeroPad(owner, 32),
+        ethers.utils.hexZeroPad(allowanceSlot, 32),
+      ]),
+    );
 
-      await process.nextTick(() => {}); // https://stackoverflow.com/questions/69169492/async-external-function-leaves-open-handles-jest-supertest-express
-      const { data } = await axios.post(
-        `https://api.tenderly.co/api/v1/account/${TENDERLY_ACCOUNT_ID}/project/${TENDERLY_PROJECT}/vnets/${this.vnetId}/transactions`,
-        {
-          callArgs: {
-            from: params.from,
-            to: params.to,
-            value:
-              params.value === '0'
-                ? '0x0'
-                : ethers.utils.hexStripZeros(
-                    ethers.utils.hexlify(BigInt(params.value)),
-                  ),
-            gas: ethers.utils.hexStripZeros(
-              ethers.utils.hexlify(BigInt(this.maxGasLimit)),
-            ),
-            data: params.data,
-          },
-          blockNumber: 'pending',
-        },
-        {
-          timeout: 30 * 1000,
-          headers: {
-            'x-access-key': TENDERLY_TOKEN!,
-          },
-        },
-      );
+    return ethers.utils.keccak256(
+      ethers.utils.concat([ethers.utils.hexZeroPad(spender, 32), slotHash]),
+    );
+  }
+
+  /**
+   *
+   * @param allowanceSlot storage slot of `allowance` mapping
+   * @param owner account's address
+   * @param spender spender's address
+   */
+  calculateVyperAddressAllowanceSlot(
+    allowanceSlot: string,
+    owner: string,
+    spender: string,
+  ) {
+    const slotHash = ethers.utils.keccak256(
+      ethers.utils.concat([
+        ethers.utils.hexZeroPad(allowanceSlot, 32),
+        ethers.utils.hexZeroPad(owner, 32),
+      ]),
+    );
+
+    return ethers.utils.keccak256(
+      ethers.utils.concat([slotHash, ethers.utils.hexZeroPad(spender, 32)]),
+    );
+  }
+
+  buildBalanceOfSimulationRequest(
+    chainId: number,
+    token: string,
+    owner: string,
+  ): SimulateTransactionRequest {
+    const iface = new ethers.utils.Interface([
+      'function balanceOf(address owner) view returns (uint)',
+    ]);
+
+    return {
+      from: ethers.constants.AddressZero,
+      to: token,
+      data: iface.encodeFunctionData('balanceOf', [owner]),
+      chainId,
+    };
+  }
+
+  buildAllowanceSimulationRequest(
+    chainId: number,
+    token: string,
+    owner: string,
+    spender: string,
+  ): SimulateTransactionRequest {
+    const iface = new ethers.utils.Interface([
+      'function allowance(address owner, address spender) view returns (uint)',
+    ]);
+
+    return {
+      from: ethers.constants.AddressZero,
+      to: token,
+      data: iface.encodeFunctionData('allowance', [owner, spender]),
+      chainId,
+    };
+  }
+
+  /**
+   * Finds the slot of the `balanceOf` mapping in given token contract's storage.
+   * Supports `Solidity` and `Vyper` contracts
+   * @param chainId token chain id
+   * @param token token address
+   */
+  async findTokenBalanceOfSlot(
+    chainId: number,
+    token: string,
+  ): Promise<{ slot: string; isVyper?: boolean }> {
+    const account = TenderlySimulator.DEFAULT_OWNER;
+
+    const balanceOfSimulationRequest = this.buildBalanceOfSimulationRequest(
+      chainId,
+      token,
+      account,
+    );
+
+    const { id: simulationId } = await this.simulateTransaction(
+      balanceOfSimulationRequest,
+    );
+
+    const simulationDetails = await this.getSimulatedTransactionDetails(
+      simulationId,
+    );
+
+    if (!simulationDetails) {
+      throw `No simulation with id ${simulationId} details found`;
+    }
 
-      if (data.status === 'success') {
-        return {
-          success: true,
-          gasUsed: data.gasUsed,
-          url: `https://dashboard.tenderly.co/${TENDERLY_ACCOUNT_ID}/${TENDERLY_PROJECT}/testnet/${
-            this.vnetId
-          }/tx/${this.chainIdToChainNameMap[this.network]}/${data.id}`,
-          transaction: data.input,
-        };
-      } else {
-        return {
-          success: false,
-          url: `https://dashboard.tenderly.co/${TENDERLY_ACCOUNT_ID}/${TENDERLY_PROJECT}/testnet/${
-            this.vnetId
-          }/tx/${this.chainIdToChainNameMap[this.network]}/${data.id}`,
-          error: `Simulation failed ${data.error_reason}`,
-        };
+    const callTrace = simulationDetails.transaction.transaction_info.call_trace;
+
+    const sloadCalls = this.getSLOADCalls(callTrace);
+    // token's storage slots that were read during the `balanceOf` call
+    const readSlots = sloadCalls
+      .map(call => call.storage_slot?.[0])
+      .filter<string>((slot): slot is string => !!slot);
+
+    const startingPoints = [
+      // regular contract
+      0n,
+      // ERC20Upgradeable
+      BigInt(
+        '0x52c63247e1f47db19d5ce0460030c497f067ca4cebf71ba98eeadabe20bace00',
+      ),
+    ];
+
+    for (const startingPoint of startingPoints) {
+      for (let i = startingPoint; i < startingPoint + 1_000n; i += 1n) {
+        const candidateSlot = ethers.utils.defaultAbiCoder.encode(
+          ['uint'],
+          [i],
+        );
+        // try solidity slot
+        const solitidyBalanceOfSlot = this.calculateSolidityAddressBalanceSlot(
+          candidateSlot,
+          account,
+        );
+        if (readSlots.includes(solitidyBalanceOfSlot)) {
+          return { slot: candidateSlot };
+        }
+        // try vyper slot
+        const vyperBalanceOfSlot = this.calculateVyperAddressBalanceSlot(
+          candidateSlot,
+          account,
+        );
+        if (readSlots.includes(vyperBalanceOfSlot)) {
+          return { slot: candidateSlot, isVyper: true };
+        }
       }
-    } catch (e) {
-      console.error('TenderlySimulation_simulate_error', e);
-      return {
-        success: false,
-      };
     }
+
+    throw new Error(
+      `Could not find a 'balanceOf' mapping slot for token ${token} on chain ${chainId}`,
+    );
   }
 
-  async executeStateOverrides(
-    stateOverridesParams: Record<string, StateSimulateApiOverride>,
-  ) {
-    if (!this.rpcURL) {
-      throw new Error(
-        `rpcURL is not defined for testnet: ${this.vnetId} (https://dashboard.tenderly.co/${TENDERLY_ACCOUNT_ID}/${TENDERLY_PROJECT}/testnet/${this.vnetId}`,
-      );
+  /**
+   * Finds the slot of the `allowance` mapping in given token contract's storage
+   * Supports `Solidity` and `Vyper` contracts
+   * @param chainId token chain id
+   * @param token token address
+   */
+  async findTokenAllowanceSlot(
+    chainId: number,
+    token: string,
+  ): Promise<{ slot: string; isVyper?: boolean }> {
+    const account = TenderlySimulator.DEFAULT_OWNER;
+    const spender = TenderlySimulator.DEFAULT_SPENDER;
+
+    const allowanceSimulationRequest = this.buildAllowanceSimulationRequest(
+      chainId,
+      token,
+      account,
+      spender,
+    );
+
+    const { id: simulationId } = await this.simulateTransaction(
+      allowanceSimulationRequest,
+    );
+
+    const simulationDetails = await this.getSimulatedTransactionDetails(
+      simulationId,
+    );
+
+    if (!simulationDetails) {
+      throw `No simulation with id ${simulationId} details found`;
     }
 
-    const testNetRPC = new StaticJsonRpcProvider(this.rpcURL);
-    await Promise.all(
-      Object.keys(stateOverridesParams).map(async addr => {
-        const storage = stateOverridesParams[addr].storage;
-
-        await Promise.all(
-          Object.keys(storage).map(async slot => {
-            const txHash = await testNetRPC!.send('tenderly_setStorageAt', [
-              addr,
-              slot,
-              storage[slot],
-            ]);
-
-            const transaction = await testNetRPC!.waitForTransaction(txHash);
-            if (!transaction.status) {
-              console.log(`Transaction failed: ${txHash}`);
-            }
-          }),
+    const callTrace = simulationDetails.transaction.transaction_info.call_trace;
+
+    const sloadCalls = this.getSLOADCalls(callTrace);
+    // token's storage slots that were read during the `allowance` call
+    const readSlots = sloadCalls
+      .map(call => call.storage_slot?.[0])
+      .filter<string>((slot): slot is string => !!slot);
+
+    const startingPoints = [
+      // regular contract
+      0n,
+      // ERC20Upgradeable
+      BigInt(
+        '0x52c63247e1f47db19d5ce0460030c497f067ca4cebf71ba98eeadabe20bace00',
+      ),
+    ];
+
+    for (const init of startingPoints) {
+      for (let i = init; i < init + 1000n; i += 1n) {
+        const candidateSlot = ethers.utils.defaultAbiCoder.encode(
+          ['uint'],
+          [i],
         );
-      }),
+        // try solidity
+        const solidityAllowanceSlot =
+          this.calculateSolidityAddressAllowanceSlot(
+            candidateSlot,
+            account,
+            spender,
+          );
+        if (readSlots.includes(solidityAllowanceSlot)) {
+          return { slot: candidateSlot };
+        }
+        // try vyper
+        const vyperAllowanceSlot = this.calculateVyperAddressAllowanceSlot(
+          candidateSlot,
+          account,
+          spender,
+        );
+        if (readSlots.includes(vyperAllowanceSlot)) {
+          return { slot: candidateSlot, isVyper: true };
+        }
+      }
+    }
+
+    throw new Error(
+      `Could not find a 'allowance' mapping slot for token ${token} on chain ${chainId}`,
+    );
+  }
+
+  /**
+   * Returns storage slots for the given token contract.
+   * @param chainId Token chain ID
+   * @param token Token address. Doesn't have to be normalized
+   */
+  async getTokenStorageSlots(
+    chainId: number,
+    token: string,
+  ): Promise<TokenStorageSlots> {
+    const normalizedToken = token.toLowerCase();
+
+    if (normalizedToken === ETHER_ADDRESS) {
+      throw new Error('Cannot provide storage slots for native token');
+    }
+
+    // check cache
+    const cachedSlots = await TokenStorageSlotsCache.getTokenStorageSlots(
+      chainId,
+      normalizedToken,
+    );
+
+    // return if cached
+    if (cachedSlots) {
+      return cachedSlots;
+    }
+
+    // find the slots
+    const [balanceSlot, allowanceSlot] = await Promise.all([
+      this.findTokenBalanceOfSlot(chainId, token),
+      this.findTokenAllowanceSlot(chainId, token),
+    ]);
+
+    // save the slots and return
+    const slots: TokenStorageSlots = {
+      balanceSlot: balanceSlot.slot,
+      allowanceSlot: allowanceSlot.slot,
+      isVyper: balanceSlot.isVyper,
+    };
+
+    // no need to await
+    void TokenStorageSlotsCache.setTokenStorageSlots(chainId, token, slots);
+
+    return slots;
+  }
+
+  /**
+   * Adds native balance override to an existing `StateOverride` object
+   * @param stateOverride object to add the override to
+   * @param account address to be given the balance
+   * @param amount token amount in wei
+   */
+  addBalanceOverride(
+    stateOverride: StateOverride,
+    account: string,
+    amount: bigint,
+  ): void {
+    // add the balance override
+    stateOverride[account] ||= {};
+    stateOverride[account].balance = amount.toString();
+  }
+
+  /**
+   * Adds token balance override to an existing `StateOverride` object
+   * @param stateOverride object to add the override to
+   * @param chainId token chain ID
+   * @param token token address
+   * @param account address to be given the balance
+   * @param amount token amount in wei
+   */
+  async addTokenBalanceOverride(
+    stateOverride: StateOverride,
+    chainId: number,
+    token: string,
+    account: string,
+    amount: bigint,
+  ): Promise<void> {
+    // get mapping slots
+    const tokenSlots = await this.getTokenStorageSlots(chainId, token);
+    // calculate balance slot
+    const slotToOverride = this.calculateAddressBalanceSlot(
+      tokenSlots.balanceSlot,
+      account,
+      tokenSlots.isVyper,
+    );
+    // add the balance override
+    stateOverride[token] ||= {};
+    stateOverride[token].storage ||= {};
+    stateOverride[token].storage[slotToOverride] =
+      ethers.utils.defaultAbiCoder.encode(['uint'], [amount]);
+  }
+
+  /**
+   * Adds token allowance override to an existing `StateOverride` object
+   * @param stateOverride object to add the override to
+   * @param chainId token chain ID
+   * @param token token address
+   * @param account owner address
+   * @param spender spender address
+   * @param amount token amount in wei
+   */
+  async addAllowanceOverride(
+    stateOverride: StateOverride,
+    chainId: number,
+    token: string,
+    account: string,
+    spender: string,
+    amount: bigint,
+  ): Promise<void> {
+    // get mapping slots
+    const tokenSlots = await this.getTokenStorageSlots(chainId, token);
+    // calculate allowance slot
+    const slotToOverride = this.calculateAddressAllowanceSlot(
+      tokenSlots.allowanceSlot,
+      account,
+      spender,
+      tokenSlots.isVyper,
     );
+    // add the allowance override
+    stateOverride[token] ||= {};
+    stateOverride[token].storage ||= {};
+    stateOverride[token].storage[slotToOverride] =
+      ethers.utils.defaultAbiCoder.encode(['uint'], [amount]);
   }
 }