fix(protocol-kit): Fix memoization issue across different chains in t…

…he predict Safe address util (safe-global#640)

* fix memoization issue across different chains

* Added chainId to the memoizedGetSafeContract fn

* fixed performance issue in the predictSafeAddress

---------

Co-authored-by: Yago Pérez Vázquez <yago@safe.global>
Co-authored-by: Daniel <25051234+dasanra@users.noreply.github.com>

packages/account-abstraction-kit/src/AccountAbstraction.test.ts

@@ -13,7 +13,8 @@ const SafeMock = Safe as jest.MockedClass<typeof Safe>
 describe('AccountAbstraction', () => {
   const ethersAdapter = {
     getSignerAddress: jest.fn(),
-    isContractDeployed: jest.fn()
+    isContractDeployed: jest.fn(),
+    getChainId: jest.fn()
   }
   const signerAddress = '0xSignerAddress'
   const predictSafeAddress = '0xPredictSafeAddressMock'

packages/account-abstraction-kit/src/AccountAbstraction.ts

@@ -40,6 +40,7 @@ class AccountAbstraction {
 
     const safeAddress = await predictSafeAddress({
       ethAdapter: this.#ethAdapter,
+      chainId: await this.#ethAdapter.getChainId(),
       safeAccountConfig
     })
 

packages/protocol-kit/src/Safe.ts

@@ -186,6 +186,7 @@ class Safe {
       const chainId = await this.#ethAdapter.getChainId()
       return predictSafeAddress({
         ethAdapter: this.#ethAdapter,
+        chainId,
         customContracts: this.#contractManager.contractNetworks?.[chainId.toString()],
         ...this.#predictedSafe
       })

packages/protocol-kit/src/contracts/safeDeploymentContracts.ts

@@ -26,13 +26,13 @@ import {
 } from '@safe-global/safe-deployments'
 import { safeDeploymentsL1ChainIds, safeDeploymentsVersions } from './config'
 
-interface GetContractInstanceProps {
+export interface GetContractInstanceProps {
   ethAdapter: EthAdapter
   safeVersion: SafeVersion
   customContracts?: ContractNetworkConfig
 }
 
-interface GetSafeContractInstanceProps extends GetContractInstanceProps {
+export interface GetSafeContractInstanceProps extends GetContractInstanceProps {
   isL1SafeSingleton?: boolean
   customSafeAddress?: string
 }

packages/protocol-kit/src/contracts/utils.ts

@@ -13,6 +13,8 @@ import { generateAddress2, keccak256, toBuffer } from 'ethereumjs-util'
 import semverSatisfies from 'semver/functions/satisfies'
 
 import {
+  GetContractInstanceProps,
+  GetSafeContractInstanceProps,
   getCompatibilityFallbackHandlerContract,
   getProxyFactoryContract,
   getSafeContract
@@ -46,6 +48,7 @@ const ZKSYNC_CREATE2_PREFIX = '0x2020dba91b30cc0006188af794c2fb30dd8520db7e2c088
 
 export interface PredictSafeAddressProps {
   ethAdapter: EthAdapter
+  chainId: bigint // required for performance
   safeAccountConfig: SafeAccountConfig
   safeDeploymentConfig?: SafeDeploymentConfig
   isL1SafeSingleton?: boolean
@@ -132,28 +135,43 @@ export async function encodeSetupCallData({
   ])
 }
 
-const memoizedGetProxyFactoryContract = createMemoizedFunction(getProxyFactoryContract)
-const memoizedGetSafeContract = createMemoizedFunction(getSafeContract)
+// we need to include the chainId as string to prevent memoization issues see: https://github.com/safe-global/safe-core-sdk/issues/598
+type MemoizedGetProxyFactoryContractProps = GetContractInstanceProps & { chainId: string }
+type MemoizedGetSafeContractInstanceProps = GetSafeContractInstanceProps & { chainId: string }
+
+const memoizedGetProxyFactoryContract = createMemoizedFunction(
+  ({ ethAdapter, safeVersion, customContracts }: MemoizedGetProxyFactoryContractProps) =>
+    getProxyFactoryContract({ ethAdapter, safeVersion, customContracts })
+)
+
 const memoizedGetProxyCreationCode = createMemoizedFunction(
   async ({
     ethAdapter,
     safeVersion,
-    customContracts
-  }: {
-    ethAdapter: EthAdapter
-    safeVersion: SafeVersion
-    customContracts?: ContractNetworkConfig
-  }) => {
+    customContracts,
+    chainId
+  }: MemoizedGetProxyFactoryContractProps) => {
     const safeProxyFactoryContract = await memoizedGetProxyFactoryContract({
       ethAdapter,
       safeVersion,
-      customContracts
+      customContracts,
+      chainId
     })
 
     return safeProxyFactoryContract.proxyCreationCode()
   }
 )
 
+const memoizedGetSafeContract = createMemoizedFunction(
+  ({
+    ethAdapter,
+    safeVersion,
+    isL1SafeSingleton,
+    customContracts
+  }: MemoizedGetSafeContractInstanceProps) =>
+    getSafeContract({ ethAdapter, safeVersion, isL1SafeSingleton, customContracts })
+)
+
 /**
  * Provides a chain-specific default salt nonce for generating unique addresses
  * for the same Safe configuration across different chains.
@@ -167,6 +185,7 @@ export function getChainSpecificDefaultSaltNonce(chainId: bigint): string {
 
 export async function predictSafeAddress({
   ethAdapter,
+  chainId,
   safeAccountConfig,
   safeDeploymentConfig = {},
   isL1SafeSingleton = false,
@@ -175,8 +194,6 @@ export async function predictSafeAddress({
   validateSafeAccountConfig(safeAccountConfig)
   validateSafeDeploymentConfig(safeDeploymentConfig)
 
-  const chainId = await ethAdapter.getChainId()
-
   const {
     safeVersion = DEFAULT_SAFE_VERSION,
     saltNonce = getChainSpecificDefaultSaltNonce(chainId)
@@ -185,20 +202,23 @@ export async function predictSafeAddress({
   const safeProxyFactoryContract = await memoizedGetProxyFactoryContract({
     ethAdapter,
     safeVersion,
-    customContracts
+    customContracts,
+    chainId: chainId.toString()
   })
 
   const proxyCreationCode = await memoizedGetProxyCreationCode({
     ethAdapter,
     safeVersion,
-    customContracts
+    customContracts,
+    chainId: chainId.toString()
   })
 
   const safeContract = await memoizedGetSafeContract({
     ethAdapter,
     safeVersion,
     isL1SafeSingleton,
-    customContracts
+    customContracts,
+    chainId: chainId.toString()
   })
 
   const initializer = await encodeSetupCallData({

packages/protocol-kit/src/safeFactory/index.ts

@@ -128,6 +128,7 @@ class SafeFactory {
 
     return predictSafeAddress({
       ethAdapter: this.#ethAdapter,
+      chainId,
       safeAccountConfig,
       safeDeploymentConfig,
       isL1SafeSingleton: this.#isL1SafeSingleton,

packages/protocol-kit/tests/e2e/utilsContracts.test.ts

@@ -72,6 +72,7 @@ describe('Contract utils', () => {
 
       const predictedSafeAddress = await predictSafeAddress({
         ethAdapter,
+        chainId,
         safeAccountConfig,
         safeDeploymentConfig,
         customContracts
@@ -115,6 +116,7 @@ describe('Contract utils', () => {
 
       const predictedSafeAddress = await predictSafeAddress({
         ethAdapter,
+        chainId,
         safeAccountConfig,
         safeDeploymentConfig,
         customContracts
@@ -158,6 +160,7 @@ describe('Contract utils', () => {
 
       const predictedSafeAddress = await predictSafeAddress({
         ethAdapter,
+        chainId,
         safeAccountConfig,
         safeDeploymentConfig,
         customContracts
@@ -201,6 +204,7 @@ describe('Contract utils', () => {
 
       const predictSafeAddressWithInvalidThreshold = {
         ethAdapter,
+        chainId,
         safeAccountConfig,
         safeDeploymentConfig,
         customContracts
@@ -234,6 +238,7 @@ describe('Contract utils', () => {
 
       const predictSafeAddressWithInvalidThreshold = {
         ethAdapter,
+        chainId,
         safeAccountConfig,
         safeDeploymentConfig,
         customContracts
@@ -267,6 +272,7 @@ describe('Contract utils', () => {
 
       const predictSafeAddressWithInvalidThreshold = {
         ethAdapter,
+        chainId,
         safeAccountConfig,
         safeDeploymentConfig,
         customContracts
@@ -299,6 +305,7 @@ describe('Contract utils', () => {
 
       const predictSafeAddressWithInvalidThreshold = {
         ethAdapter,
+        chainId,
         safeAccountConfig,
         safeDeploymentConfig,
         customContracts
@@ -331,6 +338,7 @@ describe('Contract utils', () => {
 
       const predictedSafeAddress1 = await predictSafeAddress({
         ethAdapter,
+        chainId,
         safeAccountConfig,
         safeDeploymentConfig: {
           safeVersion,
@@ -353,6 +361,7 @@ describe('Contract utils', () => {
 
       const predictedSafeAddress2 = await predictSafeAddress({
         ethAdapter,
+        chainId,
         safeAccountConfig,
         safeDeploymentConfig: {
           safeVersion,
@@ -375,6 +384,7 @@ describe('Contract utils', () => {
 
       const predictedSafeAddress3 = await predictSafeAddress({
         ethAdapter,
+        chainId,
         safeAccountConfig,
         safeDeploymentConfig: {
           safeVersion,
@@ -429,20 +439,23 @@ describe('Contract utils', () => {
 
       const firstPredictedSafeAddress = await predictSafeAddress({
         ethAdapter,
+        chainId,
         safeAccountConfig,
         safeDeploymentConfig,
         customContracts
       })
 
       const secondPredictedSafeAddress = await predictSafeAddress({
         ethAdapter,
+        chainId,
         safeAccountConfig,
         safeDeploymentConfig,
         customContracts
       })
 
       const thirdPredictedSafeAddress = await predictSafeAddress({
         ethAdapter,
+        chainId,
         safeAccountConfig,
         safeDeploymentConfig,
         customContracts
@@ -474,6 +487,7 @@ describe('Contract utils', () => {
 
         const predictedSafeAddress = await predictSafeAddress({
           ethAdapter,
+          chainId,
           safeAccountConfig,
           customContracts
         })
@@ -517,6 +531,7 @@ describe('Contract utils', () => {
 
         const firstPredictedSafeAddress = await predictSafeAddress({
           ethAdapter,
+          chainId,
           safeAccountConfig: safeAccountConfig1,
           safeDeploymentConfig: safeDeploymentConfig1,
           customContracts
@@ -538,6 +553,7 @@ describe('Contract utils', () => {
 
         const secondPredictedSafeAddress = await predictSafeAddress({
           ethAdapter,
+          chainId,
           safeAccountConfig: safeAccountConfig2,
           safeDeploymentConfig: safeDeploymentConfig2,
           customContracts
@@ -560,6 +576,7 @@ describe('Contract utils', () => {
 
         const thirdPredictedSafeAddress = await predictSafeAddress({
           ethAdapter,
+          chainId,
           safeAccountConfig: safeAccountConfig3,
           safeDeploymentConfig: safeDeploymentConfig3,
           customContracts
@@ -571,5 +588,69 @@ describe('Contract utils', () => {
         chai.expect(thirdPredictedSafeAddress).to.be.equal(expectedSafeAddress3)
       }
     )
+
+    itif(safeVersionDeployed === '1.3.0')(
+      // see: https://github.com/safe-global/safe-core-sdk/issues/598
+      'returns the correct predicted address for each chain',
+      async () => {
+        const { accounts } = await setupTests()
+        const [owner] = accounts
+        const safeVersion = safeVersionDeployed
+
+        const gnosisEthAdapter = await getEthAdapter(getNetworkProvider('gnosis'))
+        const zkSyncEthAdapter = await getEthAdapter(getNetworkProvider('zksync'))
+        const sepoliaEthAdapter = await getEthAdapter(getNetworkProvider('sepolia'))
+        const mainnetEthAdapter = await getEthAdapter(getNetworkProvider('mainnet'))
+
+        // 1/1 Safe
+        const safeAccountConfig: SafeAccountConfig = {
+          owners: [owner.address],
+          threshold: 1
+        }
+        const safeDeploymentConfig: SafeDeploymentConfig = {
+          safeVersion,
+          saltNonce: '1691490995332'
+        }
+
+        const gnosisPredictedSafeAddress = await predictSafeAddress({
+          ethAdapter: gnosisEthAdapter,
+          chainId: await gnosisEthAdapter.getChainId(),
+          safeAccountConfig: safeAccountConfig,
+          safeDeploymentConfig: safeDeploymentConfig
+        })
+
+        const zkSyncPredictedSafeAddress = await predictSafeAddress({
+          ethAdapter: zkSyncEthAdapter,
+          chainId: await zkSyncEthAdapter.getChainId(),
+          safeAccountConfig: safeAccountConfig,
+          safeDeploymentConfig: safeDeploymentConfig
+        })
+
+        const sepoliaPredictedSafeAddress = await predictSafeAddress({
+          ethAdapter: sepoliaEthAdapter,
+          chainId: await sepoliaEthAdapter.getChainId(),
+          safeAccountConfig: safeAccountConfig,
+          safeDeploymentConfig: safeDeploymentConfig
+        })
+
+        const mainnetPredictedSafeAddress = await predictSafeAddress({
+          ethAdapter: mainnetEthAdapter,
+          chainId: await mainnetEthAdapter.getChainId(),
+          safeAccountConfig: safeAccountConfig,
+          safeDeploymentConfig: safeDeploymentConfig
+        })
+
+        const expectedGnosisSafeAddress = '0x30421B2bE26942448CD6C690f21F551BF6C8A45F'
+        const expectedSkSyncSafeAddress = '0x4680B7AC23A98d5D68c21e3d6F8cBC9576A5920A'
+        const expectedSepoliaSafeAddress = '0x7f44E49C9E4C7D19fA2A704c2E66527Bd4688f99'
+        const expectedMainnetSafeAddress = '0x9C1C8c37a68242cEC6d68Ab091583c81FBF479C0'
+
+        // returns the correct predicted address for each chain
+        chai.expect(gnosisPredictedSafeAddress).to.be.equal(expectedGnosisSafeAddress)
+        chai.expect(zkSyncPredictedSafeAddress).to.be.equal(expectedSkSyncSafeAddress)
+        chai.expect(sepoliaPredictedSafeAddress).to.be.equal(expectedSepoliaSafeAddress)
+        chai.expect(mainnetPredictedSafeAddress).to.be.equal(expectedMainnetSafeAddress)
+      }
+    )
   })
 })

playground/protocol-kit/generate-safe-address.ts

@@ -29,6 +29,8 @@ async function generateSafeAddresses() {
     threshold: config.threshold
   }
 
+  const chainId = await ethAdapter.getChainId()
+
   // infinite loop to search a valid Safe addresses
   for (saltNonce; true; saltNonce++ && iteractions++) {
     // we updete the Deployment config with the current saltNonce
@@ -40,6 +42,7 @@ async function generateSafeAddresses() {
     // we predict the Safe address using the current saltNonce
     const predictedSafeAddress = await predictSafeAddress({
       ethAdapter,
+      chainId,
       safeAccountConfig,
       safeDeploymentConfig
     })