airdrop-api 0.1.0

// Merkle proof verification utilities
// Aligned with miracle-copy and escrow-copy's merkle tree construction approach
// Our approach: Single merkle tree per campaign (simpler than dual merkle tree)

use solana_program::hash::hashv;
use solana_program::pubkey::Pubkey;

/// Airdrop allocation data for merkle tree construction
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AirdropAllocation {
    /// Recipient wallet address
    pub recipient: Pubkey,
    /// Token amount to airdrop
    pub amount: u64,
}

impl AirdropAllocation {
    /// Create a new airdrop allocation
    pub fn new(recipient: Pubkey, amount: u64) -> Self {
        Self { recipient, amount }
    }

    /// Compute leaf hash for this allocation
    pub fn compute_leaf_hash(&self) -> [u8; 32] {
        compute_leaf_hash(&self.recipient, self.amount)
    }
}

/// Merkle tree for airdrop campaign
///
/// ## Construction Rules (aligned with miracle-copy and escrow-copy):
/// 1. **Single leaf case**: Duplicate the leaf before building tree
///    - Root = hash(leaf, leaf) instead of root = leaf
/// 2. **Odd nodes at level**: Hash last node with itself
///    - Ensures consistent tree structure
/// 3. **Deterministic**: Same allocations always produce same root
#[derive(Clone, Debug)]
pub struct MerkleTree {
    /// Merkle root hash (32 bytes)
    root: [u8; 32],
    /// Tree levels (for proof generation)
    /// tree[0] = leaf level, tree[last] = root level
    tree: Vec<Vec<[u8; 32]>>,
    /// Original allocations (sorted)
    allocations: Vec<AirdropAllocation>,
}

impl MerkleTree {
    /// Create a new merkle tree from airdrop allocations
    ///
    /// ## Parameters
    /// - `allocations`: Vector of (recipient, amount) allocations
    ///
    /// ## Returns
    /// - `Ok(MerkleTree)` if successful
    /// - `Err(String)` if allocations is empty
    ///
    /// ## Single Leaf Handling
    /// If only one allocation exists, the leaf is duplicated before building the tree.
    /// This matches miracle-copy and escrow-copy's approach for consistency.
    ///
    /// ## Example
    /// ```rust
    /// use airdrop_api::merkle::{MerkleTree, AirdropAllocation};
    /// use solana_program::pubkey::Pubkey;
    ///
    /// let allocations = vec![
    ///     AirdropAllocation::new(Pubkey::new_unique(), 1000),
    ///     AirdropAllocation::new(Pubkey::new_unique(), 2000),
    /// ];
    ///
    /// let tree = MerkleTree::new(allocations)?;
    /// println!("Merkle root: {:?}", tree.root());
    /// ```
    pub fn new(mut allocations: Vec<AirdropAllocation>) -> Result<Self, String> {
        if allocations.is_empty() {
            return Err("Cannot create merkle tree from empty allocations list".to_string());
        }

        // Sort allocations deterministically by recipient for consistent merkle tree construction
        // This ensures that the same allocations always produce the same merkle root regardless of input order
        allocations.sort_by(|a, b| a.recipient.cmp(&b.recipient));

        // Generate leaf hashes from sorted allocations
        let mut leaves: Vec<[u8; 32]> = allocations.iter().map(|a| a.compute_leaf_hash()).collect();

        // CRITICAL: Handle single leaf case - duplicate the leaf to force one hash operation
        // This matches miracle-copy and escrow-copy's logic exactly for consistency
        if leaves.len() == 1 {
            leaves.push(leaves[0]); // Duplicate the single leaf
        }

        // Build merkle tree bottom-up
        let mut tree = vec![leaves.clone()];

        while leaves.len() > 1 {
            let mut new_level = Vec::new();
            for chunk in leaves.chunks(2) {
                if chunk.len() == 2 {
                    // Two nodes: hash(left, right)
                    let combined = hashv(&[&chunk[0], &chunk[1]]).to_bytes();
                    new_level.push(combined);
                } else {
                    // Single element case - hash with itself for consistency
                    let combined = hashv(&[&chunk[0], &chunk[0]]).to_bytes();
                    new_level.push(combined);
                }
            }
            leaves = new_level;
            tree.push(leaves.clone());
        }

        let root = tree
            .last()
            .ok_or("Failed to get tree root: tree is empty")?
            .first()
            .ok_or("Failed to get tree root: root level is empty")?
            .clone();

        Ok(Self {
            root,
            tree,
            allocations,
        })
    }

    /// Get the merkle root hash
    pub fn root(&self) -> [u8; 32] {
        self.root
    }

    /// Generate a merkle proof for a specific allocation
    ///
    /// ## Parameters
    /// - `allocation_index`: Index of the allocation in the original data
    ///
    /// ## Returns
    /// - `Ok(MerkleProof)` if successful
    /// - `Err(String)` if index is out of bounds
    ///
    /// ## Example
    /// ```rust
    /// let tree = MerkleTree::new(allocations)?;
    /// let proof = tree.generate_proof(0)?;
    /// ```
    pub fn generate_proof(&self, allocation_index: usize) -> Result<MerkleProof, String> {
        if allocation_index >= self.allocations.len() {
            return Err(format!(
                "Allocation index {} out of bounds (max: {})",
                allocation_index,
                self.allocations.len() - 1
            ));
        }

        let mut path = Vec::new();
        let mut indices = Vec::new();
        let mut current_index = allocation_index;

        // Traverse up the tree to build the proof
        for level in 0..self.tree.len() - 1 {
            let current_level = &self.tree[level];
            let sibling_index = if current_index % 2 == 0 {
                // Current is left, sibling is right
                current_index + 1
            } else {
                // Current is right, sibling is left
                current_index - 1
            };

            // Check if sibling exists (might not exist at odd levels)
            if sibling_index < current_level.len() {
                let sibling_hash = current_level[sibling_index];
                path.push(sibling_hash);
                // is_right = true if sibling is on the right (current_index is even)
                indices.push(current_index % 2 == 0);
            } else {
                // No sibling (odd level), use current node as sibling (duplicate)
                let current_hash = current_level[current_index];
                path.push(current_hash);
                // Always treat as right sibling for consistency
                indices.push(true);
            }

            // Move to parent level
            current_index /= 2;
        }

        Ok(MerkleProof::new(path, indices))
    }

    /// Find and generate proof for a specific recipient
    ///
    /// ## Parameters
    /// - `recipient`: Recipient wallet address
    ///
    /// ## Returns
    /// - `Some(Ok(MerkleProof))` if recipient found
    /// - `Some(Err(String))` if found but proof generation failed
    /// - `None` if recipient not found
    pub fn find_and_generate_proof(
        &self,
        recipient: &Pubkey,
    ) -> Option<Result<MerkleProof, String>> {
        self.allocations
            .iter()
            .position(|a| &a.recipient == recipient)
            .map(|index| self.generate_proof(index))
    }

    /// Get the number of allocations in the tree
    pub fn allocation_count(&self) -> usize {
        self.allocations.len()
    }

    /// Get allocation data by index
    pub fn get_allocation(&self, index: usize) -> Option<&AirdropAllocation> {
        self.allocations.get(index)
    }

    /// Get all allocation data
    pub fn allocations(&self) -> &[AirdropAllocation] {
        &self.allocations
    }
}

/// Merkle proof structure for airdrop claims
#[derive(Clone, Debug)]
pub struct MerkleProof {
    /// Proof path nodes (32-byte hashes)
    pub path: Vec<[u8; 32]>,
    /// Direction indices (false = left sibling, true = right sibling)
    pub indices: Vec<bool>,
}

impl MerkleProof {
    /// Create a new merkle proof
    pub fn new(path: Vec<[u8; 32]>, indices: Vec<bool>) -> Self {
        Self { path, indices }
    }

    /// Verify merkle proof for airdrop claim
    ///
    /// ## Parameters
    /// - `recipient`: Recipient wallet address
    /// - `amount`: Token amount to claim
    /// - `merkle_root`: Campaign's merkle root
    ///
    /// ## Returns
    /// - `true` if proof is valid
    /// - `false` if proof is invalid
    ///
    /// ## Algorithm (from miracle-copy's essence)
    /// 1. Compute leaf hash: hashv(&[recipient, amount])
    /// 2. Traverse proof path
    /// 3. For each step: hash(left, right) based on direction
    /// 4. Final hash should match merkle_root
    pub fn verify(&self, recipient: &Pubkey, amount: u64, merkle_root: [u8; 32]) -> bool {
        // Validate proof structure
        if self.path.len() != self.indices.len() {
            return false;
        }

        // Compute leaf hash: hash(recipient, amount)
        let leaf_hash = compute_leaf_hash(recipient, amount);

        // Traverse proof path
        let mut current_hash = leaf_hash;
        for (i, &path_node) in self.path.iter().enumerate() {
            let is_right = self.indices[i];

            if is_right {
                // Sibling is on the right: hash(current_hash, path_node)
                current_hash = hashv(&[&current_hash, &path_node]).to_bytes();
            } else {
                // Sibling is on the left: hash(path_node, current_hash)
                current_hash = hashv(&[&path_node, &current_hash]).to_bytes();
            }
        }

        current_hash == merkle_root
    }
}

/// Compute leaf hash for airdrop allocation
///
/// ## Formula
/// leaf_hash = hashv(&[recipient.as_ref(), &amount.to_le_bytes()])
///
/// ## Deterministic
/// Same (recipient, amount) always produces same hash
pub fn compute_leaf_hash(recipient: &Pubkey, amount: u64) -> [u8; 32] {
    hashv(&[recipient.as_ref(), &amount.to_le_bytes()]).to_bytes()
}

/// Verify merkle proof for airdrop claim (standalone function)
///
/// ## Parameters
/// - `proof`: Merkle proof structure
/// - `recipient`: Recipient wallet address
/// - `amount`: Token amount to claim
/// - `merkle_root`: Campaign's merkle root
///
/// ## Returns
/// - `true` if proof is valid
/// - `false` if proof is invalid
pub fn verify_merkle_proof(
    proof: &MerkleProof,
    recipient: &Pubkey,
    amount: u64,
    merkle_root: [u8; 32],
) -> bool {
    proof.verify(recipient, amount, merkle_root)
}