use crate::errors::{MerkleTreeError, MerkleTreeErrorKind};
use crate::hasher::Arity2Hasher;
use std::marker::PhantomData;
type TreeSizeType = u64;
fn count_set_bits(mut n: TreeSizeType) -> u8 {
let mut count = 0;
while n != 0 {
n &= n - 1;
count += 1
}
count
}
fn num_bits(mut n: TreeSizeType) -> u8 {
if n == 0 {
return 0;
}
let mut index = 0;
while n != 0 {
index += 1;
n >>= 1;
}
index
}
fn least_significant_set_bit(n: TreeSizeType) -> u8 {
if n == 0 {
return 0;
}
n.trailing_zeros() as u8
}
fn largest_power_of_2_less_than(n: TreeSizeType) -> TreeSizeType {
if n < 2 {
return 0;
}
let mut cur = 1u64;
let mut largest = 1u64;
while cur < n {
largest = cur;
let r = cur.wrapping_shl(1);
if r < cur {
break;
} else {
cur = r;
}
}
largest
}
fn powers_of_2(mut n: TreeSizeType) -> Vec<TreeSizeType> {
if n == 0 {
return vec![];
}
let mut powers = vec![];
loop {
if n.is_power_of_two() {
powers.push(n);
break;
} else {
let p = largest_power_of_2_less_than(n);
n = n - p;
powers.push(p);
}
}
powers
}
pub trait HashDb<H> {
fn add_leaf(&mut self, leaf_hash: H) -> Result<(), MerkleTreeError>;
fn add_full_subtree_root(&mut self, node_hash: H) -> Result<(), MerkleTreeError>;
fn get_leaf(&self, leaf_index: TreeSizeType) -> Result<H, MerkleTreeError>;
fn get_full_subtree_root(&self, node_index: TreeSizeType) -> Result<H, MerkleTreeError>;
}
#[derive(Clone, Debug)]
pub struct InMemoryHashDb<H> {
leaves: Vec<H>,
nodes: Vec<H>,
}
impl<H: Clone> HashDb<H> for InMemoryHashDb<H> {
fn add_leaf(&mut self, leaf_hash: H) -> Result<(), MerkleTreeError> {
self.leaves.push(leaf_hash);
Ok(())
}
fn add_full_subtree_root(&mut self, node_hash: H) -> Result<(), MerkleTreeError> {
self.nodes.push(node_hash);
Ok(())
}
fn get_leaf(&self, leaf_index: TreeSizeType) -> Result<H, MerkleTreeError> {
let i = leaf_index as usize;
if i >= self.leaves.len() {
Err(MerkleTreeError::from_kind(
MerkleTreeErrorKind::LeafIndexNotFoundInDB { index: i as u64 },
))
} else {
Ok(self.leaves[i].clone())
}
}
fn get_full_subtree_root(&self, node_index: TreeSizeType) -> Result<H, MerkleTreeError> {
let i = node_index as usize;
if i >= self.nodes.len() {
Err(MerkleTreeError::from_kind(
MerkleTreeErrorKind::NodeIndexNotFoundInDB { index: i as u64 },
))
} else {
Ok(self.nodes[i].clone())
}
}
}
impl<H> InMemoryHashDb<H> {
pub fn new() -> Self {
Self {
leaves: vec![],
nodes: vec![],
}
}
pub fn leaf_count(&self) -> TreeSizeType {
self.leaves.len() as TreeSizeType
}
pub fn node_count(&self) -> TreeSizeType {
self.nodes.len() as TreeSizeType
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct CompactMerkleTree<D: Clone, H: Clone, MTH>
where
MTH: Arity2Hasher<D, H>,
{
pub size: TreeSizeType,
pub full_subtree_roots: Vec<H>,
pub hasher: MTH,
pub phantom: PhantomData<D>,
}
impl<D: Clone, H: Clone + PartialEq, MTH> CompactMerkleTree<D, H, MTH>
where
MTH: Arity2Hasher<D, H>,
{
pub fn new(hasher: MTH) -> Self {
Self {
size: 0,
full_subtree_roots: vec![],
hasher,
phantom: PhantomData,
}
}
pub fn new_from_hash_db(
hasher: MTH,
tree_size: TreeSizeType,
hash_db: &dyn HashDb<H>,
) -> Result<Self, MerkleTreeError> {
let mut full_subtree_roots = Self::get_leaf_inclusion_proof_for_tree_size(
&hasher,
tree_size,
tree_size + 1,
hash_db,
)?;
full_subtree_roots.reverse();
let mut new_tree = Self::new(hasher);
new_tree.size = tree_size;
new_tree.full_subtree_roots = full_subtree_roots;
Ok(new_tree)
}
pub fn append(
&mut self,
leaf_data: D,
hash_db: &mut dyn HashDb<H>,
) -> Result<Vec<H>, MerkleTreeError> {
let mut inclusion_proof = self.full_subtree_roots.clone();
inclusion_proof.reverse();
self.push_full_subtree(vec![leaf_data], hash_db)?;
Ok(inclusion_proof)
}
pub fn extend(
&mut self,
mut leaves: Vec<D>,
hash_db: &mut dyn HashDb<H>,
) -> Result<(), MerkleTreeError> {
if leaves.is_empty() {
return Err(MerkleTreeErrorKind::NoLeafProvided.into());
}
let mut num_leaves_remaining = leaves.len() as TreeSizeType;
if self.size > 0 {
loop {
let max_leaves_to_insert = 1 << self.smallest_subtree_height() as TreeSizeType;
if max_leaves_to_insert <= num_leaves_remaining {
let to_insert = leaves.drain(0..max_leaves_to_insert as usize).collect();
self.push_full_subtree(to_insert, hash_db)?;
num_leaves_remaining -= max_leaves_to_insert;
} else {
break;
}
}
}
if num_leaves_remaining > 0 {
let (leaf_hashes, node_hashes) = Self::hash_leaves(&self.hasher, leaves.drain(0..).collect())?;
let mut idx = 0;
for p in powers_of_2(num_leaves_remaining) {
if p == 1 {
self.added_subtree(1, leaf_hashes[(num_leaves_remaining - 1) as usize].clone());
} else {
idx += (p - 1);
self.added_subtree(p, node_hashes[(idx - 1) as usize].clone());
}
}
for l in leaf_hashes {
hash_db.add_leaf(l)?;
}
for n in node_hashes {
hash_db.add_full_subtree_root(n)?;
}
}
Ok(())
}
pub fn get_root_hash(&self) -> Result<H, MerkleTreeError> {
if self.size == 0 {
return Err(MerkleTreeErrorKind::CannotQueryEmptyTree.into());
}
let mut cur_root = self.full_subtree_roots[self.full_subtree_roots.len() - 1].clone();
for i in (0..self.full_subtree_roots.len() - 1).rev() {
cur_root = self
.hasher
.hash_tree_nodes(self.full_subtree_roots[i].clone(), cur_root)?;
}
Ok(cur_root)
}
pub fn get_leaf_inclusion_proof(
&self,
leaf_index: TreeSizeType,
hash_db: &dyn HashDb<H>,
) -> Result<Vec<H>, MerkleTreeError> {
Self::get_leaf_inclusion_proof_for_tree_size(&self.hasher, leaf_index, self.size, hash_db)
}
pub fn get_leaf_inclusion_proof_for_tree_size(
hasher: &MTH,
leaf_index: TreeSizeType,
tree_size: TreeSizeType,
hash_db: &dyn HashDb<H>,
) -> Result<Vec<H>, MerkleTreeError> {
if leaf_index >= tree_size {
return Err(MerkleTreeErrorKind::TreeSmallerThanExpected {
expected: leaf_index + 1,
given: tree_size,
}
.into());
}
Self::path(hasher, leaf_index, 0, tree_size, hash_db)
}
pub fn get_consistency_proof(
&self,
old_tree_size: TreeSizeType,
hash_db: &dyn HashDb<H>,
) -> Result<Vec<H>, MerkleTreeError> {
Self::get_consistency_proof_for_tree_size(&self.hasher, old_tree_size, self.size, hash_db)
}
pub fn get_consistency_proof_for_tree_size(
hasher: &MTH,
old_tree_size: TreeSizeType,
new_tree_size: TreeSizeType,
hash_db: &dyn HashDb<H>,
) -> Result<Vec<H>, MerkleTreeError> {
if old_tree_size > new_tree_size {
return Err(MerkleTreeErrorKind::TreeSmallerThanExpected {
expected: old_tree_size,
given: new_tree_size,
}
.into());
}
if old_tree_size == new_tree_size {
Ok(vec![])
} else {
let proof = Self::subproof(hasher, old_tree_size, 0, new_tree_size, true, hash_db)?;
Ok(proof)
}
}
pub fn verify_leaf_inclusion_proof(
hasher: &MTH,
leaf_index: TreeSizeType,
leaf_val: D,
tree_size: TreeSizeType,
root: &H,
proof: Vec<H>,
) -> Result<bool, MerkleTreeError> {
if leaf_index >= tree_size {
return Err(MerkleTreeErrorKind::TreeSmallerThanExpected {
expected: leaf_index + 1,
given: tree_size,
}
.into());
}
let cur_hash = hasher.hash_leaf_data(leaf_val)?;
let (right_border_len, inner_path_len) =
Self::get_right_border_and_inner_node_count(leaf_index, tree_size);
if (right_border_len + inner_path_len) != proof.len() as u8 {
return Err(MerkleTreeErrorKind::ShorterInclusionProof {
expected: right_border_len + inner_path_len,
given: proof.len() as u8,
}
.into());
}
Self::_verify_leaf_inclusion_proof(
hasher,
leaf_index,
inner_path_len,
cur_hash,
proof,
root,
)
}
pub fn verify_consistency_proof(
hasher: &MTH,
old_tree_size: TreeSizeType,
new_tree_size: TreeSizeType,
old_root: &H,
new_root: &H,
mut proof: Vec<H>,
) -> Result<bool, MerkleTreeError> {
if old_tree_size > new_tree_size {
return Err(MerkleTreeErrorKind::ConsistencyProofIncorrectTreeSize {
new: new_tree_size,
old: old_tree_size,
}
.into());
}
if old_tree_size == 0 || new_tree_size == 0 {
return Err(MerkleTreeErrorKind::ConsistencyProofWithEmptyTree.into());
}
if old_tree_size == new_tree_size {
return Ok(old_root == new_root);
}
let start_hash = if old_tree_size.is_power_of_two() {
old_root.clone()
} else {
proof.remove(0)
};
let (right_border_len, mut inner_path_len) =
Self::get_right_border_and_inner_node_count(old_tree_size - 1, new_tree_size);
let shift = least_significant_set_bit(old_tree_size);
inner_path_len -= shift;
if (right_border_len + inner_path_len) != proof.len() as u8 {
return Err(MerkleTreeErrorKind::ShorterConsistencyProof {
expected: right_border_len + inner_path_len,
given: proof.len() as u8,
}
.into());
}
let node_index = (old_tree_size - 1) >> shift as TreeSizeType;
let mut index_for_old = node_index.clone();
let mut expected_old_root = start_hash.clone();
for p in proof.iter().take(inner_path_len as usize) {
if index_for_old % 2 == 1 {
expected_old_root = hasher.hash_tree_nodes(p.clone(), expected_old_root)?;
}
index_for_old >>= 1;
}
for p in proof.iter().skip(inner_path_len as usize) {
expected_old_root = hasher.hash_tree_nodes(p.clone(), expected_old_root)?;
}
if expected_old_root != *old_root {
return Err(MerkleTreeErrorKind::InconsistentOldRoot.into());
}
Self::_verify_leaf_inclusion_proof(
hasher,
node_index,
inner_path_len,
start_hash,
proof,
new_root,
)
}
fn _verify_leaf_inclusion_proof(
hasher: &MTH,
mut leaf_index: TreeSizeType,
inner_path_len: u8,
mut cur_hash: H,
mut proof: Vec<H>,
expected_root: &H,
) -> Result<bool, MerkleTreeError> {
for p in proof.drain(0..inner_path_len as usize) {
if leaf_index % 2 == 1 {
cur_hash = hasher.hash_tree_nodes(p, cur_hash)?;
} else {
cur_hash = hasher.hash_tree_nodes(cur_hash, p)?;
}
leaf_index >>= 1;
}
for p in proof.drain(0..) {
cur_hash = hasher.hash_tree_nodes(p, cur_hash)?;
}
Ok(cur_hash == *expected_root)
}
fn push_full_subtree(
&mut self,
leaves: Vec<D>,
hash_db: &mut dyn HashDb<H>,
) -> Result<(), MerkleTreeError> {
if !leaves.len().is_power_of_two() {
return Err(MerkleTreeErrorKind::ErrorInsertingSubtree {
msg: format!(
"Leaves should form a full subtree but only {} leaves given",
leaves.len()
),
}
.into());
}
let new_subtree_height = leaves.len().trailing_zeros() as u8;
let min_subtree_height = self.smallest_subtree_height();
if (self.size != 0) && (new_subtree_height > min_subtree_height) {
return Err(MerkleTreeErrorKind::ErrorInsertingSubtree {msg: format!("Leaves should form full subtree of height at most {} but form subtree of height {}", min_subtree_height, new_subtree_height)}.into());
}
let (leaf_hashes, node_hashes) = Self::hash_leaves(&self.hasher, leaves)?;
let subtree_root = if node_hashes.is_empty() {
leaf_hashes.last().unwrap().clone()
} else {
node_hashes.last().unwrap().clone()
};
for l in leaf_hashes {
hash_db.add_leaf(l)?;
}
for n in node_hashes {
hash_db.add_full_subtree_root(n)?;
}
let new_roots = self.push_full_subtree_hash(new_subtree_height, subtree_root)?;
for n in new_roots {
hash_db.add_full_subtree_root(n)?;
}
Ok(())
}
fn push_full_subtree_hash(
&mut self,
subtree_height: u8,
subtree_hash: H,
) -> Result<Vec<H>, MerkleTreeError> {
let min_subtree_height = self.smallest_subtree_height();
let subtree_size = 1 << (subtree_height as TreeSizeType);
if (self.size != 0) && (subtree_height > min_subtree_height) {
return Err(MerkleTreeErrorKind::ErrorInsertingSubtree {msg: format!("Leaves should form full subtree of height at most {} but form subtree of height {}", min_subtree_height, subtree_height)}.into());
}
if self.size == 0 || subtree_height < min_subtree_height {
self.added_subtree(subtree_size, subtree_hash);
Ok(vec![])
} else {
let removed_subtree_hash = self.removed_subtree(subtree_size)?;
let next_root = self
.hasher
.hash_tree_nodes(removed_subtree_hash, subtree_hash)?;
let mut remain_roots =
self.push_full_subtree_hash(subtree_height + 1, next_root.clone())?;
remain_roots.insert(0, next_root);
Ok(remain_roots)
}
}
fn added_subtree(&mut self, subtree_size: TreeSizeType, subtree_root: H) {
self.size += subtree_size;
self.full_subtree_roots.push(subtree_root);
}
fn removed_subtree(&mut self, subtree_size: TreeSizeType) -> Result<H, MerkleTreeError> {
if self.size < subtree_size {
return Err(MerkleTreeErrorKind::TreeSmallerThanExpected {
expected: subtree_size,
given: self.size,
}
.into());
}
self.size -= subtree_size;
Ok(self.full_subtree_roots.pop().unwrap())
}
fn hash_leaves(hasher: &MTH, mut leaves: Vec<D>) -> Result<(Vec<H>, Vec<H>), MerkleTreeError> {
if leaves.is_empty() {
return Err(MerkleTreeErrorKind::NoLeafProvided.into());
}
if leaves.len() == 1 {
Ok((vec![hasher.hash_leaf_data(leaves.remove(0))?], vec![]))
} else {
let left_subtree_size = largest_power_of_2_less_than(leaves.len() as TreeSizeType);
let right_subtree_size = leaves.len() as TreeSizeType - left_subtree_size;
let right_subtree_leaves = leaves.split_off(left_subtree_size as usize);
let (mut left_leaf_hashes, mut left_node_hashes) = Self::hash_leaves(hasher, leaves)?;
let (mut right_leaf_hashes, mut right_node_hashes) =
Self::hash_leaves(hasher, right_subtree_leaves)?;
let root = if left_subtree_size == right_subtree_size {
let left_root = if left_node_hashes.is_empty() {
left_leaf_hashes.last().unwrap().clone()
} else {
left_node_hashes.last().unwrap().clone()
};
let right_root = if right_node_hashes.is_empty() {
right_leaf_hashes.last().unwrap().clone()
} else {
right_node_hashes.last().unwrap().clone()
};
let root = hasher.hash_tree_nodes(left_root, right_root)?;
Some(root)
} else {
None
};
let mut all_leaf_hashes = vec![];
all_leaf_hashes.append(&mut left_leaf_hashes);
all_leaf_hashes.append(&mut right_leaf_hashes);
let mut all_node_hashes = vec![];
all_node_hashes.append(&mut left_node_hashes);
all_node_hashes.append(&mut right_node_hashes);
if left_subtree_size == right_subtree_size {
all_node_hashes.push(root.unwrap());
}
Ok((all_leaf_hashes, all_node_hashes))
}
}
fn path(
hasher: &MTH,
leaf_index: TreeSizeType,
from: TreeSizeType,
to: TreeSizeType,
hash_db: &dyn HashDb<H>,
) -> Result<Vec<H>, MerkleTreeError> {
let subtree_size = to - from;
if subtree_size == 1 {
return Ok(vec![]);
}
let k = largest_power_of_2_less_than(subtree_size);
let (mut path, mth) = if leaf_index < k {
let path = Self::path(hasher, leaf_index, from, from + k, hash_db)?;
let mth = Self::subtree_root_hash_from_db(hasher, from + k, to, hash_db)?;
(path, mth)
} else {
let path = Self::path(hasher, leaf_index - k, from + k, to, hash_db)?;
let mth = Self::subtree_root_hash_from_db(hasher, from, from + k, hash_db)?;
(path, mth)
};
path.push(mth);
Ok(path)
}
fn subproof(
hasher: &MTH,
old_tree_size: TreeSizeType,
from: TreeSizeType,
to: TreeSizeType,
root_present_in_old_tree: bool,
hash_db: &dyn HashDb<H>,
) -> Result<Vec<H>, MerkleTreeError> {
let subtree_size = to - from;
if subtree_size == old_tree_size && root_present_in_old_tree {
Ok(vec![])
} else if subtree_size == old_tree_size && !root_present_in_old_tree {
Ok(vec![Self::subtree_root_hash_from_db(
hasher, from, to, hash_db,
)?])
} else {
let k = largest_power_of_2_less_than(subtree_size);
let (mut subproof, mth) = if old_tree_size <= k {
let subproof = Self::subproof(
hasher,
old_tree_size,
from,
from + k,
root_present_in_old_tree,
hash_db,
)?;
let mth = Self::subtree_root_hash_from_db(hasher, from + k, to, hash_db)?;
(subproof, mth)
} else {
let subproof =
Self::subproof(hasher, old_tree_size - k, from + k, to, false, hash_db)?;
let mth = Self::subtree_root_hash_from_db(hasher, from, from + k, hash_db)?;
(subproof, mth)
};
subproof.push(mth);
Ok(subproof)
}
}
fn subtree_root_hash_from_db(
hasher: &MTH,
from: TreeSizeType,
to: TreeSizeType,
hash_db: &dyn HashDb<H>,
) -> Result<H, MerkleTreeError> {
if from >= to {
return Err(MerkleTreeErrorKind::IncorrectSpan { from, to }.into());
}
let subtree_size = to - from;
if subtree_size == 1 {
hash_db.get_leaf(from)
} else if subtree_size.is_power_of_two() {
let node_index = Self::get_subtree_root_index(from, to);
hash_db.get_full_subtree_root(node_index)
} else {
let left_subtree_size = largest_power_of_2_less_than(subtree_size);
let left_root =
Self::subtree_root_hash_from_db(hasher, from, from + left_subtree_size, hash_db)?;
let right_root =
Self::subtree_root_hash_from_db(hasher, from + left_subtree_size, to, hash_db)?;
hasher.hash_tree_nodes(left_root, right_root)
}
}
fn get_subtree_root_index(from: TreeSizeType, to: TreeSizeType) -> TreeSizeType {
fn root_index(m: TreeSizeType, from: TreeSizeType, to: TreeSizeType) -> TreeSizeType {
if from == m {
to - from - 1 - 1
} else {
let d = to - from;
let k = largest_power_of_2_less_than(d);
(k - 1) + root_index(m, from + k, to)
}
}
root_index(from, 0, to)
}
fn smallest_subtree_height(&self) -> u8 {
least_significant_set_bit(self.size)
}
fn get_right_border_and_inner_node_count(
leaf_index: TreeSizeType,
size: TreeSizeType,
) -> (u8, u8) {
let last_leaf_idx = (size - 1);
let inner_path_len = num_bits(leaf_index ^ last_leaf_idx);
let right_border_path_len = count_set_bits(leaf_index >> inner_path_len as u64);
(right_border_path_len, inner_path_len)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::hasher::Sha256Hasher;
#[test]
fn test_largest_power_of_2_less_than() {
assert_eq!(largest_power_of_2_less_than(0), 0);
assert_eq!(largest_power_of_2_less_than(1), 0);
assert_eq!(largest_power_of_2_less_than(2), 1);
assert_eq!(largest_power_of_2_less_than(3), 2);
assert_eq!(largest_power_of_2_less_than(4), 2);
assert_eq!(largest_power_of_2_less_than(5), 4);
assert_eq!(largest_power_of_2_less_than(6), 4);
assert_eq!(largest_power_of_2_less_than(7), 4);
assert_eq!(largest_power_of_2_less_than(8), 4);
assert_eq!(
largest_power_of_2_less_than(u64::max_value()),
9223372036854775808
);
assert_eq!(
largest_power_of_2_less_than(u64::max_value() - 1),
9223372036854775808
);
assert_eq!(
largest_power_of_2_less_than(9223372036854775808u64),
4611686018427387904u64
);
assert_eq!(
largest_power_of_2_less_than(9223372036854775808u64 - 1),
4611686018427387904u64
);
assert_eq!(
largest_power_of_2_less_than(9223372036854775808u64 + 1),
9223372036854775808u64
);
}
#[test]
fn test_count_set_bits() {
assert_eq!(count_set_bits(0), 0);
assert_eq!(count_set_bits(1), 1);
assert_eq!(count_set_bits(2), 1);
assert_eq!(count_set_bits(3), 2);
assert_eq!(count_set_bits(4), 1);
assert_eq!(count_set_bits(u64::max_value()), 64);
assert_eq!(count_set_bits(9223372036854775808), 1);
assert_eq!(count_set_bits(9223372036854775808 - 1), 63);
}
#[test]
fn test_least_significant_set_bit() {
assert_eq!(least_significant_set_bit(0), 0);
assert_eq!(least_significant_set_bit(1), 0);
assert_eq!(least_significant_set_bit(2), 1);
assert_eq!(least_significant_set_bit(3), 0);
assert_eq!(least_significant_set_bit(4), 2);
assert_eq!(least_significant_set_bit(u64::max_value()), 0);
assert_eq!(least_significant_set_bit(9223372036854775808), 63);
assert_eq!(least_significant_set_bit(9223372036854775808 - 1), 0);
}
#[test]
fn test_num_bits() {
assert_eq!(num_bits(0), 0);
assert_eq!(num_bits(1), 1);
assert_eq!(num_bits(2), 2);
assert_eq!(num_bits(3), 2);
assert_eq!(num_bits(4), 3);
assert_eq!(num_bits(5), 3);
assert_eq!(num_bits(6), 3);
assert_eq!(num_bits(7), 3);
assert_eq!(num_bits(8), 4);
assert_eq!(num_bits(u64::max_value()), 64);
assert_eq!(num_bits(9223372036854775808), 64);
assert_eq!(num_bits(9223372036854775808 - 1), 63);
}
#[test]
fn test_break_in_powers_of_2() {
assert_eq!(powers_of_2(1), vec![1]);
assert_eq!(powers_of_2(2), vec![2]);
assert_eq!(powers_of_2(3), vec![2, 1]);
assert_eq!(powers_of_2(4), vec![4]);
assert_eq!(powers_of_2(5), vec![4, 1]);
assert_eq!(powers_of_2(6), vec![4, 2]);
assert_eq!(powers_of_2(7), vec![4, 2, 1]);
assert_eq!(powers_of_2(8), vec![8]);
assert_eq!(powers_of_2(9), vec![8, 1]);
assert_eq!(powers_of_2(10), vec![8, 2]);
assert_eq!(powers_of_2(11), vec![8, 2, 1]);
assert_eq!(powers_of_2(12), vec![8, 4]);
assert_eq!(powers_of_2(13), vec![8, 4, 1]);
assert_eq!(powers_of_2(14), vec![8, 4, 2]);
assert_eq!(powers_of_2(15), vec![8, 4, 2, 1]);
assert_eq!(powers_of_2(16), vec![16]);
}
#[test]
fn test_get_subtree_root_index() {
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_subtree_root_index(0, 2),
0
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_subtree_root_index(2, 4),
1
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_subtree_root_index(0, 4),
2
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_subtree_root_index(4, 6),
3
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_subtree_root_index(6, 8),
4
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_subtree_root_index(4, 8),
5
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_subtree_root_index(0, 8),
6
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_subtree_root_index(8, 10),
7
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_subtree_root_index(10, 12),
8
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_subtree_root_index(8, 12),
9
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_subtree_root_index(12, 14),
10
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_subtree_root_index(14, 16),
11
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_subtree_root_index(12, 16),
12
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_subtree_root_index(8, 16),
13
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_subtree_root_index(0, 16),
14
);
}
#[test]
fn test_hash_db() {
let mut db = InMemoryHashDb::<Vec<u8>>::new();
let hasher = Sha256Hasher {
leaf_data_domain_separator: 0,
node_domain_separator: 1,
};
let l_0 = "leaf_0";
let l_1 = "leaf_1";
let l_2 = "leaf_2";
let l_3 = "leaf_3";
let h_l_0 = hasher.hash_leaf(l_0).unwrap();
let h_l_1 = hasher.hash_leaf(l_1).unwrap();
let h_l_2 = hasher.hash_leaf(l_2).unwrap();
let h_l_3 = hasher.hash_leaf(l_3).unwrap();
db.add_leaf(h_l_0.clone()).unwrap();
db.add_leaf(h_l_1.clone()).unwrap();
db.add_leaf(h_l_2.clone()).unwrap();
db.add_leaf(h_l_3.clone()).unwrap();
let n_0_1 = hasher
.hash_tree_nodes(h_l_0.clone(), h_l_1.clone())
.unwrap();
let n_2_3 = hasher
.hash_tree_nodes(h_l_2.clone(), h_l_3.clone())
.unwrap();
db.add_full_subtree_root(n_0_1.clone()).unwrap();
db.add_full_subtree_root(n_2_3.clone()).unwrap();
assert_eq!(db.get_leaf(0).unwrap(), h_l_0);
assert_eq!(db.get_leaf(1).unwrap(), h_l_1);
assert_eq!(db.get_leaf(2).unwrap(), h_l_2);
assert_eq!(db.get_leaf(3).unwrap(), h_l_3);
assert_eq!(db.get_full_subtree_root(0).unwrap(), n_0_1);
assert_eq!(db.get_full_subtree_root(1).unwrap(), n_2_3);
}
#[test]
fn test_hash_leaves() {
let hasher = Sha256Hasher {
leaf_data_domain_separator: 0,
node_domain_separator: 1,
};
let tree = CompactMerkleTree::new(hasher.clone());
let l_0 = "leaf_0";
let l_1 = "leaf_1";
let l_2 = "leaf_2";
let l_3 = "leaf_3";
let l_4 = "leaf_4";
let l_5 = "leaf_5";
let l_6 = "leaf_6";
let (h_l_0_0, h_n_0_0) = CompactMerkleTree::hash_leaves(&hasher, vec![l_0]).unwrap();
assert_eq!(h_l_0_0.len(), 1);
assert!(h_n_0_0.is_empty());
let h_l_0 = hasher.hash_leaf(l_0).unwrap();
let h_l_1 = hasher.hash_leaf(l_1).unwrap();
let expected_h_l_0_1 = hasher.hash_tree_nodes(h_l_0, h_l_1).unwrap();
let (h_l_0_1, h_n_0_1) = CompactMerkleTree::hash_leaves(&hasher, vec![l_0, l_1]).unwrap();
assert_eq!(h_l_0_1.len(), 2);
assert_eq!(h_n_0_1.len(), 1);
assert_eq!(h_n_0_1[0], expected_h_l_0_1);
let h_l_2 = hasher.hash_leaf(l_2).unwrap();
let (h_l_0_2, h_n_0_2) = CompactMerkleTree::hash_leaves(&hasher, vec![l_0, l_1, l_2]).unwrap();
assert_eq!(h_l_0_2.len(), 3);
assert_eq!(h_n_0_2.len(), 1);
assert_eq!(h_n_0_2[0], expected_h_l_0_1);
let h_l_3 = hasher.hash_leaf(l_3).unwrap();
let expected_h_l_2_3 = hasher.hash_tree_nodes(h_l_2, h_l_3).unwrap();
let expected_h_l_0_3 = hasher
.hash_tree_nodes(expected_h_l_0_1.clone(), expected_h_l_2_3.clone())
.unwrap();
let (h_l_0_3, h_n_0_3) = CompactMerkleTree::hash_leaves(&hasher, vec![l_0, l_1, l_2, l_3]).unwrap();
assert_eq!(h_l_0_3.len(), 4);
assert_eq!(h_n_0_3.len(), 3);
assert_eq!(h_n_0_3[0], expected_h_l_0_1);
assert_eq!(h_n_0_3[1], expected_h_l_2_3);
assert_eq!(h_n_0_3[2], expected_h_l_0_3);
let h_l_4 = hasher.hash_leaf(l_4).unwrap();
let (h_l_0_4, h_n_0_4) = CompactMerkleTree::hash_leaves(&hasher, vec![l_0, l_1, l_2, l_3, l_4]).unwrap();
assert_eq!(h_l_0_4.len(), 5);
assert_eq!(h_n_0_4.len(), 3);
assert_eq!(h_n_0_4[0], expected_h_l_0_1);
assert_eq!(h_n_0_4[1], expected_h_l_2_3);
assert_eq!(h_n_0_4[2], expected_h_l_0_3);
let h_l_5 = hasher.hash_leaf(l_5).unwrap();
let expected_h_l_4_5 = hasher.hash_tree_nodes(h_l_4, h_l_5).unwrap();
let (h_l_0_5, h_n_0_5) = CompactMerkleTree::hash_leaves(&hasher, vec![l_0, l_1, l_2, l_3, l_4, l_5])
.unwrap();
assert_eq!(h_l_0_5.len(), 6);
assert_eq!(h_n_0_5.len(), 4);
assert_eq!(h_n_0_5[0], expected_h_l_0_1);
assert_eq!(h_n_0_5[1], expected_h_l_2_3);
assert_eq!(h_n_0_5[2], expected_h_l_0_3);
assert_eq!(h_n_0_5[3], expected_h_l_4_5);
let h_l_6 = hasher.hash_leaf(l_6).unwrap();
let (h_l_0_6, h_n_0_6) = CompactMerkleTree::hash_leaves(&hasher, vec![l_0, l_1, l_2, l_3, l_4, l_5, l_6])
.unwrap();
assert_eq!(h_l_0_6.len(), 7);
assert_eq!(h_n_0_6.len(), 4);
assert_eq!(h_n_0_6[0], expected_h_l_0_1);
assert_eq!(h_n_0_6[1], expected_h_l_2_3);
assert_eq!(h_n_0_6[2], expected_h_l_0_3);
assert_eq!(h_n_0_6[3], expected_h_l_4_5);
}
#[test]
fn test_append() {
let mut db = InMemoryHashDb::<Vec<u8>>::new();
let hasher = Sha256Hasher {
leaf_data_domain_separator: 0,
node_domain_separator: 1,
};
let mut tree = CompactMerkleTree::new(hasher.clone());
let l_0 = "leaf_0";
let l_1 = "leaf_1";
let l_2 = "leaf_2";
let l_3 = "leaf_3";
let l_4 = "leaf_4";
let l_5 = "leaf_5";
let path_0 = tree.append(l_0, &mut db).unwrap();
assert_eq!(tree.size, 1);
assert_eq!(tree.full_subtree_roots.len(), 1);
assert_eq!(
hasher.hash_leaf(l_0).unwrap(),
tree.get_root_hash().unwrap()
);
assert_eq!(path_0.len(), 0);
let path_1 = tree.append(l_1, &mut db).unwrap();
assert_eq!(tree.size, 2);
assert_eq!(tree.full_subtree_roots.len(), 1);
let rh_2 = tree.get_root_hash().unwrap();
let h_l_1 = hasher.hash_leaf(l_1).unwrap();
assert_eq!(
tree.hasher
.hash_tree_nodes(hasher.hash_leaf(l_0).unwrap(), h_l_1.clone())
.unwrap(),
rh_2
);
assert_eq!(path_1.len(), 1);
assert_eq!(
tree.hasher
.hash_tree_nodes(path_1[0].clone(), h_l_1)
.unwrap(),
rh_2
);
let path_2 = tree.append(l_2, &mut db).unwrap();
assert_eq!(tree.size, 3);
assert_eq!(tree.full_subtree_roots.len(), 2);
let rh_3 = tree.get_root_hash().unwrap();
let h_l_2 = hasher.hash_leaf(l_2).unwrap();
assert_eq!(
tree.hasher
.hash_tree_nodes(rh_2.clone(), h_l_2.clone())
.unwrap(),
rh_3
);
assert_eq!(path_2.len(), 1);
assert_eq!(
tree.hasher
.hash_tree_nodes(path_2[0].clone(), h_l_2.clone())
.unwrap(),
rh_3
);
let path_3 = tree.append(l_3, &mut db).unwrap();
assert_eq!(tree.size, 4);
assert_eq!(tree.full_subtree_roots.len(), 1);
let rh_4 = tree.get_root_hash().unwrap();
assert_eq!(
tree.hasher
.hash_tree_nodes(
rh_2.clone(),
tree.hasher
.hash_tree_nodes(
hasher.hash_leaf(l_2).unwrap(),
hasher.hash_leaf(l_3).unwrap()
)
.unwrap()
)
.unwrap(),
rh_4
);
assert_eq!(path_3.len(), 2);
assert_eq!(path_3[0], h_l_2);
assert_eq!(path_3[1], rh_2);
let path_4 = tree.append(l_4, &mut db).unwrap();
assert_eq!(tree.size, 5);
assert_eq!(tree.full_subtree_roots.len(), 2);
let rh_5 = tree.get_root_hash().unwrap();
assert_eq!(
tree.hasher
.hash_tree_nodes(rh_4.clone(), hasher.hash_leaf(l_4).unwrap())
.unwrap(),
rh_5
);
let path_5 = tree.append(l_5, &mut db).unwrap();
assert_eq!(tree.size, 6);
assert_eq!(tree.full_subtree_roots.len(), 2);
let rh_6 = tree.get_root_hash().unwrap();
assert_eq!(
tree.hasher
.hash_tree_nodes(
rh_4,
tree.hasher
.hash_tree_nodes(
hasher.hash_leaf(l_4).unwrap(),
hasher.hash_leaf(l_5).unwrap()
)
.unwrap()
)
.unwrap(),
rh_6
);
}
#[test]
fn test_inner_border_split() {
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_right_border_and_inner_node_count(
0, 4
),
(0, 2)
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_right_border_and_inner_node_count(
1, 4
),
(0, 2)
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_right_border_and_inner_node_count(
2, 4
),
(1, 1)
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_right_border_and_inner_node_count(
3, 4
),
(2, 0)
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_right_border_and_inner_node_count(
0, 5
),
(0, 3)
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_right_border_and_inner_node_count(
1, 5
),
(0, 3)
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_right_border_and_inner_node_count(
2, 5
),
(0, 3)
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_right_border_and_inner_node_count(
3, 5
),
(0, 3)
);
assert_eq!(
CompactMerkleTree::<&str, Vec<u8>, Sha256Hasher>::get_right_border_and_inner_node_count(
4, 5
),
(1, 0)
);
}
#[test]
fn test_verify_leaf_inclusion_path() {
let mut db = InMemoryHashDb::<Vec<u8>>::new();
let hasher = Sha256Hasher {
leaf_data_domain_separator: 0,
node_domain_separator: 1,
};
let mut tree = CompactMerkleTree::new(hasher.clone());
let l_0 = "leaf_0";
let l_1 = "leaf_1";
let l_2 = "leaf_2";
let l_3 = "leaf_3";
let l_4 = "leaf_4";
let l_5 = "leaf_5";
let l_6 = "leaf_6";
let l_7 = "leaf_7";
let path_0 = tree.append(l_0, &mut db).unwrap();
let rh_1 = tree.get_root_hash().unwrap();
assert!(CompactMerkleTree::verify_leaf_inclusion_proof(
&hasher, 0, l_0, tree.size, &rh_1, path_0
)
.unwrap());
let path_1 = tree.append(l_1, &mut db).unwrap();
let rh_2 = tree.get_root_hash().unwrap();
assert!(CompactMerkleTree::verify_leaf_inclusion_proof(
&hasher, 1, l_1, tree.size, &rh_2, path_1
)
.unwrap());
let path_2 = tree.append(l_2, &mut db).unwrap();
let rh_3 = tree.get_root_hash().unwrap();
assert!(CompactMerkleTree::verify_leaf_inclusion_proof(
&hasher, 2, l_2, tree.size, &rh_3, path_2
)
.unwrap());
let path_3 = tree.append(l_3, &mut db).unwrap();
let rh_4 = tree.get_root_hash().unwrap();
assert!(CompactMerkleTree::verify_leaf_inclusion_proof(
&hasher, 3, l_3, tree.size, &rh_4, path_3
)
.unwrap());
let path_4 = tree.append(l_4, &mut db).unwrap();
let rh_5 = tree.get_root_hash().unwrap();
assert!(CompactMerkleTree::verify_leaf_inclusion_proof(
&hasher, 4, l_4, tree.size, &rh_5, path_4
)
.unwrap());
let path_5 = tree.append(l_5, &mut db).unwrap();
let rh_6 = tree.get_root_hash().unwrap();
assert!(CompactMerkleTree::verify_leaf_inclusion_proof(
&hasher, 5, l_5, tree.size, &rh_6, path_5
)
.unwrap());
let path_6 = tree.append(l_6, &mut db).unwrap();
let rh_7 = tree.get_root_hash().unwrap();
assert!(CompactMerkleTree::verify_leaf_inclusion_proof(
&hasher, 6, l_6, tree.size, &rh_7, path_6
)
.unwrap());
let path_7 = tree.append(l_7, &mut db).unwrap();
let rh_8 = tree.get_root_hash().unwrap();
assert!(CompactMerkleTree::verify_leaf_inclusion_proof(
&hasher, 7, l_7, tree.size, &rh_8, path_7
)
.unwrap());
}
fn verify_inclusion_proof_for_all(
test_cases: usize,
hasher: &Sha256Hasher,
tree: &CompactMerkleTree<&str, Vec<u8>, Sha256Hasher>,
leaf_data: &Vec<String>,
db: &InMemoryHashDb<Vec<u8>>,
) {
let rh = tree.get_root_hash().unwrap();
for i in 0..test_cases {
let path = CompactMerkleTree::get_leaf_inclusion_proof_for_tree_size(
hasher,
i as TreeSizeType,
tree.size,
db,
)
.unwrap();
assert!(CompactMerkleTree::verify_leaf_inclusion_proof(
hasher,
i as TreeSizeType,
&leaf_data[i],
tree.size,
&rh,
path
)
.unwrap());
}
}
#[test]
fn test_append_and_verify_proof() {
let mut db = InMemoryHashDb::<Vec<u8>>::new();
let hasher = Sha256Hasher {
leaf_data_domain_separator: 0,
node_domain_separator: 1,
};
let mut tree = CompactMerkleTree::new(hasher.clone());
let test_cases = 100;
let mut leaf_data = vec![];
for i in 0..test_cases {
leaf_data.push(i.to_string())
}
for i in 0..test_cases {
let path = tree.append(&leaf_data[i], &mut db).unwrap();
let rh = tree.get_root_hash().unwrap();
assert!(CompactMerkleTree::verify_leaf_inclusion_proof(
&hasher,
i as TreeSizeType,
&leaf_data[i],
tree.size,
&rh,
path
)
.unwrap());
verify_inclusion_proof_for_all(i, &hasher, &tree, &leaf_data, &db)
}
let rh = tree.get_root_hash().unwrap();
for i in 0..test_cases {
let path = CompactMerkleTree::get_leaf_inclusion_proof_for_tree_size(
&hasher,
i as TreeSizeType,
tree.size,
&db,
)
.unwrap();
assert!(CompactMerkleTree::verify_leaf_inclusion_proof(
&hasher,
i as TreeSizeType,
&leaf_data[i],
tree.size,
&rh,
path
)
.unwrap());
}
}
#[test]
fn test_extend_and_verify_proof() {
let mut db = InMemoryHashDb::<Vec<u8>>::new();
let hasher = Sha256Hasher {
leaf_data_domain_separator: 0,
node_domain_separator: 1,
};
let mut tree = CompactMerkleTree::new(hasher.clone());
let test_cases = 100;
let mut leaf_data = vec![];
for i in 0..test_cases {
leaf_data.push(i.to_string())
}
tree.extend(vec![&leaf_data[0], &leaf_data[1]], &mut db)
.unwrap();
assert_eq!(tree.size, 2);
verify_inclusion_proof_for_all(tree.size as usize, &hasher, &tree, &leaf_data, &db);
tree.extend(vec![&leaf_data[2]], &mut db).unwrap();
assert_eq!(tree.size, 3);
verify_inclusion_proof_for_all(tree.size as usize, &hasher, &tree, &leaf_data, &db);
tree.extend(vec![&leaf_data[3]], &mut db).unwrap();
assert_eq!(tree.size, 4);
verify_inclusion_proof_for_all(tree.size as usize, &hasher, &tree, &leaf_data, &db);
tree.extend(vec![&leaf_data[4], &leaf_data[5], &leaf_data[6]], &mut db)
.unwrap();
assert_eq!(tree.size, 7);
verify_inclusion_proof_for_all(tree.size as usize, &hasher, &tree, &leaf_data, &db);
tree.extend(
leaf_data[7..15].iter().map(|s| s.as_str()).collect(),
&mut db,
)
.unwrap();
assert_eq!(tree.size, 15);
verify_inclusion_proof_for_all(tree.size as usize, &hasher, &tree, &leaf_data, &db);
tree.extend(
leaf_data[15..32].iter().map(|s| s.as_str()).collect(),
&mut db,
)
.unwrap();
assert_eq!(tree.size, 32);
verify_inclusion_proof_for_all(tree.size as usize, &hasher, &tree, &leaf_data, &db);
tree.extend(
leaf_data[32..64].iter().map(|s| s.as_str()).collect(),
&mut db,
)
.unwrap();
assert_eq!(tree.size, 64);
verify_inclusion_proof_for_all(tree.size as usize, &hasher, &tree, &leaf_data, &db);
tree.extend(
leaf_data[64..70].iter().map(|s| s.as_str()).collect(),
&mut db,
)
.unwrap();
assert_eq!(tree.size, 70);
verify_inclusion_proof_for_all(tree.size as usize, &hasher, &tree, &leaf_data, &db);
tree.extend(
leaf_data[70..80].iter().map(|s| s.as_str()).collect(),
&mut db,
)
.unwrap();
assert_eq!(tree.size, 80);
verify_inclusion_proof_for_all(tree.size as usize, &hasher, &tree, &leaf_data, &db);
tree.extend(
leaf_data[80..100].iter().map(|s| s.as_str()).collect(),
&mut db,
)
.unwrap();
assert_eq!(tree.size, 100);
verify_inclusion_proof_for_all(tree.size as usize, &hasher, &tree, &leaf_data, &db);
}
#[test]
fn test_consistency_proof() {
let mut db = InMemoryHashDb::<Vec<u8>>::new();
let hasher = Sha256Hasher {
leaf_data_domain_separator: 0,
node_domain_separator: 1,
};
let mut tree = CompactMerkleTree::new(hasher.clone());
let test_cases = 100;
let mut leaf_data = vec![];
for i in 0..test_cases {
leaf_data.push(i.to_string())
}
let mut roots = vec![];
for i in 0..test_cases {
tree.append(&leaf_data[i], &mut db).unwrap();
roots.push(tree.get_root_hash().unwrap());
if i > 0 {
for j in 0..i {
let consistency_proof = tree
.get_consistency_proof((j + 1) as TreeSizeType, &db)
.unwrap();
assert!(CompactMerkleTree::verify_consistency_proof(
&hasher,
(j + 1) as TreeSizeType,
(i + 1) as TreeSizeType,
&roots[j],
&roots[i],
consistency_proof
)
.unwrap());
}
}
}
}
#[test]
fn test_create_tree_from_hash_db() {
let mut db = InMemoryHashDb::<Vec<u8>>::new();
let hasher = Sha256Hasher {
leaf_data_domain_separator: 0,
node_domain_separator: 1,
};
let mut tree = CompactMerkleTree::new(hasher.clone());
let test_cases = 100;
let mut leaf_data = vec![];
for i in 0..test_cases {
leaf_data.push(i.to_string())
}
for i in 0..test_cases {
tree.append(&leaf_data[i], &mut db).unwrap();
let new_tree =
CompactMerkleTree::new_from_hash_db(hasher.clone(), tree.size, &db).unwrap();
assert_eq!(new_tree.size, tree.size);
assert_eq!(new_tree.full_subtree_roots, tree.full_subtree_roots);
}
}
}