sp_domains/
merkle_tree.rs

1#[cfg(not(feature = "std"))]
2extern crate alloc;
3
4use crate::OperatorPublicKey;
5#[cfg(not(feature = "std"))]
6use alloc::vec::Vec;
7use blake2::digest::FixedOutput;
8use blake2::{Blake2b, Digest};
9use parity_scale_codec::{Decode, Encode};
10use rs_merkle::Hasher;
11use scale_info::TypeInfo;
12use sp_runtime::traits::{BlakeTwo256, Hash};
13
14/// Merkle tree using [`Blake2b256Algorithm`].
15pub type MerkleTree = rs_merkle::MerkleTree<Blake2b256Algorithm>;
16
17/// Merkle proof using [`Blake2b256Algorithm`].
18pub type MerkleProof = rs_merkle::MerkleProof<Blake2b256Algorithm>;
19
20/// Merke proof based Witness.
21#[derive(Debug, Decode, Encode, TypeInfo, PartialEq, Eq, Clone, Default)]
22pub struct Witness {
23    /// Index of the leaf the proof is for.
24    pub leaf_index: u32,
25    /// Merkle proof in bytes.
26    pub proof: Vec<u8>,
27    /// Number of leaves in the original tree.
28    pub number_of_leaves: u32,
29}
30
31#[derive(Clone)]
32pub struct Blake2b256Algorithm;
33
34impl Default for Blake2b256Algorithm {
35    #[inline]
36    fn default() -> Self {
37        Self
38    }
39}
40
41impl Hasher for Blake2b256Algorithm {
42    type Hash = [u8; 32];
43
44    fn hash(data: &[u8]) -> Self::Hash {
45        let mut hasher = Blake2b::new();
46        hasher.update(data);
47        hasher.finalize_fixed().into()
48    }
49}
50
51/// Constructs a merkle tree from given authorities.
52pub fn authorities_merkle_tree<StakeWeight: Encode>(
53    authorities: &[(OperatorPublicKey, StakeWeight)],
54) -> MerkleTree {
55    let leaves = authorities
56        .iter()
57        .map(|x| BlakeTwo256::hash_of(&x.encode()).to_fixed_bytes())
58        .collect::<Vec<_>>();
59    MerkleTree::from_leaves(&leaves)
60}
61
62#[cfg(test)]
63mod tests {
64    use super::MerkleTree;
65
66    #[test]
67    fn test_merkle_tree() {
68        let leaves = [[1u8; 32], [2u8; 32], [3u8; 32], [4u8; 32], [5u8; 32]];
69
70        let merkle_tree = MerkleTree::from_leaves(&leaves);
71
72        let indices_to_prove = (0..leaves.len()).collect::<Vec<_>>();
73        let leaves_to_prove = leaves
74            .get(0..leaves.len())
75            .expect("can't get leaves to prove");
76
77        let proof = merkle_tree.proof(&indices_to_prove);
78        let root = merkle_tree.root().expect("couldn't get the merkle root");
79
80        assert!(proof.verify(root, &indices_to_prove, leaves_to_prove, leaves.len()));
81    }
82}