jito-distributor-sdk/src/distributor.ts

TypeScript SDK for JITO Merkle Distributor with production-ready versioning and double-hashing support

import { 
  PublicKey, 
  SystemProgram,
  SYSVAR_RENT_PUBKEY,
} from '@solana/web3.js';
import { 
  Program, 
  AnchorProvider, 
  BN 
} from '@coral-xyz/anchor';
import { 
  TOKEN_PROGRAM_ID, 
  ASSOCIATED_TOKEN_PROGRAM_ID,
  getAssociatedTokenAddress 
} from '@solana/spl-token';

import idl from '../idl/merkle_distributor.json';
import { 
  PROGRAM_ID, 
  MerkleDistributor as MerkleDistributorAccount,
  ClaimStatus,
  CreateDistributorArgs,
  ClaimArgs,
  ClaimLockedArgs 
} from './types';
import { 
  getDistributorPDA, 
  getClaimStatusPDA,
  bigintToBN,
  validateTimestamps,
  validateMerkleProof 
} from './utils';

/**
 * MerkleDistributor SDK class providing a clean interface to the Anchor program
 */
export class MerkleDistributor {
  public readonly program: Program;
  public readonly provider: AnchorProvider;
  public readonly programId: PublicKey;

  constructor(provider: AnchorProvider, programId?: PublicKey) {
    this.provider = provider;
    this.programId = programId || PROGRAM_ID;
    this.program = new Program(idl as any, this.programId, provider);
  }

  /**
   * Creates a new merkle distributor
   * @param args CreateDistributorArgs
   * @returns Transaction signature
   */
  async createDistributor(args: CreateDistributorArgs): Promise<string> {
    // Validate timestamps
    const timestampValidation = validateTimestamps(
      args.startVestingTs,
      args.endVestingTs,
      args.clawbackStartTs
    );
    
    if (!timestampValidation.valid) {
      throw new Error(`Invalid timestamps: ${timestampValidation.error}`);
    }

    // Derive PDAs using the custom program ID
    const versionBuffer = Buffer.alloc(8);
    versionBuffer.writeBigUInt64LE(args.version);
    const [distributorPDA] = PublicKey.findProgramAddressSync(
      [
        Buffer.from('MerkleDistributor'),
        args.mint.toBuffer(),
        versionBuffer
      ],
      this.programId
    );

    const distributorTokenAccount = await getAssociatedTokenAddress(
      args.mint,
      distributorPDA,
      true // allowOwnerOffCurve
    );

    // Convert arguments to the format expected by Anchor
    const anchorArgs = {
      version: bigintToBN(args.version),
      root: Array.from(args.root),
      maxTotalClaim: bigintToBN(args.maxTotalClaim),
      maxNumNodes: bigintToBN(args.maxNumNodes),
      startVestingTs: bigintToBN(args.startVestingTs),
      endVestingTs: bigintToBN(args.endVestingTs),
      clawbackStartTs: bigintToBN(args.clawbackStartTs),
    };

    const signature = await this.program.methods
      .newDistributor(
        anchorArgs.version,
        anchorArgs.root,
        anchorArgs.maxTotalClaim,
        anchorArgs.maxNumNodes,
        anchorArgs.startVestingTs,
        anchorArgs.endVestingTs,
        anchorArgs.clawbackStartTs
      )
      .accounts({
        distributor: distributorPDA,
        clawbackReceiver: args.clawbackReceiver,
        mint: args.mint,
        tokenVault: distributorTokenAccount,
        admin: args.admin,
        systemProgram: SystemProgram.programId,
        associatedTokenProgram: ASSOCIATED_TOKEN_PROGRAM_ID,
        tokenProgram: TOKEN_PROGRAM_ID,
      })
      .rpc();

    return signature;
  }

  /**
   * Claims tokens from the distributor
   * @param args ClaimArgs
   * @returns Transaction signature
   */
  async claim(args: ClaimArgs): Promise<string> {
    // Validate proof format
    if (!validateMerkleProof(args.proof)) {
      throw new Error('Invalid merkle proof format');
    }

    // Get distributor info to derive accounts
    const distributorInfo = await this.getDistributor(args.distributor);
    
    // Derive PDAs using the custom program ID
    const [claimStatusPDA] = PublicKey.findProgramAddressSync(
      [
        Buffer.from('ClaimStatus'),
        args.claimant.toBuffer(),
        args.distributor.toBuffer()
      ],
      this.programId
    );
    
    const distributorTokenAccount = await getAssociatedTokenAddress(
      distributorInfo.mint,
      args.distributor,
      true // allowOwnerOffCurve
    );

    // Convert proof to the format expected by Anchor
    const anchorProof = args.proof.map(p => Array.from(p));

    const signature = await this.program.methods
      .newClaim(
        bigintToBN(args.amountUnlocked),
        bigintToBN(args.amountLocked),
        anchorProof
      )
      .accounts({
        distributor: args.distributor,
        claimStatus: claimStatusPDA,
        from: distributorTokenAccount,
        to: args.claimantTokenAccount,
        claimant: args.claimant,
        tokenProgram: TOKEN_PROGRAM_ID,
        systemProgram: SystemProgram.programId,
      })
      .rpc();

    return signature;
  }

  /**
   * Claims locked tokens after vesting period
   * @param args ClaimLockedArgs
   * @returns Transaction signature
   */
  async claimLocked(args: ClaimLockedArgs): Promise<string> {
    // Get distributor info to derive accounts
    const distributorInfo = await this.getDistributor(args.distributor);
    
    // Derive PDAs using the custom program ID
    const [claimStatusPDA] = PublicKey.findProgramAddressSync(
      [
        Buffer.from('ClaimStatus'),
        args.claimant.toBuffer(),
        args.distributor.toBuffer()
      ],
      this.programId
    );
    
    const distributorTokenAccount = await getAssociatedTokenAddress(
      distributorInfo.mint,
      args.distributor,
      true // allowOwnerOffCurve
    );

    const signature = await this.program.methods
      .claimLocked()
      .accounts({
        distributor: args.distributor,
        claimStatus: claimStatusPDA,
        from: distributorTokenAccount,
        to: args.claimantTokenAccount,
        claimant: args.claimant,
        tokenProgram: TOKEN_PROGRAM_ID,
      })
      .rpc();

    return signature;
  }

  /**
   * Claws back remaining tokens to the clawback receiver
   * @param distributor Distributor public key
   * @param claimant Claimant public key (can be anyone after clawback period)
   * @returns Transaction signature
   */
  async clawback(distributor: PublicKey, claimant: PublicKey): Promise<string> {
    // Get distributor info
    const distributorInfo = await this.getDistributor(distributor);
    
    const distributorTokenAccount = await getAssociatedTokenAddress(
      distributorInfo.mint,
      distributor,
      true // allowOwnerOffCurve
    );

    const clawbackTokenAccount = await getAssociatedTokenAddress(
      distributorInfo.mint,
      distributorInfo.clawbackReceiver
    );

    const signature = await this.program.methods
      .clawback()
      .accounts({
        distributor,
        from: distributorTokenAccount,
        to: clawbackTokenAccount,
        claimant,
        systemProgram: SystemProgram.programId,
        tokenProgram: TOKEN_PROGRAM_ID,
      })
      .rpc();

    return signature;
  }

  /**
   * Sets a new admin for the distributor
   * @param distributor Distributor public key
   * @param currentAdmin Current admin public key
   * @param newAdmin New admin public key
   * @returns Transaction signature
   */
  async setAdmin(distributor: PublicKey, currentAdmin: PublicKey, newAdmin: PublicKey): Promise<string> {
    const signature = await this.program.methods
      .setAdmin()
      .accounts({
        distributor,
        admin: currentAdmin,
        newAdmin,
      })
      .rpc();

    return signature;
  }

  /**
   * Sets a new clawback receiver for the distributor
   * @param distributor Distributor public key
   * @param newClawbackReceiver New clawback receiver public key
   * @param admin Admin public key
   * @returns Transaction signature
   */
  async setClawbackReceiver(
    distributor: PublicKey, 
    newClawbackReceiver: PublicKey, 
    admin: PublicKey
  ): Promise<string> {
    const signature = await this.program.methods
      .setClawbackReceiver()
      .accounts({
        distributor,
        newClawbackAccount: newClawbackReceiver,
        admin,
      })
      .rpc();

    return signature;
  }

  /**
   * Fetches a distributor account
   * @param distributor Distributor public key
   * @returns MerkleDistributor account data
   */
  async getDistributor(distributor: PublicKey): Promise<MerkleDistributorAccount> {
    const account = await this.program.account.merkleDistributor.fetch(distributor);
    return account as unknown as MerkleDistributorAccount;
  }

  /**
   * Fetches a claim status account
   * @param claimStatus Claim status public key
   * @returns ClaimStatus account data
   */
  async getClaimStatus(claimStatus: PublicKey): Promise<ClaimStatus> {
    const account = await this.program.account.claimStatus.fetch(claimStatus);
    return account as unknown as ClaimStatus;
  }

  /**
   * Fetches claim status for a specific claimant and distributor
   * @param claimant Claimant public key
   * @param distributor Distributor public key
   * @returns ClaimStatus account data or null if not found
   */
  async getClaimStatusForClaimant(claimant: PublicKey, distributor: PublicKey): Promise<ClaimStatus | null> {
    try {
      const [claimStatusPDA] = PublicKey.findProgramAddressSync(
        [
          Buffer.from('ClaimStatus'),
          claimant.toBuffer(),
          distributor.toBuffer()
        ],
        this.programId
      );
      return await this.getClaimStatus(claimStatusPDA);
    } catch (error) {
      // Account doesn't exist yet
      return null;
    }
  }

  /**
   * Checks if a claimant has already claimed
   * @param claimant Claimant public key
   * @param distributor Distributor public key
   * @returns Boolean indicating if tokens have been claimed
   */
  async hasClaimed(claimant: PublicKey, distributor: PublicKey): Promise<boolean> {
    const claimStatus = await this.getClaimStatusForClaimant(claimant, distributor);
    return claimStatus !== null;
  }

  /**
   * Queries all existing distributors for a given mint
   * @param mint The mint to query distributors for
   * @param maxVersion Maximum version to check (default: 100)
   * @returns Map of version to distributor info
   */
  async queryDistributorsForMint(mint: PublicKey, maxVersion: number = 100): Promise<Map<bigint, {
    pda: PublicKey;
    account: MerkleDistributorAccount;
    version: bigint;
  }>> {
    const distributors = new Map();
    
    for (let version = 0n; version <= BigInt(maxVersion); version++) {
      try {
        const [pda] = getDistributorPDA(mint, version);
        const account = await this.getDistributor(pda);
        
        distributors.set(version, {
          pda,
          account,
          version
        });
      } catch (error) {
        // Distributor doesn't exist for this version, continue
        continue;
      }
    }
    
    return distributors;
  }

  /**
   * Finds the next available version for a mint
   * @param mint The mint to find next version for
   * @param startFrom Starting version to check from (default: 0)
   * @param maxCheck Maximum version to check (default: 1000)
   * @returns Next available version number
   */
  async findNextAvailableVersion(mint: PublicKey, startFrom: bigint = 0n, maxCheck: number = 1000): Promise<bigint> {
    for (let version = startFrom; version <= BigInt(maxCheck); version++) {
      try {
        const [pda] = getDistributorPDA(mint, version);
        await this.getDistributor(pda);
        // If we get here, distributor exists, continue to next version
        continue;
      } catch (error) {
        // Distributor doesn't exist, this version is available
        return version;
      }
    }
    
    throw new Error(`No available version found between ${startFrom} and ${maxCheck}`);
  }

  /**
   * Gets a comprehensive overview of distributions for a mint
   * @param mint The mint to get overview for
   * @param maxVersion Maximum version to check (default: 100)
   * @returns Distribution overview with used versions, next available, and stats
   */
  async getDistributionOverview(mint: PublicKey, maxVersion: number = 100): Promise<{
    mint: PublicKey;
    usedVersions: bigint[];
    nextAvailableVersion: bigint;
    totalDistributors: number;
    totalClaimed: bigint;
    totalUnclaimed: bigint;
    distributors: Map<bigint, {
      pda: PublicKey;
      account: MerkleDistributorAccount;
      version: bigint;
      claimedAmount: bigint;
      remainingAmount: bigint;
    }>;
  }> {
    const distributors = await this.queryDistributorsForMint(mint, maxVersion);
    const usedVersions = Array.from(distributors.keys()).sort((a, b) => Number(a - b));
    const nextAvailableVersion = await this.findNextAvailableVersion(mint, 0n, maxVersion + 100);
    
    let totalClaimed = 0n;
    let totalUnclaimed = 0n;
    
    const distributorMap = new Map();
    
    for (const [version, info] of distributors) {
      const claimedAmount = BigInt(info.account.totalAmountClaimed.toString());
      const maxClaim = BigInt(info.account.maxTotalClaim.toString());
      const remainingAmount = maxClaim - claimedAmount;
      
      totalClaimed += claimedAmount;
      totalUnclaimed += remainingAmount;
      
      distributorMap.set(version, {
        ...info,
        claimedAmount,
        remainingAmount
      });
    }
    
    return {
      mint,
      usedVersions,
      nextAvailableVersion,
      totalDistributors: distributors.size,
      totalClaimed,
      totalUnclaimed,
      distributors: distributorMap
    };
  }

  /**
   * Checks if a version is available for a mint
   * @param mint The mint to check
   * @param version The version to check
   * @returns Boolean indicating if version is available
   */
  async isVersionAvailable(mint: PublicKey, version: bigint): Promise<boolean> {
    try {
      const [pda] = getDistributorPDA(mint, version);
      await this.getDistributor(pda);
      return false; // Distributor exists, version not available
    } catch (error) {
      return true; // Distributor doesn't exist, version available
    }
  }

  /**
   * Gets the PDA for a specific mint and version
   * @param mint The mint
   * @param version The version
   * @returns [PDA, bump] tuple
   */
  getDistributorPDA(mint: PublicKey, version: bigint): [PublicKey, number] {
    return getDistributorPDA(mint, version);
  }

  /**
   * Batch check multiple versions for availability
   * @param mint The mint to check versions for
   * @param versions Array of versions to check
   * @returns Map of version to availability status
   */
  async batchCheckVersions(mint: PublicKey, versions: bigint[]): Promise<Map<bigint, boolean>> {
    const results = new Map<bigint, boolean>();
    
    const checks = versions.map(async (version) => {
      const available = await this.isVersionAvailable(mint, version);
      results.set(version, available);
    });
    
    await Promise.all(checks);
    return results;
  }
}