ghost-node/pallets/helpers/src/merkle_tree.rs

use ghost_traits::hashing::GhostHasher;
use sp_arithmetic::traits::AtLeast8BitUnsigned;
use sp_runtime::traits::UniqueSaturatedInto;
use sp_std::vec::Vec;

const LEAF_PREFIX: u8 = 0;
const NODE_PREFIX: u8 = 1;

fn leaf_prefixed<H>(preimage: &[u8]) -> H::Hash
where
	H: GhostHasher,
{
    let mut prefixed_data = sp_std::vec![LEAF_PREFIX; 1 + preimage.len()];
    prefixed_data[1..].copy_from_slice(preimage);
    H::hash(&prefixed_data)
}

fn node_prefixed<H>(preimage1: &[u8], preimage2: &[u8]) -> H::Hash
where
	H: GhostHasher,
{
    let total_len = preimage1.len() + preimage2.len();
    let mut prefixed_data = sp_std::vec![NODE_PREFIX; 1 + total_len];
    prefixed_data[1..1 + preimage1.len()].copy_from_slice(preimage1);
    prefixed_data[1 + preimage1.len()..].copy_from_slice(preimage2);
    H::hash(&prefixed_data)
}

pub fn generate_tree<H, I, S, F, E>(
    max_index: I,
    raw_values: S,
    generate_preimage: F,
) -> Result<Vec<H::Hash>, E>
where
    H: GhostHasher,
    S: IntoIterator,
    I: AtLeast8BitUnsigned + UniqueSaturatedInto<usize>,
    F: Fn(S::Item) -> Result<(usize, Vec<u8>), E>,
{
    let num_of_leaves: usize = (max_index + I::one()).unique_saturated_into();
    let padded_leaves = num_of_leaves.next_power_of_two();
    let total_capacity = (2 * padded_leaves) - 1;

    let empty_hash = H::empty();
    let empty_leaf_hash = leaf_prefixed::<H>(empty_hash.as_ref());
    let mut merkle_tree = sp_std::vec![empty_leaf_hash; total_capacity];

    for item in raw_values.into_iter() {
        let (index, preimage) = generate_preimage(item)?;
        if index >= padded_leaves { return Ok(sp_std::vec![]); }
        merkle_tree[index] = leaf_prefixed::<H>(&preimage)
    }

    let mut layer_start = 0;
    let mut current_layer_len = padded_leaves;
    let mut write_ptr = padded_leaves;

    while current_layer_len > 1 {
        for i in (0..current_layer_len).step_by(2) {
            let layer_index = layer_start + i;
            let left_bytes = merkle_tree[layer_index].as_ref();
            let right_bytes = merkle_tree[layer_index | 1].as_ref();

            merkle_tree[write_ptr] = node_prefixed::<H>(left_bytes, right_bytes);
            write_ptr += 1;
        }
        layer_start += current_layer_len;
        current_layer_len >>= 1;
    }

    Ok(merkle_tree)
}

pub fn generate_proof<H, I>(
    merkle_tree: &[H::Hash],
    max_index: I,
    index: I,
) -> Vec<H::Hash>
where
    H: GhostHasher,
    I: AtLeast8BitUnsigned + UniqueSaturatedInto<usize>,
{
    if index > max_index { return sp_std::vec![]; }

    let num_of_leaves: usize = (max_index + I::one()).unique_saturated_into();
    let mut current_index: usize = index.unique_saturated_into();
    let mut current_layer_len = num_of_leaves.next_power_of_two();

    let tree_height = current_layer_len.trailing_zeros() as usize;
    let mut proof = Vec::with_capacity(tree_height);

    let mut layer_start = 0;

    while current_layer_len > 1 {
        let sibling_idx_in_layer = current_index ^ 1;
        let sibling_hash = merkle_tree[layer_start + sibling_idx_in_layer];
        proof.push(sibling_hash);

        layer_start += current_layer_len;
        current_layer_len >>= 1;
        current_index >>= 1;
    }

    proof
}

pub fn verify_tree_proof<H, I>(
    preimage: &[u8],
    merkle_proof: &[H::Hash],
    merkle_root: H::Hash,
    index: I,
) -> bool
where
    H: GhostHasher,
    I: PartialOrd + UniqueSaturatedInto<usize>,
{
    let mut current_index: usize = index.unique_saturated_into();
    if current_index >= (1 << merkle_proof.len()) { return false; }

    let mut current_hash = leaf_prefixed::<H>(preimage);

    for sibling in merkle_proof.iter() {
        let sibling_bytes = sibling.as_ref();
        let hash_bytes = current_hash.as_ref();

        if current_index % 2 == 0 {
            current_hash = node_prefixed::<H>(hash_bytes, sibling_bytes);
        } else {
            current_hash = node_prefixed::<H>(sibling_bytes, hash_bytes);
        }

        current_index >>= 1;
    }

    current_hash == merkle_root
}

#[cfg(test)]
mod tests {
    use super::*;
    use sp_core::H256;
    use sp_io::hashing::blake2_256;

    pub struct TestBlake2Hasher;
    impl GhostHasher for TestBlake2Hasher {
        type Hash = H256;
        type HashBytes = [u8; 32];

        fn hash(data: &[u8]) -> Self::Hash {
            H256::from(blake2_256(data))
        }
        fn empty() -> Self::Hash {
            H256::zero()
        }
        fn hash_len() -> usize {
            H256::len_bytes()
        }
    }

    #[test]
    fn test_generate_tree_perfect_power_of_two() {
        let max_index = 3u32;
        let raw_data = sp_std::vec![
            (0, sp_std::vec![1]),
            (1, sp_std::vec![2]),
            (2, sp_std::vec![3]),
            (3, sp_std::vec![4]),
        ];

        let tree = generate_tree::<TestBlake2Hasher, u32, _, _, ()>(
            max_index,
            raw_data,
            |(index, bytes)| Ok((index, bytes)),
        )
        .unwrap();

        assert_eq!(tree.len(), 7);

        let h0 = TestBlake2Hasher::hash(&[ &[0u8], [1].as_slice() ].concat());
        let h1 = TestBlake2Hasher::hash(&[ &[0u8], [2].as_slice() ].concat());
        let h2 = TestBlake2Hasher::hash(&[ &[0u8], [3].as_slice() ].concat());
        let h3 = TestBlake2Hasher::hash(&[ &[0u8], [4].as_slice() ].concat());

        let mut c1 = [0u8; 65];
        c1[0] = NODE_PREFIX;
        c1[1..33].copy_from_slice(h0.as_ref());
        c1[33..].copy_from_slice(h1.as_ref());
        let parent_left = TestBlake2Hasher::hash(&c1);

        let mut c2 = [0u8; 65];
        c2[0] = NODE_PREFIX;
        c2[1..33].copy_from_slice(h2.as_ref());
        c2[33..].copy_from_slice(h3.as_ref());
        let parent_right = TestBlake2Hasher::hash(&c2);

        let mut c_root = [0u8; 65];
        c_root[0] = NODE_PREFIX;
        c_root[1..33].copy_from_slice(parent_left.as_ref());
        c_root[33..].copy_from_slice(parent_right.as_ref());
        let expected_root = TestBlake2Hasher::hash(&c_root);

        assert_eq!(tree.last().unwrap(), &expected_root);
    }

    #[test]
    fn test_generate_tree_with_padding_hashing() {
        let max_index = 2u32;
        let raw_data = sp_std::vec![
            (0, sp_std::vec![1]),
            (1, sp_std::vec![2]),
            (2, sp_std::vec![3]),
        ];

        let tree = generate_tree::<TestBlake2Hasher, u32, _, _, ()>(
            max_index,
            raw_data,
            |(index, bytes)| Ok((index, bytes)),
        )
        .unwrap();

        let mut empty_data = sp_std::vec![LEAF_PREFIX];
        empty_data.extend_from_slice(&H256::zero().to_fixed_bytes());
        let empty_hash = TestBlake2Hasher::hash(&empty_data);

        assert_eq!(tree.len(), 7);
        assert_eq!(tree[3], empty_hash);

        let mut leaf2_data = sp_std::vec![LEAF_PREFIX];
        leaf2_data.extend_from_slice(&[3]);
        let h2 = TestBlake2Hasher::hash(&leaf2_data);

        let mut combined_right = [0u8; 65];
        combined_right[0] = NODE_PREFIX;
        combined_right[1..33].copy_from_slice(h2.as_ref());
        combined_right[33..].copy_from_slice(empty_hash.as_ref());

        let expected_parent_right = TestBlake2Hasher::hash(&combined_right);
        assert_eq!(tree[5], expected_parent_right);
    }

    #[test]
    fn test_generate_tree_determinism() {
        let max_index = 1u32;

        let data_1 = sp_std::vec![(0, sp_std::vec![100]), (1, sp_std::vec![200])];
        let data_2 = sp_std::vec![(0, sp_std::vec![100]), (1, sp_std::vec![200])];
        let data_3 = sp_std::vec![(0, sp_std::vec![101]), (1, sp_std::vec![200])];

        let tree_1 =
            generate_tree::<TestBlake2Hasher, u32, _, _, ()>(max_index, data_1, |(i, b)| {
                Ok((i, b))
            })
            .unwrap();
        let tree_2 =
            generate_tree::<TestBlake2Hasher, u32, _, _, ()>(max_index, data_2, |(i, b)| {
                Ok((i, b))
            })
            .unwrap();
        let tree_3 =
            generate_tree::<TestBlake2Hasher, u32, _, _, ()>(max_index, data_3, |(i, b)| {
                Ok((i, b))
            })
            .unwrap();

        assert_eq!(tree_1.last().unwrap(), tree_2.last().unwrap());
        assert_ne!(tree_1.last().unwrap(), tree_3.last().unwrap());
    }

    #[test]
    fn test_generate_tree_single_element() {
        let max_index = 0u32;
        let raw_data = sp_std::vec![(0, sp_std::vec![42])];

        let tree = generate_tree::<TestBlake2Hasher, u32, _, _, ()>(
            max_index,
            raw_data,
            |(index, bytes)| Ok((index, bytes)),
        )
        .unwrap();

        assert_eq!(tree.len(), 1);
        assert_eq!(tree[0], TestBlake2Hasher::hash(&vec![LEAF_PREFIX, 42]));
    }

    #[test]
    fn test_proof_and_verification_end_to_end() {
        let max_index = 3u32;
        let raw_data = sp_std::vec![
            (0, sp_std::vec![1]),
            (1, sp_std::vec![2]),
            (2, sp_std::vec![3]),
            (3, sp_std::vec![4]),
        ];

        let tree = generate_tree::<TestBlake2Hasher, u32, _, _, ()>(
            max_index,
            raw_data,
            |(index, bytes)| Ok((index, bytes)),
        )
        .unwrap();

        let root = *tree.last().unwrap();

        let target_index = 2u32;
        let proof = generate_proof::<TestBlake2Hasher, u32>(&tree, max_index, target_index);

        assert_eq!(proof.len(), 2);

        let preimage = sp_std::vec![3];
        let is_valid =
            verify_tree_proof::<TestBlake2Hasher, u32>(&preimage, &proof, root, target_index);

        assert!(is_valid);
    }

    #[test]
    fn test_proof_index_out_of_bounds() {
        let max_index = 3u32;
        let raw_data = sp_std::vec![
            (0, sp_std::vec![1]),
            (1, sp_std::vec![2]),
            (2, sp_std::vec![3]),
            (3, sp_std::vec![4]),
        ];

        let tree = generate_tree::<TestBlake2Hasher, u32, _, _, ()>(
            max_index,
            raw_data,
            |(index, bytes)| Ok((index, bytes)),
        )
        .unwrap();

        let root = *tree.last().unwrap();

        let target_index = 2u32;
        let proof = generate_proof::<TestBlake2Hasher, u32>(&tree, max_index, target_index);

        assert_eq!(proof.len(), 2);

        let preimage = sp_std::vec![3];
        let is_valid =
            verify_tree_proof::<TestBlake2Hasher, u32>(&preimage, &proof, root, target_index + 1);

        assert!(!is_valid);
    }

    #[test]
    fn test_proof_contains_correct_siblings() {
        let max_index = 3u32;
        let raw_data = sp_std::vec![
            (0, sp_std::vec![10]),
            (1, sp_std::vec![20]),
            (2, sp_std::vec![30]),
            (3, sp_std::vec![40]),
        ];

        let tree = generate_tree::<TestBlake2Hasher, u32, _, _, ()>(
            max_index,
            raw_data,
            |(index, bytes)| Ok((index, bytes)),
        )
        .unwrap();

        let h0 = tree[0];
        let h1 = tree[1];
        let h2 = tree[2];
        let h3 = tree[3];

        let parent_left = tree[4]; // hash(h0 + h1)
        let parent_right = tree[5]; // hash(h2 + h3)

        let proof_for_0 = generate_proof::<TestBlake2Hasher, u32>(&tree, max_index, 0);
        assert_eq!(proof_for_0[0], h1);
        assert_eq!(proof_for_0[1], parent_right);

        let proof_for_1 = generate_proof::<TestBlake2Hasher, u32>(&tree, max_index, 1);
        assert_eq!(proof_for_1[0], h0);
        assert_eq!(proof_for_1[1], parent_right);

        let proof_for_2 = generate_proof::<TestBlake2Hasher, u32>(&tree, max_index, 2);
        assert_eq!(proof_for_2[0], h3);
        assert_eq!(proof_for_2[1], parent_left);

        let proof_for_3 = generate_proof::<TestBlake2Hasher, u32>(&tree, max_index, 3);
        assert_eq!(proof_for_3[0], h2);
        assert_eq!(proof_for_3[1], parent_left);
    }

    #[test]
    fn test_proof_for_padded_tree_leaf() {
        let max_index = 2u32;
        let raw_data = sp_std::vec![
            (0, sp_std::vec![1]),
            (1, sp_std::vec![2]),
            (2, sp_std::vec![3]),
        ];

        let tree = generate_tree::<TestBlake2Hasher, u32, _, _, ()>(
            max_index,
            raw_data,
            |(index, bytes)| Ok((index, bytes)),
        )
            .unwrap();

        let proof = generate_proof::<TestBlake2Hasher, u32>(&tree, max_index, 2);

        let mut empty_data = sp_std::vec![LEAF_PREFIX];
        empty_data.extend_from_slice(&H256::zero().to_fixed_bytes());
        let empty_hash = TestBlake2Hasher::hash(&empty_data);
        assert_eq!(proof[0], empty_hash);

        let root = *tree.last().unwrap();
        let is_valid =
            verify_tree_proof::<TestBlake2Hasher, u32>(&sp_std::vec![3], &proof, root, 2);
        assert!(is_valid);
    }

    #[test]
    fn test_proof_for_single_element_tree() {
        let max_index = 0u32;
        let raw_data = sp_std::vec![(0, sp_std::vec![99])];

        let tree = generate_tree::<TestBlake2Hasher, u32, _, _, ()>(
            max_index,
            raw_data,
            |(index, bytes)| Ok((index, bytes)),
        )
        .unwrap();

        let proof = generate_proof::<TestBlake2Hasher, u32>(&tree, max_index, 0);
        assert!(proof.is_empty());
    }

    #[test]
    fn test_verify_tree_proof_with_h0_h3_and_random_values() {
        let max_index = 3u32;
        let leaf_0 = sp_std::vec![10u8];
        let leaf_1 = sp_std::vec![20u8];
        let leaf_2 = sp_std::vec![30u8];
        let leaf_3 = sp_std::vec![40u8];

        let raw_data = sp_std::vec![
            (0, leaf_0.clone()),
            (1, leaf_1.clone()),
            (2, leaf_2.clone()),
            (3, leaf_3.clone())
        ];

        let tree = generate_tree::<TestBlake2Hasher, u32, _, _, ()>(
            max_index,
            raw_data,
            |(index, bytes)| Ok((index, bytes)),
        )
        .unwrap();

        let root = *tree.last().unwrap();

        let h0 = tree[0];
        let h1 = tree[1];
        let h2 = tree[2];
        let h3 = tree[3];

        let proof_for_1 = generate_proof::<TestBlake2Hasher, u32>(&tree, max_index, 1);
        assert_eq!(proof_for_1[0], h0);

        let is_valid_1 = verify_tree_proof::<TestBlake2Hasher, u32>(&leaf_1, &proof_for_1, root, 1);
        assert!(is_valid_1);

        let proof_for_2 = generate_proof::<TestBlake2Hasher, u32>(&tree, max_index, 2);
        assert_eq!(proof_for_2[0], h3);

        let is_valid_2 = verify_tree_proof::<TestBlake2Hasher, u32>(&leaf_2, &proof_for_2, root, 2);
        assert!(is_valid_2);

        let proof_for_0 = generate_proof::<TestBlake2Hasher, u32>(&tree, max_index, 0);
        assert_eq!(proof_for_0[0], h1);
        assert!(verify_tree_proof::<TestBlake2Hasher, u32>(
            &leaf_0,
            &proof_for_0,
            root,
            0
        ));

        let proof_for_3 = generate_proof::<TestBlake2Hasher, u32>(&tree, max_index, 3);
        assert_eq!(proof_for_3[0], h2);
        assert!(verify_tree_proof::<TestBlake2Hasher, u32>(
            &leaf_3,
            &proof_for_3,
            root,
            3
        ));

        let mut bad_proof = proof_for_1.clone();
        bad_proof[0] = H256::random();
        let is_valid_bad_proof =
            verify_tree_proof::<TestBlake2Hasher, u32>(&leaf_1, &bad_proof, root, 1);
        assert!(!is_valid_bad_proof);

        let bad_preimage = vec![99u8];
        let is_valid_bad_preimage =
            verify_tree_proof::<TestBlake2Hasher, u32>(&bad_preimage, &proof_for_1, root, 1);
        assert!(!is_valid_bad_preimage);
    }

    #[test]
    fn test_verify_tree_proof_deep_tree_depth_4() {
        let max_index = 15u32;

        let mut raw_data = sp_std::vec![];
        for i in 0..16 {
            raw_data.push((i, sp_std::vec![i as u8]));
        }

        let tree = generate_tree::<TestBlake2Hasher, u32, _, _, ()>(
            max_index,
            raw_data,
            |(index, bytes)| Ok((index, bytes)),
        )
        .unwrap();

        let root = *tree.last().unwrap();

        assert_eq!(tree.len(), 31);

        let target_index = 11u32;
        let preimage = vec![target_index as u8];

        let proof = generate_proof::<TestBlake2Hasher, u32>(&tree, max_index, target_index);
        assert_eq!(proof.len(), 4);

        let h10 = tree[10];
        assert_eq!(proof[0], h10);

        let is_valid =
            verify_tree_proof::<TestBlake2Hasher, u32>(&preimage, &proof, root, target_index);
        assert!(is_valid);

        let is_valid_wrong_index =
            verify_tree_proof::<TestBlake2Hasher, u32>(&preimage, &proof, root, 4);
        assert!(!is_valid_wrong_index);
    }

    #[test]
    fn test_verification_index_at_upper_bound() {
        let max_index = 3u32;
        let raw_data = sp_std::vec![
            (0, sp_std::vec![1]),
            (1, sp_std::vec![2]),
            (2, sp_std::vec![3]),
            (3, sp_std::vec![4]),
        ];

        let tree = generate_tree::<TestBlake2Hasher, u32, _, _, ()>(
            max_index,
            raw_data,
            |(index, bytes)| Ok((index, bytes)),
        )
        .unwrap();

        let root = *tree.last().unwrap();

        let target_index = 3u32;
        let proof = generate_proof::<TestBlake2Hasher, u32>(&tree, max_index, target_index);

        assert_eq!(proof.len(), 2);

        let preimage = sp_std::vec![4];
        let is_valid =
            verify_tree_proof::<TestBlake2Hasher, u32>(&preimage, &proof, root, target_index);

        assert!(is_valid);
    }

    #[test]
    fn test_verification_fails_when_index_out_of_bounds() {
        let max_index = 3u32;
        let raw_data = sp_std::vec![
            (0, sp_std::vec![1]),
            (1, sp_std::vec![2]),
            (2, sp_std::vec![3]),
            (3, sp_std::vec![4]),
        ];

        let tree = generate_tree::<TestBlake2Hasher, u32, _, _, ()>(
            max_index,
            raw_data,
            |(index, bytes)| Ok((index, bytes)),
        )
        .unwrap();

        let root = *tree.last().unwrap();

        let proof = generate_proof::<TestBlake2Hasher, u32>(&tree, max_index, 3u32);

        let target_index = 4u32;
        let preimage = sp_std::vec![4];

        let is_valid =
            verify_tree_proof::<TestBlake2Hasher, u32>(&preimage, &proof, root, target_index);

        assert!(!is_valid);
    }
}