light-sdk-types 0.24.0

use light_hasher::{Hasher, Poseidon};

use crate::lca::{CompressedAccountError, Pubkey, TreeType, Vec};

#[repr(C)]
#[cfg_attr(
    all(feature = "std", feature = "anchor"),
    derive(anchor_lang::AnchorDeserialize, anchor_lang::AnchorSerialize)
)]
#[cfg_attr(
    not(feature = "anchor"),
    derive(borsh::BorshDeserialize, borsh::BorshSerialize)
)]
#[derive(Debug, PartialEq, Default, Clone)]
pub struct PackedCompressedAccountWithMerkleContext {
    pub compressed_account: CompressedAccount,
    pub merkle_context: PackedMerkleContext,
    /// Index of root used in inclusion validity proof.
    pub root_index: u16,
    pub read_only: bool,
}

#[cfg_attr(
    all(feature = "std", feature = "anchor"),
    derive(anchor_lang::AnchorDeserialize, anchor_lang::AnchorSerialize)
)]
#[cfg_attr(
    not(feature = "anchor"),
    derive(borsh::BorshDeserialize, borsh::BorshSerialize)
)]
#[derive(Debug, PartialEq, Default, Clone)]
pub struct InCompressedAccountWithMerkleContext {
    pub compressed_account: InCompressedAccount,
    pub merkle_context: MerkleContext,
}

impl From<CompressedAccount> for InCompressedAccount {
    fn from(value: CompressedAccount) -> Self {
        let data = value.data.unwrap_or_default();
        InCompressedAccount {
            owner: value.owner,
            lamports: value.lamports,
            address: value.address,
            discriminator: data.discriminator,
            data_hash: data.data_hash,
        }
    }
}

#[cfg_attr(
    all(feature = "std", feature = "anchor"),
    derive(anchor_lang::AnchorDeserialize, anchor_lang::AnchorSerialize)
)]
#[cfg_attr(
    not(feature = "anchor"),
    derive(borsh::BorshDeserialize, borsh::BorshSerialize)
)]
#[derive(Debug, PartialEq, Default, Clone)]
pub struct PackedInCompressedAccountWithMerkleContext {
    pub compressed_account: InCompressedAccount,
    pub merkle_context: PackedMerkleContext,
    /// Index of root used in inclusion validity proof.
    pub root_index: u16,
}

impl From<PackedCompressedAccountWithMerkleContext> for PackedInCompressedAccountWithMerkleContext {
    fn from(value: PackedCompressedAccountWithMerkleContext) -> Self {
        Self {
            compressed_account: value.compressed_account.into(),
            merkle_context: value.merkle_context,
            root_index: value.root_index,
        }
    }
}

impl From<CompressedAccountWithMerkleContext> for InCompressedAccountWithMerkleContext {
    fn from(value: CompressedAccountWithMerkleContext) -> Self {
        Self {
            compressed_account: value.compressed_account.into(),
            merkle_context: value.merkle_context,
        }
    }
}

#[cfg_attr(
    all(feature = "std", feature = "anchor"),
    derive(anchor_lang::AnchorDeserialize, anchor_lang::AnchorSerialize)
)]
#[cfg_attr(
    not(feature = "anchor"),
    derive(borsh::BorshDeserialize, borsh::BorshSerialize)
)]
#[derive(Debug, PartialEq, Default, Clone)]
pub struct CompressedAccountWithMerkleContext {
    pub compressed_account: CompressedAccount,
    pub merkle_context: MerkleContext,
}

impl CompressedAccountWithMerkleContext {
    pub fn hash(&self) -> Result<[u8; 32], CompressedAccountError> {
        self.compressed_account.hash(
            &self.merkle_context.merkle_tree_pubkey,
            &self.merkle_context.leaf_index,
            self.merkle_context.tree_type == TreeType::StateV2,
        )
    }
}

impl CompressedAccountWithMerkleContext {
    pub fn into_read_only(
        &self,
        root_index: Option<u16>,
    ) -> Result<ReadOnlyCompressedAccount, CompressedAccountError> {
        let account_hash = self.hash()?;
        let merkle_context = if root_index.is_none() {
            let mut merkle_context = self.merkle_context;
            merkle_context.prove_by_index = true;
            merkle_context
        } else {
            self.merkle_context
        };
        Ok(ReadOnlyCompressedAccount {
            account_hash,
            merkle_context,
            root_index: root_index.unwrap_or_default(),
        })
    }
}

#[cfg_attr(
    all(feature = "std", feature = "anchor"),
    derive(anchor_lang::AnchorDeserialize, anchor_lang::AnchorSerialize)
)]
#[cfg_attr(
    not(feature = "anchor"),
    derive(borsh::BorshDeserialize, borsh::BorshSerialize)
)]
#[derive(Debug, PartialEq, Default, Clone)]
pub struct ReadOnlyCompressedAccount {
    pub account_hash: [u8; 32],
    pub merkle_context: MerkleContext,
    pub root_index: u16,
}

#[repr(C)]
#[cfg_attr(
    all(feature = "std", feature = "anchor"),
    derive(anchor_lang::AnchorDeserialize, anchor_lang::AnchorSerialize)
)]
#[cfg_attr(
    not(feature = "anchor"),
    derive(borsh::BorshDeserialize, borsh::BorshSerialize)
)]
#[derive(Debug, PartialEq, Default, Clone)]
pub struct PackedReadOnlyCompressedAccount {
    pub account_hash: [u8; 32],
    pub merkle_context: PackedMerkleContext,
    pub root_index: u16,
}

#[cfg_attr(
    all(feature = "std", feature = "anchor"),
    derive(anchor_lang::AnchorDeserialize, anchor_lang::AnchorSerialize)
)]
#[cfg_attr(
    not(feature = "anchor"),
    derive(borsh::BorshDeserialize, borsh::BorshSerialize)
)]
#[derive(Debug, Clone, Copy, PartialEq, Default)]
pub struct MerkleContext {
    pub merkle_tree_pubkey: Pubkey,
    pub queue_pubkey: Pubkey,
    pub leaf_index: u32,
    pub prove_by_index: bool,
    pub tree_type: TreeType,
}

#[repr(C)]
#[cfg_attr(
    all(feature = "std", feature = "anchor"),
    derive(anchor_lang::AnchorDeserialize, anchor_lang::AnchorSerialize)
)]
#[cfg_attr(
    not(feature = "anchor"),
    derive(borsh::BorshDeserialize, borsh::BorshSerialize)
)]
#[derive(Debug, Clone, Copy, PartialEq, Default)]
pub struct PackedMerkleContext {
    pub merkle_tree_pubkey_index: u8,
    pub queue_pubkey_index: u8,
    pub leaf_index: u32,
    pub prove_by_index: bool,
}

#[repr(C)]
#[cfg_attr(
    all(feature = "std", feature = "anchor"),
    derive(anchor_lang::AnchorDeserialize, anchor_lang::AnchorSerialize)
)]
#[cfg_attr(
    not(feature = "anchor"),
    derive(borsh::BorshDeserialize, borsh::BorshSerialize)
)]
#[derive(Debug, PartialEq, Default, Clone)]
pub struct CompressedAccount {
    pub owner: Pubkey,
    pub lamports: u64,
    pub address: Option<[u8; 32]>,
    pub data: Option<CompressedAccountData>,
}

#[cfg_attr(
    all(feature = "std", feature = "anchor"),
    derive(anchor_lang::AnchorDeserialize, anchor_lang::AnchorSerialize)
)]
#[cfg_attr(
    not(feature = "anchor"),
    derive(borsh::BorshDeserialize, borsh::BorshSerialize)
)]
#[derive(Debug, PartialEq, Default, Clone)]
pub struct InCompressedAccount {
    pub owner: Pubkey,
    pub lamports: u64,
    pub discriminator: [u8; 8],
    pub data_hash: [u8; 32],
    pub address: Option<[u8; 32]>,
}

#[repr(C)]
#[cfg_attr(
    all(feature = "std", feature = "anchor"),
    derive(anchor_lang::AnchorDeserialize, anchor_lang::AnchorSerialize)
)]
#[cfg_attr(
    not(feature = "anchor"),
    derive(borsh::BorshDeserialize, borsh::BorshSerialize)
)]
#[derive(Debug, PartialEq, Default, Clone)]
pub struct CompressedAccountData {
    pub discriminator: [u8; 8],
    pub data: Vec<u8>,
    pub data_hash: [u8; 32],
}

pub fn hash_with_hashed_values(
    lamports: &u64,
    address: Option<&[u8]>,
    data: Option<(&[u8], &[u8])>,
    owner_hashed: &[u8; 32],
    merkle_tree_hashed: &[u8; 32],
    leaf_index: &u32,
    is_batched: bool,
) -> Result<[u8; 32], CompressedAccountError> {
    // Use ArrayVec with max capacity of 7 elements:
    // owner_hashed + leaf_index + merkle_tree_hashed + lamports + address + discriminator + data_hash
    let mut vec: tinyvec::ArrayVec<[&[u8]; 7]> = tinyvec::ArrayVec::new();
    vec.push(owner_hashed.as_slice());

    // leaf index and merkle tree pubkey are used to make every compressed account hash unique
    let mut leaf_index_bytes = [0u8; 32];
    if is_batched {
        leaf_index_bytes[28..].copy_from_slice(&leaf_index.to_be_bytes());
    } else {
        leaf_index_bytes[28..].copy_from_slice(&leaf_index.to_le_bytes());
    };
    vec.push(leaf_index_bytes.as_slice());

    vec.push(merkle_tree_hashed.as_slice());

    // Lamports are only hashed if non-zero to save CU.
    // For safety, we prefix lamports with 1 in 1 byte.
    // Thus, even if the discriminator has the same value as the lamports, the hash will be different.
    let mut lamports_bytes = [0u8; 32];
    if *lamports != 0 {
        if is_batched {
            lamports_bytes[24..].copy_from_slice(&lamports.to_be_bytes());
        } else {
            lamports_bytes[24..].copy_from_slice(&lamports.to_le_bytes());
        };
        lamports_bytes[23] = 1;

        vec.push(lamports_bytes.as_slice());
    }
    if let Some(address) = address {
        vec.push(address);
    }

    let mut discriminator_bytes = [0u8; 32];
    if let Some((discriminator, data_hash)) = data {
        discriminator_bytes[24..].copy_from_slice(discriminator);
        discriminator_bytes[23] = 2;
        vec.push(&discriminator_bytes);
        vec.push(data_hash);
    }
    Ok(Poseidon::hashv(&vec)?)
}

/// Hashing scheme:
/// H(owner || leaf_index || merkle_tree_pubkey || lamports || address || data.discriminator || data.data_hash)
impl CompressedAccount {
    pub fn hash_with_hashed_values(
        &self,
        owner_hashed: &[u8; 32],
        merkle_tree_hashed: &[u8; 32],
        leaf_index: &u32,
        is_batched: bool,
    ) -> Result<[u8; 32], CompressedAccountError> {
        hash_with_hashed_values(
            &self.lamports,
            self.address.as_ref().map(|x| x.as_slice()),
            self.data
                .as_ref()
                .map(|x| (x.discriminator.as_slice(), x.data_hash.as_slice())),
            owner_hashed,
            merkle_tree_hashed,
            leaf_index,
            is_batched,
        )
    }

    pub fn hash(
        &self,
        &merkle_tree_pubkey: &Pubkey,
        leaf_index: &u32,
        is_batched: bool,
    ) -> Result<[u8; 32], CompressedAccountError> {
        use light_hasher::hash_to_field_size::hash_to_bn254_field_size_be;
        let hashed_mt = hash_to_bn254_field_size_be(merkle_tree_pubkey.as_ref());

        self.hash_with_hashed_values(
            &hash_to_bn254_field_size_be(self.owner.as_ref()),
            &hashed_mt,
            leaf_index,
            is_batched,
        )
    }
}