@arcium-hq/client 0.9.2 → 0.9.3

package/src/onchain.ts

@@ -0,0 +1,1020 @@
+import { Program, AnchorProvider } from '@coral-xyz/anchor';
+import * as anchor from '@coral-xyz/anchor';
+import { ConfirmOptions, PublicKey } from '@solana/web3.js';
+import { ARCIUM_IDL, ArciumIdlType } from './idl/index.js';
+import { sha256 } from './cryptography/cryptography.js';
+import { Arcium } from './idl/arcium.js';
+import {
+    COMP_DEF_ACC_SEED,
+    COMP_DEF_OFFSET_SIZE,
+    DISCRIMINATOR_SIZE,
+    MAX_ACCOUNT_SIZE,
+    MAX_EMBIGGEN_IX_PER_TX,
+    MAX_REALLOC_PER_IX,
+    MAX_UPLOAD_PER_TX_BYTES,
+    OFFSET_BUFFER_SIZE,
+} from './constants.js';
+import {
+    getArciumProgramId,
+    getMXEAccAddress,
+    getComputationAccAddress,
+    getMempoolAccAddress,
+    getExecutingPoolAccAddress,
+    getFeePoolAccAddress,
+    getRawCircuitAccAddress,
+    getLookupTableAddress,
+} from './pda.js';
+import { ComputationReference, isNullRef, optionalLog } from './utils.js';
+
+const TINY_MEMPOOL_ACC_NAME: Arcium['accounts']['21']['name'] = 'tinyMempool';
+const TINY_MEMPOOL_DISCRIMINATOR: Arcium['accounts']['21']['discriminator'] = [
+    176,
+    33,
+    67,
+    108,
+    73,
+    135,
+    110,
+    166,
+];
+
+const TINY_EXECPOOL_ACC_NAME: Arcium['accounts']['20']['name'] = 'tinyExecPool';
+const TINY_EXECPOOL_DISCRIMINATOR: Arcium['accounts']['20']['discriminator'] = [
+    80,
+    245,
+    5,
+    90,
+    154,
+    189,
+    190,
+    172,
+];
+
+const SMALL_MEMPOOL_ACC_NAME: Arcium['accounts']['19']['name'] = 'smallMempool';
+const SMALL_MEMPOOL_DISCRIMINATOR: Arcium['accounts']['19']['discriminator'] = [
+    123,
+    153,
+    151,
+    118,
+    126,
+    71,
+    73,
+    92,
+];
+
+const SMALL_EXECPOOL_ACC_NAME: Arcium['accounts']['18']['name'] = 'smallExecPool';
+const SMALL_EXECPOOL_DISCRIMINATOR: Arcium['accounts']['18']['discriminator'] = [
+    37,
+    147,
+    249,
+    253,
+    217,
+    136,
+    3,
+    87,
+];
+const MEDIUM_MEMPOOL_ACC_NAME: Arcium['accounts']['12']['name'] = 'mediumMempool';
+const MEDIUM_MEMPOOL_DISCRIMINATOR: Arcium['accounts']['12']['discriminator'] = [
+    10,
+    249,
+    58,
+    39,
+    255,
+    231,
+    199,
+    168,
+];
+
+const MEDIUM_EXECPOOL_ACC_NAME: Arcium['accounts']['11']['name'] = 'mediumExecPool';
+const MEDIUM_EXECPOOL_DISCRIMINATOR: Arcium['accounts']['11']['discriminator'] = [
+    97,
+    117,
+    128,
+    202,
+    213,
+    76,
+    5,
+    163,
+];
+
+const LARGE_MEMPOOL_ACC_NAME: Arcium['accounts']['9']['name'] = 'largeMempool';
+const LARGE_MEMPOOL_DISCRIMINATOR: Arcium['accounts']['9']['discriminator'] = [
+    16,
+    168,
+    90,
+    235,
+    249,
+    207,
+    73,
+    223,
+];
+
+const LARGE_EXECPOOL_ACC_NAME: Arcium['accounts']['8']['name'] = 'largeExecPool';
+const LARGE_EXECPOOL_DISCRIMINATOR: Arcium['accounts']['8']['discriminator'] = [
+    147,
+    145,
+    148,
+    170,
+    30,
+    13,
+    43,
+    216,
+];
+
+const MEMPOOL_DISCRIMINATOR_MAP = {
+    [TINY_MEMPOOL_DISCRIMINATOR.toString()]: TINY_MEMPOOL_ACC_NAME,
+    [SMALL_MEMPOOL_DISCRIMINATOR.toString()]: SMALL_MEMPOOL_ACC_NAME,
+    [MEDIUM_MEMPOOL_DISCRIMINATOR.toString()]: MEDIUM_MEMPOOL_ACC_NAME,
+    [LARGE_MEMPOOL_DISCRIMINATOR.toString()]: LARGE_MEMPOOL_ACC_NAME,
+} as const;
+
+const EXECPOOL_DISCRIMINATOR_MAP = {
+    [TINY_EXECPOOL_DISCRIMINATOR.toString()]: TINY_EXECPOOL_ACC_NAME,
+    [SMALL_EXECPOOL_DISCRIMINATOR.toString()]: SMALL_EXECPOOL_ACC_NAME,
+    [MEDIUM_EXECPOOL_DISCRIMINATOR.toString()]: MEDIUM_EXECPOOL_ACC_NAME,
+    [LARGE_EXECPOOL_DISCRIMINATOR.toString()]: LARGE_EXECPOOL_ACC_NAME,
+} as const;
+
+/**
+ * Represents a mempool account of any size (tiny, small, medium, or large).
+ * Mempools store pending computations prioritized by fee, with a time-to-live of 180 slots.
+ * Each size supports different maximum heap capacities:
+ * - Tiny: 1 computation
+ * - Small: 3 computations
+ * - Medium: 10 computations
+ * - Large: 100 computations
+ */
+export type MempoolAccount = anchor.IdlTypes<ArciumIdlType>['tinyMempool'] | anchor.IdlTypes<ArciumIdlType>['smallMempool'] | anchor.IdlTypes<ArciumIdlType>['mediumMempool'] | anchor.IdlTypes<ArciumIdlType>['largeMempool'];
+
+/**
+ * Represents an executing pool account of any size (tiny, small, medium, or large).
+ * Executing pools manage parallel computation execution with account locking.
+ * Each size supports different maximum parallel computations:
+ * - Tiny: 1 parallel computation
+ * - Small: 3 parallel computations
+ * - Medium: 10 parallel computations
+ * - Large: 100 parallel computations
+ */
+export type ExecutingPoolAccount = anchor.IdlTypes<ArciumIdlType>['tinyExecPool'] | anchor.IdlTypes<ArciumIdlType>['smallExecPool'] | anchor.IdlTypes<ArciumIdlType>['mediumExecPool'] | anchor.IdlTypes<ArciumIdlType>['largeExecPool'];
+
+/**
+ * IDL-generated type for circuit source location.
+ * Discriminated union with variants: local, onChain, offChain.
+ */
+export type CircuitSource = anchor.IdlTypes<ArciumIdlType>['circuitSource'];
+
+/**
+ * Fetch and decode the mempool account info for any mempool account size.
+ * @param provider - Anchor provider to use for fetching accounts.
+ * @param mempoolAccPubkey - Public key of the mempool account.
+ * @returns Decoded mempool account info.
+ * @throws Error if the account cannot be fetched or the discriminator is unknown.
+ */
+export async function getMempoolAccInfo(provider: AnchorProvider, mempoolAccPubkey: anchor.web3.PublicKey): Promise<MempoolAccount> {
+    const accData = await provider.connection.getAccountInfo(mempoolAccPubkey);
+    if (accData === null) {
+        throw new Error(`Failed to fetch mempool account ${mempoolAccPubkey.toBase58()}`);
+    }
+
+    const discriminator = Array.from(accData.data.subarray(0, DISCRIMINATOR_SIZE)).toString();
+    const accName = MEMPOOL_DISCRIMINATOR_MAP[discriminator];
+
+    if (accName === undefined) {
+        throw new Error(`Unknown mempool account discriminator: ${discriminator}`);
+    }
+    const program = getArciumProgram(provider);
+    return program.coder.accounts.decode(accName, accData.data);
+}
+
+/**
+ * Fetch and decode the executing pool account info for any pool size.
+ * @param provider - Anchor provider to use for fetching accounts.
+ * @param executingPoolAccPubkey - Public key of the executing pool account.
+ * @returns Decoded executing pool account info.
+ * @throws Error if the account cannot be fetched or the discriminator is unknown.
+ */
+export async function getExecutingPoolAccInfo(provider: AnchorProvider, executingPoolAccPubkey: anchor.web3.PublicKey): Promise<ExecutingPoolAccount> {
+    const accData = await provider.connection.getAccountInfo(executingPoolAccPubkey);
+    if (accData === null) {
+        throw new Error(`Failed to fetch executing pool account ${executingPoolAccPubkey.toBase58()}`);
+    }
+
+    const discriminator = Array.from(accData.data.subarray(0, DISCRIMINATOR_SIZE)).toString();
+    const accName = EXECPOOL_DISCRIMINATOR_MAP[discriminator];
+
+    if (accName === undefined) {
+        throw new Error(`Unknown executing pool account discriminator: ${discriminator}`);
+    }
+    const program = getArciumProgram(provider);
+    return program.coder.accounts.decode(accName, accData.data);
+}
+
+/**
+ * Return all computation references in the mempool for a given account.
+ * Only non-stake computations are included.
+ * @param arciumProgram - Anchor program instance.
+ * @param address - Public key of the mempool account.
+ * @returns Array of ComputationReference objects.
+ */
+export async function getComputationsInMempool(
+    arciumProgram: anchor.Program<ArciumIdlType>,
+    address: PublicKey,
+): Promise<ComputationReference[]> {
+    const mempool = await getMempoolAccInfo(
+        arciumProgram.provider as anchor.AnchorProvider,
+        address,
+    );
+    const startIndex = mempool.inner.computations.startIndex;
+    const length = mempool.inner.computations.length;
+    const elems = mempool.inner.computations.elems;
+
+    function isValid(validBits: number[], idx: number): boolean {
+        const byte = idx >>> 3;
+        const bit = idx & 7;
+
+        if (byte >= validBits.length) {
+            // This should never happen, so we'll want to know about it
+            throw new Error(
+                `isValid: byte ${byte} >= validBits.length ${validBits.length}`,
+            );
+        }
+
+        return (validBits[byte] & (1 << bit)) !== 0;
+    }
+
+    // Handle circular buffer wraparound
+    const refs = [];
+    for (let i = 0; i < length; i++) {
+        const idx = (startIndex + i) % elems.length;
+        // Only save non-stake computations
+        if (isValid(mempool.inner.computations.validBits, idx)) {
+            refs.push(...elems[idx].entries);
+        }
+    }
+
+    return refs
+        .flat()
+        .filter(
+            (ref) => !isNullRef(ref as ComputationReference),
+        ) as ComputationReference[];
+}
+
+/**
+ * Statistics about priority fees for computations in a mempool.
+ * @property mean - Average priority fee across all computations.
+ * @property median - Middle value of priority fees when sorted.
+ * @property min - Lowest priority fee in the mempool.
+ * @property max - Highest priority fee in the mempool.
+ * @property count - Total number of computations in the mempool.
+ */
+export interface MempoolPriorityFeeStats {
+    mean: anchor.BN;
+    median: anchor.BN;
+    min: anchor.BN;
+    max: anchor.BN;
+    count: number;
+}
+
+/**
+ * Calculate priority fee statistics for computations in a mempool.
+ * @param arciumProgram - Anchor program instance.
+ * @param mempoolAddress - Public key of the mempool account.
+ * @returns Priority fee statistics (mean, median, min, max, count).
+ */
+export async function getMempoolPriorityFeeStats(
+    arciumProgram: anchor.Program<ArciumIdlType>,
+    mempoolAddress: PublicKey,
+): Promise<MempoolPriorityFeeStats> {
+    const refs = await getComputationsInMempool(arciumProgram, mempoolAddress);
+
+    if (refs.length === 0) {
+        const zero = new anchor.BN(0);
+        return { mean: zero, median: zero, min: zero, max: zero, count: 0 };
+    }
+
+    const fees = refs.map(ref => ref.priorityFee).sort((a, b) => a.cmp(b));
+
+    const sum = fees.reduce((acc, fee) => acc.add(fee), new anchor.BN(0));
+    const mean = sum.div(new anchor.BN(fees.length));
+
+    const mid = Math.floor(fees.length / 2);
+    const median = fees.length % 2 === 0
+        ? fees[mid - 1].add(fees[mid]).div(new anchor.BN(2))
+        : fees[mid];
+
+    return {
+        mean,
+        median,
+        min: fees[0],
+        max: fees[fees.length - 1],
+        count: fees.length,
+    };
+}
+
+/**
+ * Fetch a specific utility key from the MXE account.
+ * @param provider - Anchor provider to use for fetching accounts.
+ * @param mxeProgramId - Public key of the MXE program.
+ * @param field - Field name to extract ('x25519Pubkey' or 'ed25519VerifyingKey').
+ * @returns Utility key as a Uint8Array, or null if not set.
+ */
+async function getMXEUtilityKey(
+    provider: AnchorProvider,
+    mxeProgramId: anchor.web3.PublicKey,
+    field: 'x25519Pubkey' | 'ed25519VerifyingKey',
+): Promise<Uint8Array | null> {
+    const program = getArciumProgram(provider);
+    const mxeAccAddress = getMXEAccAddress(mxeProgramId);
+    const mxeAccInfo = await program.account.mxeAccount.fetch(mxeAccAddress);
+
+    if ('set' in mxeAccInfo.utilityPubkeys) {
+        const setData = mxeAccInfo.utilityPubkeys.set as {
+            0:
+            {
+                x25519Pubkey: number[],
+                ed25519VerifyingKey: number[],
+                elGamalPubkey: number[],
+                pubkeyValidityProof: number[]
+            }
+        };
+        return new Uint8Array(setData[0][field]);
+    }
+    else if ('unset' in mxeAccInfo.utilityPubkeys) {
+        const unsetData = mxeAccInfo.utilityPubkeys.unset as {
+            0: {
+                x25519Pubkey: number[],
+                ed25519VerifyingKey: number[],
+                elGamalPubkey: number[],
+                pubkeyValidityProof: number[]
+            }; 1: boolean[]
+        };
+        if (unsetData[1].every(Boolean)) {
+            return new Uint8Array(unsetData[0][field]);
+        }
+    }
+    return null;
+}
+
+/**
+ * Fetch and extract the MXE X25519 public key from the MXE account.
+ * @param provider - Anchor provider to use for fetching accounts.
+ * @param mxeProgramId - Public key of the MXE program.
+ * @returns MXE's X25519 public key as a Uint8Array, or null if not set.
+ */
+export async function getMXEPublicKey(
+    provider: AnchorProvider,
+    mxeProgramId: anchor.web3.PublicKey,
+): Promise<Uint8Array | null> {
+    return getMXEUtilityKey(provider, mxeProgramId, 'x25519Pubkey');
+}
+
+/**
+ * Fetch and extract the MXE arcis Ed25519 verifying key from the MXE account.
+ * @param provider - Anchor provider to use for fetching accounts.
+ * @param mxeProgramId - Public key of the MXE program.
+ * @returns MXE's arcis Ed25519 verifying key as a Uint8Array, or null if not set.
+ */
+export async function getMXEArcisEd25519VerifyingKey(
+    provider: AnchorProvider,
+    mxeProgramId: anchor.web3.PublicKey,
+): Promise<Uint8Array | null> {
+    return getMXEUtilityKey(provider, mxeProgramId, 'ed25519VerifyingKey');
+}
+
+/**
+ * Represents the current state of a circuit's on-chain presence.
+ * @internal Used by `uploadCircuit` - most users don't need this directly.
+ */
+export type CircuitState = 'OnchainPending' | 'OnchainFinalized' | 'Offchain';
+
+/**
+ * Determine the current state of a circuit based on its on-chain configuration.
+ * @internal Called internally by `uploadCircuit` - most users don't need this directly.
+ * @param circuitSource - circuitSource field from ComputationDefinitionAccount.
+ * @returns Current state of the circuit.
+ */
+export function getCircuitState(circuitSource: CircuitSource): CircuitState {
+    if (!('onChain' in circuitSource) || !circuitSource.onChain) {
+        return 'Offchain';
+    }
+    if (circuitSource.onChain[0].isCompleted) {
+        return 'OnchainFinalized';
+    }
+    return 'OnchainPending';
+}
+
+/**
+ * Upload a circuit to the blockchain, splitting it into multiple accounts if necessary.
+ * @param provider - Anchor provider to use for transactions.
+ * @param circuitName - Name of the circuit.
+ * @param mxeProgramId - Public key of the MXE program.
+ * @param rawCircuit - Raw circuit data as a Uint8Array.
+ * @param logging - Whether to log progress (default: true).
+ * @param chunkSize - Number of upload transactions to send in parallel (default: 500).
+ * @param confirmOptions - Transaction confirmation options.
+ * @returns Array of transaction signatures for all upload and finalize transactions.
+ */
+export async function uploadCircuit(
+    provider: AnchorProvider,
+    circuitName: string,
+    mxeProgramId: anchor.web3.PublicKey,
+    rawCircuit: Uint8Array,
+    logging: boolean = true,
+    chunkSize: number = 500,
+    confirmOptions?: ConfirmOptions,
+): Promise<string[]> {
+    const compDefAccInfo = getCompDefAccInfo(circuitName, mxeProgramId);
+    const program = getArciumProgram(provider);
+    const compDefAcc = await program.account.computationDefinitionAccount.fetch(compDefAccInfo.pubkey);
+
+    // Type assertion at boundary where Anchor's fetch() returns the account data
+    const state = getCircuitState(compDefAcc.circuitSource as CircuitSource);
+    if (state !== 'OnchainPending') {
+        optionalLog(logging, `Circuit ${circuitName} skipped: ${state}`);
+        return [];
+    }
+
+    // 9 = 8-byte discriminator + 1-byte bump for ComputationDefinitionRaw account
+    const numAccs = Math.ceil(rawCircuit.length / (MAX_ACCOUNT_SIZE - 9));
+    const sigs: string[] = [];
+    const uploadPromises = [];
+    for (let i = 0; i < numAccs; i++) {
+        uploadPromises.push(
+            uploadToCircuitAcc(
+                provider,
+                program,
+                rawCircuit.subarray(
+                    i * (MAX_ACCOUNT_SIZE - 9),
+                    (i + 1) * (MAX_ACCOUNT_SIZE - 9),
+                ),
+                i,
+                compDefAccInfo,
+                mxeProgramId,
+                logging,
+                chunkSize,
+                confirmOptions,
+            ),
+        );
+    }
+
+    sigs.push(...(await Promise.all(uploadPromises)).flat());
+
+    const finalizeCompDefTx = await buildFinalizeCompDefTx(provider, compDefAccInfo.offset, mxeProgramId);
+    sigs.push(await signAndSendWithBlockhash(provider, finalizeCompDefTx, await provider.connection.getLatestBlockhash({ commitment: confirmOptions?.commitment || 'confirmed' }), confirmOptions));
+
+    return sigs;
+}
+
+/**
+ * Queue a key recovery initialization for an MXE on a given cluster.
+ * @param provider - Anchor provider for signing and sending.
+ * @param clusterOffset - Cluster offset to recover keys from.
+ * @param mxeProgramId - Public key of the MXE program.
+ * @param confirmOptions - Transaction confirmation options.
+ * @returns Array of transaction signatures.
+ */
+export async function queueKeyRecoveryInit(
+    provider: AnchorProvider,
+    clusterOffset: number,
+    mxeProgramId: anchor.web3.PublicKey,
+    confirmOptions?: ConfirmOptions,
+): Promise<string[]> {
+    const program = getArciumProgram(provider);
+    const sigs: string[] = [];
+    const queueKeyRecoveryInitTx = await program.methods
+        .queueKeyRecoveryInit(clusterOffset, mxeProgramId)
+        .accounts({
+            signer: provider.publicKey,
+            mxeProgram: mxeProgramId,
+        })
+        .transaction();
+    sigs.push(await signAndSendWithBlockhash(
+        provider,
+        queueKeyRecoveryInitTx,
+        await provider.connection.getLatestBlockhash({ commitment: confirmOptions?.commitment || "confirmed" }),
+        confirmOptions,
+    ));
+    return sigs;
+}
+
+/**
+ * Build a transaction to finalize a computation definition.
+ * @param provider - Anchor provider to use for transactions.
+ * @param compDefOffset - Offset of the computation definition.
+ * @param mxeProgramId - Public key of the MXE program.
+ * @returns Transaction to finalize the computation definition.
+ */
+export async function buildFinalizeCompDefTx(
+    provider: AnchorProvider,
+    compDefOffset: number,
+    mxeProgramId: anchor.web3.PublicKey,
+): Promise<anchor.web3.Transaction> {
+    const program = getArciumProgram(provider);
+    const compDefOffsetBuffer = Buffer.alloc(OFFSET_BUFFER_SIZE);
+    compDefOffsetBuffer.writeUInt32LE(compDefOffset, 0);
+    return program.methods
+        .finalizeComputationDefinition(compDefOffset, mxeProgramId)
+        .accounts({
+            signer: provider.publicKey,
+        })
+        .transaction();
+}
+
+async function uploadToCircuitAcc(
+    provider: AnchorProvider,
+    program: Program<Arcium>,
+    rawCircuitPart: Uint8Array,
+    rawCircuitIndex: number,
+    compDefAccInfo: { pubkey: anchor.web3.PublicKey; offset: number },
+    mxeProgramId: anchor.web3.PublicKey,
+    shouldLog: boolean = true,
+    chunkSize: number = 500,
+    confirmOptions?: ConfirmOptions,
+): Promise<string[]> {
+    const rawCircuitPda = getRawCircuitAccAddress(compDefAccInfo.pubkey, rawCircuitIndex);
+    const existingAcc = await provider.connection.getAccountInfo(rawCircuitPda);
+
+    // Skip entirely if account exists with correct size (pre-seeded at genesis or already uploaded)
+    // Account layout: 8-byte discriminator + 1-byte bump + circuit data
+    const requiredAccountSize = rawCircuitPart.length + 9;
+    if (existingAcc !== null && existingAcc.data.length >= requiredAccountSize) {
+        optionalLog(
+            shouldLog,
+            `Raw circuit acc ${rawCircuitIndex} already exists with sufficient size, skipping`,
+        );
+        return [];
+    }
+
+    const sigs: string[] = [];
+
+    if (existingAcc === null) {
+        const initTx = await program.methods
+            .initRawCircuitAcc(compDefAccInfo.offset, mxeProgramId, rawCircuitIndex)
+            .accounts({
+                signer: provider.publicKey,
+            })
+            .rpc();
+        optionalLog(
+            shouldLog,
+            `Initiated raw circuit acc with raw circuit index ${rawCircuitIndex}`,
+        );
+        sigs.push(initTx);
+    }
+
+    if (rawCircuitPart.length > MAX_REALLOC_PER_IX) {
+        // We only need to add size if the init didn't already make us big enough
+        const nonAsyncTxCount = Math.ceil(
+            rawCircuitPart.length / (MAX_REALLOC_PER_IX * MAX_EMBIGGEN_IX_PER_TX),
+        );
+
+        for (let i = 0; i < nonAsyncTxCount; i++) {
+            optionalLog(shouldLog, `Sending resize tx ${i} of ${nonAsyncTxCount}`);
+            // eslint-disable-next-line no-await-in-loop
+            const tx = await buildResizeTx(
+                program,
+                provider.publicKey,
+                compDefAccInfo,
+                mxeProgramId,
+                rawCircuitIndex,
+                MAX_REALLOC_PER_IX
+                + i * (MAX_REALLOC_PER_IX * MAX_EMBIGGEN_IX_PER_TX),
+                rawCircuitPart.length,
+            );
+            // eslint-disable-next-line no-await-in-loop
+            const blockInfo = await provider.connection.getLatestBlockhash({ commitment: confirmOptions?.commitment || 'confirmed' });
+            sigs.push(
+                // eslint-disable-next-line no-await-in-loop
+                await signAndSendWithBlockhash(provider, tx, blockInfo, confirmOptions),
+            );
+            optionalLog(shouldLog, `Sent resize tx ${i} of ${nonAsyncTxCount}`);
+        }
+    }
+
+    optionalLog(shouldLog, 'Done sending resize txs');
+
+    // The previous txs have allocated all the space we need now, so we can send the remaining upload txs async to one another now
+    const remainingTxCount = Math.ceil(
+        rawCircuitPart.length / MAX_UPLOAD_PER_TX_BYTES,
+    );
+
+    optionalLog(shouldLog, `Sending ${remainingTxCount} upload txs`);
+
+    const remainingTxs: string[] = [];
+    const blockInfo = await provider.connection.getLatestBlockhash({ commitment: confirmOptions?.commitment || 'confirmed' });
+    for (let i = 0; i < remainingTxCount; i += chunkSize) {
+        optionalLog(
+            shouldLog,
+            `Sending chunk ${i / chunkSize + 1
+            } of ${Math.ceil(remainingTxCount / chunkSize)}`,
+        );
+        // Array to hold promises for the current chunk
+        const chunkPromises = [];
+        // Handle the last chunk
+        const currentChunkSize = Math.min(chunkSize, remainingTxCount - i);
+
+        for (let j = 0; j < currentChunkSize; j++) {
+            const offset = MAX_UPLOAD_PER_TX_BYTES * (i + j);
+            // eslint-disable-next-line no-await-in-loop
+            const tx = await buildUploadCircuitTx(
+                program,
+                provider.publicKey,
+                compDefAccInfo,
+                mxeProgramId,
+                Buffer.copyBytesFrom(rawCircuitPart, offset, MAX_UPLOAD_PER_TX_BYTES),
+                offset,
+                rawCircuitIndex,
+            );
+            chunkPromises.push(
+                signAndSendWithBlockhash(provider, tx, blockInfo, confirmOptions),
+            );
+        }
+
+        // Wait for the current chunk to complete before proceeding
+        // eslint-disable-next-line no-await-in-loop
+        const chunkResults = await Promise.all(chunkPromises);
+        // Add the results of the current chunk to the main array
+        remainingTxs.push(...chunkResults);
+        optionalLog(
+            shouldLog,
+            `Done sending chunk ${i / chunkSize + 1
+            } of ${Math.ceil(remainingTxCount / chunkSize)}`,
+        );
+    }
+
+    return sigs.concat(remainingTxs);
+}
+
+async function buildResizeTx(
+    program: anchor.Program<Arcium>,
+    signerPubkey: anchor.web3.PublicKey,
+    compDefAccInfo: { pubkey: anchor.web3.PublicKey; offset: number },
+    mxeProgramId: anchor.web3.PublicKey,
+    rawCircuitIndex: number,
+    currentSize: number,
+    requiredSize: number,
+): Promise<anchor.web3.Transaction> {
+    const ix = await program.methods
+        .embiggenRawCircuitAcc(compDefAccInfo.offset, mxeProgramId, rawCircuitIndex)
+        .accounts({
+            signer: signerPubkey,
+        })
+        .instruction();
+    const resizeSize = Math.min(
+        requiredSize - currentSize,
+        MAX_EMBIGGEN_IX_PER_TX * MAX_REALLOC_PER_IX,
+    );
+    // We might alloc a bit too much here, but that's ok because we fix it in the upload tx
+    const ixCount = Math.ceil(resizeSize / MAX_REALLOC_PER_IX);
+    const tx = new anchor.web3.Transaction();
+    for (let i = 0; i < ixCount; i++) {
+        tx.add(ix);
+    }
+
+    return tx;
+}
+
+async function buildUploadCircuitTx(
+    program: anchor.Program<Arcium>,
+    signerPubkey: anchor.web3.PublicKey,
+    compDefAccInfo: { pubkey: anchor.web3.PublicKey; offset: number },
+    mxeProgramId: anchor.web3.PublicKey,
+    bytes: Buffer,
+    circuitOffset: number,
+    rawCircuitIndex: number,
+): Promise<anchor.web3.Transaction> {
+    if (bytes.length > MAX_UPLOAD_PER_TX_BYTES) {
+        throw new Error(`Upload circuit bytes must be ${MAX_UPLOAD_PER_TX_BYTES} bytes or less per tx`);
+    }
+
+    let bytesInner = bytes;
+
+    if (bytesInner.length < MAX_UPLOAD_PER_TX_BYTES) {
+        const paddedBytes = Buffer.allocUnsafe(MAX_UPLOAD_PER_TX_BYTES);
+        paddedBytes.set(bytesInner);
+        bytesInner = paddedBytes;
+    }
+
+    return program.methods
+        .uploadCircuit(
+            compDefAccInfo.offset,
+            mxeProgramId,
+            rawCircuitIndex,
+            Array.from(bytesInner),
+            circuitOffset,
+        )
+        .accounts({
+            signer: signerPubkey,
+        })
+        .transaction();
+}
+
+async function signAndSendWithBlockhash(
+    provider: AnchorProvider,
+    tx: anchor.web3.Transaction,
+    block: anchor.web3.BlockhashWithExpiryBlockHeight,
+    confirmOptions?: ConfirmOptions,
+): Promise<string> {
+    // eslint-disable-next-line no-param-reassign
+    tx.recentBlockhash = block.blockhash;
+    // eslint-disable-next-line no-param-reassign
+    tx.lastValidBlockHeight = block.lastValidBlockHeight;
+    return provider.sendAndConfirm(tx, [], confirmOptions || { commitment: 'confirmed' });
+}
+
+/**
+ * Return the base seed for an Arcium account, given its name.
+ * @param accName - Name of the account.
+ * @returns Base seed as a Uint8Array.
+ */
+export function getArciumAccountBaseSeed(accName: string): Uint8Array {
+    return Buffer.from(accName, 'utf-8');
+}
+
+/**
+ * Compute the offset for a computation definition account, based on the circuit name.
+ * @param circuitName - Name of the circuit.
+ * @returns Offset as a 4-byte Uint8Array.
+ */
+export function getCompDefAccOffset(circuitName: string): Uint8Array {
+    const hash = new Uint8Array(sha256([Buffer.from(circuitName, 'utf-8')]));
+    return hash.slice(0, COMP_DEF_OFFSET_SIZE);
+}
+
+/**
+ * Return an Anchor program instance for the Arcium program.
+ * @param provider - Anchor provider to use.
+ * @returns Anchor program instance for Arcium.
+ */
+export function getArciumProgram(provider: AnchorProvider): Program<Arcium> {
+    return new Program(ARCIUM_IDL as unknown as Arcium, provider);
+}
+
+/**
+ * Return the public key and offset for a computation definition account, given the circuit name and MXE program ID.
+ * @param circuitName - Name of the circuit.
+ * @param mxeProgramId - Public key of the MXE program.
+ * @returns Object containing the public key and offset for the computation definition account.
+ */
+function getCompDefAccInfo(circuitName: string, mxeProgramId: anchor.web3.PublicKey): { pubkey: anchor.web3.PublicKey; offset: number } {
+    const offset = getCompDefAccOffset(circuitName);
+    const pda = getCompDefAccPDA(getArciumProgramId(), mxeProgramId, offset);
+    return { pubkey: pda, offset: Buffer.from(offset).readUInt32LE(0) };
+}
+
+/**
+ * Return the PDA for a computation definition account, given the program ID, MXE program ID, and offset.
+ * @param arciumProgramId - Public key of the Arcium program.
+ * @param mxeProgramId - Public key of the MXE program.
+ * @param offset - Offset as a Uint8Array.
+ * @returns PDA for the computation definition account.
+ */
+function getCompDefAccPDA(arciumProgramId: anchor.web3.PublicKey, mxeProgramId: anchor.web3.PublicKey, offset: Uint8Array): anchor.web3.PublicKey {
+    return anchor.web3.PublicKey.findProgramAddressSync([Buffer.from(COMP_DEF_ACC_SEED, 'utf-8'), mxeProgramId.toBuffer(), offset], arciumProgramId)[0];
+}
+
+/**
+ * Set an MXE to Recovery status, initiating the key recovery process.
+ * @param provider - Anchor provider to use for transactions.
+ * @param mxeProgramId - Public key of the MXE program to recover.
+ * @param confirmOptions - Transaction confirmation options.
+ * @returns Transaction signature.
+ */
+export async function recoverMxe(
+    provider: AnchorProvider,
+    mxeProgramId: anchor.web3.PublicKey,
+    confirmOptions?: ConfirmOptions,
+): Promise<string> {
+    const program = getArciumProgram(provider);
+    const tx = await program.methods
+        .recoverMxe(mxeProgramId)
+        .accountsPartial({
+            authority: provider.publicKey,
+            mxeProgram: mxeProgramId,
+        })
+        .transaction();
+    return signAndSendWithBlockhash(provider, tx, await provider.connection.getLatestBlockhash({ commitment: confirmOptions?.commitment || 'confirmed' }), confirmOptions);
+}
+
+/**
+ * Initialize key recovery execution by creating the MxeRecoveryAccount and
+ * registering the key_recovery_final computation definition on the backup MXE.
+ * This is split into two parts due to Solana's 10KB per-instruction allocation limit.
+ * @param provider - Anchor provider to use for transactions.
+ * @param originalMxeProgramId - Public key of the original MXE program being recovered.
+ * @param backupMxeProgramId - Public key of the backup MXE program that will take over.
+ * @param confirmOptions - Transaction confirmation options.
+ * @returns Transaction signature from part2.
+ */
+export async function initKeyRecoveryExecution(
+    provider: AnchorProvider,
+    originalMxeProgramId: anchor.web3.PublicKey,
+    backupMxeProgramId: anchor.web3.PublicKey,
+    confirmOptions?: ConfirmOptions,
+): Promise<string> {
+    // Part 1: Create MxeRecoveryAccount with partial size
+    const program = getArciumProgram(provider);
+    const tx1 = await program.methods
+        .initKeyRecoveryExecutionPart1(originalMxeProgramId, backupMxeProgramId)
+        .accountsPartial({
+            payer: provider.publicKey,
+            originalMxeProgram: originalMxeProgramId,
+            backupMxeProgram: backupMxeProgramId,
+        })
+        .transaction();
+    await signAndSendWithBlockhash(provider, tx1, await provider.connection.getLatestBlockhash({ commitment: confirmOptions?.commitment || 'confirmed' }), confirmOptions);
+
+    // Part 2: Reallocate to full size and create computation definition
+    const tx2 = await program.methods
+        .initKeyRecoveryExecutionPart2(originalMxeProgramId, backupMxeProgramId)
+        .accountsPartial({
+            authority: provider.publicKey,
+            payer: provider.publicKey,
+            originalMxeProgram: originalMxeProgramId,
+            backupMxeProgram: backupMxeProgramId,
+        })
+        .transaction();
+    return signAndSendWithBlockhash(provider, tx2, await provider.connection.getLatestBlockhash({ commitment: confirmOptions?.commitment || 'confirmed' }), confirmOptions);
+}
+
+/**
+ * Submit a re-encrypted key recovery share from a recovery peer.
+ * Recovery peers must decrypt shares using their x25519 private key and re-encrypt
+ * them for the backup MXE before submission.
+ * @param provider - Anchor provider to use for transactions.
+ * @param originalMxeProgramId - Public key of the original MXE program being recovered.
+ * @param backupMxeProgramId - Public key of the backup MXE program.
+ * @param peerOffset - Offset of the recovery peer.
+ * @param peerIndex - Index of this peer in the recovery peers list.
+ * @param share - Re-encrypted share: 5 field elements of 32 bytes each (160 bytes total).
+ * @param confirmOptions - Transaction confirmation options.
+ * @returns Transaction signature.
+ */
+export async function submitKeyRecoveryShare(
+    provider: AnchorProvider,
+    originalMxeProgramId: anchor.web3.PublicKey,
+    backupMxeProgramId: anchor.web3.PublicKey,
+    peerOffset: number,
+    peerIndex: number,
+    share: number[][] | Uint8Array[],
+    confirmOptions?: ConfirmOptions,
+): Promise<string> {
+    const program = getArciumProgram(provider);
+    // Convert to array of 5 elements, each 32 bytes
+    const shareArrays: number[][] = share.map(elem => Array.from(elem));
+    if (shareArrays.length !== 5) {
+        throw new Error(`Share must contain exactly 5 elements, got ${shareArrays.length}`);
+    }
+    for (let i = 0; i < 5; i++) {
+        if (shareArrays[i].length !== 32) {
+            throw new Error(`Share element ${i} must be exactly 32 bytes, got ${shareArrays[i].length}`);
+        }
+    }
+    const tx = await program.methods
+        .submitKeyRecoveryShare(
+            originalMxeProgramId,
+            backupMxeProgramId,
+            peerOffset,
+            peerIndex,
+            shareArrays as number[][],
+        )
+        .accountsPartial({
+            signer: provider.publicKey,
+            originalMxeProgram: originalMxeProgramId,
+            backupMxeProgram: backupMxeProgramId,
+        })
+        .transaction();
+    return signAndSendWithBlockhash(provider, tx, await provider.connection.getLatestBlockhash({ commitment: confirmOptions?.commitment || 'confirmed' }), confirmOptions);
+}
+
+/**
+ * Finalize key recovery execution after the submission threshold is met.
+ * This queues the key_recovery_finalize MPC computation on the backup cluster.
+ * @param provider - Anchor provider to use for transactions.
+ * @param originalMxeProgramId - Public key of the original MXE program being recovered.
+ * @param backupMxeProgramId - Public key of the backup MXE program.
+ * @param clusterOffset - Cluster offset where the backup MXE is deployed.
+ * @param keyRecoveryFinalizeOffset - Computation offset for the key_recovery_finalize computation.
+ * @param confirmOptions - Transaction confirmation options.
+ * @returns Transaction signature.
+ */
+export async function finalizeKeyRecoveryExecution(
+    provider: AnchorProvider,
+    originalMxeProgramId: anchor.web3.PublicKey,
+    backupMxeProgramId: anchor.web3.PublicKey,
+    clusterOffset: number,
+    keyRecoveryFinalizeOffset: anchor.BN,
+    confirmOptions?: ConfirmOptions,
+): Promise<string> {
+    const program = getArciumProgram(provider);
+    const tx = await program.methods
+        .finalizeKeyRecoveryExecution(originalMxeProgramId, backupMxeProgramId, clusterOffset)
+        .accountsPartial({
+            authority: provider.publicKey,
+            payer: provider.publicKey,
+            keyRecoveryFinalizeComputation: getComputationAccAddress(clusterOffset, keyRecoveryFinalizeOffset),
+            executingPool: getExecutingPoolAccAddress(clusterOffset),
+            mempool: getMempoolAccAddress(clusterOffset),
+            originalMxeProgram: originalMxeProgramId,
+            backupMxeProgram: backupMxeProgramId,
+        })
+        .transaction();
+    return signAndSendWithBlockhash(provider, tx, await provider.connection.getLatestBlockhash({ commitment: confirmOptions?.commitment || 'confirmed' }), confirmOptions);
+}
+
+/**
+ * Initialize an MXE (part 1). Due to Solana's 10KB per-instruction allocation limit,
+ * this only partially allocates recovery_cluster_acc.
+ * Call initMxePart2 afterwards to finish allocation and add keygen to mempool.
+ * @param provider - Anchor provider to use for transactions.
+ * @param mxeProgramId - Public key to use as the MXE program ID.
+ * @param confirmOptions - Transaction confirmation options.
+ * @returns Transaction signature.
+ */
+export async function initMxePart1(
+    provider: AnchorProvider,
+    mxeProgramId: anchor.web3.PublicKey,
+    confirmOptions?: ConfirmOptions,
+): Promise<string> {
+    const program = getArciumProgram(provider);
+    const tx = await program.methods
+        .initMxePart1()
+        .accountsPartial({
+            signer: provider.publicKey,
+            mxeProgram: mxeProgramId,
+        })
+        .transaction();
+    return signAndSendWithBlockhash(provider, tx, await provider.connection.getLatestBlockhash({ commitment: confirmOptions?.commitment || 'confirmed' }), confirmOptions);
+}
+
+/**
+ * Finish MXE initialization (part 2).
+ * Reallocate recovery_cluster_acc to full size, initialize recovery_peers,
+ * and add the keygen computation to the mempool.
+ * @param provider - Anchor provider to use for transactions.
+ * @param clusterOffset - Cluster offset to associate with the MXE.
+ * @param mxeProgramId - Public key to use as the MXE program ID.
+ * @param recoveryPeers - Array of 100 node offsets for recovery peers (0 for unused slots).
+ * @param keygenOffset - Computation offset for the keygen computation.
+ * @param keyRecoveryInitOffset - Computation offset for the key_recovery_init computation.
+ * @param lutOffset - Lookup table offset for the MXE's address lookup table.
+ * @param mxeAuthority - Optional authority for the MXE (defaults to provider.publicKey).
+ * @param confirmOptions - Transaction confirmation options.
+ * @returns Transaction signature.
+ */
+export async function initMxePart2(
+    provider: AnchorProvider,
+    clusterOffset: number,
+    mxeProgramId: anchor.web3.PublicKey,
+    recoveryPeers: number[],
+    keygenOffset: anchor.BN,
+    keyRecoveryInitOffset: anchor.BN,
+    lutOffset: anchor.BN,
+    mxeAuthority?: anchor.web3.PublicKey,
+    confirmOptions?: ConfirmOptions,
+): Promise<string> {
+    const program = getArciumProgram(provider);
+    // Ensure recoveryPeers has exactly 100 elements
+    const paddedRecoveryPeers = [...recoveryPeers];
+    while (paddedRecoveryPeers.length < 100) {
+        paddedRecoveryPeers.push(0);
+    }
+    if (paddedRecoveryPeers.length > 100) {
+        throw new Error('recoveryPeers must have at most 100 elements');
+    }
+
+    const tx = await program.methods
+        .initMxePart2(
+            clusterOffset,
+            mxeProgramId,
+            paddedRecoveryPeers,
+            keygenOffset,
+            keyRecoveryInitOffset,
+            lutOffset,
+        )
+        .accountsPartial({
+            signer: provider.publicKey,
+            mxeProgram: mxeProgramId,
+            poolAccount: getFeePoolAccAddress(),
+            mxeAuthority: mxeAuthority ?? provider.publicKey,
+            addressLookupTable: getLookupTableAddress(mxeProgramId, lutOffset),
+        })
+        .transaction();
+    return signAndSendWithBlockhash(provider, tx, await provider.connection.getLatestBlockhash({ commitment: confirmOptions?.commitment || 'confirmed' }), confirmOptions);
+}
+
+export async function claimComputationRent(
+    provider: AnchorProvider,
+    clusterOffset: number,
+    computationOffset: anchor.BN,
+    confirmOptions?: ConfirmOptions,
+): Promise<string> {
+    const program = getArciumProgram(provider);
+    const tx = await program.methods
+        .claimComputationRent(computationOffset, clusterOffset)
+        .accountsPartial({
+            signer: provider.publicKey,
+            comp: getComputationAccAddress(clusterOffset, computationOffset),
+        })
+        .transaction();
+    return signAndSendWithBlockhash(provider, tx, await provider.connection.getLatestBlockhash({ commitment: confirmOptions?.commitment || 'confirmed' }), confirmOptions);
+}