use sha2::{Digest, Sha256};
pub const LEAF_HASH_PREFIX: u8 = 0x00;
pub const NODE_HASH_PREFIX: u8 = 0x01;
pub fn leaf_hash(leaf_bytes: &[u8]) -> [u8; 32] {
let mut h = Sha256::new();
h.update([LEAF_HASH_PREFIX]);
h.update(leaf_bytes);
h.finalize().into()
}
pub fn node_hash(left: &[u8; 32], right: &[u8; 32]) -> [u8; 32] {
let mut h = Sha256::new();
h.update([NODE_HASH_PREFIX]);
h.update(left);
h.update(right);
h.finalize().into()
}
fn largest_pow2_lt(n: usize) -> usize {
debug_assert!(n >= 2);
let mut k = 1usize;
while k * 2 < n {
k *= 2;
}
k
}
pub fn merkle_tree_hash(leaf_hashes: &[[u8; 32]]) -> [u8; 32] {
match leaf_hashes.len() {
0 => Sha256::digest(b"").into(),
1 => leaf_hashes[0],
n => {
let k = largest_pow2_lt(n);
node_hash(
&merkle_tree_hash(&leaf_hashes[..k]),
&merkle_tree_hash(&leaf_hashes[k..]),
)
}
}
}
pub fn inclusion_path(leaf_index: usize, leaf_hashes: &[[u8; 32]]) -> Option<Vec<[u8; 32]>> {
let n = leaf_hashes.len();
if leaf_index >= n {
return None;
}
if n == 1 {
return Some(Vec::new());
}
let k = largest_pow2_lt(n);
let mut path = if leaf_index < k {
let mut p = inclusion_path(leaf_index, &leaf_hashes[..k])?;
p.push(merkle_tree_hash(&leaf_hashes[k..]));
p
} else {
let mut p = inclusion_path(leaf_index - k, &leaf_hashes[k..])?;
p.push(merkle_tree_hash(&leaf_hashes[..k]));
p
};
path.shrink_to_fit();
Some(path)
}
pub fn consistency_proof(first: usize, leaf_hashes: &[[u8; 32]]) -> Option<Vec<[u8; 32]>> {
if first == 0 || first > leaf_hashes.len() {
return None;
}
fn subproof(m: usize, d: &[[u8; 32]], b: bool) -> Vec<[u8; 32]> {
let n = d.len();
if m == n {
return if b {
Vec::new()
} else {
vec![merkle_tree_hash(d)]
};
}
let k = largest_pow2_lt(n);
let mut p = if m <= k {
let mut p = subproof(m, &d[..k], b);
p.push(merkle_tree_hash(&d[k..]));
p
} else {
let mut p = subproof(m - k, &d[k..], false);
p.push(merkle_tree_hash(&d[..k]));
p
};
p.shrink_to_fit();
p
}
Some(subproof(first, leaf_hashes, true))
}
#[must_use]
pub fn verify_inclusion(
leaf_hash: &[u8; 32],
leaf_index: u64,
tree_size: u64,
path: &[[u8; 32]],
root: &[u8; 32],
) -> bool {
if leaf_index >= tree_size {
return false;
}
let mut fnode = leaf_index;
let mut snode = tree_size - 1;
let mut r = *leaf_hash;
for p in path {
if snode == 0 {
return false;
}
if fnode % 2 == 1 || fnode == snode {
r = node_hash(p, &r);
if fnode % 2 == 0 {
while fnode % 2 == 0 && fnode != 0 {
fnode >>= 1;
snode >>= 1;
}
}
} else {
r = node_hash(&r, p);
}
fnode >>= 1;
snode >>= 1;
}
snode == 0 && r == *root
}
#[must_use]
pub fn verify_consistency(
first: u64,
second: u64,
path: &[[u8; 32]],
first_root: &[u8; 32],
second_root: &[u8; 32],
) -> bool {
if first == second {
return path.is_empty() && first_root == second_root;
}
if first == 0 || first > second || path.is_empty() {
return false;
}
let mut elements: Vec<&[u8; 32]> = Vec::with_capacity(path.len() + 1);
if first & (first - 1) == 0 {
elements.push(first_root);
}
elements.extend(path.iter());
let mut elements = elements.into_iter();
let mut fnode = first - 1;
let mut snode = second - 1;
while fnode % 2 == 1 {
fnode >>= 1;
snode >>= 1;
}
let Some(first_el) = elements.next() else {
return false;
};
let mut fr = *first_el;
let mut sr = *first_el;
for c in elements {
if snode == 0 {
return false;
}
if fnode % 2 == 1 || fnode == snode {
fr = node_hash(c, &fr);
sr = node_hash(c, &sr);
if fnode % 2 == 0 {
while fnode % 2 == 0 && fnode != 0 {
fnode >>= 1;
snode >>= 1;
}
}
} else {
sr = node_hash(&sr, c);
}
fnode >>= 1;
snode >>= 1;
}
fr == *first_root && sr == *second_root && snode == 0
}
#[cfg(test)]
mod tests {
use super::*;
fn leaves(n: usize) -> Vec<[u8; 32]> {
(0..n)
.map(|i| leaf_hash(format!("leaf-{i}").as_bytes()))
.collect()
}
#[test]
fn empty_tree_root_is_sha256_of_empty_string() {
assert_eq!(
hex::encode(merkle_tree_hash(&[])),
"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
);
}
#[test]
fn prefixes_domain_separate() {
let l = [7u8; 32];
let r = [9u8; 32];
let mut concat = Vec::new();
concat.extend_from_slice(&l);
concat.extend_from_slice(&r);
assert_ne!(node_hash(&l, &r), leaf_hash(&concat));
}
#[test]
fn inclusion_and_consistency_exhaustive_to_8() {
for n in 1..=8usize {
let d = leaves(n);
let root = merkle_tree_hash(&d);
for m in 0..n {
let path = inclusion_path(m, &d).expect("path exists");
assert!(
verify_inclusion(&d[m], m as u64, n as u64, &path, &root),
"PATH({m}, D[{n}]) must verify"
);
let mut bad_leaf = d[m];
bad_leaf[0] ^= 1;
assert!(!verify_inclusion(
&bad_leaf, m as u64, n as u64, &path, &root
));
for i in 0..path.len() {
let mut bad = path.clone();
bad[i][0] ^= 1;
assert!(!verify_inclusion(&d[m], m as u64, n as u64, &bad, &root));
}
if !path.is_empty() {
assert!(!verify_inclusion(
&d[m],
m as u64,
n as u64,
&path[..path.len() - 1],
&root
));
}
let mut long = path.clone();
long.push([0u8; 32]);
assert!(!verify_inclusion(&d[m], m as u64, n as u64, &long, &root));
}
for m in 1..=n {
let first_root = merkle_tree_hash(&d[..m]);
let proof = consistency_proof(m, &d).expect("proof exists");
assert!(
verify_consistency(m as u64, n as u64, &proof, &first_root, &root),
"PROOF({m}, D[{n}]) must verify"
);
if m == n {
assert!(proof.is_empty());
}
for i in 0..proof.len() {
let mut bad = proof.clone();
bad[i][0] ^= 1;
assert!(!verify_consistency(
m as u64,
n as u64,
&bad,
&first_root,
&root
));
}
if m < n {
assert!(!verify_consistency(
m as u64,
n as u64,
&proof,
&root,
&first_root
));
}
}
}
}
#[test]
fn consistency_degenerate_cases() {
let d = leaves(4);
let root = merkle_tree_hash(&d);
let empty_root = merkle_tree_hash(&[]);
assert!(!verify_consistency(0, 4, &[], &empty_root, &root));
assert!(!verify_consistency(0, 4, &[[0u8; 32]], &empty_root, &root));
assert!(!verify_consistency(5, 4, &[[0u8; 32]], &root, &root));
assert!(!verify_consistency(
2,
4,
&[],
&merkle_tree_hash(&d[..2]),
&root
));
assert!(verify_consistency(4, 4, &[], &root, &root));
assert!(!verify_consistency(4, 4, &[[0u8; 32]], &root, &root));
assert!(!verify_consistency(
4,
4,
&[],
&merkle_tree_hash(&d[..2]),
&root
));
assert!(inclusion_path(4, &d).is_none());
assert!(inclusion_path(0, &[]).is_none());
assert!(consistency_proof(0, &d).is_none());
assert!(consistency_proof(5, &d).is_none());
}
}