use bytes::Bytes;
use itertools::Itertools;
use crate::{
hash::{hash_data, hash_node},
Hashed,
};
pub struct DenseMerkleTree {
datablocks: Vec<Bytes>,
bottom_to_top: Vec<Hashed>,
}
impl DenseMerkleTree {
pub fn new<R: AsRef<[u8]>>(datablocks: &[R]) -> Self {
let mut btt = vec![];
for blk in datablocks {
let hash = hash_data(blk.as_ref());
btt.push(hash);
}
let mut npp = btt.len().next_power_of_two();
while btt.len() < npp {
btt.push(Hashed::default())
}
while npp > 1 {
let index_range = btt.len() - npp..btt.len();
index_range.tuples().for_each(|(a, b)| {
let a = btt[a];
let b = btt[b];
let combined_hash = hash_node(a, b);
btt.push(combined_hash)
});
npp /= 2
}
let datablocks_vec: Vec<Bytes> = datablocks
.iter()
.map(|s| s.as_ref().to_vec().into())
.collect();
Self {
datablocks: datablocks_vec,
bottom_to_top: btt,
}
}
pub fn root_hash(&self) -> Hashed {
*self.bottom_to_top.last().unwrap()
}
pub fn data(&self) -> &[Bytes] {
&self.datablocks
}
pub fn proof(&self, mut idx: usize) -> Vec<Hashed> {
let mut accum = Vec::new();
let mut ptr = &self.bottom_to_top[..];
while ptr.len() > 1 {
let (left_half, right_half) = ptr.split_at(ptr.len() / 2 + 1);
accum.push(left_half[idx ^ 1]);
idx >>= 1;
ptr = right_half
}
accum
}
}
pub fn verify_dense(proof: &[Hashed], root: Hashed, mut idx: usize, leaf: Hashed) -> bool {
let mut ptr = leaf;
for elem in proof {
let bit = idx & 1;
idx >>= 1;
if bit > 0 {
ptr = hash_node(*elem, ptr)
} else {
ptr = hash_node(ptr, *elem)
}
}
ptr == root
}
#[cfg(test)]
mod tests {
use crate::{dense::verify_dense, hash::hash_data};
use super::DenseMerkleTree;
#[test]
fn dense_simple() {
let vals = [b"hello", b"world", b"fdfef"];
let mt = DenseMerkleTree::new(&vals);
dbg!(mt.proof(1));
assert!(verify_dense(
&mt.proof(1),
mt.root_hash(),
1,
hash_data(b"world")
))
}
#[test]
fn get_datablock() {
let vals = [b"hello", b"world", b"fdfef"];
let mt = DenseMerkleTree::new(&vals);
let blk = &mt.data()[1];
assert!(blk[..] == b"world"[..])
}
}