feat(deal): faster deal via access lists

Rubilmax · Oct 14, 2024 · d915990 · d915990
1 parent 8031567
commit d915990
Show file tree

Hide file tree

Showing 3 changed files with 110 additions and 118 deletions.
diff --git a/src/actions/test/deal.ts b/src/actions/test/deal.ts
@@ -1,49 +1,36 @@
 import type { TestClientMode } from "node_modules/viem/_types/clients/createTestClient.js";
 import {
+  type AccessList,
   type Account,
   type Address,
+  type BlockNumber,
+  type BlockTag,
   type Chain,
+  type ExactPartial,
   type GetStorageAtErrorType,
+  type Hash,
+  type Quantity,
   type ReadContractErrorType,
   type SetStorageAtErrorType,
   type TestClient,
+  type TransactionRequest,
   type Transport,
-  encodeAbiParameters,
+  encodeFunctionData,
   erc20Abi,
-  keccak256,
   numberToHex,
 } from "viem";
 import { getStorageAt, readContract, setStorageAt } from "viem/actions";
 
-let cache:
-  | Record<
-      Address,
-      {
-        type: StorageLayoutType;
-        slot: number;
-      }
-    >
-  | undefined;
-let cachePath: string | undefined;
-
-if (typeof process !== "undefined") {
-  const { homedir } = await import("node:os");
-  const { join } = await import("node:path");
-
-  cachePath = join(homedir(), ".foundry", "cache", "deal");
-
-  try {
-    const { readFileSync } = await import("node:fs");
-
-    cache = JSON.parse(await readFileSync(cachePath, "utf-8"));
-  } catch (error) {
-    console.debug(`Could not load cache: ${error}, re-initializing.`);
-
-    cache = {};
-  }
-}
-
-export type StorageLayoutType = "solidity" | "vyper";
+export type CreateAccessListRpcSchema = {
+  Method: "eth_createAccessList";
+  Parameters:
+    | [transaction: ExactPartial<TransactionRequest>]
+    | [transaction: ExactPartial<TransactionRequest>, block: BlockNumber | BlockTag | Hash];
+  ReturnType: {
+    accessList: AccessList;
+    gasUsed: Quantity;
+  };
+};
 
 export type DealParameters = {
   /* The address of the ERC20 token to deal. */
@@ -52,23 +39,10 @@ export type DealParameters = {
   recipient: Address;
   /* The amount of tokens to deal. */
   amount: bigint;
-  /* The storage slot of the `balanceOf` mapping, if known. */
-  slot?: number;
-  /* The type of storage layout used by the ERC20 token. */
-  storageType?: StorageLayoutType;
-  /* The maximum storage slot to brute-forcefully look for a `balanceOf` mapping. */
-  maxSlot?: number;
 };
 
 export type DealErrorType = GetStorageAtErrorType | SetStorageAtErrorType | ReadContractErrorType;
 
-export function getBalanceOfSlot(type: StorageLayoutType, slot: bigint, recipient: Address) {
-  if (type === "vyper")
-    return keccak256(encodeAbiParameters([{ type: "uint256" }, { type: "address" }], [slot, recipient]));
-
-  return keccak256(encodeAbiParameters([{ type: "address" }, { type: "uint256" }], [recipient, slot]));
-}
-
 /**
  * Deals ERC20 tokens to a recipient, by overriding the storage of `balanceOf(recipient)`.
  *
@@ -95,71 +69,65 @@ export function getBalanceOfSlot(type: StorageLayoutType, slot: bigint, recipien
  */
 export async function deal<chain extends Chain | undefined, account extends Account | undefined>(
   client: TestClient<TestClientMode, Transport, chain, account, false>,
-  { erc20, recipient, amount, slot, storageType, maxSlot = 256 }: DealParameters,
+  { erc20, recipient, amount }: DealParameters,
 ) {
-  const trySlot = async ({
-    type,
-    slot,
-  }: {
-    type: StorageLayoutType;
-    slot: number;
-  }) => {
-    const balanceOfSlot = getBalanceOfSlot(type, BigInt(slot), recipient);
-    const storageBefore = await getStorageAt(client, { address: erc20, slot: balanceOfSlot });
-
-    await setStorageAt(client, { address: erc20, index: balanceOfSlot, value: numberToHex(amount, { size: 32 }) });
-
-    const balance = await readContract(client, {
-      abi: erc20Abi,
-      address: erc20,
-      functionName: "balanceOf",
-      args: [recipient],
-    });
-
-    if (balance === amount) return true;
-
-    if (storageBefore != null)
-      await setStorageAt(client, { address: erc20, index: balanceOfSlot, value: storageBefore });
+  const value = numberToHex(amount, { size: 32 });
 
-    return false;
-  };
-
-  const cached = cache?.[erc20];
-  if (cached != null && (await trySlot(cached))) return;
-
-  const switchStorageType = storageType == null;
-
-  slot ??= 0;
-  storageType ??= "solidity";
-
-  let success = await trySlot({ type: storageType, slot });
-
-  while (!success && slot <= maxSlot) {
-    if (switchStorageType) {
-      if (storageType === "solidity") storageType = "vyper";
-      else {
-        ++slot;
-        storageType = "solidity";
-      }
-    } else ++slot;
-
-    success = await trySlot({ type: storageType, slot });
-  }
-
-  if (!success) throw Error(`Could not deal ERC20 tokens: cannot brute-force "balanceOf" storage slot at "${erc20}"`);
-
-  if (cache != null && cachePath != null) {
-    cache[erc20] = { type: storageType, slot };
-
-    try {
-      const { dirname } = await import("node:path");
-      const { mkdirSync, writeFileSync } = await import("node:fs");
-
-      await mkdirSync(dirname(cachePath), { recursive: true });
-
-      await writeFileSync(cachePath, JSON.stringify(cache));
-    } catch (error) {
-      console.error(`Could not save cache: ${error}`);
+  const { accessList } = await client.request<CreateAccessListRpcSchema>({
+    method: "eth_createAccessList",
+    params: [
+      {
+        to: erc20,
+        data: encodeFunctionData({
+          abi: erc20Abi,
+          functionName: "balanceOf",
+          args: [recipient],
+        }),
+      },
+    ],
+  });
+
+  for (const { address: address_, storageKeys } of reverse(accessList)) {
+    // Address needs to be lower-case to work with setStorageAt.
+    const address = address_.toLowerCase() as Address;
+
+    for (const slot of reverse(storageKeys)) {
+      const storageBefore = await getStorageAt(client, { address, slot });
+
+      await setStorageAt(client, { address, index: slot, value });
+
+      try {
+        const balance = await readContract(client, {
+          abi: erc20Abi,
+          address: erc20,
+          functionName: "balanceOf",
+          args: [recipient],
+        });
+
+        if (balance === amount) return;
+      } catch {}
+
+      if (storageBefore != null) await setStorageAt(client, { address, index: slot, value: storageBefore });
     }
   }
+
+  throw Error(`Could not deal ERC20 tokens: cannot find valid "balanceOf" storage slot for "${erc20}"`);
 }
+
+const reverse = <T>(arr: readonly T[]) => {
+  let index = arr.length;
+
+  return {
+    next() {
+      index--;
+
+      return {
+        done: index < 0,
+        value: arr[index]!,
+      };
+    },
+    [Symbol.iterator]() {
+      return this;
+    },
+  };
+};
diff --git a/src/dealActions.ts b/src/dealActions.ts
@@ -6,9 +6,7 @@ export type DealActions = {
   /**
    * Deals ERC20 tokens to a recipient, by overriding the storage of `balanceOf(recipient)`.
    *
-   * - Docs: https://viem.sh/docs/actions/test/deal
-   *
-   * @param args - {@link DropTransactionParameters}
+   * @param args - {@link DealParameters}
    *
    * @example
    * import { createTestClient, http } from 'viem'

diff --git a/test/deal.test.ts b/test/deal.test.ts
@@ -7,6 +7,8 @@ const aave = "0x7fc66500c84a76ad7e9c93437bfc5ac33e2ddae9";
 const crv = "0xd533a949740bb3306d119cc777fa900ba034cd52";
 const cbEth = "0xBe9895146f7AF43049ca1c1AE358B0541Ea49704";
 const maDai = "0x36F8d0D0573ae92326827C4a82Fe4CE4C244cAb6";
+const usd0 = "0x35D8949372D46B7a3D5A56006AE77B215fc69bC0";
+const stEth = "0xae7ab96520DE3A18E5e111B5EaAb095312D7fE84";
 
 describe("deal", () => {
   test("should deal USDC", async ({ client }) => {
@@ -109,7 +111,6 @@ describe("deal", () => {
       erc20: cbEth,
       recipient: client.account.address,
       amount: expected,
-      storageType: "solidity",
     });
 
     const balance = await client.readContract({
@@ -150,16 +151,41 @@ describe("deal", () => {
     expect(balance).toEqual(expected);
   });
 
-  test("should not deal USD0", async ({ client }) => {
-    await expect(
+  test("should deal USD0", async ({ client }) => {
+    const expected = parseUnits("1000", 18);
+
+    expect(
+      await client.readContract({
+        abi: erc20Abi,
+        address: usd0,
+        functionName: "balanceOf",
+        args: [client.account.address],
+      }),
+    ).not.toEqual(expected);
+
+    await client.deal({
+      erc20: usd0,
+      recipient: client.account.address,
+      amount: expected,
+    });
+
+    const balance = await client.readContract({
+      abi: erc20Abi,
+      address: usd0,
+      functionName: "balanceOf",
+      args: [client.account.address],
+    });
+
+    expect(balance).toEqual(expected);
+  });
+
+  test("should not deal stETH", async ({ client }) => {
+    expect(
       client.deal({
-        erc20: "0x35D8949372D46B7a3D5A56006AE77B215fc69bC0",
+        erc20: stEth,
         recipient: client.account.address,
         amount: 1n,
-        maxSlot: 10,
       }),
-    ).rejects.toThrow(
-      `Could not deal ERC20 tokens: cannot brute-force "balanceOf" storage slot at "0x35D8949372D46B7a3D5A56006AE77B215fc69bC0"`,
-    );
+    ).rejects.toBe(`Could not deal ERC20 tokens: cannot find valid "balanceOf" storage slot for "${stEth}"`);
   });
 });