feat(protocol-kit): Predict Safe address improvements

packages/protocol-kit/src/Safe.ts

@@ -21,7 +21,7 @@ import {
   getPredictedSafeAddressInitCode,
   predictSafeAddress
 } from './contracts/utils'
-import { DEFAULT_SAFE_VERSION } from './contracts/config'
+import { ContractInfo, DEFAULT_SAFE_VERSION, getContractInfo } from './contracts/config'
 import ContractManager from './managers/contractManager'
 import FallbackHandlerManager from './managers/fallbackHandlerManager'
 import GuardManager from './managers/guardManager'
@@ -124,12 +124,12 @@ class Safe {
   async #initializeProtocolKit(config: SafeConfig) {
     const { provider, signer, isL1SafeSingleton, contractNetworks } = config
 
-    this.#safeProvider = await SafeProvider.init(
+    this.#safeProvider = await SafeProvider.init({
       provider,
       signer,
-      DEFAULT_SAFE_VERSION,
+      safeVersion: DEFAULT_SAFE_VERSION,
       contractNetworks
-    )
+    })
 
     if (isSafeConfigWithPredictedSafe(config)) {
       this.#predictedSafe = config.predictedSafe
@@ -155,7 +155,12 @@ class Safe {
     }
 
     const safeVersion = this.getContractVersion()
-    this.#safeProvider = await SafeProvider.init(provider, signer, safeVersion, contractNetworks)
+    this.#safeProvider = await SafeProvider.init({
+      provider,
+      signer,
+      safeVersion,
+      contractNetworks
+    })
 
     this.#ownerManager = new OwnerManager(this.#safeProvider, this.#contractManager.safeContract)
     this.#moduleManager = new ModuleManager(this.#safeProvider, this.#contractManager.safeContract)
@@ -169,14 +174,14 @@ class Safe {
     if (isPasskeySigner) {
       const safeAddress = await this.getAddress()
       const owners = await this.getOwners()
-      this.#safeProvider = await SafeProvider.init(
+      this.#safeProvider = await SafeProvider.init({
         provider,
         signer,
         safeVersion,
         contractNetworks,
         safeAddress,
         owners
-      )
+      })
     }
   }
 
@@ -1636,6 +1641,14 @@ class Safe {
       return false
     }
   }
+
+  getContractInfo = ({
+    contractAddress
+  }: {
+    contractAddress: string
+  }): ContractInfo | undefined => {
+    return getContractInfo(contractAddress)
+  }
 }
 
 export default Safe

packages/protocol-kit/src/SafeFactory.ts

@@ -62,7 +62,12 @@ class SafeFactory {
   }: SafeFactoryInitConfig) {
     this.#provider = provider
     this.#signer = signer
-    this.#safeProvider = await SafeProvider.init(provider, signer, safeVersion, contractNetworks)
+    this.#safeProvider = await SafeProvider.init({
+      provider,
+      signer,
+      safeVersion,
+      contractNetworks
+    })
     this.#safeVersion = safeVersion
     this.#isL1SafeSingleton = isL1SafeSingleton
     this.#contractNetworks = contractNetworks

packages/protocol-kit/src/SafeProvider.ts

@@ -17,20 +17,18 @@
   EIP712TypedDataMessage,
   EIP712TypedDataTx,
   Eip3770Address,
-  SafeEIP712Args,
-  SafeVersion
+  SafeEIP712Args
 } from '@safe-global/types-kit'
 import {
   SafeProviderTransaction,
   SafeProviderConfig,
+  SafeProviderInitOptions,
   ExternalClient,
   ExternalSigner,
   Eip1193Provider,
   HttpTransport,
   SocketTransport,
   SafeSigner,
-  SafeConfig,
-  ContractNetworksConfig,
   PasskeyArgType,
   PasskeyClient
 } from '@safe-global/protocol-kit/types'
@@ -102,14 +100,14 @@
     return this.#externalProvider
   }
 
-  static async init(
-    provider: SafeConfig['provider'],
-    signer?: SafeConfig['signer'],
-    safeVersion: SafeVersion = DEFAULT_SAFE_VERSION,
-    contractNetworks?: ContractNetworksConfig,
-    safeAddress?: string,
-    owners?: string[]
-  ): Promise<SafeProvider> {
+  static async init({
+    provider,
+    signer,
+    safeVersion = DEFAULT_SAFE_VERSION,
+    contractNetworks,
+    safeAddress,
+    owners
+  }: SafeProviderInitOptions): Promise<SafeProvider> {
     const isPasskeySigner = signer && typeof signer !== 'string'
 
     if (isPasskeySigner) {
@@ -365,11 +363,11 @@
   }
 
   // TODO: fix anys
   encodeParameters(types: string, values: any[]): string {
     return encodeAbiParameters(parseAbiParameters(types), values)
   }
 
   decodeParameters(types: string, values: string): { [key: string]: any } {
     return decodeAbiParameters(parseAbiParameters(types), asHex(values))
   }
 

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

@@ -9,7 +9,9 @@
   Chain
 } from 'viem'
 import { estimateContractGas, getTransactionReceipt } from 'viem/actions'
+import { SingletonDeployment } from '@safe-global/safe-deployments'
 import { contractName, getContractDeployment } from '@safe-global/protocol-kit/contracts/config'
+import { DeploymentType } from '@safe-global/protocol-kit/types'
 import SafeProvider from '@safe-global/protocol-kit/SafeProvider'
 import {
   EncodeFunction,
@@ -62,6 +64,7 @@
    * @param safeVersion - The version of the Safe contract.
    * @param customContractAddress - Optional custom address for the contract. If not provided, the address is derived from the Safe deployments based on the chainId and safeVersion.
    * @param customContractAbi - Optional custom ABI for the contract. If not provided, the ABI is derived from the Safe deployments or the defaultAbi is used.
+   * @param deploymentType - Optional deployment type for the contract. If not provided, the first deployment retrieved from the safe-deployments array will be used.
    */
   constructor(
     contractName: contractName,
@@ -70,24 +73,27 @@
     defaultAbi: ContractAbiType,
     safeVersion: SafeVersion,
     customContractAddress?: string,
-    customContractAbi?: ContractAbiType
+    customContractAbi?: ContractAbiType,
+    deploymentType?: DeploymentType
   ) {
     const deployment = getContractDeployment(safeVersion, chainId, contractName)
 
-    const contractAddress =
-      customContractAddress || deployment?.networkAddresses[chainId.toString()]
+    const resolvedAddress =
+      customContractAddress ??
+      this.#resolveAddress(
+        deployment?.networkAddresses[chainId.toString()],
+        deployment,
+        deploymentType
+      )
 
-    if (!contractAddress) {
+    if (!resolvedAddress) {
       throw new Error(`Invalid ${contractName.replace('Version', '')} contract address`)
     }
 
     this.chainId = chainId
     this.contractName = contractName
     this.safeVersion = safeVersion
-    this.contractAddress =
-      Array.isArray(contractAddress) && contractAddress.length
-        ? contractAddress[0]
-        : contractAddress.toString()
+    this.contractAddress = resolvedAddress
     this.contractAbi =
       customContractAbi ||
       (deployment?.abi as unknown as ContractAbiType) || // this cast is required because abi is set as any[] in safe-deployments
@@ -97,6 +103,31 @@
     this.safeProvider = safeProvider
   }
 
+  #resolveAddress(
+    networkAddresses: string | string[] | undefined,
+    deployment: SingletonDeployment,
+    deploymentType?: DeploymentType
+  ): string | undefined {
+    if (!networkAddresses) {
+      return undefined
+    }
+
+    if (typeof networkAddresses === 'string') {
+      return networkAddresses
+    }
+
+    if (deploymentType) {
+      const customDeploymentTypeAddress = deployment.deployments[deploymentType]?.address
+
+      return (
+        networkAddresses.find((address) => address === customDeploymentTypeAddress) ??
+        networkAddresses[0]
+      )
+    }
+
+    return networkAddresses[0]
+  }
+
   async init() {
     this.wallet = await this.safeProvider.getExternalSigner()
   }
@@ -175,7 +206,7 @@
       functionName
     >
   >(functionName: functionName, args: functionArgs, options?: TransactionOptions) {
     const converted = this.convertOptions(options) as any
 
     return await this.getWallet().writeContract({
       address: this.contractAddress,

packages/protocol-kit/src/contracts/CompatibilityFallbackHandler/CompatibilityFallbackHandlerBaseContract.ts

@@ -2,6 +2,7 @@ import { Abi } from 'abitype'
 
 import SafeProvider from '@safe-global/protocol-kit/SafeProvider'
 import BaseContract from '@safe-global/protocol-kit/contracts/BaseContract'
+import { DeploymentType } from '@safe-global/protocol-kit/types'
 import { SafeVersion } from '@safe-global/types-kit'
 import { contractName } from '@safe-global/protocol-kit/contracts/config'
 
@@ -33,14 +34,16 @@ abstract class CompatibilityFallbackHandlerBaseContract<
    * @param safeVersion - The version of the Safe contract.
    * @param customContractAddress - Optional custom address for the contract. If not provided, the address is derived from the Safe deployments based on the chainId and safeVersion.
    * @param customContractAbi - Optional custom ABI for the contract. If not provided, the ABI is derived from the Safe deployments or the defaultAbi is used.
+   * @param deploymentType - Optional deployment type for the contract. If not provided, the first deployment retrieved from the safe-deployments array will be used.
    */
   constructor(
     chainId: bigint,
     safeProvider: SafeProvider,
     defaultAbi: CompatibilityFallbackHandlerContractAbiType,
     safeVersion: SafeVersion,
     customContractAddress?: string,
-    customContractAbi?: CompatibilityFallbackHandlerContractAbiType
+    customContractAbi?: CompatibilityFallbackHandlerContractAbiType,
+    deploymentType?: DeploymentType
   ) {
     const contractName = 'compatibilityFallbackHandler'
 
@@ -51,7 +54,8 @@ abstract class CompatibilityFallbackHandlerBaseContract<
       defaultAbi,
       safeVersion,
       customContractAddress,
-      customContractAbi
+      customContractAbi,
+      deploymentType
     )
 
     this.contractName = contractName

packages/protocol-kit/src/contracts/CompatibilityFallbackHandler/v1.3.0/CompatibilityFallbackHandlerContract_v1_3_0.ts

@@ -1,5 +1,6 @@
 import CompatibilityFallbackHandlerBaseContract from '@safe-global/protocol-kit/contracts/CompatibilityFallbackHandler/CompatibilityFallbackHandlerBaseContract'
 import SafeProvider from '@safe-global/protocol-kit/SafeProvider'
+import { DeploymentType } from '@safe-global/protocol-kit/types'
 import {
   CompatibilityFallbackHandlerContract_v1_3_0_Abi,
   CompatibilityFallbackHandlerContract_v1_3_0_Contract,
@@ -25,17 +26,27 @@ class CompatibilityFallbackHandlerContract_v1_3_0
    * @param safeProvider - An instance of SafeProvider.
    * @param customContractAddress - Optional custom address for the contract. If not provided, the address is derived from the CompatibilityFallbackHandler deployments based on the chainId and safeVersion.
    * @param customContractAbi - Optional custom ABI for the contract. If not provided, the default ABI for version 1.3.0 is used.
+   * @param deploymentType - Optional deployment type for the contract. If not provided, the first deployment retrieved from the safe-deployments array will be used.
    */
   constructor(
     chainId: bigint,
     safeProvider: SafeProvider,
     customContractAddress?: string,
-    customContractAbi?: CompatibilityFallbackHandlerContract_v1_3_0_Abi
+    customContractAbi?: CompatibilityFallbackHandlerContract_v1_3_0_Abi,
+    deploymentType?: DeploymentType
   ) {
     const safeVersion = '1.3.0'
     const defaultAbi = compatibilityFallbackHandler_1_3_0_ContractArtifacts.abi
 
-    super(chainId, safeProvider, defaultAbi, safeVersion, customContractAddress, customContractAbi)
+    super(
+      chainId,
+      safeProvider,
+      defaultAbi,
+      safeVersion,
+      customContractAddress,
+      customContractAbi,
+      deploymentType
+    )
   }
 }
 

packages/protocol-kit/src/contracts/CompatibilityFallbackHandler/v1.4.1/CompatibilityFallbackHandlerContract_v1_4_1.ts

@@ -1,5 +1,6 @@
 import CompatibilityFallbackHandlerBaseContract from '@safe-global/protocol-kit/contracts/CompatibilityFallbackHandler/CompatibilityFallbackHandlerBaseContract'
 import SafeProvider from '@safe-global/protocol-kit/SafeProvider'
+import { DeploymentType } from '@safe-global/protocol-kit/types'
 import {
   compatibilityFallbackHandler_1_4_1_ContractArtifacts,
   CompatibilityFallbackHandlerContract_v1_4_1_Abi,
@@ -25,17 +26,27 @@ class CompatibilityFallbackHandlerContract_v1_4_1
    * @param safeProvider - An instance of SafeProvider.
    * @param customContractAddress - Optional custom address for the contract. If not provided, the address is derived from the CompatibilityFallbackHandler deployments based on the chainId and safeVersion.
    * @param customContractAbi - Optional custom ABI for the contract. If not provided, the default ABI for version 1.4.1 is used.
+   * @param deploymentType - Optional deployment type for the contract. If not provided, the first deployment retrieved from the safe-deployments array will be used.
    */
   constructor(
     chainId: bigint,
     safeProvider: SafeProvider,
     customContractAddress?: string,
-    customContractAbi?: CompatibilityFallbackHandlerContract_v1_4_1_Abi
+    customContractAbi?: CompatibilityFallbackHandlerContract_v1_4_1_Abi,
+    deploymentType?: DeploymentType
   ) {
     const safeVersion = '1.4.1'
     const defaultAbi = compatibilityFallbackHandler_1_4_1_ContractArtifacts.abi
 
-    super(chainId, safeProvider, defaultAbi, safeVersion, customContractAddress, customContractAbi)
+    super(
+      chainId,
+      safeProvider,
+      defaultAbi,
+      safeVersion,
+      customContractAddress,
+      customContractAbi,
+      deploymentType
+    )
   }
 }
 

packages/protocol-kit/src/contracts/CreateCall/CreateCallBaseContract.ts

@@ -1,6 +1,7 @@
 import { Abi } from 'abitype'
 
 import SafeProvider from '@safe-global/protocol-kit/SafeProvider'
+import { DeploymentType } from '@safe-global/protocol-kit/types'
 import BaseContract from '@safe-global/protocol-kit/contracts/BaseContract'
 import { SafeVersion } from '@safe-global/types-kit'
 import { contractName } from '@safe-global/protocol-kit/contracts/config'
@@ -33,14 +34,16 @@ abstract class CreateCallBaseContract<
    * @param safeVersion - The version of the Safe contract.
    * @param customContractAddress - Optional custom address for the contract. If not provided, the address is derived from the Safe deployments based on the chainId and safeVersion.
    * @param customContractAbi - Optional custom ABI for the contract. If not provided, the ABI is derived from the Safe deployments or the defaultAbi is used.
+   * @param deploymentType - Optional deployment type for the contract. If not provided, the first deployment retrieved from the safe-deployments array will be used.
    */
   constructor(
     chainId: bigint,
     safeProvider: SafeProvider,
     defaultAbi: CreateCallContractAbiType,
     safeVersion: SafeVersion,
     customContractAddress?: string,
-    customContractAbi?: CreateCallContractAbiType
+    customContractAbi?: CreateCallContractAbiType,
+    deploymentType?: DeploymentType
   ) {
     const contractName = 'createCallVersion'
 
@@ -51,7 +54,8 @@ abstract class CreateCallBaseContract<
       defaultAbi,
       safeVersion,
       customContractAddress,
-      customContractAbi
+      customContractAbi,
+      deploymentType
     )
 
     this.contractName = contractName