use metamorphic_crypto::hash::sha256;
pub const HASH_LEN: usize = 32;
pub type Hash = [u8; HASH_LEN];
const LEAF_PREFIX: u8 = 0x00;
const NODE_PREFIX: u8 = 0x01;
#[inline]
#[must_use]
pub fn empty_root() -> Hash {
sha256(&[])
}
#[inline]
#[must_use]
pub fn hash_leaf(leaf_bytes: &[u8]) -> Hash {
let mut buf = Vec::with_capacity(1 + leaf_bytes.len());
buf.push(LEAF_PREFIX);
buf.extend_from_slice(leaf_bytes);
sha256(&buf)
}
#[inline]
#[must_use]
pub fn hash_children(left: &Hash, right: &Hash) -> Hash {
let mut buf = [0u8; 1 + 2 * HASH_LEN];
buf[0] = NODE_PREFIX;
buf[1..1 + HASH_LEN].copy_from_slice(left);
buf[1 + HASH_LEN..].copy_from_slice(right);
sha256(&buf)
}
#[inline]
fn largest_power_of_two_below(n: u64) -> u64 {
debug_assert!(n > 1);
let bits = u64::BITS - (n - 1).leading_zeros();
1u64 << (bits - 1)
}
#[derive(Debug, Clone, Default)]
pub struct MerkleTree {
leaves: Vec<Hash>,
}
impl MerkleTree {
#[must_use]
pub fn new() -> Self {
Self { leaves: Vec::new() }
}
pub fn push(&mut self, leaf_bytes: &[u8]) -> u64 {
let index = self.leaves.len() as u64;
self.leaves.push(hash_leaf(leaf_bytes));
index
}
pub fn push_leaf_hash(&mut self, leaf_hash: Hash) -> u64 {
let index = self.leaves.len() as u64;
self.leaves.push(leaf_hash);
index
}
#[must_use]
pub fn size(&self) -> u64 {
self.leaves.len() as u64
}
#[must_use]
pub fn leaf_hash(&self, index: u64) -> Option<Hash> {
self.leaves.get(index as usize).copied()
}
#[must_use]
pub fn root_at(&self, size: u64) -> Hash {
assert!(size <= self.size(), "root_at: size exceeds tree size");
mth(&self.leaves[..size as usize])
}
#[must_use]
pub fn root(&self) -> Hash {
self.root_at(self.size())
}
#[must_use]
pub fn inclusion_proof(&self, index: u64, size: u64) -> Vec<Hash> {
assert!(index < size, "inclusion_proof: index beyond size");
assert!(
size <= self.size(),
"inclusion_proof: size exceeds tree size"
);
inclusion_path(index, &self.leaves[..size as usize])
}
#[must_use]
pub fn consistency_proof(&self, size1: u64, size2: u64) -> Vec<Hash> {
assert!(size1 > 0, "consistency_proof: size1 must be > 0");
assert!(size1 <= size2, "consistency_proof: size1 > size2");
assert!(
size2 <= self.size(),
"consistency_proof: size2 exceeds tree size"
);
consistency_path(size1, &self.leaves[..size2 as usize])
}
}
#[inline]
#[must_use]
pub fn merkle_tree_hash(node_hashes: &[Hash]) -> Hash {
mth(node_hashes)
}
fn mth(leaves: &[Hash]) -> Hash {
match leaves.len() {
0 => empty_root(),
1 => leaves[0],
n => {
let k = largest_power_of_two_below(n as u64) as usize;
let left = mth(&leaves[..k]);
let right = mth(&leaves[k..]);
hash_children(&left, &right)
}
}
}
fn inclusion_path(m: u64, leaves: &[Hash]) -> Vec<Hash> {
let n = leaves.len();
if n <= 1 {
return Vec::new();
}
let k = largest_power_of_two_below(n as u64);
if m < k {
let mut path = inclusion_path(m, &leaves[..k as usize]);
path.push(mth(&leaves[k as usize..]));
path
} else {
let mut path = inclusion_path(m - k, &leaves[k as usize..]);
path.push(mth(&leaves[..k as usize]));
path
}
}
fn consistency_path(m: u64, leaves: &[Hash]) -> Vec<Hash> {
subproof(m, leaves, true)
}
fn subproof(m: u64, leaves: &[Hash], b: bool) -> Vec<Hash> {
let n = leaves.len() as u64;
if m == n {
if b {
return Vec::new();
}
return vec![mth(leaves)];
}
let k = largest_power_of_two_below(n);
if m <= k {
let mut proof = subproof(m, &leaves[..k as usize], b);
proof.push(mth(&leaves[k as usize..]));
proof
} else {
let mut proof = subproof(m - k, &leaves[k as usize..], false);
proof.push(mth(&leaves[..k as usize]));
proof
}
}
#[cfg(all(test, not(target_arch = "wasm32")))]
mod tests {
use super::*;
#[test]
fn empty_root_is_sha256_of_empty() {
let expected = hex32("e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855");
assert_eq!(empty_root(), expected);
}
#[test]
fn single_empty_leaf_hash() {
let expected = hex32("6e340b9cffb37a989ca544e6bb780a2c78901d3fb33738768511a30617afa01d");
assert_eq!(hash_leaf(b""), expected);
}
#[test]
fn largest_power_of_two_below_cases() {
assert_eq!(largest_power_of_two_below(2), 1);
assert_eq!(largest_power_of_two_below(3), 2);
assert_eq!(largest_power_of_two_below(4), 2);
assert_eq!(largest_power_of_two_below(5), 4);
assert_eq!(largest_power_of_two_below(8), 4);
assert_eq!(largest_power_of_two_below(9), 8);
}
fn hex32(s: &str) -> Hash {
assert_eq!(s.len(), 64, "expected 64 hex chars");
let mut out = [0u8; 32];
for (i, b) in out.iter_mut().enumerate() {
*b = u8::from_str_radix(&s[2 * i..2 * i + 2], 16).unwrap();
}
out
}
}