From c25cf4116bcba033930f11b2823d0cbd4697e274 Mon Sep 17 00:00:00 2001
From: Christophe <contact@christophedeveaux.com>
Date: Mon, 15 Jul 2024 15:45:20 +0000
Subject: [PATCH] Add upgradeExecutor flag

---
 src/actions/addExecutor.ts    | 61 ++++++++++++++---------------------
 src/actions/removeExecutor.ts | 61 ++++++++++++++---------------------
 src/withUpgradeExecutor.ts    |  4 +--
 3 files changed, 52 insertions(+), 74 deletions(-)

diff --git a/src/actions/addExecutor.ts b/src/actions/addExecutor.ts
index f4f5ca37..e9cc5c83 100644
--- a/src/actions/addExecutor.ts
+++ b/src/actions/addExecutor.ts
@@ -1,50 +1,39 @@
-import {
-  Address,
-  Chain,
-  PrepareTransactionRequestParameters,
-  PrepareTransactionRequestReturnType,
-  PublicClient,
-  Transport,
-  encodeFunctionData,
-} from 'viem';
+import { Address, Chain, PrepareTransactionRequestParameters, PublicClient, Transport } from 'viem';
 import { upgradeExecutor } from '../contracts';
-import { ActionParameters, WithAccount } from '../types/Actions';
+import { PrepareTransactionRequestReturnTypeWithChainId, WithAccount } from '../types/Actions';
 import { Prettify } from '../types/utils';
 import { UPGRADE_EXECUTOR_ROLE_EXECUTOR } from '../upgradeExecutorEncodeFunctionData';
+import { withUpgradeExecutor } from '../withUpgradeExecutor';
 
-export type AddExecutorParameters<Curried extends boolean = false> = Prettify<
-  WithAccount<
-    ActionParameters<
-      {
-        address: Address;
-      },
-      'upgradeExecutor',
-      Curried
-    >
-  >
+export type AddExecutorParameters = Prettify<
+  WithAccount<{
+    address: Address;
+    upgradeExecutor: Address;
+  }>
 >;
 
-export type AddExecutorReturnType = PrepareTransactionRequestReturnType;
-
-function upgradeExecutorFunctionData({ address }: AddExecutorParameters) {
-  return encodeFunctionData({
-    abi: upgradeExecutor.abi,
-    functionName: 'grantRole',
-    args: [UPGRADE_EXECUTOR_ROLE_EXECUTOR, address],
-  });
-}
+export type AddExecutorReturnType = PrepareTransactionRequestReturnTypeWithChainId;
 
 export async function addExecutor<TChain extends Chain | undefined>(
   client: PublicClient<Transport, TChain>,
-  args: AddExecutorParameters,
+  params: AddExecutorParameters,
 ): Promise<AddExecutorReturnType> {
-  const data = upgradeExecutorFunctionData(args);
+  if (!client.chain) {
+    throw new Error("[addExecutor] client doesn't have a chain property");
+  }
+  const { account, upgradeExecutor: upgradeExecutorAddress, address } = params;
 
-  return client.prepareTransactionRequest({
-    to: args.upgradeExecutor,
-    value: BigInt(0),
+  const request = await client.prepareTransactionRequest({
     chain: client.chain,
-    data,
-    account: args.account,
+    account,
+    ...withUpgradeExecutor({
+      to: upgradeExecutorAddress,
+      upgradeExecutor: upgradeExecutorAddress,
+      args: [UPGRADE_EXECUTOR_ROLE_EXECUTOR, address],
+      abi: upgradeExecutor.abi,
+      functionName: 'grantRole',
+    }),
   } satisfies PrepareTransactionRequestParameters);
+
+  return { ...request, chainId: client.chain.id };
 }
diff --git a/src/actions/removeExecutor.ts b/src/actions/removeExecutor.ts
index 71249d26..0feadeb8 100644
--- a/src/actions/removeExecutor.ts
+++ b/src/actions/removeExecutor.ts
@@ -1,50 +1,39 @@
-import {
-  Address,
-  Chain,
-  PrepareTransactionRequestParameters,
-  PrepareTransactionRequestReturnType,
-  PublicClient,
-  Transport,
-  encodeFunctionData,
-} from 'viem';
+import { Address, Chain, PrepareTransactionRequestParameters, PublicClient, Transport } from 'viem';
 import { upgradeExecutor } from '../contracts';
-import { ActionParameters, WithAccount } from '../types/Actions';
+import { PrepareTransactionRequestReturnTypeWithChainId, WithAccount } from '../types/Actions';
 import { Prettify } from '../types/utils';
 import { UPGRADE_EXECUTOR_ROLE_EXECUTOR } from '../upgradeExecutorEncodeFunctionData';
+import { withUpgradeExecutor } from '../withUpgradeExecutor';
 
-export type RemoveExecutorParameters<Curried extends boolean = false> = Prettify<
-  WithAccount<
-    ActionParameters<
-      {
-        address: Address;
-      },
-      'upgradeExecutor',
-      Curried
-    >
-  >
+export type RemoveExecutorParameters = Prettify<
+  WithAccount<{
+    upgradeExecutor: Address;
+    address: Address;
+  }>
 >;
 
-export type RemoveExecutorReturnType = PrepareTransactionRequestReturnType;
-
-function upgradeExecutorFunctionData({ address }: RemoveExecutorParameters) {
-  return encodeFunctionData({
-    abi: upgradeExecutor.abi,
-    functionName: 'revokeRole',
-    args: [UPGRADE_EXECUTOR_ROLE_EXECUTOR, address],
-  });
-}
+export type RemoveExecutorReturnType = PrepareTransactionRequestReturnTypeWithChainId;
 
 export async function removeExecutor<TChain extends Chain | undefined>(
   client: PublicClient<Transport, TChain>,
-  args: RemoveExecutorParameters,
+  params: RemoveExecutorParameters,
 ): Promise<RemoveExecutorReturnType> {
-  const data = upgradeExecutorFunctionData(args);
+  if (!client.chain) {
+    throw new Error("[removeExecutor] client doesn't have a chain property");
+  }
+  const { account, upgradeExecutor: upgradeExecutorAddress, address } = params;
 
-  return client.prepareTransactionRequest({
-    to: args.upgradeExecutor,
-    value: BigInt(0),
+  const request = await client.prepareTransactionRequest({
     chain: client.chain,
-    data,
-    account: args.account,
+    account,
+    ...withUpgradeExecutor({
+      to: upgradeExecutorAddress,
+      upgradeExecutor: upgradeExecutorAddress,
+      args: [UPGRADE_EXECUTOR_ROLE_EXECUTOR, address],
+      abi: upgradeExecutor.abi,
+      functionName: 'revokeRole',
+    }),
   } satisfies PrepareTransactionRequestParameters);
+
+  return { ...request, chainId: client.chain.id };
 }
diff --git a/src/withUpgradeExecutor.ts b/src/withUpgradeExecutor.ts
index 901ea589..ee50ffe7 100644
--- a/src/withUpgradeExecutor.ts
+++ b/src/withUpgradeExecutor.ts
@@ -4,10 +4,10 @@ import {
   EncodeFunctionDataParameters as ViemEncodeFunctionDataParameters,
 } from 'viem';
 import { GetFunctionName } from './types/utils';
-import { arbOwner, sequencerInbox } from './contracts';
+import { arbOwner, sequencerInbox, upgradeExecutor } from './contracts';
 import { upgradeExecutorEncodeFunctionData } from './upgradeExecutorEncodeFunctionData';
 
-type ABIs = typeof sequencerInbox.abi | typeof arbOwner.abi;
+type ABIs = typeof sequencerInbox.abi | typeof arbOwner.abi | typeof upgradeExecutor.abi;
 type FunctionName<TAbi extends ABIs> = GetFunctionName<TAbi>;
 
 type EncodeFunctionDataParameters<