feat(oft-solana): introduce addComputeUnitInstructions (#1091)

Co-authored-by: nazreen <nazreen.mohamad@layerzerolabs.org>

examples/oft-solana/tasks/solana/createOFT.ts

@@ -26,7 +26,7 @@ import { OFT_DECIMALS as DEFAULT_SHARED_DECIMALS, oft, types } from '@layerzerol
 import { checkMultisigSigners, createMintAuthorityMultisig } from './multisig'
 import { assertAccountInitialized } from './utils'
 
-import { deriveConnection, deriveKeys, getExplorerTxLink, output } from './index'
+import { addComputeUnitInstructions, deriveConnection, deriveKeys, getExplorerTxLink, output } from './index'
 
 const DEFAULT_LOCAL_DECIMALS = 9
 
@@ -102,6 +102,8 @@ interface CreateOFTTaskArgs {
      * The URI for the token metadata.
      */
     uri: string
+
+    computeUnitPriceScaleFactor: number
 }
 
 // Define a Hardhat task for creating OFT on Solana
@@ -135,6 +137,7 @@ task('lz:oft:solana:create', 'Mints new SPL Token and creates new OFT Store acco
         devtoolsTypes.string
     )
     .addParam('uri', 'URI for token metadata', '', devtoolsTypes.string)
+    .addParam('computeUnitPriceScaleFactor', 'The compute unit price scale factor', 4, devtoolsTypes.float, true)
     .setAction(
         async ({
             amount,
@@ -151,6 +154,7 @@ task('lz:oft:solana:create', 'Mints new SPL Token and creates new OFT Store acco
             onlyOftStore,
             tokenProgram: tokenProgramStr,
             uri,
+            computeUnitPriceScaleFactor,
         }: CreateOFTTaskArgs) => {
             const isMABA = !!mintStr
             if (tokenProgramStr !== TOKEN_PROGRAM_ID.toBase58() && !isMABA) {
@@ -221,34 +225,49 @@ task('lz:oft:solana:create', 'Mints new SPL Token and creates new OFT Store acco
                             })
                         )
                 }
+                txBuilder = await addComputeUnitInstructions(
+                    connection,
+                    umi,
+                    eid,
+                    txBuilder,
+                    umiWalletSigner,
+                    computeUnitPriceScaleFactor
+                )
                 const createTokenTx = await txBuilder.sendAndConfirm(umi)
                 await assertAccountInitialized(connection, toWeb3JsPublicKey(mint.publicKey))
                 console.log(`createTokenTx: ${getExplorerTxLink(bs58.encode(createTokenTx.signature), isTestnet)}`)
             }
 
             const lockboxSigner = createSignerFromKeypair({ eddsa: eddsa }, lockBox)
-            const { signature } = await transactionBuilder()
-                .add(
-                    oft.initOft(
-                        {
-                            payer: umiWalletSigner,
-                            admin: umiWalletKeyPair.publicKey,
-                            mint: mint.publicKey,
-                            escrow: lockboxSigner,
-                        },
-                        types.OFTType.Native,
-                        sharedDecimals,
-                        {
-                            oft: programId,
-                            token: tokenProgramId,
-                        }
-                    )
+            let txBuilder = transactionBuilder().add(
+                oft.initOft(
+                    {
+                        payer: umiWalletSigner,
+                        admin: umiWalletKeyPair.publicKey,
+                        mint: mint.publicKey,
+                        escrow: lockboxSigner,
+                    },
+                    types.OFTType.Native,
+                    sharedDecimals,
+                    {
+                        oft: programId,
+                        token: tokenProgramId,
+                    }
                 )
-                .sendAndConfirm(umi)
+            )
+            txBuilder = await addComputeUnitInstructions(
+                connection,
+                umi,
+                eid,
+                txBuilder,
+                umiWalletSigner,
+                computeUnitPriceScaleFactor
+            )
+            const { signature } = await txBuilder.sendAndConfirm(umi)
             console.log(`initOftTx: ${getExplorerTxLink(bs58.encode(signature), isTestnet)}`)
 
             if (!isMABA) {
-                const { signature } = await transactionBuilder()
+                let txBuilder = transactionBuilder()
                     .add(
                         setAuthority(umi, {
                             owned: mint.publicKey,
@@ -265,7 +284,15 @@ task('lz:oft:solana:create', 'Mints new SPL Token and creates new OFT Store acco
                             authorityType: 1,
                         })
                     )
-                    .sendAndConfirm(umi)
+                txBuilder = await addComputeUnitInstructions(
+                    connection,
+                    umi,
+                    eid,
+                    txBuilder,
+                    umiWalletSigner,
+                    computeUnitPriceScaleFactor
+                )
+                const { signature } = await txBuilder.sendAndConfirm(umi)
                 console.log(`setAuthorityTx: ${getExplorerTxLink(bs58.encode(signature), isTestnet)}`)
             }
             output(eid, programIdStr, mint.publicKey, mintAuthorityPublicKey.toBase58(), escrowPK, oftStorePda)

examples/oft-solana/tasks/solana/createOFTAdapter.ts

@@ -8,7 +8,7 @@ import { types as devtoolsTypes } from '@layerzerolabs/devtools-evm-hardhat'
 import { EndpointId } from '@layerzerolabs/lz-definitions'
 import { OFT_DECIMALS, oft, types } from '@layerzerolabs/oft-v2-solana-sdk'
 
-import { deriveConnection, deriveKeys, getExplorerTxLink, output } from './index'
+import { addComputeUnitInstructions, deriveConnection, deriveKeys, getExplorerTxLink, output } from './index'
 
 interface CreateOFTAdapterTaskArgs {
     /**
@@ -30,6 +30,8 @@ interface CreateOFTAdapterTaskArgs {
      * The Token Program public key.
      */
     tokenProgram: string
+
+    computeUnitPriceScaleFactor: number
 }
 
 // Define a Hardhat task for creating OFTAdapter on Solana
@@ -38,14 +40,16 @@ task('lz:oft-adapter:solana:create', 'Creates new OFT Adapter (OFT Store PDA)')
     .addParam('programId', 'The OFT program ID')
     .addParam('eid', 'Solana mainnet or testnet', undefined, devtoolsTypes.eid)
     .addParam('tokenProgram', 'The Token Program public key', TOKEN_PROGRAM_ID.toBase58(), devtoolsTypes.string, true)
+    .addParam('computeUnitPriceScaleFactor', 'The compute unit price scale factor', 4, devtoolsTypes.float, true)
     .setAction(
         async ({
             eid,
             mint: mintStr,
             programId: programIdStr,
             tokenProgram: tokenProgramStr,
+            computeUnitPriceScaleFactor,
         }: CreateOFTAdapterTaskArgs) => {
-            const { connection, umi, umiWalletKeyPair } = await deriveConnection(eid)
+            const { connection, umi, umiWalletKeyPair, umiWalletSigner } = await deriveConnection(eid)
             const { programId, lockBox, escrowPK, oftStorePda, eddsa } = deriveKeys(programIdStr)
 
             const tokenProgram = publicKey(tokenProgramStr)
@@ -55,21 +59,31 @@ task('lz:oft-adapter:solana:create', 'Creates new OFT Adapter (OFT Store PDA)')
 
             const mintAuthority = mintPDA.mintAuthority
 
-            const initOftIx = oft.initOft(
-                {
-                    payer: createSignerFromKeypair({ eddsa: eddsa }, umiWalletKeyPair),
-                    admin: umiWalletKeyPair.publicKey,
-                    mint,
-                    escrow: createSignerFromKeypair({ eddsa: eddsa }, lockBox),
-                },
-                types.OFTType.Adapter,
-                OFT_DECIMALS,
-                {
-                    oft: programId,
-                    token: tokenProgram ? publicKey(tokenProgram) : undefined,
-                }
+            let txBuilder = transactionBuilder().add(
+                oft.initOft(
+                    {
+                        payer: createSignerFromKeypair({ eddsa: eddsa }, umiWalletKeyPair),
+                        admin: umiWalletKeyPair.publicKey,
+                        mint: mint,
+                        escrow: createSignerFromKeypair({ eddsa: eddsa }, lockBox),
+                    },
+                    types.OFTType.Adapter,
+                    OFT_DECIMALS,
+                    {
+                        oft: programId,
+                        token: tokenProgram ? publicKey(tokenProgram) : undefined,
+                    }
+                )
+            )
+            txBuilder = await addComputeUnitInstructions(
+                connection,
+                umi,
+                eid,
+                txBuilder,
+                umiWalletSigner,
+                computeUnitPriceScaleFactor
             )
-            const { signature } = await transactionBuilder().add(initOftIx).sendAndConfirm(umi)
+            const { signature } = await txBuilder.sendAndConfirm(umi)
             console.log(`initOftTx: ${getExplorerTxLink(bs58.encode(signature), eid == EndpointId.SOLANA_V2_TESTNET)}`)
 
             output(eid, programIdStr, mint, mintAuthority ? mintAuthority.toBase58() : '', escrowPK, oftStorePda)

examples/oft-solana/tasks/solana/getPrioFees.ts

@@ -20,7 +20,7 @@ interface GetPrioFeesTaskArgs {
 
 task('lz:solana:get-priority-fees', 'Fetches prioritization fees from the Solana network')
     .addParam('eid', 'The endpoint ID for the Solana network', undefined, devtoolsTypes.eid)
-    .addParam(
+    .addOptionalParam(
         'address',
         'The address (program ID or account address) that will be written to',
         undefined,

examples/oft-solana/tasks/solana/index.ts

@@ -1,15 +1,43 @@
+import assert from 'assert'
 import { existsSync, mkdirSync, writeFileSync } from 'node:fs'
 
-import { mplToolbox } from '@metaplex-foundation/mpl-toolbox'
-import { EddsaInterface, createSignerFromKeypair, publicKey, signerIdentity } from '@metaplex-foundation/umi'
+import {
+    fetchAddressLookupTable,
+    mplToolbox,
+    setComputeUnitLimit,
+    setComputeUnitPrice,
+} from '@metaplex-foundation/mpl-toolbox'
+import {
+    AddressLookupTableInput,
+    EddsaInterface,
+    Instruction,
+    KeypairSigner,
+    PublicKey,
+    TransactionBuilder,
+    Umi,
+    createSignerFromKeypair,
+    publicKey,
+    signerIdentity,
+    transactionBuilder,
+} from '@metaplex-foundation/umi'
 import { createUmi } from '@metaplex-foundation/umi-bundle-defaults'
 import { createWeb3JsEddsa } from '@metaplex-foundation/umi-eddsa-web3js'
+import { toWeb3JsInstruction, toWeb3JsPublicKey } from '@metaplex-foundation/umi-web3js-adapters'
+import { AddressLookupTableAccount, Connection } from '@solana/web3.js'
+import { getSimulationComputeUnits } from '@solana-developers/helpers'
 import bs58 from 'bs58'
 
+import { formatEid } from '@layerzerolabs/devtools'
 import { EndpointId, endpointIdToNetwork } from '@layerzerolabs/lz-definitions'
 import { OftPDA } from '@layerzerolabs/oft-v2-solana-sdk'
 
 import { createSolanaConnectionFactory } from '../common/utils'
+import getFee from '../utils/getFee'
+
+const LOOKUP_TABLE_ADDRESS: Partial<Record<EndpointId, PublicKey>> = {
+    [EndpointId.SOLANA_V2_MAINNET]: publicKey('AokBxha6VMLLgf97B5VYHEtqztamWmYERBmmFvjuTzJB'),
+    [EndpointId.SOLANA_V2_TESTNET]: publicKey('9thqPdbR27A1yLWw2spwJLySemiGMXxPnEvfmXVk4KuK'),
+}
 
 const getFromEnv = (key: string): string => {
     const value = process.env[key]
@@ -107,3 +135,79 @@ export const getLayerZeroScanLink = (hash: string, isTestnet = false) =>
 
 export const getExplorerTxLink = (hash: string, isTestnet = false) =>
     `https://explorer.solana.com/tx/${hash}?cluster=${isTestnet ? 'devnet' : 'mainnet-beta'}`
+
+export const getAddressLookupTable = async (connection: Connection, umi: Umi, fromEid: EndpointId) => {
+    // Lookup Table Address and Priority Fee Calculation
+    const lookupTableAddress = LOOKUP_TABLE_ADDRESS[fromEid]
+    assert(lookupTableAddress != null, `No lookup table found for ${formatEid(fromEid)}`)
+    const addressLookupTableInput: AddressLookupTableInput = await fetchAddressLookupTable(umi, lookupTableAddress)
+    if (!addressLookupTableInput) {
+        throw new Error(`No address lookup table found for ${lookupTableAddress}`)
+    }
+    const { value: lookupTableAccount } = await connection.getAddressLookupTable(toWeb3JsPublicKey(lookupTableAddress))
+    if (!lookupTableAccount) {
+        throw new Error(`No address lookup table account found for ${lookupTableAddress}`)
+    }
+    return {
+        lookupTableAddress,
+        addressLookupTableInput,
+        lookupTableAccount,
+    }
+}
+
+export const getComputeUnitPriceAndLimit = async (
+    connection: Connection,
+    ixs: Instruction[],
+    wallet: KeypairSigner,
+    lookupTableAccount: AddressLookupTableAccount
+) => {
+    const { averageFeeExcludingZeros } = await getFee(connection)
+    const priorityFee = Math.round(averageFeeExcludingZeros)
+    const computeUnitPrice = BigInt(priorityFee)
+
+    const computeUnits = await getSimulationComputeUnits(
+        connection,
+        ixs.map((ix) => toWeb3JsInstruction(ix)),
+        toWeb3JsPublicKey(wallet.publicKey),
+        [lookupTableAccount]
+    )
+
+    if (!computeUnits) {
+        throw new Error('Unable to compute units')
+    }
+
+    return {
+        computeUnitPrice,
+        computeUnits,
+    }
+}
+
+export const addComputeUnitInstructions = async (
+    connection: Connection,
+    umi: Umi,
+    eid: EndpointId,
+    txBuilder: TransactionBuilder,
+    umiWalletSigner: KeypairSigner,
+    computeUnitPriceScaleFactor: number
+) => {
+    const computeUnitLimitScaleFactor = 1.1 // hardcoded to 1.1 as the estimations are not perfect and can fall slightly short of the actual CU usage on-chain
+    const { addressLookupTableInput, lookupTableAccount } = await getAddressLookupTable(connection, umi, eid)
+    const { computeUnitPrice, computeUnits } = await getComputeUnitPriceAndLimit(
+        connection,
+        txBuilder.getInstructions(),
+        umiWalletSigner,
+        lookupTableAccount
+    )
+    // Since transaction builders are immutable, we must be careful to always assign the result of the add and prepend
+    // methods to a new variable.
+    const newTxBuilder = transactionBuilder()
+        .add(
+            setComputeUnitPrice(umi, {
+                microLamports: computeUnitPrice * BigInt(Math.floor(computeUnitPriceScaleFactor)),
+            })
+        )
+        .add(setComputeUnitLimit(umi, { units: computeUnits * computeUnitLimitScaleFactor }))
+        .setAddressLookupTables([addressLookupTableInput])
+        .add(txBuilder)
+    return newTxBuilder
+}