use std::any::Any;
use std::fs::File;
use std::io::Write;
use std::mem::size_of;
use std::path::PathBuf;
use anyhow::ensure;
use filecoin_hashers::{Domain, Hasher, PoseidonArity};
use generic_array::typenum::{Unsigned, U0};
use log::trace;
use merkletree::{
merkle::{
get_merkle_tree_leafs, is_merkle_tree_size_valid, FromIndexedParallelIterator, MerkleTree,
},
store::{DiskStore, ExternalReader, LevelCacheStore, ReplicaConfig, Store, StoreConfig},
};
use rand::Rng;
use rayon::prelude::{IntoParallelIterator, ParallelIterator};
use crate::{
error::{Error, Result},
merkle::{DiskTree, LCTree, MerkleTreeTrait, MerkleTreeWrapper},
util::{data_at_node, default_rows_to_discard, NODE_SIZE},
};
type LCMerkleTree<H, U> = LCTree<H, U, U0, U0>;
pub fn create_disk_tree<Tree: MerkleTreeTrait>(
base_tree_len: usize,
configs: &[StoreConfig],
) -> Result<DiskTree<Tree::Hasher, Tree::Arity, Tree::SubTreeArity, Tree::TopTreeArity>> {
let base_tree_leafs = get_merkle_tree_leafs(base_tree_len, Tree::Arity::to_usize())?;
if Tree::TopTreeArity::to_usize() > 0 {
ensure!(
Tree::SubTreeArity::to_usize() > 0,
"Invalid top arity specified without sub arity"
);
DiskTree::from_sub_tree_store_configs(base_tree_leafs, configs)
} else if Tree::SubTreeArity::to_usize() > 0 {
ensure!(
!configs.is_empty(),
"Cannot create sub-tree with a single tree config"
);
DiskTree::from_store_configs(base_tree_leafs, configs)
} else {
ensure!(configs.len() == 1, "Invalid tree-shape specified");
let store = DiskStore::new_from_disk(base_tree_len, Tree::Arity::to_usize(), &configs[0])?;
DiskTree::from_data_store(store, base_tree_leafs)
}
}
pub fn create_lc_tree<Tree: MerkleTreeTrait>(
base_tree_len: usize,
configs: &[StoreConfig],
replica_config: &ReplicaConfig,
) -> Result<LCTree<Tree::Hasher, Tree::Arity, Tree::SubTreeArity, Tree::TopTreeArity>> {
let base_tree_leafs = get_merkle_tree_leafs(base_tree_len, Tree::Arity::to_usize())?;
if Tree::TopTreeArity::to_usize() > 0 {
ensure!(
Tree::SubTreeArity::to_usize() > 0,
"Invalid top arity specified without sub arity"
);
LCTree::from_sub_tree_store_configs_and_replica(base_tree_leafs, configs, replica_config)
} else if Tree::SubTreeArity::to_usize() > 0 {
ensure!(
!configs.is_empty(),
"Cannot create sub-tree with a single tree config"
);
LCTree::from_store_configs_and_replica(base_tree_leafs, configs, replica_config)
} else {
ensure!(configs.len() == 1, "Invalid tree-shape specified");
let store = LevelCacheStore::new_from_disk_with_reader(
base_tree_len,
Tree::Arity::to_usize(),
&configs[0],
ExternalReader::new_from_path(&replica_config.path)?,
)?;
LCTree::from_data_store(store, base_tree_leafs)
}
}
pub fn create_tree<Tree: MerkleTreeTrait>(
base_tree_len: usize,
configs: &[StoreConfig],
replica_config: Option<&ReplicaConfig>,
) -> Result<
MerkleTreeWrapper<
<Tree as MerkleTreeTrait>::Hasher,
<Tree as MerkleTreeTrait>::Store,
<Tree as MerkleTreeTrait>::Arity,
<Tree as MerkleTreeTrait>::SubTreeArity,
<Tree as MerkleTreeTrait>::TopTreeArity,
>,
>
where
Tree::Store: 'static,
{
let base_tree_leafs = get_base_tree_leafs::<Tree>(base_tree_len)?;
let mut trees = Vec::with_capacity(configs.len());
for i in 0..configs.len() {
let mut store = Tree::Store::new_with_config(
base_tree_len,
Tree::Arity::to_usize(),
configs[i].clone(),
)?;
if let Some(lc_store) = <dyn Any>::downcast_mut::<
LevelCacheStore<<Tree::Hasher as Hasher>::Domain, File>,
>(&mut store)
{
ensure!(
replica_config.is_some(),
"Cannot create LCTree without replica paths"
);
let replica_config = replica_config.expect("replica config failure");
lc_store.set_external_reader(ExternalReader::new_from_config(replica_config, i)?)?;
}
if configs.len() == 1 {
return MerkleTreeWrapper::<
Tree::Hasher,
Tree::Store,
Tree::Arity,
Tree::SubTreeArity,
Tree::TopTreeArity,
>::from_data_store(store, base_tree_leafs);
} else {
trees.push(MerkleTreeWrapper::<
Tree::Hasher,
Tree::Store,
Tree::Arity,
U0,
U0,
>::from_data_store(store, base_tree_leafs)?);
}
}
ensure!(
Tree::TopTreeArity::to_usize() > 0 || Tree::SubTreeArity::to_usize() > 0,
"Cannot have a sub/top tree without more than 1 config"
);
if Tree::TopTreeArity::to_usize() > 0 {
ensure!(
Tree::SubTreeArity::to_usize() > 0,
"Invalid top arity specified without sub arity"
);
MerkleTreeWrapper::<
Tree::Hasher,
Tree::Store,
Tree::Arity,
Tree::SubTreeArity,
Tree::TopTreeArity,
>::from_sub_trees_as_trees(trees)
} else {
ensure!(
!configs.is_empty(),
"Cannot create sub-tree with a single tree config"
);
MerkleTreeWrapper::from_trees(trees)
}
}
pub fn create_base_merkle_tree<Tree: MerkleTreeTrait>(
config: Option<StoreConfig>,
size: usize,
data: &[u8],
) -> Result<Tree> {
ensure!(
data.len() == (NODE_SIZE * size),
Error::InvalidMerkleTreeArgs(data.len(), NODE_SIZE, size)
);
trace!("create_merkle_tree called with size {}", size);
trace!(
"is_merkle_tree_size_valid({}, arity {}) = {}",
size,
Tree::Arity::to_usize(),
is_merkle_tree_size_valid(size, Tree::Arity::to_usize())
);
ensure!(
is_merkle_tree_size_valid(size, Tree::Arity::to_usize()),
"Invalid merkle tree size given the arity"
);
let f = |i| {
let d = data_at_node(data, i).expect("data_at_node math failed");
<Tree::Hasher as Hasher>::Domain::try_from_bytes(d)
.expect("failed to convert node data to domain element")
};
let tree = match config {
Some(x) => MerkleTree::<
<Tree::Hasher as Hasher>::Domain,
<Tree::Hasher as Hasher>::Function,
Tree::Store,
Tree::Arity,
Tree::SubTreeArity,
Tree::TopTreeArity,
>::from_par_iter_with_config((0..size).into_par_iter().map(f), x),
None => MerkleTree::<
<Tree::Hasher as Hasher>::Domain,
<Tree::Hasher as Hasher>::Function,
Tree::Store,
Tree::Arity,
Tree::SubTreeArity,
Tree::TopTreeArity,
>::from_par_iter((0..size).into_par_iter().map(f)),
}?;
Ok(Tree::from_merkle(tree))
}
pub fn create_base_lcmerkle_tree<H: Hasher, BaseTreeArity: 'static + PoseidonArity>(
config: StoreConfig,
size: usize,
data: &[u8],
replica_config: &ReplicaConfig,
) -> Result<LCMerkleTree<H, BaseTreeArity>> {
trace!("create_base_lcmerkle_tree called with size {}", size);
trace!(
"is_merkle_tree_size_valid({}, arity {}) = {}",
size,
BaseTreeArity::to_usize(),
is_merkle_tree_size_valid(size, BaseTreeArity::to_usize())
);
ensure!(
is_merkle_tree_size_valid(size, BaseTreeArity::to_usize()),
"Invalid merkle tree size given the arity"
);
ensure!(
data.len() == size * size_of::<H::Domain>(),
"Invalid data length for merkle tree"
);
let f = |i| {
let d = data_at_node(data, i)?;
H::Domain::try_from_bytes(d)
};
let mut lc_tree: LCMerkleTree<H, BaseTreeArity> =
LCMerkleTree::<H, BaseTreeArity>::try_from_iter_with_config((0..size).map(f), config)?;
lc_tree.set_external_reader_path(&replica_config.path)?;
Ok(lc_tree)
}
pub fn split_config(config: StoreConfig, count: usize) -> Result<Vec<StoreConfig>> {
if count == 1 {
return Ok(vec![config]);
}
let configs = (0..count)
.map(|i| StoreConfig {
path: config.path.clone(),
id: format!("{}-{}", config.id, i),
size: config.size,
rows_to_discard: config.rows_to_discard,
})
.collect();
Ok(configs)
}
pub fn split_config_and_replica(
config: StoreConfig,
replica_path: PathBuf,
sub_tree_width: usize, count: usize,
) -> Result<(Vec<StoreConfig>, ReplicaConfig)> {
if count == 1 {
return Ok((
vec![config],
ReplicaConfig {
path: replica_path,
offsets: vec![0],
},
));
}
let (configs, replica_offsets) = (0..count)
.map(|i| {
(
StoreConfig {
path: config.path.clone(),
id: format!("{}-{}", config.id, i),
size: config.size,
rows_to_discard: config.rows_to_discard,
},
i * sub_tree_width * NODE_SIZE,
)
})
.unzip();
Ok((
configs,
ReplicaConfig {
path: replica_path,
offsets: replica_offsets,
},
))
}
pub fn get_base_tree_count<Tree: MerkleTreeTrait>() -> usize {
if Tree::TopTreeArity::to_usize() == 0 && Tree::SubTreeArity::to_usize() == 0 {
return 1;
}
if Tree::TopTreeArity::to_usize() > 0 {
assert!(Tree::SubTreeArity::to_usize() != 0);
Tree::TopTreeArity::to_usize() * Tree::SubTreeArity::to_usize()
} else {
Tree::SubTreeArity::to_usize()
}
}
pub fn get_base_tree_leafs<Tree: MerkleTreeTrait>(base_tree_size: usize) -> Result<usize> {
get_merkle_tree_leafs(base_tree_size, Tree::Arity::to_usize())
}
pub type ResTree<Tree> = MerkleTreeWrapper<
<Tree as MerkleTreeTrait>::Hasher,
<Tree as MerkleTreeTrait>::Store,
<Tree as MerkleTreeTrait>::Arity,
<Tree as MerkleTreeTrait>::SubTreeArity,
<Tree as MerkleTreeTrait>::TopTreeArity,
>;
fn generate_base_tree<R: Rng, Tree: MerkleTreeTrait>(
rng: &mut R,
nodes: usize,
temp_path: Option<PathBuf>,
) -> (Vec<u8>, ResTree<Tree>)
where
Tree::Store: 'static,
{
let elements = (0..nodes)
.map(|_| <Tree::Hasher as Hasher>::Domain::random(rng))
.collect::<Vec<_>>();
let mut data = Vec::new();
for el in &elements {
data.extend_from_slice(AsRef::<[u8]>::as_ref(el));
}
if let Some(ref temp_path) = temp_path {
let id: u64 = rng.gen();
let replica_path = temp_path.join(format!("replica-path-{}", id));
let config = StoreConfig::new(
temp_path,
format!("test-lc-tree-{}", id),
default_rows_to_discard(nodes, Tree::Arity::to_usize()),
);
let mut tree =
MerkleTreeWrapper::try_from_iter_with_config(elements.iter().map(|v| (Ok(*v))), config)
.expect("try from iter with config failure");
let mut f = File::create(&replica_path).expect("replica file create failure");
f.write_all(&data).expect("replica file write failure");
{
if let Some(lc_tree) = <dyn Any>::downcast_mut::<
MerkleTree<
<Tree::Hasher as Hasher>::Domain,
<Tree::Hasher as Hasher>::Function,
LevelCacheStore<<Tree::Hasher as Hasher>::Domain, File>,
Tree::Arity,
Tree::SubTreeArity,
Tree::TopTreeArity,
>,
>(&mut tree.inner)
{
lc_tree
.set_external_reader_path(&replica_path)
.expect("lc tree set external reader failure");
}
}
(data, tree)
} else {
(
data,
MerkleTreeWrapper::try_from_iter(elements.iter().map(|v| Ok(*v)))
.expect("try from iter map failure"),
)
}
}
fn generate_sub_tree<R: Rng, Tree: MerkleTreeTrait>(
rng: &mut R,
nodes: usize,
temp_path: Option<PathBuf>,
) -> (Vec<u8>, ResTree<Tree>)
where
Tree::Store: 'static,
{
let base_tree_count = Tree::SubTreeArity::to_usize();
let base_tree_size = nodes / base_tree_count;
let mut trees = Vec::with_capacity(base_tree_count);
let mut data = Vec::new();
for _ in 0..base_tree_count {
let (inner_data, tree) = generate_base_tree::<
R,
MerkleTreeWrapper<Tree::Hasher, Tree::Store, Tree::Arity>,
>(rng, base_tree_size, temp_path.clone());
trees.push(tree);
data.extend(inner_data);
}
(
data,
MerkleTreeWrapper::from_trees(trees).expect("from trees failure"),
)
}
pub fn generate_tree<Tree: MerkleTreeTrait, R: Rng>(
rng: &mut R,
nodes: usize,
temp_path: Option<PathBuf>,
) -> (Vec<u8>, ResTree<Tree>)
where
Tree::Store: 'static,
{
let sub_tree_arity = Tree::SubTreeArity::to_usize();
let top_tree_arity = Tree::TopTreeArity::to_usize();
if top_tree_arity > 0 {
assert!(
sub_tree_arity != 0,
"malformed tree with TopTreeArity > 0 and SubTreeARity == 0"
);
let mut sub_trees = Vec::with_capacity(top_tree_arity);
let mut data = Vec::new();
for _i in 0..top_tree_arity {
let (inner_data, tree) = generate_sub_tree::<
R,
MerkleTreeWrapper<Tree::Hasher, Tree::Store, Tree::Arity, Tree::SubTreeArity, U0>,
>(rng, nodes / top_tree_arity, temp_path.clone());
sub_trees.push(tree);
data.extend(inner_data);
}
(
data,
MerkleTreeWrapper::from_sub_trees(sub_trees).expect("from sub trees failure"),
)
} else if sub_tree_arity > 0 {
generate_sub_tree::<R, Tree>(rng, nodes, temp_path)
} else {
generate_base_tree::<R, Tree>(rng, nodes, temp_path)
}
}