Fix typing for prepareTransactionRequest

Author: chrstph-dvx

src/arbOwnerPrepareTransactionRequest.ts

@@ -5,45 +5,30 @@ import {
   Address,
   Chain,
   Transport,
+  GetFunctionArgs,
 } from 'viem';
 
-import { arbOwner, ArbOSVersions } from './contracts';
+import { arbOwner, ArbOSVersions, ArbOwnerABIs } from './contracts';
 import { upgradeExecutorEncodeFunctionData } from './upgradeExecutor';
-import { GetFunctionName } from './types/utils';
+import { ArbOwnerPublicAbi, ArbOwnerPublicFunctionName } from './arbOwnerReadContract';
 
-type ArbOwnerAbi = typeof arbOwner.abi;
-export type ArbOwnerPrepareTransactionRequestFunctionName = GetFunctionName<ArbOwnerAbi>;
 export type ArbOwnerEncodeFunctionDataParameters<
-  TFunctionName extends ArbOwnerPrepareTransactionRequestFunctionName,
-> = EncodeFunctionDataParameters<ArbOwnerAbi, TFunctionName>;
-
-function arbOwnerEncodeFunctionData<
-  TFunctionName extends ArbOwnerPrepareTransactionRequestFunctionName,
->({ functionName, abi, args }: ArbOwnerEncodeFunctionDataParameters<TFunctionName>) {
-  return encodeFunctionData({
-    abi,
-    functionName,
-    args,
-  });
-}
-
-type ArbOwnerPrepareFunctionDataParameters<
-  TFunctionName extends ArbOwnerPrepareTransactionRequestFunctionName,
-> = ArbOwnerEncodeFunctionDataParameters<TFunctionName> & {
+  TArbOsVersion extends ArbOSVersions,
+  TFunctionName extends ArbOwnerPublicFunctionName<TArbOsVersion>,
+> = EncodeFunctionDataParameters<ArbOwnerPublicAbi<TArbOsVersion>, TFunctionName> & {
   upgradeExecutor: Address | false;
-  abi: ArbOwnerAbi;
 };
+
 function arbOwnerPrepareFunctionData<
-  TFunctionName extends ArbOwnerPrepareTransactionRequestFunctionName,
->(params: ArbOwnerPrepareFunctionDataParameters<TFunctionName>) {
+  TArbOsVersion extends ArbOSVersions,
+  TFunctionName extends ArbOwnerPublicFunctionName<TArbOsVersion>,
+>(params: ArbOwnerEncodeFunctionDataParameters<TArbOsVersion, TFunctionName>) {
   const { upgradeExecutor } = params;
 
   if (!upgradeExecutor) {
     return {
       to: arbOwner.address,
-      data: arbOwnerEncodeFunctionData(
-        params as ArbOwnerEncodeFunctionDataParameters<TFunctionName>,
-      ),
+      data: encodeFunctionData(params as EncodeFunctionDataParameters<ArbOwnerPublicAbi<TArbOsVersion>, TFunctionName>),
       value: BigInt(0),
     };
   }
@@ -54,37 +39,40 @@ function arbOwnerPrepareFunctionData<
       functionName: 'executeCall',
       args: [
         arbOwner.address, // target
-        arbOwnerEncodeFunctionData(params as ArbOwnerEncodeFunctionDataParameters<TFunctionName>), // targetCallData
+        encodeFunctionData(params as EncodeFunctionDataParameters<ArbOwnerPublicAbi<TArbOsVersion>, TFunctionName>), // targetCallData
       ],
     }),
     value: BigInt(0),
   };
 }
 
-export type ArbOwnerPrepareTransactionRequestParameters<TFunctionName extends ArbOwnerPrepareTransactionRequestFunctionName> = ArbOwnerEncodeFunctionDataParameters<TFunctionName> & {
-  upgradeExecutor: Address | false;
+export type ArbOwnerPrepareTransactionRequestParameters<
+  TArbOsVersion extends ArbOSVersions,
+  TFunctionName extends ArbOwnerPublicFunctionName<TArbOsVersion>,
+> = {
+  functionName: TFunctionName;
   account: Address;
-}
+  upgradeExecutor: Address | false;
+} & GetFunctionArgs<ArbOwnerPublicAbi<TArbOsVersion>, TFunctionName>;
 
 export async function arbOwnerPrepareTransactionRequest<
   TArbOsVersion extends ArbOSVersions,
-  TFunctionName extends ArbOwnerPrepareTransactionRequestFunctionName,
   TChain extends Chain | undefined,
+  TFunctionName extends ArbOwnerPublicFunctionName<TArbOsVersion>,
 >(
   client: PublicClient<Transport, TChain>,
-  params: ArbOwnerPrepareTransactionRequestParameters<TFunctionName> & {
-    arbOSVersion: TArbOsVersion;
+  params: ArbOwnerPrepareTransactionRequestParameters<TArbOsVersion, TFunctionName> & {
+    arbOsVersion: TArbOsVersion;
   },
 ) {
   if (typeof client.chain === 'undefined') {
     throw new Error('[arbOwnerPrepareTransactionRequest] client.chain is undefined');
   }
 
-  // params is extending ArbOwnerPrepareFunctionDataParameters, it's safe to cast
   const { to, data, value } = arbOwnerPrepareFunctionData({
     ...params,
-    abi: arbOwner.abi,
-  } as unknown as ArbOwnerPrepareFunctionDataParameters<TFunctionName>);
+    abi: ArbOwnerABIs[params.arbOsVersion],
+  } as unknown as ArbOwnerEncodeFunctionDataParameters<TArbOsVersion, TFunctionName>);
 
   // @ts-ignore (todo: fix viem type issue)
   const request = await client.prepareTransactionRequest({

src/arbOwnerReadContract.ts

@@ -29,13 +29,13 @@ export function arbOwnerReadContract<
 >(
   client: PublicClient<Transport, TChain>,
   params: ArbOwnerReadContractParameters<TArbOsVersion, TFunctionName> & {
-    arbOSVersion: TArbOsVersion;
+    arbOsVersion: TArbOsVersion;
   },
 ): Promise<ArbOwnerReadContractReturnType<TArbOsVersion, TFunctionName>> {
   // @ts-ignore (todo: fix viem type issue)
   return client.readContract({
     address: arbOwnerPublic.address,
-    abi: ArbOwnerABIs[params.arbOSVersion],
+    abi: ArbOwnerABIs[params.arbOsVersion],
     functionName: params.functionName,
     args: params.args,
   });

src/contracts.ts

@@ -216,6 +216,13 @@ export const ArbOwnerABIs = {
       name: 'onlyOnArbOS10',
       outputs: [{ name: '', internalType: 'address[]', type: 'address[]' }],
     },
+    {
+      stateMutability: 'nonpayable',
+      type: 'function',
+      inputs: [{ name: 'recipient', internalType: 'address[]', type: 'address[]' }],
+      name: 'setL1PricingRewardRecipient',
+      outputs: [],
+    },
   ],
   11: [
     {
@@ -232,6 +239,13 @@ export const ArbOwnerABIs = {
       name: 'onlyOnArbOS11',
       outputs: [{ name: '', internalType: 'address[]', type: 'address[]' }],
     },
+    {
+      stateMutability: 'nonpayable',
+      type: 'function',
+      inputs: [{ name: 'recipient', internalType: 'uint64', type: 'uint64' }],
+      name: 'setL1PricingRewardRecipient',
+      outputs: [],
+    },
   ],
   20: arbOwnerConfig.abi,
 } as const;

src/decorators/arbOwnerPrepareTransactionRequest.unit.test.ts

@@ -14,7 +14,7 @@ import { generatePrivateKey, privateKeyToAccount } from 'viem/accounts';
 const client = createPublicClient({
   chain: nitroTestnodeL2,
   transport: http(),
-}).extend(arbOwnerPublicActions);
+}).extend(arbOwnerPublicActions({ arbOsVersion: 20 }));
 const randomAccount = privateKeyToAccount(generatePrivateKey());
 
 it('Infer parameters based on function name', async () => {

src/decorators/arbOwnerPublicActions.integration.test.ts

@@ -15,7 +15,7 @@ const randomAccount = privateKeyToAccount(generatePrivateKey());
 const client = createPublicClient({
   chain: nitroTestnodeL2,
   transport: http(),
-}).extend(arbOwnerPublicActions({ arbOSVersion: 20 }));
+}).extend(arbOwnerPublicActions({ arbOsVersion: 20 }));
 
 it('successfully fetches network fee receiver', async () => {
   const result = await client.arbOwnerReadContract({

src/decorators/arbOwnerPublicActions.ts

@@ -8,7 +8,6 @@ import {
 } from '../arbOwnerReadContract';
 import {
   arbOwnerPrepareTransactionRequest,
-  ArbOwnerPrepareTransactionRequestFunctionName,
   ArbOwnerPrepareTransactionRequestParameters,
 } from '../arbOwnerPrepareTransactionRequest';
 import { ArbOSVersions } from '../contracts';
@@ -22,25 +21,25 @@ export type ArbOwnerPublicActions<
   ) => Promise<ArbOwnerReadContractReturnType<TArbOsVersion, TFunctionName>>;
 
   arbOwnerPrepareTransactionRequest: <
-    TFunctionName extends ArbOwnerPrepareTransactionRequestFunctionName,
+    TFunctionName extends ArbOwnerPublicFunctionName<TArbOsVersion>,
   >(
-    args: ArbOwnerPrepareTransactionRequestParameters<TFunctionName>,
+    args: ArbOwnerPrepareTransactionRequestParameters<TArbOsVersion, TFunctionName>,
   ) => Promise<PrepareTransactionRequestReturnType<TChain> & { chainId: number }>;
 };
 
 export function arbOwnerPublicActions<
   TArbOsVersion extends ArbOSVersions,
   TTransport extends Transport = Transport,
   TChain extends Chain | undefined = Chain | undefined,
->({ arbOSVersion }: { arbOSVersion: TArbOsVersion }) {
+>({ arbOsVersion }: { arbOsVersion: TArbOsVersion }) {
   return (
     client: PublicClient<TTransport, TChain>,
   ): ArbOwnerPublicActions<TArbOsVersion, TChain> => {
     return {
-      arbOwnerReadContract: (args) => arbOwnerReadContract(client, { ...args, arbOSVersion }),
+      arbOwnerReadContract: (args) => arbOwnerReadContract(client, { ...args, arbOsVersion }),
 
       arbOwnerPrepareTransactionRequest: (args) =>
-        arbOwnerPrepareTransactionRequest(client, { ...args, arbOSVersion }),
+        arbOwnerPrepareTransactionRequest(client, { ...args, arbOsVersion }),
     };
   };
 }

src/decorators/arbOwnerPublicActions.unit.test.ts

@@ -1,81 +1,112 @@
-import { it, expect, expectTypeOf } from 'vitest';
+import { describe, it, expect, expectTypeOf } from 'vitest';
 
 import { AbiFunctionNotFoundError, createPublicClient, http } from 'viem';
 import { nitroTestnodeL2 } from '../chains';
 import { arbOwnerPublicActions } from './arbOwnerPublicActions';
+import { generatePrivateKey, privateKeyToAccount } from 'viem/accounts';
 
 const client10 = createPublicClient({
   chain: nitroTestnodeL2,
   transport: http(),
-}).extend(arbOwnerPublicActions({ arbOSVersion: 10 }));
+}).extend(arbOwnerPublicActions({ arbOsVersion: 10 }));
 const client11 = createPublicClient({
   chain: nitroTestnodeL2,
   transport: http(),
-}).extend(arbOwnerPublicActions({ arbOSVersion: 11 }));
+}).extend(arbOwnerPublicActions({ arbOsVersion: 11 }));
 const client20 = createPublicClient({
   chain: nitroTestnodeL2,
   transport: http(),
-}).extend(arbOwnerPublicActions({ arbOSVersion: 20 }));
+}).extend(arbOwnerPublicActions({ arbOsVersion: 20 }));
+const randomAccount = privateKeyToAccount(generatePrivateKey());
+const upgradeExecutorAddress = '0x24198F8A339cd3C47AEa3A764A20d2dDaB4D1b5b';
 
-it('Accept function name based on arbOSVersion', async () => {
-  // Version 10
-  client10.arbOwnerReadContract({
-    functionName: 'onlyOnArbOS10',
-  });
+describe('Accept function name based on arbOSVersion', async () => {
+  it('Version 10', () => {
+    expectTypeOf<typeof client10.arbOwnerReadContract<'onlyOnArbOS10'>>().toBeCallableWith({
+      functionName: 'onlyOnArbOS10'
+    })
 
-  expect(
-    client10.arbOwnerReadContract({
-      // @ts-expect-error Not available for version 10
-      functionName: 'onlyOnArbOS20',
-    }),
-  ).rejects.toThrowError(AbiFunctionNotFoundError);
+    expectTypeOf<typeof client10.arbOwnerPrepareTransactionRequest<'setL1PricingRewardRecipient'>>().toBeCallableWith({
+      functionName: 'setL1PricingRewardRecipient',
+      account: randomAccount.address,
+      upgradeExecutor: upgradeExecutorAddress,
+      args: [[randomAccount.address, randomAccount.address]]
+    })
 
-  // Version 11
-  client11.arbOwnerReadContract({
-    functionName: 'onlyOnArbOS11',
-  });
+    expect(
+      client10.arbOwnerReadContract({
+        // @ts-expect-error Not available for version 10
+        functionName: 'onlyOnArbOS20',
+      }),
+    ).rejects.toThrowError(AbiFunctionNotFoundError);
 
-  expect(
-    client11.arbOwnerReadContract({
-      // @ts-expect-error Not available for version 11
-      functionName: 'onlyOnArbOS20',
-    }),
-  ).rejects.toThrowError(AbiFunctionNotFoundError);
 
-  // Version 20
-  client20.arbOwnerReadContract({
-    functionName: 'getInfraFeeAccount',
-  });
+  })
 
-  expect(
-    client20.arbOwnerReadContract({
-      // @ts-expect-error Not available for version 20
-      functionName: 'onlyOnArbOS10',
-    }),
-  ).rejects.toThrowError(AbiFunctionNotFoundError);
+  it('Version 11', () => {
+    expectTypeOf<typeof client11.arbOwnerReadContract<'onlyOnArbOS11'>>().toBeCallableWith({
+      functionName: 'onlyOnArbOS11'
+    })
+
+    expectTypeOf<typeof client11.arbOwnerPrepareTransactionRequest<'setL1PricingRewardRecipient'>>().toBeCallableWith({
+      functionName: 'setL1PricingRewardRecipient',
+      account: randomAccount.address,
+      upgradeExecutor: upgradeExecutorAddress,
+      args: [100n]
+    })
+
+    expect(
+      client11.arbOwnerReadContract({
+        // @ts-expect-error Not available for version 11
+        functionName: 'onlyOnArbOS20',
+      }),
+    ).rejects.toThrowError(AbiFunctionNotFoundError);
+  })
+
+  it('Version 20', () => {
+    expectTypeOf<typeof client20.arbOwnerReadContract<'getInfraFeeAccount'>>().toBeCallableWith({
+      functionName: 'getInfraFeeAccount'
+    })
+
+    expectTypeOf<typeof client20.arbOwnerPrepareTransactionRequest<'setL1PricingRewardRecipient'>>().toBeCallableWith({
+      functionName: 'setL1PricingRewardRecipient',
+      account: randomAccount.address,
+      upgradeExecutor: upgradeExecutorAddress,
+      args: [randomAccount.address]
+    })
+
+    expect(
+      client20.arbOwnerReadContract({
+        // @ts-expect-error Not available for version 20
+        functionName: 'onlyOnArbOS10',
+      }),
+    ).rejects.toThrowError(AbiFunctionNotFoundError);
+  })
 });
 
 // Those tests won't fail if the return type is wrong
 // But they will display an error in the IDE
-it('Type return values for function in multiple versions', () => {
-  // Version 10
-  expectTypeOf(
-    client10.arbOwnerReadContract({
-      functionName: 'getAllChainOwners',
-    }),
-  ).resolves.toEqualTypeOf<`0x${string}`>();
-
-  // Version 11
-  expectTypeOf(
-    client11.arbOwnerReadContract({
-      functionName: 'getAllChainOwners',
-    }),
-  ).resolves.toEqualTypeOf<bigint>();
+describe('Type return values for function in multiple versions', () => {
+  it('Version 10', () => {
+    expectTypeOf(
+      client10.arbOwnerReadContract({
+        functionName: 'getAllChainOwners',
+      }),
+    ).resolves.toEqualTypeOf<`0x${string}`>();
+  })
 
-  // Version 20
-  expectTypeOf(
-    client20.arbOwnerReadContract({
-      functionName: 'getAllChainOwners',
-    }),
-  ).resolves.toEqualTypeOf<readonly `0x${string}`[]>();
-});
+  it('Version 11', () => {
+    expectTypeOf(
+      client11.arbOwnerReadContract({
+        functionName: 'getAllChainOwners',
+      }),
+    ).resolves.toEqualTypeOf<bigint>();
+  })
+  it('Version 11', () => {
+    expectTypeOf(
+      client20.arbOwnerReadContract({
+        functionName: 'getAllChainOwners',
+      }),
+    ).resolves.toEqualTypeOf<readonly `0x${string}`[]>();
+  })
+});