From 221ecbf1656098e48794aa05e4d6cd23221b1e94 Mon Sep 17 00:00:00 2001
From: Christophe <contact@christophedeveaux.com>
Date: Tue, 2 Jul 2024 10:22:32 +0000
Subject: [PATCH] Fetch rollupAddress from sequencerInbox in getters

---
 src/actions/getConfirmPeriodBlocks.ts      |  8 ++++--
 src/actions/getExtraChallengeTimeBlocks.ts |  8 ++++--
 src/actions/getMinimumAssertionPeriod.ts   | 13 ++++-----
 src/actions/getWasmModuleRoot.ts           |  8 ++++--
 src/decorators/publicActionsParentChain.ts |  8 ++++--
 src/getRollupAddress.ts                    | 27 ++++++++++++++++++
 src/types/Actions.ts                       | 32 ++++++++++++++++++++++
 7 files changed, 86 insertions(+), 18 deletions(-)
 create mode 100644 src/getRollupAddress.ts

diff --git a/src/actions/getConfirmPeriodBlocks.ts b/src/actions/getConfirmPeriodBlocks.ts
index 327aa571..120db86a 100644
--- a/src/actions/getConfirmPeriodBlocks.ts
+++ b/src/actions/getConfirmPeriodBlocks.ts
@@ -1,8 +1,9 @@
 import { Chain, PublicClient, ReadContractReturnType, Transport } from 'viem';
 import { rollupAdminLogic } from '../contracts';
-import { ActionParameters } from '../types/Actions';
+import { WithContractAddress } from '../types/Actions';
+import { getRollupAddress } from '../getRollupAddress';
 
-export type GetConfirmPeriodBlocksParameters<Curried extends boolean = false> = ActionParameters<
+export type GetConfirmPeriodBlocksParameters<Curried extends boolean = false> = WithContractAddress<
   {},
   'rollupAdminLogic',
   Curried
@@ -17,9 +18,10 @@ export async function getConfirmPeriodBlocks<TChain extends Chain | undefined>(
   client: PublicClient<Transport, TChain>,
   args: GetConfirmPeriodBlocksParameters,
 ): Promise<GetConfirmPeriodBlocksReturnType> {
+  const rollupAdminLogicAddresss = await getRollupAddress(client, args);
   return client.readContract({
     abi: rollupAdminLogic.abi,
     functionName: 'confirmPeriodBlocks',
-    address: args.rollupAdminLogic,
+    address: rollupAdminLogicAddresss,
   });
 }
diff --git a/src/actions/getExtraChallengeTimeBlocks.ts b/src/actions/getExtraChallengeTimeBlocks.ts
index ede8f8b2..0bce1c4f 100644
--- a/src/actions/getExtraChallengeTimeBlocks.ts
+++ b/src/actions/getExtraChallengeTimeBlocks.ts
@@ -1,9 +1,10 @@
 import { Chain, PublicClient, ReadContractReturnType, Transport } from 'viem';
 import { rollupAdminLogic } from '../contracts';
-import { ActionParameters } from '../types/Actions';
+import { WithContractAddress } from '../types/Actions';
+import { getRollupAddress } from '../getRollupAddress';
 
 export type GetExtraChallengeTimeBlocksParameters<Curried extends boolean = false> =
-  ActionParameters<{}, 'rollupAdminLogic', Curried>;
+  WithContractAddress<{}, 'rollupAdminLogic', Curried>;
 
 export type GetExtraChallengeTimeBlocksReturnType = ReadContractReturnType<
   typeof rollupAdminLogic.abi,
@@ -14,9 +15,10 @@ export async function getExtraChallengeTimeBlocks<TChain extends Chain | undefin
   client: PublicClient<Transport, TChain>,
   args: GetExtraChallengeTimeBlocksParameters,
 ): Promise<GetExtraChallengeTimeBlocksReturnType> {
+  const rollupAdminLogicAddresss = await getRollupAddress(client, args);
   return client.readContract({
     abi: rollupAdminLogic.abi,
     functionName: 'extraChallengeTimeBlocks',
-    address: args.rollupAdminLogic,
+    address: rollupAdminLogicAddresss,
   });
 }
diff --git a/src/actions/getMinimumAssertionPeriod.ts b/src/actions/getMinimumAssertionPeriod.ts
index c250b2ce..0ff81960 100644
--- a/src/actions/getMinimumAssertionPeriod.ts
+++ b/src/actions/getMinimumAssertionPeriod.ts
@@ -1,12 +1,10 @@
 import { Chain, PublicClient, ReadContractReturnType, Transport } from 'viem';
 import { rollupAdminLogic } from '../contracts';
-import { ActionParameters } from '../types/Actions';
+import { WithContractAddress } from '../types/Actions';
+import { getRollupAddress } from '../getRollupAddress';
 
-export type GetMinimumAssertionPeriodParameters<Curried extends boolean = false> = ActionParameters<
-  {},
-  'rollupAdminLogic',
-  Curried
->;
+export type GetMinimumAssertionPeriodParameters<Curried extends boolean = false> =
+  WithContractAddress<{}, 'rollupAdminLogic', Curried>;
 
 export type GetMinimumAssertionPeriodReturnType = ReadContractReturnType<
   typeof rollupAdminLogic.abi,
@@ -17,9 +15,10 @@ export async function getMinimumAssertionPeriod<TChain extends Chain | undefined
   client: PublicClient<Transport, TChain>,
   args: GetMinimumAssertionPeriodParameters,
 ): Promise<GetMinimumAssertionPeriodReturnType> {
+  const rollupAdminLogicAddresss = await getRollupAddress(client, args);
   return client.readContract({
     abi: rollupAdminLogic.abi,
     functionName: 'minimumAssertionPeriod',
-    address: args.rollupAdminLogic,
+    address: rollupAdminLogicAddresss,
   });
 }
diff --git a/src/actions/getWasmModuleRoot.ts b/src/actions/getWasmModuleRoot.ts
index ba6b23cd..8798b32d 100644
--- a/src/actions/getWasmModuleRoot.ts
+++ b/src/actions/getWasmModuleRoot.ts
@@ -1,8 +1,9 @@
 import { Chain, PublicClient, ReadContractReturnType, Transport } from 'viem';
 import { rollupAdminLogic } from '../contracts';
-import { ActionParameters } from '../types/Actions';
+import { WithContractAddress } from '../types/Actions';
+import { getRollupAddress } from '../getRollupAddress';
 
-export type GetWasmModuleRootParameters<Curried extends boolean = false> = ActionParameters<
+export type GetWasmModuleRootParameters<Curried extends boolean = false> = WithContractAddress<
   {},
   'rollupAdminLogic',
   Curried
@@ -17,9 +18,10 @@ export async function getWasmModuleRoot<TChain extends Chain | undefined>(
   client: PublicClient<Transport, TChain>,
   args: GetWasmModuleRootParameters,
 ): Promise<GetWasmModuleRootReturnType> {
+  const rollupAdminLogicAddresss = await getRollupAddress(client, args);
   return client.readContract({
     abi: rollupAdminLogic.abi,
     functionName: 'wasmModuleRoot',
-    address: args.rollupAdminLogic,
+    address: rollupAdminLogicAddresss,
   });
 }
diff --git a/src/decorators/publicActionsParentChain.ts b/src/decorators/publicActionsParentChain.ts
index b4995147..4b9a6d13 100644
--- a/src/decorators/publicActionsParentChain.ts
+++ b/src/decorators/publicActionsParentChain.ts
@@ -72,9 +72,13 @@ export type PublicActionsParentChain<Curried extends boolean> = {
 
 function getSequencerInboxAddress(
   params: Params,
-  args: { sequencerInbox?: Address } | void,
+  args: { sequencerInbox?: Address } | {} | void,
 ): Address {
-  return ((args && args.sequencerInbox) ?? (params && params.sequencerInbox)) as unknown as Address;
+  if (args && 'sequencerInbox' in args && args.sequencerInbox) {
+    return args.sequencerInbox;
+  }
+
+  return (params && params.sequencerInbox) as Address;
 }
 
 /**
diff --git a/src/getRollupAddress.ts b/src/getRollupAddress.ts
new file mode 100644
index 00000000..532d0e8a
--- /dev/null
+++ b/src/getRollupAddress.ts
@@ -0,0 +1,27 @@
+import { Address, Chain, PublicClient, Transport } from 'viem';
+import { sequencerInboxABI } from './abi';
+
+const cache: Record<Address, Address> = {};
+export async function getRollupAddress<TChain extends Chain | undefined>(
+  publicClient: PublicClient<Transport, TChain>,
+  params: { sequencerInbox: Address } | { rollupAdminLogic: Address },
+): Promise<Address> {
+  // rollupAdminLogic was passed as an override, return directly
+  if ('rollupAdminLogic' in params) {
+    return params.rollupAdminLogic;
+  }
+
+  const addressFromCache = cache[params.sequencerInbox];
+  if (addressFromCache) {
+    return addressFromCache;
+  }
+
+  // Otherwise, fetch the rollup address from sequencerInbox contract
+  const rollupAddress = await publicClient.readContract({
+    functionName: 'rollup',
+    address: params.sequencerInbox,
+    abi: sequencerInboxABI,
+  });
+  cache[params.sequencerInbox] = rollupAddress;
+  return rollupAddress;
+}
diff --git a/src/types/Actions.ts b/src/types/Actions.ts
index e61648f0..fb7717f0 100644
--- a/src/types/Actions.ts
+++ b/src/types/Actions.ts
@@ -24,3 +24,35 @@ export type ActionParameters<Args, ContractName extends string, Curried extends
 export type WithAccount<Args> = Args & {
   account: Address;
 };
+
+/**
+ * Some actions require a different contract than sequencerInbox.
+ * We either accept sequencerInbox to fetch the contract address
+ * or an override
+ */
+export type WithContractAddress<
+  Args,
+  ContractName extends string,
+  Curried extends boolean,
+> = Curried extends true
+  ? /**
+     * If sequencerInbox was passed to the decorator. We accept the contract address,
+     * an sequencerInbox override, or no parameters
+     */
+    | (Args & {
+          sequencerInbox?: Address;
+        })
+      | (Args & {
+          [key in ContractName]?: Address;
+        })
+      | void
+  : /**
+     * If sequencerInbox wasn't passed to the decorator. We need one of the address to be passed
+     * We either accept the contract address or an sequencerInbox override
+     */
+    | (Args & {
+          sequencerInbox: Address;
+        })
+      | (Args & {
+          [key in ContractName]: Address;
+        });