use std::marker::PhantomData;
use std::path::PathBuf;
use std::sync::{Arc, RwLock};
use anyhow::{Context, Result};
use log::debug;
use rayon::prelude::*;
use typenum::marker_traits::Unsigned;
use typenum::{U0, U2};
use crate::hash::{Algorithm, Hashable};
use crate::proof::Proof;
use crate::store::{
ExternalReader, LevelCacheStore, Store, StoreConfig, VecStore, BUILD_CHUNK_NODES,
};
pub const BUILD_DATA_BLOCK_SIZE: usize = 64 * BUILD_CHUNK_NODES;
#[derive(Clone, Eq, PartialEq)]
#[allow(clippy::enum_variant_names)]
enum Data<E: Element, A: Algorithm<E>, S: Store<E>, BaseTreeArity: Unsigned, SubTreeArity: Unsigned>
{
BaseTree(S),
SubTree(Vec<MerkleTree<E, A, S, BaseTreeArity>>),
TopTree(Vec<MerkleTree<E, A, S, BaseTreeArity, SubTreeArity>>),
}
impl<E: Element, A: Algorithm<E>, S: Store<E>, BaseTreeArity: Unsigned, SubTreeArity: Unsigned>
Data<E, A, S, BaseTreeArity, SubTreeArity>
{
fn store(&self) -> Option<&S> {
match self {
Data::BaseTree(s) => Some(s),
_ => None,
}
}
fn store_mut(&mut self) -> Option<&mut S> {
match self {
Data::BaseTree(s) => Some(s),
_ => None,
}
}
#[allow(clippy::type_complexity)]
fn sub_trees(&self) -> Option<&Vec<MerkleTree<E, A, S, BaseTreeArity, SubTreeArity>>> {
match self {
Data::TopTree(s) => Some(s),
_ => None,
}
}
fn base_trees(&self) -> Option<&Vec<MerkleTree<E, A, S, BaseTreeArity>>> {
match self {
Data::SubTree(s) => Some(s),
_ => None,
}
}
}
impl<E: Element, A: Algorithm<E>, S: Store<E>, BaseTreeArity: Unsigned, SubTreeArity: Unsigned>
std::fmt::Debug for Data<E, A, S, BaseTreeArity, SubTreeArity>
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("enum Data").finish()
}
}
#[allow(clippy::type_complexity)]
#[derive(Clone, Eq, PartialEq)]
pub struct MerkleTree<E, A, S, BaseTreeArity = U2, SubTreeArity = U0, TopTreeArity = U0>
where
E: Element,
A: Algorithm<E>,
S: Store<E>,
BaseTreeArity: Unsigned,
SubTreeArity: Unsigned,
TopTreeArity: Unsigned,
{
data: Data<E, A, S, BaseTreeArity, SubTreeArity>,
leafs: usize,
len: usize,
height: usize,
root: E,
_a: PhantomData<A>,
_e: PhantomData<E>,
_bta: PhantomData<BaseTreeArity>,
_sta: PhantomData<SubTreeArity>,
_tta: PhantomData<TopTreeArity>,
}
impl<
E: Element,
A: Algorithm<E>,
S: Store<E>,
BaseTreeArity: Unsigned,
SubTreeArity: Unsigned,
TopTreeArity: Unsigned,
> std::fmt::Debug for MerkleTree<E, A, S, BaseTreeArity, SubTreeArity, TopTreeArity>
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("MerkleTree")
.field("data", &self.data)
.field("leafs", &self.leafs)
.field("len", &self.len)
.field("height", &self.height)
.field("root", &self.root)
.finish()
}
}
pub trait Element: Ord + Clone + AsRef<[u8]> + Sync + Send + Default + std::fmt::Debug {
fn byte_len() -> usize;
fn from_slice(bytes: &[u8]) -> Self;
fn copy_to_slice(&self, bytes: &mut [u8]);
}
impl<
E: Element,
A: Algorithm<E>,
BaseTreeArity: Unsigned,
SubTreeArity: Unsigned,
TopTreeArity: Unsigned,
>
MerkleTree<E, A, LevelCacheStore<E, std::fs::File>, BaseTreeArity, SubTreeArity, TopTreeArity>
{
pub fn set_external_reader_path(&mut self, path: &PathBuf) -> Result<()> {
ensure!(self.data.store_mut().is_some(), "store data required");
self.data
.store_mut()
.unwrap()
.set_external_reader(ExternalReader::new_from_path(path)?)
}
#[allow(clippy::type_complexity)]
pub fn from_store_configs_and_replicas(
leafs: usize,
configs: &[StoreConfig],
replica_paths: &[PathBuf],
) -> Result<
MerkleTree<
E,
A,
LevelCacheStore<E, std::fs::File>,
BaseTreeArity,
SubTreeArity,
TopTreeArity,
>,
> {
let branches = BaseTreeArity::to_usize();
let mut trees = Vec::with_capacity(configs.len());
ensure!(
configs.len() == replica_paths.len(),
"Config and Replica list lengths are invalid"
);
for i in 0..configs.len() {
let config = &configs[i];
let replica_path = &replica_paths[i];
let data = LevelCacheStore::new_from_disk_with_reader(
get_merkle_tree_len(leafs, branches)?,
branches,
config,
ExternalReader::new_from_path(replica_path)?,
)
.context("failed to instantiate levelcache store")?;
trees.push(
MerkleTree::<E, A, LevelCacheStore<_, _>, BaseTreeArity>::from_data_store(
data, leafs,
)?,
);
}
Self::from_trees(trees)
}
#[allow(clippy::type_complexity)]
pub fn from_sub_tree_store_configs_and_replicas(
leafs: usize,
configs: &[StoreConfig],
replica_paths: &[PathBuf],
) -> Result<
MerkleTree<
E,
A,
LevelCacheStore<E, std::fs::File>,
BaseTreeArity,
SubTreeArity,
TopTreeArity,
>,
> {
ensure!(
configs.len() == replica_paths.len(),
"Config and Replica list lengths are invalid"
);
let sub_tree_count = TopTreeArity::to_usize();
let mut start = 0;
let mut end = configs.len() / sub_tree_count;
let mut trees = Vec::with_capacity(sub_tree_count);
for _ in 0..sub_tree_count {
trees.push(MerkleTree::<
E,
A,
LevelCacheStore<_, _>,
BaseTreeArity,
SubTreeArity,
>::from_store_configs_and_replicas(
leafs,
&configs[start..end],
&replica_paths[start..end],
)?);
start = end;
end += configs.len() / sub_tree_count;
}
Self::from_sub_trees(trees)
}
}
impl<
E: Element,
A: Algorithm<E>,
S: Store<E>,
BaseTreeArity: Unsigned,
SubTreeArity: Unsigned,
TopTreeArity: Unsigned,
> MerkleTree<E, A, S, BaseTreeArity, SubTreeArity, TopTreeArity>
{
pub fn new<I: IntoIterator<Item = E>>(
data: I,
) -> Result<MerkleTree<E, A, S, BaseTreeArity, SubTreeArity, TopTreeArity>> {
Self::try_from_iter(data.into_iter().map(Ok))
}
pub fn new_with_config<I: IntoIterator<Item = E>>(
data: I,
config: StoreConfig,
) -> Result<MerkleTree<E, A, S, BaseTreeArity, SubTreeArity, TopTreeArity>> {
Self::try_from_iter_with_config(data.into_iter().map(Ok), config)
}
pub fn from_data<O: Hashable<A>, I: IntoIterator<Item = O>>(
data: I,
) -> Result<MerkleTree<E, A, S, BaseTreeArity, SubTreeArity, TopTreeArity>> {
let mut a = A::default();
Self::try_from_iter(data.into_iter().map(|x| {
a.reset();
x.hash(&mut a);
Ok(a.hash())
}))
}
pub fn from_data_with_config<O: Hashable<A>, I: IntoIterator<Item = O>>(
data: I,
config: StoreConfig,
) -> Result<MerkleTree<E, A, S, BaseTreeArity, SubTreeArity, TopTreeArity>> {
let mut a = A::default();
Self::try_from_iter_with_config(
data.into_iter().map(|x| {
a.reset();
x.hash(&mut a);
Ok(a.hash())
}),
config,
)
}
pub fn from_data_store(
data: S,
leafs: usize,
) -> Result<MerkleTree<E, A, S, BaseTreeArity, SubTreeArity, TopTreeArity>> {
ensure!(
SubTreeArity::to_usize() == 0,
"Data stores must not have sub-tree layers"
);
ensure!(
TopTreeArity::to_usize() == 0,
"Data stores must not have a top layer"
);
let branches = BaseTreeArity::to_usize();
ensure!(next_pow2(leafs) == leafs, "leafs MUST be a power of 2");
ensure!(
next_pow2(branches) == branches,
"branches MUST be a power of 2"
);
let tree_len = get_merkle_tree_len(leafs, branches)?;
ensure!(tree_len == data.len(), "Inconsistent tree data");
ensure!(
is_merkle_tree_size_valid(leafs, branches),
"MerkleTree size is invalid given the arity"
);
let height = get_merkle_tree_height(leafs, branches);
let root = data.read_at(data.len() - 1)?;
Ok(MerkleTree {
data: Data::BaseTree(data),
leafs,
len: tree_len,
height,
root,
_a: PhantomData,
_e: PhantomData,
_bta: PhantomData,
_sta: PhantomData,
_tta: PhantomData,
})
}
pub fn from_tree_slice(
data: &[u8],
leafs: usize,
) -> Result<MerkleTree<E, A, S, BaseTreeArity, SubTreeArity, TopTreeArity>> {
ensure!(
SubTreeArity::to_usize() == 0,
"Data slice must not have sub-tree layers"
);
ensure!(
TopTreeArity::to_usize() == 0,
"Data slice must not have a top layer"
);
let branches = BaseTreeArity::to_usize();
let height = get_merkle_tree_height(leafs, branches);
let tree_len = get_merkle_tree_len(leafs, branches)?;
ensure!(
tree_len == data.len() / E::byte_len(),
"Inconsistent tree data"
);
ensure!(
is_merkle_tree_size_valid(leafs, branches),
"MerkleTree size is invalid given the arity"
);
let store = S::new_from_slice(tree_len, &data).context("failed to create data store")?;
let root = store.read_at(data.len() - 1)?;
Ok(MerkleTree {
data: Data::BaseTree(store),
leafs,
len: tree_len,
height,
root,
_a: PhantomData,
_e: PhantomData,
_bta: PhantomData,
_sta: PhantomData,
_tta: PhantomData,
})
}
pub fn from_tree_slice_with_config(
data: &[u8],
leafs: usize,
config: StoreConfig,
) -> Result<MerkleTree<E, A, S, BaseTreeArity, SubTreeArity, TopTreeArity>> {
ensure!(
SubTreeArity::to_usize() == 0,
"Data slice must not have sub-tree layers"
);
ensure!(
TopTreeArity::to_usize() == 0,
"Data slice must not have a top layer"
);
let branches = BaseTreeArity::to_usize();
let height = get_merkle_tree_height(leafs, branches);
let tree_len = get_merkle_tree_len(leafs, branches)?;
ensure!(
tree_len == data.len() / E::byte_len(),
"Inconsistent tree data"
);
ensure!(
is_merkle_tree_size_valid(leafs, branches),
"MerkleTree size is invalid given the arity"
);
let store = S::new_from_slice_with_config(tree_len, branches, &data, config)
.context("failed to create data store")?;
let root = store.read_at(data.len() - 1)?;
Ok(MerkleTree {
data: Data::BaseTree(store),
leafs,
len: tree_len,
height,
root,
_a: PhantomData,
_e: PhantomData,
_bta: PhantomData,
_sta: PhantomData,
_tta: PhantomData,
})
}
pub fn from_trees(
trees: Vec<MerkleTree<E, A, S, BaseTreeArity>>,
) -> Result<MerkleTree<E, A, S, BaseTreeArity, SubTreeArity, TopTreeArity>> {
ensure!(
SubTreeArity::to_usize() > 0,
"Cannot use from_trees if not constructing a structure with sub-trees"
);
ensure!(
trees.iter().all(|ref mt| mt.height() == trees[0].height()),
"All passed in trees must have the same height"
);
ensure!(
trees.iter().all(|ref mt| mt.len() == trees[0].len()),
"All passed in trees must have the same length"
);
let sub_tree_layer_nodes = SubTreeArity::to_usize();
ensure!(
trees.len() == sub_tree_layer_nodes,
"Length of trees MUST equal the number of sub tree layer nodes"
);
let (leafs, len, height, root) = if sub_tree_layer_nodes == 0 {
(
trees[0].leafs(),
trees[0].len(),
trees[0].height(),
trees[0].root(),
)
} else {
let leafs = trees.iter().fold(0, |leafs, mt| leafs + mt.leafs());
let len = trees.iter().fold(0, |len, mt| len + mt.len()) + 1;
let height = trees[0].height() + 1;
let roots: Vec<E> = trees.iter().map(|x| x.root()).collect();
let root = A::default().multi_node(&roots, 1);
(leafs, len, height, root)
};
Ok(MerkleTree {
data: Data::SubTree(trees),
leafs,
len,
height,
root,
_a: PhantomData,
_e: PhantomData,
_bta: PhantomData,
_sta: PhantomData,
_tta: PhantomData,
})
}
pub fn from_sub_trees(
trees: Vec<MerkleTree<E, A, S, BaseTreeArity, SubTreeArity>>,
) -> Result<MerkleTree<E, A, S, BaseTreeArity, SubTreeArity, TopTreeArity>> {
ensure!(
TopTreeArity::to_usize() > 0,
"Cannot use from_sub_trees if not constructing a structure with sub-trees"
);
ensure!(
trees.iter().all(|ref mt| mt.height() == trees[0].height()),
"All passed in trees must have the same height"
);
ensure!(
trees.iter().all(|ref mt| mt.len() == trees[0].len()),
"All passed in trees must have the same length"
);
let top_layer_nodes = TopTreeArity::to_usize();
ensure!(
trees.len() == top_layer_nodes,
"Length of trees MUST equal the number of top layer nodes"
);
let (leafs, len, height, root) = {
let leafs = trees.iter().fold(0, |leafs, mt| leafs + mt.leafs());
let len = trees.iter().fold(0, |len, mt| len + mt.len()) + 1;
let height = trees[0].height() + 1;
let roots: Vec<E> = trees.iter().map(|x| x.root()).collect();
let root = A::default().multi_node(&roots, 1);
(leafs, len, height, root)
};
Ok(MerkleTree {
data: Data::TopTree(trees),
leafs,
len,
height,
root,
_a: PhantomData,
_e: PhantomData,
_bta: PhantomData,
_sta: PhantomData,
_tta: PhantomData,
})
}
pub fn from_sub_trees_as_trees(
mut trees: Vec<MerkleTree<E, A, S, BaseTreeArity>>,
) -> Result<MerkleTree<E, A, S, BaseTreeArity, SubTreeArity, TopTreeArity>> {
ensure!(
TopTreeArity::to_usize() > 0,
"Cannot use from_sub_trees if not constructing a structure with sub-trees"
);
ensure!(
trees.iter().all(|ref mt| mt.height() == trees[0].height()),
"All passed in trees must have the same height"
);
ensure!(
trees.iter().all(|ref mt| mt.len() == trees[0].len()),
"All passed in trees must have the same length"
);
let sub_tree_count = TopTreeArity::to_usize();
let top_layer_nodes = sub_tree_count * SubTreeArity::to_usize();
ensure!(
trees.len() == top_layer_nodes,
"Length of trees MUST equal the number of top layer nodes"
);
let mut grouped_trees = Vec::with_capacity(sub_tree_count);
for _ in (0..sub_tree_count).step_by(trees.len() / sub_tree_count) {
grouped_trees.push(trees.split_off(trees.len() / sub_tree_count));
}
grouped_trees.insert(0, trees);
let mut sub_trees: Vec<MerkleTree<E, A, S, BaseTreeArity, SubTreeArity>> =
Vec::with_capacity(sub_tree_count);
for group in grouped_trees {
sub_trees.push(MerkleTree::from_trees(group)?);
}
let (leafs, len, height, root) = {
let leafs = sub_trees.iter().fold(0, |leafs, mt| leafs + mt.leafs());
let len = sub_trees.iter().fold(0, |len, mt| len + mt.len()) + 1;
let height = sub_trees[0].height() + 1;
let roots: Vec<E> = sub_trees.iter().map(|x| x.root()).collect();
let root = A::default().multi_node(&roots, 1);
(leafs, len, height, root)
};
Ok(MerkleTree {
data: Data::TopTree(sub_trees),
leafs,
len,
height,
root,
_a: PhantomData,
_e: PhantomData,
_bta: PhantomData,
_sta: PhantomData,
_tta: PhantomData,
})
}
pub fn from_slices(
tree_data: &[&[u8]],
leafs: usize,
) -> Result<MerkleTree<E, A, S, BaseTreeArity, SubTreeArity>> {
let mut trees = Vec::with_capacity(tree_data.len());
for data in tree_data {
trees.push(MerkleTree::<E, A, S, BaseTreeArity>::from_tree_slice(
data, leafs,
)?);
}
MerkleTree::from_trees(trees)
}
pub fn from_slices_with_configs(
tree_data: &[&[u8]],
leafs: usize,
configs: &[StoreConfig],
) -> Result<MerkleTree<E, A, S, BaseTreeArity, SubTreeArity, TopTreeArity>> {
let mut trees = Vec::with_capacity(tree_data.len());
for i in 0..tree_data.len() {
trees.push(
MerkleTree::<E, A, S, BaseTreeArity>::from_tree_slice_with_config(
tree_data[i],
leafs,
configs[i].clone(),
)?,
);
}
MerkleTree::from_trees(trees)
}
pub fn from_stores(
leafs: usize,
stores: Vec<S>,
) -> Result<MerkleTree<E, A, S, BaseTreeArity, SubTreeArity, TopTreeArity>> {
let mut trees = Vec::with_capacity(stores.len());
for store in stores {
trees.push(MerkleTree::<E, A, S, BaseTreeArity>::from_data_store(
store, leafs,
)?);
}
MerkleTree::from_trees(trees)
}
pub fn from_store_configs(
leafs: usize,
configs: &[StoreConfig],
) -> Result<MerkleTree<E, A, S, BaseTreeArity, SubTreeArity, TopTreeArity>> {
let branches = BaseTreeArity::to_usize();
let mut trees = Vec::with_capacity(configs.len());
for config in configs {
let data = S::new_with_config(
get_merkle_tree_len(leafs, branches)?,
branches,
config.clone(),
)
.context("failed to create data store")?;
trees.push(MerkleTree::<E, A, S, BaseTreeArity>::from_data_store(
data, leafs,
)?);
}
MerkleTree::from_trees(trees)
}
#[allow(clippy::type_complexity)]
pub fn from_sub_tree_store_configs(
leafs: usize,
configs: &[StoreConfig],
) -> Result<MerkleTree<E, A, S, BaseTreeArity, SubTreeArity, TopTreeArity>> {
let tree_count = TopTreeArity::to_usize();
let mut start = 0;
let mut end = configs.len() / tree_count;
let mut trees = Vec::with_capacity(tree_count);
for _ in 0..tree_count {
trees.push(
MerkleTree::<E, A, S, BaseTreeArity, SubTreeArity>::from_store_configs(
leafs,
&configs[start..end],
)?,
);
start = end;
end += configs.len() / tree_count;
}
Self::from_sub_trees(trees)
}
#[inline]
fn build_partial_tree(
mut data: VecStore<E>,
leafs: usize,
height: usize,
) -> Result<MerkleTree<E, A, VecStore<E>, BaseTreeArity>> {
let root = VecStore::build::<A, BaseTreeArity>(&mut data, leafs, height, None)?;
let branches = BaseTreeArity::to_usize();
let tree_len = get_merkle_tree_len(leafs, branches)?;
ensure!(tree_len == Store::len(&data), "Inconsistent tree data");
ensure!(
is_merkle_tree_size_valid(leafs, branches),
"MerkleTree size is invalid given the arity"
);
Ok(MerkleTree {
data: Data::BaseTree(data),
leafs,
len: tree_len,
height,
root,
_a: PhantomData,
_e: PhantomData,
_bta: PhantomData,
_sta: PhantomData,
_tta: PhantomData,
})
}
#[inline]
fn gen_sub_tree_proof(
&self,
i: usize,
top_layer: bool,
arity: usize,
) -> Result<Proof<E, BaseTreeArity>> {
ensure!(arity != 0, "Invalid sub-tree arity");
let tree_index = i / (self.leafs / arity);
let sub_tree_proof = if top_layer {
ensure!(self.data.sub_trees().is_some(), "sub trees required");
let sub_trees = self.data.sub_trees().unwrap();
ensure!(arity == sub_trees.len(), "Top layer tree shape mis-match");
let tree = &sub_trees[tree_index];
let leaf_index = i % tree.leafs();
tree.gen_proof(leaf_index)
} else {
ensure!(self.data.base_trees().is_some(), "base trees required");
let base_trees = self.data.base_trees().unwrap();
ensure!(arity == base_trees.len(), "Sub tree layer shape mis-match");
let tree = &base_trees[tree_index];
let leaf_index = i % tree.leafs();
tree.gen_proof(leaf_index)
}?;
let mut path: Vec<usize> = Vec::with_capacity(1);
let mut lemma: Vec<E> = Vec::with_capacity(arity);
for i in 0..arity {
if i != tree_index {
lemma.push(if top_layer {
ensure!(self.data.sub_trees().is_some(), "sub trees required");
let sub_trees = self.data.sub_trees().unwrap();
sub_trees[i].root()
} else {
ensure!(self.data.base_trees().is_some(), "base trees required");
let base_trees = self.data.base_trees().unwrap();
base_trees[i].root()
});
}
}
lemma.push(self.root());
path.push(tree_index);
Proof::new::<TopTreeArity, SubTreeArity>(Some(Box::new(sub_tree_proof)), lemma, path)
}
#[inline]
pub fn gen_proof(&self, i: usize) -> Result<Proof<E, BaseTreeArity>> {
match &self.data {
Data::TopTree(_) => self.gen_sub_tree_proof(i, true, TopTreeArity::to_usize()),
Data::SubTree(_) => self.gen_sub_tree_proof(i, false, SubTreeArity::to_usize()),
Data::BaseTree(_) => {
ensure!(
i < self.leafs,
"{} is out of bounds (max: {})",
i,
self.leafs
);
let mut base = 0;
let mut j = i;
let mut width = self.leafs;
let branches = BaseTreeArity::to_usize();
ensure!(width == next_pow2(width), "Must be a power of 2 tree");
ensure!(
branches == next_pow2(branches),
"branches must be a power of 2"
);
let shift = log2_pow2(branches);
let mut lemma: Vec<E> =
Vec::with_capacity(get_merkle_proof_lemma_len(self.height, branches));
let mut path: Vec<usize> = Vec::with_capacity(self.height - 1);
ensure!(
SubTreeArity::to_usize() == 0,
"Data slice must not have sub-tree layers"
);
ensure!(
TopTreeArity::to_usize() == 0,
"Data slice must not have a top layer"
);
lemma.push(self.read_at(j)?);
while base + 1 < self.len() {
let hash_index = (j / branches) * branches;
for k in hash_index..hash_index + branches {
if k != j {
lemma.push(self.read_at(base + k)?)
}
}
path.push(j % branches);
base += width;
width >>= shift;
j >>= shift;
}
lemma.push(self.root());
ensure!(
lemma.len() == get_merkle_proof_lemma_len(self.height, branches),
"Invalid proof lemma length"
);
ensure!(path.len() == self.height - 1, "Invalid proof path length");
Proof::new::<U0, U0>(None, lemma, path)
}
}
}
fn gen_cached_top_tree_proof<Arity: Unsigned>(
&self,
i: usize,
levels: usize,
) -> Result<Proof<E, BaseTreeArity>> {
ensure!(Arity::to_usize() != 0, "Invalid top-tree arity");
ensure!(
i < self.leafs,
"{} is out of bounds (max: {})",
i,
self.leafs
);
ensure!(self.data.sub_trees().is_some(), "sub trees required");
let trees = &self.data.sub_trees().unwrap();
let tree_index = i / (self.leafs / Arity::to_usize());
let tree = &trees[tree_index];
let tree_leafs = tree.leafs();
let leaf_index = i % tree_leafs;
let sub_tree_proof = tree.gen_cached_proof(leaf_index, levels)?;
let mut path: Vec<usize> = Vec::with_capacity(1);
let mut lemma: Vec<E> = Vec::with_capacity(Arity::to_usize());
for i in 0..Arity::to_usize() {
if i != tree_index {
lemma.push(trees[i].root())
}
}
lemma.push(self.root());
path.push(tree_index);
Proof::new::<TopTreeArity, SubTreeArity>(Some(Box::new(sub_tree_proof)), lemma, path)
}
fn gen_cached_sub_tree_proof<Arity: Unsigned>(
&self,
i: usize,
levels: usize,
) -> Result<Proof<E, BaseTreeArity>> {
ensure!(Arity::to_usize() != 0, "Invalid sub-tree arity");
ensure!(
i < self.leafs,
"{} is out of bounds (max: {})",
i,
self.leafs
);
ensure!(self.data.base_trees().is_some(), "base trees required");
let trees = &self.data.base_trees().unwrap();
let tree_index = i / (self.leafs / Arity::to_usize());
let tree = &trees[tree_index];
let tree_leafs = tree.leafs();
let leaf_index = i % tree_leafs;
let sub_tree_proof = tree.gen_cached_proof(leaf_index, levels)?;
let mut path: Vec<usize> = Vec::with_capacity(1);
let mut lemma: Vec<E> = Vec::with_capacity(Arity::to_usize());
for i in 0..Arity::to_usize() {
if i != tree_index {
lemma.push(trees[i].root())
}
}
lemma.push(self.root());
path.push(tree_index);
Proof::new::<TopTreeArity, SubTreeArity>(Some(Box::new(sub_tree_proof)), lemma, path)
}
#[allow(clippy::type_complexity)]
pub fn gen_cached_proof(&self, i: usize, levels: usize) -> Result<Proof<E, BaseTreeArity>> {
match &self.data {
Data::TopTree(_) => self.gen_cached_top_tree_proof::<TopTreeArity>(i, levels),
Data::SubTree(_) => self.gen_cached_sub_tree_proof::<SubTreeArity>(i, levels),
Data::BaseTree(_) => {
ensure!(
i < self.leafs,
"{} is out of bounds (max: {})",
i,
self.leafs
);
ensure!(
self.leafs == next_pow2(self.leafs),
"The size of the data layer must be a power of 2"
);
let branches = BaseTreeArity::to_usize();
let total_size = get_merkle_tree_len(self.leafs, branches)?;
let cache_size = get_merkle_tree_cache_size(self.leafs, branches, levels)?;
ensure!(
cache_size < total_size,
"Generate a partial proof with all data available?"
);
let cached_leafs = get_merkle_tree_leafs(cache_size, branches)?;
ensure!(
cached_leafs == next_pow2(cached_leafs),
"The size of the cached leafs must be a power of 2"
);
let cache_height = get_merkle_tree_height(cached_leafs, branches);
let partial_height = self.height - cache_height + 1;
let segment_width = self.leafs / cached_leafs;
let segment_start = (i / segment_width) * segment_width;
let segment_end = segment_start + segment_width;
debug!("leafs {}, branches {}, total size {}, total height {}, cache_size {}, cached levels above base {}, \
partial_height {}, cached_leafs {}, segment_width {}, segment range {}-{} for {}",
self.leafs, branches, total_size, self.height, cache_size, levels, partial_height,
cached_leafs, segment_width, segment_start, segment_end, i);
let mut data_copy = vec![0; segment_width * E::byte_len()];
ensure!(self.data.store().is_some(), "store data required");
self.data.store().unwrap().read_range_into(
segment_start,
segment_end,
&mut data_copy,
)?;
let partial_store = VecStore::new_from_slice(segment_width, &data_copy)?;
ensure!(
Store::len(&partial_store) == segment_width,
"Inconsistent store length"
);
data_copy.resize(
get_merkle_tree_len(segment_width, branches)? * E::byte_len(),
0,
);
let partial_tree: MerkleTree<E, A, VecStore<E>, BaseTreeArity> =
Self::build_partial_tree(partial_store, segment_width, partial_height)?;
ensure!(
partial_height == partial_tree.height(),
"Inconsistent partial tree height"
);
let proof = self.gen_proof_with_partial_tree(i, levels, &partial_tree)?;
debug!(
"generated partial_tree of height {} and len {} with {} branches for proof at {}",
partial_tree.height,
partial_tree.len(),
branches,
i
);
Ok(proof)
}
}
}
fn gen_proof_with_partial_tree(
&self,
i: usize,
levels: usize,
partial_tree: &MerkleTree<E, A, VecStore<E>, BaseTreeArity>,
) -> Result<Proof<E, BaseTreeArity>> {
ensure!(
i < self.leafs,
"{} is out of bounds (max: {})",
i,
self.leafs
);
let mut width = self.leafs;
let branches = BaseTreeArity::to_usize();
ensure!(width == next_pow2(width), "Must be a power of 2 tree");
ensure!(
branches == next_pow2(branches),
"branches must be a power of 2"
);
let data_width = width;
let total_size = get_merkle_tree_len(data_width, branches)?;
let cache_size = get_merkle_tree_cache_size(self.leafs, branches, levels)?;
let cache_index_start = total_size - cache_size;
let cached_leafs = get_merkle_tree_leafs(cache_size, branches)?;
ensure!(
cached_leafs == next_pow2(cached_leafs),
"Cached leafs size must be a power of 2"
);
let mut segment_width = width / cached_leafs;
let segment_start = (i / segment_width) * segment_width;
let shift = log2_pow2(branches);
let mut segment_shift = segment_start;
let mut j = i;
let mut base = 0;
let mut partial_base = 0;
let mut lemma: Vec<E> =
Vec::with_capacity(get_merkle_proof_lemma_len(self.height, branches));
let mut path: Vec<usize> = Vec::with_capacity(self.height - 1);
ensure!(
SubTreeArity::to_usize() == 0,
"Data slice must not have sub-tree layers"
);
ensure!(
TopTreeArity::to_usize() == 0,
"Data slice must not have a top layer"
);
lemma.push(self.read_at(j)?);
while base + 1 < self.len() {
let hash_index = (j / branches) * branches;
for k in hash_index..hash_index + branches {
if k != j {
let read_index = base + k;
lemma.push(
if read_index < data_width || read_index >= cache_index_start {
self.read_at(base + k)?
} else {
let read_index = partial_base + k - segment_shift;
partial_tree.read_at(read_index)?
},
);
}
}
path.push(j % branches);
base += width;
width >>= shift;
partial_base += segment_width;
segment_width >>= shift;
segment_shift >>= shift;
j >>= shift;
}
lemma.push(self.root());
ensure!(
lemma.len() == get_merkle_proof_lemma_len(self.height, branches),
"Invalid proof lemma length"
);
ensure!(path.len() == self.height - 1, "Invalid proof path length");
Proof::new::<U0, U0>(None, lemma, path)
}
#[inline]
pub fn root(&self) -> E {
self.root.clone()
}
#[inline]
pub fn len(&self) -> usize {
match &self.data {
Data::TopTree(_) => self.len,
Data::SubTree(_) => self.len,
Data::BaseTree(store) => store.len(),
}
}
#[inline]
pub fn compact(&mut self, config: StoreConfig, store_version: u32) -> Result<bool> {
let branches = BaseTreeArity::to_usize();
ensure!(self.data.store_mut().is_some(), "store data required");
self.data
.store_mut()
.unwrap()
.compact(branches, config, store_version)
}
#[inline]
pub fn reinit(&mut self) -> Result<()> {
ensure!(self.data.store_mut().is_some(), "store data required");
self.data.store_mut().unwrap().reinit()
}
#[inline]
pub fn delete(&self, config: StoreConfig) -> Result<()> {
S::delete(config)
}
#[inline]
pub fn is_empty(&self) -> bool {
match &self.data {
Data::TopTree(_) => true,
Data::SubTree(_) => true,
Data::BaseTree(store) => store.is_empty(),
}
}
#[inline]
pub fn height(&self) -> usize {
self.height
}
#[inline]
pub fn leafs(&self) -> usize {
self.leafs
}
#[inline]
pub fn data(&self) -> Option<&S> {
match &self.data {
Data::TopTree(_) => None,
Data::SubTree(_) => None,
Data::BaseTree(store) => Some(store),
}
}
#[inline]
pub fn read_at(&self, i: usize) -> Result<E> {
match &self.data {
Data::TopTree(sub_trees) => {
ensure!(
TopTreeArity::to_usize() == sub_trees.len(),
"Top layer tree shape mis-match"
);
let tree_index = i / (self.leafs / TopTreeArity::to_usize());
let tree = &sub_trees[tree_index];
let tree_leafs = tree.leafs();
let leaf_index = i % tree_leafs;
tree.read_at(leaf_index)
}
Data::SubTree(base_trees) => {
ensure!(
SubTreeArity::to_usize() == base_trees.len(),
"Sub-tree shape mis-match"
);
let tree_index = i / (self.leafs / SubTreeArity::to_usize());
let tree = &base_trees[tree_index];
let tree_leafs = tree.leafs();
let leaf_index = i % tree_leafs;
tree.read_at(leaf_index)
}
Data::BaseTree(data) => {
data.read_at(i)
}
}
}
pub fn read_range(&self, start: usize, end: usize) -> Result<Vec<E>> {
ensure!(start < end, "start must be less than end");
ensure!(self.data.store().is_some(), "store data required");
self.data.store().unwrap().read_range(start..end)
}
pub fn read_range_into(&self, start: usize, end: usize, buf: &mut [u8]) -> Result<()> {
ensure!(start < end, "start must be less than end");
ensure!(self.data.store().is_some(), "store data required");
self.data.store().unwrap().read_range_into(start, end, buf)
}
pub fn read_into(&self, pos: usize, buf: &mut [u8]) -> Result<()> {
ensure!(self.data.store().is_some(), "store data required");
self.data.store().unwrap().read_into(pos, buf)
}
pub fn from_byte_slice_with_config(leafs: &[u8], config: StoreConfig) -> Result<Self> {
ensure!(
leafs.len() % E::byte_len() == 0,
"{} ist not a multiple of {}",
leafs.len(),
E::byte_len()
);
let leafs_count = leafs.len() / E::byte_len();
let branches = BaseTreeArity::to_usize();
ensure!(leafs_count > 1, "not enough leaves");
ensure!(
next_pow2(leafs_count) == leafs_count,
"size MUST be a power of 2"
);
ensure!(
next_pow2(branches) == branches,
"branches MUST be a power of 2"
);
let size = get_merkle_tree_len(leafs_count, branches)?;
let height = get_merkle_tree_height(leafs_count, branches);
let mut data = S::new_from_slice_with_config(size, branches, leafs, config.clone())
.context("failed to create data store")?;
let root = S::build::<A, BaseTreeArity>(&mut data, leafs_count, height, Some(config))?;
Ok(MerkleTree {
data: Data::BaseTree(data),
leafs: leafs_count,
len: size,
height,
root,
_a: PhantomData,
_e: PhantomData,
_bta: PhantomData,
_sta: PhantomData,
_tta: PhantomData,
})
}
pub fn from_byte_slice(leafs: &[u8]) -> Result<Self> {
ensure!(
leafs.len() % E::byte_len() == 0,
"{} is not a multiple of {}",
leafs.len(),
E::byte_len()
);
let leafs_count = leafs.len() / E::byte_len();
let branches = BaseTreeArity::to_usize();
ensure!(leafs_count > 1, "not enough leaves");
ensure!(
next_pow2(leafs_count) == leafs_count,
"size MUST be a power of 2"
);
ensure!(
next_pow2(branches) == branches,
"branches MUST be a power of 2"
);
let size = get_merkle_tree_len(leafs_count, branches)?;
let height = get_merkle_tree_height(leafs_count, branches);
let mut data = S::new_from_slice(size, leafs).context("failed to create data store")?;
let root = S::build::<A, BaseTreeArity>(&mut data, leafs_count, height, None)?;
Ok(MerkleTree {
data: Data::BaseTree(data),
leafs: leafs_count,
len: size,
height,
root,
_a: PhantomData,
_e: PhantomData,
_bta: PhantomData,
_sta: PhantomData,
_tta: PhantomData,
})
}
}
pub trait FromIndexedParallelIterator<E, BaseTreeArity>: Sized
where
E: Send,
{
fn from_par_iter<I>(par_iter: I) -> Result<Self>
where
BaseTreeArity: Unsigned,
I: IntoParallelIterator<Item = E>,
I::Iter: IndexedParallelIterator;
fn from_par_iter_with_config<I>(par_iter: I, config: StoreConfig) -> Result<Self>
where
I: IntoParallelIterator<Item = E>,
I::Iter: IndexedParallelIterator,
BaseTreeArity: Unsigned;
}
impl<
E: Element,
A: Algorithm<E>,
S: Store<E>,
BaseTreeArity: Unsigned,
SubTreeArity: Unsigned,
TopTreeArity: Unsigned,
> FromIndexedParallelIterator<E, BaseTreeArity>
for MerkleTree<E, A, S, BaseTreeArity, SubTreeArity, TopTreeArity>
{
fn from_par_iter<I>(into: I) -> Result<Self>
where
I: IntoParallelIterator<Item = E>,
I::Iter: IndexedParallelIterator,
{
let iter = into.into_par_iter();
let leafs = iter.opt_len().expect("must be sized");
let branches = BaseTreeArity::to_usize();
ensure!(leafs > 1, "not enough leaves");
ensure!(next_pow2(leafs) == leafs, "size MUST be a power of 2");
ensure!(
next_pow2(branches) == branches,
"branches MUST be a power of 2"
);
let size = get_merkle_tree_len(leafs, branches)?;
let height = get_merkle_tree_height(leafs, branches);
let mut data = S::new(size).expect("failed to create data store");
populate_data_par::<E, A, S, BaseTreeArity, _>(&mut data, iter)?;
let root = S::build::<A, BaseTreeArity>(&mut data, leafs, height, None)?;
Ok(MerkleTree {
data: Data::BaseTree(data),
leafs,
len: size,
height,
root,
_a: PhantomData,
_e: PhantomData,
_bta: PhantomData,
_sta: PhantomData,
_tta: PhantomData,
})
}
fn from_par_iter_with_config<I>(into: I, config: StoreConfig) -> Result<Self>
where
BaseTreeArity: Unsigned,
I: IntoParallelIterator<Item = E>,
I::Iter: IndexedParallelIterator,
{
let iter = into.into_par_iter();
let leafs = iter.opt_len().expect("must be sized");
let branches = BaseTreeArity::to_usize();
ensure!(leafs > 1, "not enough leaves");
ensure!(next_pow2(leafs) == leafs, "size MUST be a power of 2");
ensure!(
next_pow2(branches) == branches,
"branches MUST be a power of 2"
);
let size = get_merkle_tree_len(leafs, branches)?;
let height = get_merkle_tree_height(leafs, branches);
let mut data = S::new_with_config(size, branches, config.clone())
.context("failed to create data store")?;
if data.loaded_from_disk() {
let root = data.last().context("failed to read root")?;
return Ok(MerkleTree {
data: Data::BaseTree(data),
leafs,
len: size,
height,
root,
_a: PhantomData,
_e: PhantomData,
_bta: PhantomData,
_sta: PhantomData,
_tta: PhantomData,
});
}
populate_data_par::<E, A, S, BaseTreeArity, _>(&mut data, iter)?;
let root = S::build::<A, BaseTreeArity>(&mut data, leafs, height, Some(config))?;
Ok(MerkleTree {
data: Data::BaseTree(data),
leafs,
len: size,
height,
root,
_a: PhantomData,
_e: PhantomData,
_bta: PhantomData,
_sta: PhantomData,
_tta: PhantomData,
})
}
}
impl<
E: Element,
A: Algorithm<E>,
S: Store<E>,
BaseTreeArity: Unsigned,
SubTreeArity: Unsigned,
TopTreeArity: Unsigned,
> MerkleTree<E, A, S, BaseTreeArity, SubTreeArity, TopTreeArity>
{
pub fn try_from_iter<I: IntoIterator<Item = Result<E>>>(into: I) -> Result<Self> {
let iter = into.into_iter();
let (_, n) = iter.size_hint();
let leafs = n.ok_or_else(|| anyhow!("could not get size hint from iterator"))?;
let branches = BaseTreeArity::to_usize();
ensure!(leafs > 1, "not enough leaves");
ensure!(next_pow2(leafs) == leafs, "size MUST be a power of 2");
ensure!(
next_pow2(branches) == branches,
"branches MUST be a power of 2"
);
let size = get_merkle_tree_len(leafs, branches)?;
let height = get_merkle_tree_height(leafs, branches);
let mut data = S::new(size).context("failed to create data store")?;
populate_data::<E, A, S, BaseTreeArity, I>(&mut data, iter)
.context("failed to populate data")?;
let root = S::build::<A, BaseTreeArity>(&mut data, leafs, height, None)?;
Ok(MerkleTree {
data: Data::BaseTree(data),
leafs,
len: size,
height,
root,
_a: PhantomData,
_e: PhantomData,
_bta: PhantomData,
_sta: PhantomData,
_tta: PhantomData,
})
}
pub fn try_from_iter_with_config<I: IntoIterator<Item = Result<E>>>(
into: I,
config: StoreConfig,
) -> Result<Self> {
let iter = into.into_iter();
let (_, n) = iter.size_hint();
let leafs = n.ok_or_else(|| anyhow!("could not get size hint from iterator"))?;
let branches = BaseTreeArity::to_usize();
ensure!(leafs > 1, "not enough leaves");
ensure!(next_pow2(leafs) == leafs, "size MUST be a power of 2");
ensure!(
next_pow2(branches) == branches,
"branches MUST be a power of 2"
);
let size = get_merkle_tree_len(leafs, branches)?;
let height = get_merkle_tree_height(leafs, branches);
let mut data = S::new_with_config(size, branches, config.clone())
.context("failed to create data store")?;
if data.loaded_from_disk() {
let root = data.last().context("failed to read root")?;
return Ok(MerkleTree {
data: Data::BaseTree(data),
leafs,
len: size,
height,
root,
_a: PhantomData,
_e: PhantomData,
_bta: PhantomData,
_sta: PhantomData,
_tta: PhantomData,
});
}
populate_data::<E, A, S, BaseTreeArity, I>(&mut data, iter)
.expect("failed to populate data");
let root = S::build::<A, BaseTreeArity>(&mut data, leafs, height, Some(config))?;
Ok(MerkleTree {
data: Data::BaseTree(data),
leafs,
len: size,
height,
root,
_a: PhantomData,
_e: PhantomData,
_bta: PhantomData,
_sta: PhantomData,
_tta: PhantomData,
})
}
}
impl Element for [u8; 32] {
fn byte_len() -> usize {
32
}
fn from_slice(bytes: &[u8]) -> Self {
if bytes.len() != 32 {
panic!("invalid length {}, expected 32", bytes.len());
}
*array_ref!(bytes, 0, 32)
}
fn copy_to_slice(&self, bytes: &mut [u8]) {
bytes.copy_from_slice(self);
}
}
pub fn get_merkle_tree_len(leafs: usize, branches: usize) -> Result<usize> {
ensure!(leafs >= branches, "leaf and branch mis-match");
ensure!(
branches == next_pow2(branches),
"branches must be a power of 2"
);
if branches == 2 {
ensure!(leafs == next_pow2(leafs), "leafs must be a power of 2");
return Ok(2 * leafs - 1);
}
let mut len = leafs;
let mut cur = leafs;
let shift = log2_pow2(branches);
if shift == 0 {
return Ok(len);
}
while cur > 0 {
cur >>= shift;
ensure!(cur < leafs, "invalid input provided");
len += cur;
}
Ok(len)
}
pub fn get_merkle_tree_cache_size(leafs: usize, branches: usize, levels: usize) -> Result<usize> {
let shift = log2_pow2(branches);
let len = get_merkle_tree_len(leafs, branches)?;
let mut height = get_merkle_tree_height(leafs, branches);
let stop_height = height - levels;
let mut cache_size = len;
let mut cur_leafs = leafs;
while height > stop_height {
cache_size -= cur_leafs;
cur_leafs >>= shift;
height -= 1;
}
Ok(cache_size)
}
pub fn is_merkle_tree_size_valid(leafs: usize, branches: usize) -> bool {
if branches == 0 || leafs != next_pow2(leafs) || branches != next_pow2(branches) {
return false;
}
let mut cur = leafs;
let shift = log2_pow2(branches);
while cur != 1 {
cur >>= shift;
if cur > leafs || cur == 0 {
return false;
}
}
true
}
pub fn get_merkle_tree_height(leafs: usize, branches: usize) -> usize {
if branches == 2 {
(leafs * branches).trailing_zeros() as usize
} else {
(branches as f64 * leafs as f64).log(branches as f64) as usize
}
}
pub fn get_merkle_proof_lemma_len(height: usize, branches: usize) -> usize {
2 + ((branches - 1) * (height - 1))
}
pub fn get_merkle_tree_leafs(len: usize, branches: usize) -> Result<usize> {
ensure!(
branches == next_pow2(branches),
"branches must be a power of 2"
);
let leafs = {
if branches == 2 {
(len >> 1) + 1
} else {
let mut leafs = 1;
let mut cur = len;
let shift = log2_pow2(branches);
while cur != 1 {
leafs <<= shift;
ensure!(
cur > leafs,
"Invalid tree length provided for the specified arity"
);
cur -= leafs;
ensure!(
cur < len,
"Invalid tree length provided for the specified arity"
);
}
leafs
}
};
ensure!(
leafs == next_pow2(leafs),
"Invalid tree length provided for the specified arity"
);
Ok(leafs)
}
pub fn next_pow2(n: usize) -> usize {
n.next_power_of_two()
}
pub fn log2_pow2(n: usize) -> usize {
n.trailing_zeros() as usize
}
pub fn populate_data<
E: Element,
A: Algorithm<E>,
S: Store<E>,
BaseTreeArity: Unsigned,
I: IntoIterator<Item = Result<E>>,
>(
data: &mut S,
iter: <I as std::iter::IntoIterator>::IntoIter,
) -> Result<()> {
if !data.is_empty() {
return Ok(());
}
let mut buf = Vec::with_capacity(BUILD_DATA_BLOCK_SIZE * E::byte_len());
let mut a = A::default();
for item in iter {
let item = item?;
a.reset();
buf.extend(a.leaf(item).as_ref());
if buf.len() >= BUILD_DATA_BLOCK_SIZE * E::byte_len() {
let data_len = data.len();
data.copy_from_slice(&buf, data_len)?;
buf.clear();
}
}
let data_len = data.len();
data.copy_from_slice(&buf, data_len)?;
data.sync()?;
Ok(())
}
fn populate_data_par<E, A, S, BaseTreeArity, I>(data: &mut S, iter: I) -> Result<()>
where
E: Element,
A: Algorithm<E>,
S: Store<E>,
BaseTreeArity: Unsigned,
I: ParallelIterator<Item = E> + IndexedParallelIterator,
{
if !data.is_empty() {
return Ok(());
}
let store = Arc::new(RwLock::new(data));
iter.chunks(BUILD_DATA_BLOCK_SIZE)
.enumerate()
.try_for_each(|(index, chunk)| {
let mut a = A::default();
let mut buf = Vec::with_capacity(BUILD_DATA_BLOCK_SIZE * E::byte_len());
for item in chunk {
a.reset();
buf.extend(a.leaf(item).as_ref());
}
store
.write()
.unwrap()
.copy_from_slice(&buf[..], BUILD_DATA_BLOCK_SIZE * index)
})?;
store.write().unwrap().sync()?;
Ok(())
}
#[test]
fn test_get_merkle_tree_methods() {
assert!(get_merkle_tree_len(16, 4).is_ok());
assert!(get_merkle_tree_len(3, 1).is_ok());
assert!(get_merkle_tree_len(0, 0).is_err());
assert!(get_merkle_tree_len(1, 0).is_err());
assert!(get_merkle_tree_len(1, 2).is_err());
assert!(get_merkle_tree_len(4, 16).is_err());
assert!(get_merkle_tree_len(1024, 11).is_err());
assert!(get_merkle_tree_leafs(31, 2).is_ok());
assert!(get_merkle_tree_leafs(15, 2).is_ok());
assert!(get_merkle_tree_leafs(127, 2).is_ok());
assert!(get_merkle_tree_leafs(1398101, 4).is_ok());
assert!(get_merkle_tree_leafs(299593, 8).is_ok());
assert!(get_merkle_tree_leafs(32, 2).is_err());
assert!(get_merkle_tree_leafs(16, 2).is_err());
assert!(get_merkle_tree_leafs(128, 2).is_err());
assert!(get_merkle_tree_leafs(32, 8).is_err());
assert!(get_merkle_tree_leafs(16, 8).is_err());
assert!(get_merkle_tree_leafs(128, 8).is_err());
assert!(get_merkle_tree_leafs(1398102, 4).is_err());
assert!(get_merkle_tree_leafs(299594, 8).is_err());
}