use std::fmt;
use std::marker::PhantomData;
use ct_merkle::mem_backed_tree::MemoryBackedTree;
use serde::Serialize;
use sha2::{Digest, Sha256};
use crate::error::{TrustError, TrustResult};
use crate::merkle::leaf::{SerializedLeaf, serialize_leaf};
use crate::merkle::proof::{ConsistencyProof, InclusionProof, RootHash};
mod validation;
use validation::{check_consistency_sizes, check_inclusion_index};
pub struct AppendOnlyTree<L> {
inner: MemoryBackedTree<Sha256, SerializedLeaf>,
_marker: PhantomData<fn(L)>,
}
impl<L> AppendOnlyTree<L> {
#[must_use]
pub fn new() -> Self {
Self {
inner: MemoryBackedTree::new(),
_marker: PhantomData,
}
}
#[must_use]
pub fn len(&self) -> u64 {
self.inner.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
#[must_use]
pub fn root(&self) -> RootHash {
RootHash::from_inner(self.inner.root())
}
pub fn prove_inclusion(&self, leaf_index: u64) -> TrustResult<InclusionProof> {
check_inclusion_index(self.inner.len(), leaf_index)?;
let idx = usize::try_from(leaf_index).map_err(|_err| TrustError::MerkleTree {
reason: format!("inclusion proof leaf index {leaf_index} does not fit in usize"),
})?;
Ok(InclusionProof::from_inner(self.inner.prove_inclusion(idx)))
}
pub fn prove_consistency(&self, old_size: u64, new_size: u64) -> TrustResult<ConsistencyProof> {
check_consistency_sizes(self.inner.len(), old_size, new_size)?;
let num_additions = new_size - old_size;
let num_additions_usize =
usize::try_from(num_additions).map_err(|_err| TrustError::MerkleTree {
reason: format!(
"consistency proof num_additions {num_additions} does not fit in usize"
),
})?;
Ok(ConsistencyProof::from_inner(
self.inner.prove_consistency(num_additions_usize),
))
}
}
pub enum RawLeaf {}
impl AppendOnlyTree<RawLeaf> {
pub fn append_raw(&mut self, leaf_bytes: &[u8]) -> u64 {
self.inner
.push(SerializedLeaf::from_raw_bytes(leaf_bytes.to_vec()));
self.inner.len()
}
pub fn reconstruct_from_raw_leaves<I>(leaves: I) -> Self
where
I: IntoIterator,
I::Item: AsRef<[u8]>,
{
let mut tree = Self::new();
for leaf in leaves {
tree.append_raw(leaf.as_ref());
}
tree
}
}
impl<L: Serialize> AppendOnlyTree<L> {
pub fn append(&mut self, leaf: L) -> TrustResult<u64> {
let serialized = serialize_leaf(&leaf)?;
self.inner.push(serialized);
Ok(self.inner.len())
}
pub fn reconstruct_from_leaves(leaves: Vec<L>) -> TrustResult<Self> {
let mut tree = Self::new();
for leaf in leaves {
tree.append(leaf)?;
}
Ok(tree)
}
}
impl<L> Default for AppendOnlyTree<L> {
fn default() -> Self {
Self::new()
}
}
impl<L> fmt::Debug for AppendOnlyTree<L> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("AppendOnlyTree")
.field("num_leaves", &self.inner.len())
.finish()
}
}
const NODE_PREFIX: u8 = 0x01;
fn node_hash(left: &[u8; 32], right: &[u8; 32]) -> [u8; 32] {
let mut hasher = Sha256::new();
hasher.update([NODE_PREFIX]);
hasher.update(left);
hasher.update(right);
hasher.finalize().into()
}
fn split_point(size: u64) -> u64 {
1 << (u64::BITS - 1 - (size - 1).leading_zeros())
}
#[derive(Clone, Default)]
pub struct HashTree {
size: u64,
levels: Vec<Vec<[u8; 32]>>,
}
impl HashTree {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_capacity(leaves: u64) -> Self {
let mut levels = Vec::new();
let mut width = leaves;
while width > 0 {
levels.push(Vec::with_capacity(
usize::try_from(width).unwrap_or_default(),
));
width >>= 1;
}
Self { size: 0, levels }
}
#[must_use]
pub fn len(&self) -> u64 {
self.size
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.size == 0
}
#[must_use]
pub fn retained_bytes(&self) -> usize {
self.levels
.iter()
.map(|level| level.capacity() * std::mem::size_of::<[u8; 32]>())
.sum()
}
pub fn push_leaf_hash(&mut self, leaf_hash: [u8; 32]) -> u64 {
let mut node = leaf_hash;
let mut height = 0;
loop {
if self.levels.len() == height {
self.levels.push(Vec::new());
}
let Some(level) = self.levels.get_mut(height) else {
break;
};
level.push(node);
if level.len() % 2 == 1 {
break;
}
let [.., left, right] = level.as_slice() else {
break;
};
node = node_hash(left, right);
height += 1;
}
self.size = self.size.saturating_add(1);
self.size
}
pub fn root(&self) -> TrustResult<RootHash> {
let root = if self.size == 0 {
Sha256::digest([]).into()
} else {
self.subtree(0, self.size)?
};
Ok(RootHash::from_parts(root, self.size))
}
pub fn prove_inclusion(&self, leaf_index: u64) -> TrustResult<InclusionProof> {
check_inclusion_index(self.size, leaf_index)?;
let mut path = Vec::new();
let (mut start, mut size, mut index) = (0, self.size, leaf_index);
while size > 1 {
let split = split_point(size);
if index < split {
path.push(self.subtree(start + split, size - split)?);
size = split;
} else {
path.push(self.subtree(start, split)?);
start += split;
index -= split;
size -= split;
}
}
InclusionProof::try_from_bytes(path.iter().rev().flatten().copied().collect())
}
pub fn prove_consistency(&self, old_size: u64, new_size: u64) -> TrustResult<ConsistencyProof> {
check_consistency_sizes(self.size, old_size, new_size)?;
let mut path = Vec::new();
let (mut start, mut size, mut old, mut whole) = (0, new_size, old_size, true);
while old != size {
let split = split_point(size);
if old <= split {
path.push(self.subtree(start + split, size - split)?);
size = split;
} else {
path.push(self.subtree(start, split)?);
start += split;
old -= split;
size -= split;
whole = false;
}
}
if !whole {
path.push(self.subtree(start, size)?);
}
ConsistencyProof::try_from_bytes(path.iter().rev().flatten().copied().collect())
}
fn subtree(&self, start: u64, size: u64) -> TrustResult<[u8; 32]> {
if size.is_power_of_two() {
let height = size.trailing_zeros();
let node = if start % size == 0 {
self.stored(height, start >> height)
} else {
None
};
return node.ok_or_else(|| TrustError::MerkleTree {
reason: format!(
"hash tree of {} leaves holds no node for leaves {start}..{}",
self.size,
start.saturating_add(size)
),
});
}
if size == 0 {
return Err(TrustError::MerkleTree {
reason: format!("hash tree asked for an empty subtree at leaf {start}"),
});
}
let split = split_point(size);
let left = self.subtree(start, split)?;
let right = self.subtree(start + split, size - split)?;
Ok(node_hash(&left, &right))
}
fn stored(&self, height: u32, index: u64) -> Option<[u8; 32]> {
let level = self.levels.get(usize::try_from(height).ok()?)?;
level.get(usize::try_from(index).ok()?).copied()
}
}
impl fmt::Debug for HashTree {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("HashTree")
.field("num_leaves", &self.size)
.finish_non_exhaustive()
}
}
#[cfg(test)]
#[path = "tree_tests.rs"]
mod tests;
#[cfg(test)]
#[path = "hash_tree_tests.rs"]
mod hash_tree_tests;