Rust merkle proof POC

merkle_proof.rs

use core::{mem::MaybeUninit, ptr};


pub trait MerkleHash<const N: usize> {
    fn hash(bytes: &[u8]) -> [u8;N];
    fn merkle_hashv(l: &[u8], r: &[u8], swap: bool) -> [u8;N];
}

#[inline(always)]
unsafe fn trunc32_to_n<const N: usize>(src32: *const u8) -> [u8; N] {
    debug_assert!(N <= 32);
    let mut out = MaybeUninit::<[u8; N]>::uninit();
    unsafe {
        ptr::copy_nonoverlapping(src32, out.as_mut_ptr() as *mut u8, N);
        out.assume_init()
    }
}

#[inline(always)]
unsafe fn first_n<const N: usize>(src: &[u8]) -> [u8; N] {
    debug_assert!(src.len() >= N);
    let mut out = MaybeUninit::<[u8; N]>::uninit();
    unsafe {
        ptr::copy_nonoverlapping(src.as_ptr(), out.as_mut_ptr() as *mut u8, N);
        out.assume_init()
    }
}

pub struct Sha256;
pub struct Sha256d;

impl<const N: usize> MerkleHash<N> for Sha256 {
    #[inline(always)]
    fn hash(bytes: &[u8]) -> [u8; N] {
        let h: [u8; 32] = solana_nostd_sha256::hash(bytes);
        unsafe { trunc32_to_n::<N>(h.as_ptr()) }
    }

    #[inline(always)]
    fn merkle_hashv(l: &[u8], r: &[u8], swap: bool) -> [u8; N] {
        let (a, b) = if swap { (r, l) } else { (l, r) };

        // Grab first N bytes of each child
        let left_n: [u8; N]  = unsafe { first_n::<N>(a) };
        let right_n: [u8; N] = unsafe { first_n::<N>(b) };

        // Hash concatenated slices using hashv
        let h: [u8; 32] = solana_nostd_sha256::hashv(&[&left_n, &right_n]);
        unsafe { trunc32_to_n::<N>(h.as_ptr()) }
    }
}

impl<const N: usize> MerkleHash<N> for Sha256d {
    #[inline(always)]
    fn hash(bytes: &[u8]) -> [u8; N] {
        let h1: [u8; 32] = solana_nostd_sha256::hash(bytes);
        let h2: [u8; 32] = solana_nostd_sha256::hash(&h1);
        unsafe { trunc32_to_n::<N>(h2.as_ptr()) }
    }

    #[inline(always)]
    fn merkle_hashv(l: &[u8], r: &[u8], swap: bool) -> [u8; N] {
        let (a, b) = if swap { (r, l) } else { (l, r) };

        let left_n: [u8; N]  = unsafe { first_n::<N>(a) };
        let right_n: [u8; N] = unsafe { first_n::<N>(b) };

        // Double SHA-256
        let h1: [u8; 32] = solana_nostd_sha256::hashv(&[&left_n, &right_n]);
        let h2: [u8; 32] = solana_nostd_sha256::hash(&h1);
        unsafe { trunc32_to_n::<N>(h2.as_ptr()) }
    }
}

pub struct MerkleProof<const N: usize>;

impl<const N: usize> MerkleProof<N> {
    /// Bitcoin-style Merkle proof verification: compute root from a raw leaf *bytes* and its siblings.
    /// - `leaf`   = RAW leaf bytes (NOT prehashed),
    /// - `hashes` = siblings from bottom to top (each length >= N),
    /// - `index`  = leaf index at the bottom level.
    ///
    /// Notes:
    /// - We first hash the raw `leaf` to a fixed-size digest, then combine with siblings.
    /// - If your construction uses domain separation (e.g. prefixing leaves vs. inner nodes),
    ///   implement that inside `H::hash_leaf` and `H::hashv`.
    #[inline(always)]
    pub fn merklize<H: MerkleHash<N>>(leaf: &[u8], hashes: &[&[u8]], mut index: u32) -> [u8; N] {
        debug_assert!(N <= 32);
        for s in hashes {
            debug_assert!(s.len() >= N);
        }

        // 1) Hash the raw leaf first (domain-separated inside H::hash_leaf if desired).
        let mut current = H::hash(&leaf);

        // 2) Fold over siblings; parity of `index` decides concatenation order.
        for &sib in hashes {
            let current_bytes: &[u8] = &current;
            let swap = (index & 1) == 1; // odd => sibling || current ; even => current || sibling
            current = H::merkle_hashv(current_bytes, sib, swap);
            index >>= 1;
        }

        current
    }
}

#[cfg(test)]
mod tests {
    use crate::{MerkleProof, Sha256d};

    #[test]
    fn test_merkle_proof() {
        let leaf = [0x01, 0x00, 0x00, 0x00, 0x01, 0x71, 0x25, 0xb0, 0x44, 0x67, 0xdc, 0x2e, 0x3e, 0x76, 0x6a, 0x0d, 0xae, 0x2b, 0x7a, 0x2f, 0x74, 0x21, 0x1c, 0x7a, 0xa7, 0xbf, 0x79, 0x6d, 0x47, 0xfb, 0xf4, 0x4c, 0x25, 0x9b, 0xe2, 0x34, 0x62, 0x66, 0x11, 0x00, 0x00, 0x6b, 0x48, 0x30, 0x45, 0x02, 0x21, 0x00, 0xf1, 0xe5, 0xfd, 0xfd, 0x36, 0x83, 0x7a, 0x2e, 0x84, 0x22, 0x5e, 0x15, 0x7d, 0x25, 0xf4, 0xd3, 0x41, 0xca, 0xd4, 0x9b, 0xfd, 0xc9, 0x09, 0xe0, 0x33, 0x2e, 0x5e, 0x5a, 0x58, 0xe8, 0x49, 0xa1, 0x02, 0x20, 0x3b, 0x5c, 0x59, 0xd2, 0xf5, 0xcf, 0x4c, 0x6f, 0x84, 0xb2, 0xbc, 0x18, 0x9a, 0x03, 0xed, 0x80, 0x2d, 0x48, 0x78, 0x4f, 0x33, 0x5b, 0x71, 0x2f, 0x73, 0xe8, 0x0f, 0x80, 0x7d, 0x4c, 0xdd, 0x71, 0x41, 0x21, 0x03, 0x7d, 0x53, 0x43, 0x07, 0x15, 0xb2, 0xbc, 0x84, 0x63, 0x84, 0x7e, 0x79, 0xd7, 0xe2, 0x59, 0xc1, 0x1a, 0x7d, 0x81, 0xbf, 0x7d, 0x61, 0x66, 0xe0, 0x03, 0xe1, 0xb1, 0x03, 0xb6, 0x57, 0x31, 0xff, 0xff, 0xff, 0xff, 0xff, 0x01, 0x23, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x19, 0x76, 0xa9, 0x14, 0x0e, 0xc5, 0x69, 0x60, 0xe8, 0x3c, 0xd3, 0xc0, 0x3c, 0x88, 0x82, 0xe0, 0xfd, 0x34, 0xd4, 0x62, 0xa3, 0x4c, 0x65, 0x38, 0x88, 0xac, 0x00, 0x00, 0x00, 0x00];
        let proof: &[&[u8]] = &[
            &[0x78, 0x0f, 0x39, 0x00, 0x9c, 0x90, 0xbe, 0x8c, 0xc6, 0x2b, 0x17, 0x29, 0x56, 0x9b, 0x5c, 0x8d, 0x50, 0xb5, 0x9b, 0xad, 0x44, 0x89, 0xe9, 0x5a, 0xc7, 0xa8, 0x39, 0x55, 0x5c, 0x1e, 0xe7, 0x95],
            &[0x55, 0x03, 0xc7, 0x02, 0xbc, 0x9a, 0xc9, 0x72, 0xd2, 0xdb, 0xb4, 0x6d, 0x25, 0x34, 0x55, 0x9d, 0x2a, 0x26, 0x8a, 0x7b, 0x06, 0x87, 0x2b, 0x27, 0x90, 0x15, 0x80, 0x9a, 0x32, 0x3f, 0x3d, 0x53],
            &[0xc7, 0xdf, 0xdb, 0x01, 0x53, 0x70, 0x35, 0xa4, 0xef, 0x48, 0x52, 0x75, 0xa7, 0x27, 0xae, 0xb2, 0x93, 0x28, 0xe4, 0x9c, 0x1a, 0xfc, 0x4f, 0xed, 0xac, 0x87, 0xd9, 0x33, 0x3f, 0x05, 0x99, 0x63],
            &[0xf5, 0x4a, 0x33, 0x3d, 0x21, 0xb4, 0xfa, 0xa9, 0xf9, 0x12, 0xab, 0x2a, 0xf2, 0xe7, 0x2f, 0x69, 0x41, 0xc3, 0xd1, 0xc9, 0x79, 0x33, 0x1e, 0x34, 0xf6, 0x2f, 0xa7, 0xda, 0xc2, 0x3d, 0xaf, 0x29],
            &[0xf2, 0x4c, 0x04, 0x9d, 0x33, 0x95, 0xfe, 0xa0, 0xe7, 0x7a, 0xd4, 0xdf, 0xdf, 0x49, 0x68, 0x44, 0xcf, 0x28, 0x93, 0x90, 0x36, 0x4c, 0xe2, 0xde, 0xfd, 0xab, 0x1e, 0x2b, 0x9b, 0xc4, 0xa9, 0x35],
            &[0xd0, 0x30, 0xd9, 0xcf, 0x05, 0xfb, 0x2c, 0x79, 0x32, 0x5c, 0x80, 0xab, 0x42, 0x6d, 0x6e, 0xd2, 0x74, 0xd5, 0x44, 0x04, 0x64, 0xd7, 0x07, 0x4b, 0xdb, 0xbf, 0x47, 0x55, 0x1d, 0x91, 0xdc, 0x99],
            &[0x93, 0xa8, 0x93, 0x5e, 0xb4, 0xf4, 0xec, 0x93, 0x78, 0x10, 0x2a, 0x8a, 0x32, 0xd1, 0x76, 0x82, 0x21, 0xee, 0x7b, 0xa3, 0x5c, 0xef, 0xcd, 0x91, 0xd1, 0xe9, 0x0d, 0x39, 0xae, 0x8d, 0xe5, 0x41],
            &[0xcd, 0xdb, 0x8c, 0xab, 0x43, 0xa0, 0x1a, 0x96, 0x0e, 0x09, 0xd8, 0x20, 0xfc, 0x6c, 0x1a, 0x35, 0xcb, 0xe3, 0x72, 0xd0, 0x12, 0x1d, 0x22, 0xfd, 0x19, 0x38, 0xcd, 0xa5, 0xf8, 0x98, 0x61, 0xbe],
            &[0x9a, 0x58, 0x71, 0x73, 0xdb, 0xd8, 0xc5, 0xe2, 0x32, 0xce, 0x76, 0x6e, 0xa4, 0x3e, 0x1d, 0xc4, 0x67, 0x54, 0xb1, 0x8e, 0xcc, 0x7e, 0xce, 0xb3, 0x4f, 0x3e, 0x87, 0xd4, 0x1d, 0xb9, 0xa4, 0x8b],
            &[0xef, 0xef, 0xd4, 0x61, 0x34, 0xb1, 0x14, 0x49, 0x63, 0xed, 0xe9, 0xb1, 0x17, 0xef, 0x75, 0x56, 0x85, 0x7f, 0xb9, 0xcd, 0x4a, 0xd6, 0xf0, 0xa6, 0xa0, 0x50, 0xa4, 0x5b, 0xc7, 0xe3, 0x69, 0x8c],
            &[0x16, 0x91, 0x49, 0xca, 0xcf, 0x76, 0x1b, 0x55, 0x33, 0xd7, 0x4e, 0x9c, 0x7f, 0x2c, 0x6e, 0x01, 0xec, 0x0f, 0xba, 0x86, 0x46, 0x6c, 0x86, 0xe6, 0x2c, 0x3f, 0xd9, 0xfb, 0xd1, 0x46, 0x2f, 0xdd],
            &[0x55, 0x8d, 0x9e, 0x2a, 0xed, 0xfc, 0x65, 0xfc, 0x8d, 0x58, 0x3d, 0xdc, 0xa8, 0x68, 0xfa, 0x29, 0xac, 0xeb, 0xaa, 0x19, 0x0e, 0x42, 0xf3, 0x5b, 0x89, 0x12, 0xfe, 0x35, 0x3d, 0x33, 0x1e, 0x64],
            &[0xfd, 0x0a, 0xe5, 0x31, 0x0b, 0x7e, 0xf3, 0xe5, 0x2b, 0x1c, 0x4d, 0x2c, 0xd7, 0xe9, 0x9d, 0xd0, 0x55, 0x32, 0x59, 0x39, 0x4f, 0x44, 0xb8, 0xb7, 0x1a, 0x0a, 0xd3, 0x69, 0x18, 0xc6, 0xb2, 0xd9],
            &[0x94, 0x1e, 0x1c, 0x18, 0x72, 0xda, 0xf9, 0x35, 0x16, 0x06, 0xed, 0xd6, 0x8b, 0xf1, 0x25, 0x24, 0x6b, 0x99, 0xd2, 0xcf, 0x44, 0xec, 0xd1, 0xc5, 0x26, 0xfd, 0xc0, 0x6f, 0xbf, 0xa3, 0xd9, 0xc9]
        ];

        
        let root = MerkleProof::<32>::merklize::<Sha256d>(&leaf, proof, 47);
        assert_eq!(
            [0xe4, 0x3a, 0x1d, 0xe4, 0xdd, 0x9c, 0x52, 0x62, 0x74, 0xb8, 0xea, 0x9e, 0x4b, 0xf0, 0x1f, 0xe8, 0x92, 0x86, 0x49, 0xf8, 0xe5, 0xb9, 0x4a, 0xbc, 0x4e, 0x05, 0xd8, 0x3b, 0x8a, 0xbe, 0xb9, 0x24],
            root
        );
    }
}