hpl-toolkit 0.0.5

use anchor_lang::{
    accounts::program::Program, context::CpiContext, solana_program::account_info::AccountInfo,
    Result, ToAccountInfo,
};
use spl_account_compression::{program::SplAccountCompression, ConcurrentMerkleTree, Node, Noop};

pub trait MerkleTreeUtils {
    fn get_latest_leaf_index(&self) -> u32;
    fn tree_capacity(&self) -> u32;
    fn tree_capacity_is_zero(&self) -> bool;
    fn tree_capacity_is_more_than(&self, val: u32) -> bool;
    fn get_hpl_info(&self) -> (u32, u32, u64);
}

impl<const MAX_DEPTH: usize, const MAX_BUFFER_SIZE: usize> MerkleTreeUtils
    for ConcurrentMerkleTree<MAX_DEPTH, MAX_BUFFER_SIZE>
{
    fn get_latest_leaf_index(&self) -> u32 {
        self.rightmost_proof.index
    }

    fn tree_capacity(&self) -> u32 {
        let max_capacity = 2u32.pow(MAX_DEPTH as u32) - 1;
        max_capacity - self.rightmost_proof.index
    }

    fn tree_capacity_is_zero(&self) -> bool {
        self.tree_capacity() == 0
    }

    fn tree_capacity_is_more_than(&self, val: u32) -> bool {
        self.tree_capacity() > val
    }
    fn get_hpl_info(&self) -> (u32, u32, u64) {
        (
            self.tree_capacity(),
            self.get_latest_leaf_index(),
            self.get_seq(),
        )
    }
}

const HEADER_V1_SIZE: usize = 4 + 4 + 32 + 8 + 6;

pub fn calculate_concurrent_merkle_tree_byte_size(
    max_depth: usize,
    max_buffer_size: usize,
) -> usize {
    let u64_byte_size = 8; // Size of u64
    let u32_byte_size = 4; // Size of u32
    let public_key_byte_size = 32; // Size of public key

    // ChangeLogInternal structure
    let change_log_byte_size =
        public_key_byte_size + max_depth * public_key_byte_size + u32_byte_size + u32_byte_size;

    // Path structure
    let path_byte_size =
        max_depth * public_key_byte_size + public_key_byte_size + u32_byte_size + u32_byte_size;

    // ConcurrentMerkleTree structure
    let concurrent_merkle_tree_byte_size = u64_byte_size * 3 + // sequenceNumber, activeIndex, bufferSize
        change_log_byte_size * max_buffer_size +
        path_byte_size;

    concurrent_merkle_tree_byte_size
}

pub fn calculate_canopy_depth_header_v1(
    max_depth: usize,
    max_buffer_size: usize,
    tree_size: usize,
) -> u8 {
    let canopy_byte_size = tree_size
        - 2
        - HEADER_V1_SIZE
        - calculate_concurrent_merkle_tree_byte_size(max_depth, max_buffer_size);
    let depth = ((canopy_byte_size / 32) as f64 + 2.0).log2() - 1.0;
    depth as u8
}

pub fn init_tree<'info>(
    max_depth: u32,
    max_buffer_size: u32,
    authority: &AccountInfo<'info>,
    merkle_tree: &AccountInfo<'info>,
    compression_program: &Program<'info, SplAccountCompression>,
    noop: &Program<'info, Noop>,
    signer_seeds: Option<&[&[&[u8]]; 1]>,
) -> Result<()> {
    let program = compression_program.to_account_info();
    let accounts = spl_account_compression::cpi::accounts::Initialize {
        authority: authority.to_account_info(),
        merkle_tree: merkle_tree.to_account_info(),
        noop: noop.to_account_info(),
    };
    let mut ctx = CpiContext::new(program, accounts);
    if let Some(signer_seeds) = signer_seeds {
        ctx = ctx.with_signer(signer_seeds);
    }
    spl_account_compression::cpi::init_empty_merkle_tree(ctx, max_depth, max_buffer_size)
}

pub fn close_tree<'info>(
    recipient: &AccountInfo<'info>,
    authority: &AccountInfo<'info>,
    merkle_tree: &AccountInfo<'info>,
    compression_program: &Program<'info, SplAccountCompression>,
    signer_seeds: Option<&[&[&[u8]]; 1]>,
) -> Result<()> {
    let program = compression_program.to_account_info();
    let accounts = spl_account_compression::cpi::accounts::CloseTree {
        recipient: recipient.to_account_info(),
        authority: authority.to_account_info(),
        merkle_tree: merkle_tree.to_account_info(),
    };

    let mut ctx = CpiContext::new(program, accounts);
    if let Some(signer_seeds) = signer_seeds {
        ctx = ctx.with_signer(signer_seeds);
    }

    spl_account_compression::cpi::close_empty_tree(ctx)
}

pub fn append_leaf<'info>(
    leaf: Node,
    authority: &AccountInfo<'info>,
    merkle_tree: &AccountInfo<'info>,
    compression_program: &Program<'info, SplAccountCompression>,
    noop: &Program<'info, Noop>,
    signer_seeds: Option<&[&[&[u8]]; 1]>,
) -> Result<()> {
    let program = compression_program.to_account_info();
    let accounts = spl_account_compression::cpi::accounts::Modify {
        authority: authority.to_account_info(),
        merkle_tree: merkle_tree.to_account_info(),
        noop: noop.to_account_info(),
    };

    let mut ctx = CpiContext::new(program, accounts);
    if let Some(signer_seeds) = signer_seeds {
        ctx = ctx.with_signer(signer_seeds);
    }

    spl_account_compression::cpi::append(ctx, leaf)
}

pub fn replace_leaf<'info>(
    root: Node,
    previous_leaf: Node,
    new_leaf: Node,
    index: u32,
    authority: &AccountInfo<'info>,
    merkle_tree: &AccountInfo<'info>,
    compression_program: &Program<'info, SplAccountCompression>,
    noop: &Program<'info, Noop>,
    remaining_accounts: Vec<AccountInfo<'info>>,
    signer_seeds: Option<&[&[&[u8]]; 1]>,
) -> Result<()> {
    let program = compression_program.to_account_info();
    let accounts = spl_account_compression::cpi::accounts::Modify {
        authority: authority.to_account_info(),
        merkle_tree: merkle_tree.to_account_info(),
        noop: noop.to_account_info(),
    };

    let mut ctx = CpiContext::new(program, accounts).with_remaining_accounts(remaining_accounts);
    if let Some(signer_seeds) = signer_seeds {
        ctx = ctx.with_signer(signer_seeds);
    }

    spl_account_compression::cpi::replace_leaf(ctx, root, previous_leaf, new_leaf, index)
}

pub fn verify_leaf<'info>(
    root: Node,
    leaf: Node,
    index: u32,
    merkle_tree: &AccountInfo<'info>,
    compression_program: &Program<'info, SplAccountCompression>,
    remaining_accounts: Vec<AccountInfo<'info>>,
) -> Result<()> {
    let program = compression_program.to_account_info();
    let accounts = spl_account_compression::cpi::accounts::VerifyLeaf {
        merkle_tree: merkle_tree.to_account_info(),
    };

    let ctx = CpiContext::new(program, accounts).with_remaining_accounts(remaining_accounts);

    spl_account_compression::cpi::verify_leaf(ctx, root, leaf, index)
}