Add some changes to allow to choose the deploymentType

safe-global · Sep 20, 2024 · 4e0cea9 · 4e0cea9
1 parent d9e9761
commit 4e0cea9
Show file tree

Hide file tree

Showing 6 changed files with 22 additions and 3 deletions.
diff --git a/packages/protocol-kit/src/Safe.ts b/packages/protocol-kit/src/Safe.ts
@@ -122,7 +122,7 @@ class Safe {
    * @throws "MultiSendCallOnly contract is not deployed on the current network"
    */
   async #initializeProtocolKit(config: SafeConfig) {
-    const { provider, signer, isL1SafeSingleton, contractNetworks } = config
+    const { provider, signer, isL1SafeSingleton, contractNetworks, predictedSafe } = config
 
     this.#safeProvider = await SafeProvider.init(
       provider,
@@ -131,6 +131,8 @@ class Safe {
       contractNetworks
     )
 
+    this.#safeProvider.deploymentType = predictedSafe?.safeDeploymentConfig?.deploymentType
+
     if (isSafeConfigWithPredictedSafe(config)) {
       this.#predictedSafe = config.predictedSafe
       this.#contractManager = await ContractManager.init(

diff --git a/packages/protocol-kit/src/SafeProvider.ts b/packages/protocol-kit/src/SafeProvider.ts
@@ -94,6 +94,7 @@ class SafeProvider {
   #externalProvider: ExternalClient
   signer?: SafeSigner
   provider: Eip1193Provider | HttpTransport | SocketTransport
+  #deploymentType?: 'canonical' | 'eip155' | 'zksync'
 
   constructor({
     provider,
@@ -181,6 +182,14 @@ class SafeProvider {
     }
   }
 
+  get deploymentType(): 'canonical' | 'eip155' | 'zksync' | undefined {
+    return this.#deploymentType
+  }
+
+  set deploymentType(deploymentType: 'canonical' | 'eip155' | 'zksync' | undefined) {
+    this.#deploymentType = deploymentType
+  }
+
   async getExternalSigner(): Promise<ExternalSigner | undefined> {
     const { transport, chain = await this.#getChain() } = this.getExternalProvider()
 

diff --git a/packages/protocol-kit/src/contracts/BaseContract.ts b/packages/protocol-kit/src/contracts/BaseContract.ts
@@ -75,12 +75,16 @@ class BaseContract<ContractAbiType extends Abi> {
       throw new Error(`Invalid ${contractName.replace('Version', '')} contract address`)
     }
 
+    const customDeploymentTypeAddress = safeProvider.deploymentType
+      ? deployment?.deployments[safeProvider.deploymentType]?.address
+      : undefined
+
     this.chainId = chainId
     this.contractName = contractName
     this.safeVersion = safeVersion
     this.contractAddress =
       Array.isArray(contractAddress) && contractAddress.length
-        ? contractAddress[0]
+        ? contractAddress.find((a) => a === customDeploymentTypeAddress) || contractAddress[0]
         : contractAddress.toString()
     this.contractAbi =
       customContractAbi ||

diff --git a/packages/protocol-kit/src/contracts/config.ts b/packages/protocol-kit/src/contracts/config.ts
@@ -160,6 +160,7 @@ export function getContractInfo(contractAddress: string): ContractInfo | undefin
       const deployment = contractFunctions[contractName as contractName](
         filters
       ) as SingletonDeployment
+
       if (deployment && deployment.networkAddresses) {
         for (const [, address] of Object.entries(deployment.networkAddresses)) {
           if (address.toLowerCase() === contractAddress.toLowerCase()) {
@@ -168,10 +169,12 @@ export function getContractInfo(contractAddress: string): ContractInfo | undefin
               | 'eip155'
               | 'zksync'
             )[]
+
             const type = types.find(
               (t) =>
                 deployment.deployments[t]?.address.toLowerCase() === contractAddress.toLowerCase()
             )
+
             if (type) {
               return {
                 version: safeVersion,

diff --git a/packages/protocol-kit/src/types/safeConfig.ts b/packages/protocol-kit/src/types/safeConfig.ts
@@ -17,6 +17,7 @@ export type SafeAccountConfig = {
 export type SafeDeploymentConfig = {
   saltNonce?: string
   safeVersion?: SafeVersion
+  deploymentType?: 'canonical' | 'eip155' | 'zksync'
 }
 
 export type PredictedSafeProps = {

diff --git a/packages/protocol-kit/tests/e2e/contractInfo.test.ts b/packages/protocol-kit/tests/e2e/contractInfo.test.ts
@@ -20,7 +20,7 @@ describe('Contract Info', () => {
     })
   })
 
-  describe.only('create', async () => {
+  describe('create', async () => {
     it('should return undefined for a contract address not related to Safe', async () => {
       chai.expect(
         protocolKit.getContractInfo({