use std::{convert::Infallible, fmt};
use proptest::proptest;
use rand::{prelude::*, thread_rng};
use crate::{MMR, Merge, util::MemStore};
#[derive(Eq, PartialEq, Clone, Default)]
struct NumberRange {
start: u32,
end: u32,
}
struct MergeNumberRange;
impl fmt::Debug for NumberRange {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "NumberRange({}, {})", self.start, self.end)
}
}
impl fmt::Debug for MergeNumberRange {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "MergeNumberRange")
}
}
impl From<u32> for NumberRange {
fn from(num: u32) -> Self {
Self {
start: num,
end: num,
}
}
}
impl NumberRange {
fn is_normalized(&self) -> bool {
self.start <= self.end
}
}
impl Merge for MergeNumberRange {
type Item = NumberRange;
type Error = Infallible;
fn leaf_hash(data: &[u8]) -> Result<Self::Item, Self::Error> {
let num = data
.get(..4)
.and_then(|b| b.try_into().ok())
.map_or(0, u32::from_le_bytes);
Ok(Self::Item {
start: num,
end: num,
})
}
fn merge_pos(
_pos: u64,
left: &Self::Item,
right: &Self::Item,
) -> Result<Self::Item, Self::Error> {
Ok(Self::Item {
start: left.start,
end: right.end,
})
}
fn merge_peaks(
left: &Self::Item,
right: &Self::Item,
) -> Result<Self::Item, Self::Error> {
Self::merge_pos(0, right, left)
}
}
async fn test_sequence_sub_func(count: u32, proof_elem: Vec<u32>) {
let store = MemStore::default();
let mut mmr = MMR::<MergeNumberRange, _>::new(0, store);
let mut positions: Vec<u64> = Vec::new();
for i in 0..count {
let pos = mmr.push(&i.to_le_bytes()).await.expect("push");
positions.push(pos);
}
let root = mmr.get_root().await.expect("get_root");
assert!(root.is_normalized());
let proof = mmr
.gen_proof(
proof_elem
.iter()
.map(|elem| positions[*elem as usize])
.collect(),
)
.await
.expect("gen_proof");
for item in proof.proof_items() {
assert!(item.is_normalized());
}
mmr.commit().await.expect("commit");
let result = proof
.verify(
&root,
proof_elem
.iter()
.map(|elem| {
(
positions[*elem as usize],
MergeNumberRange::leaf_hash(&elem.to_le_bytes()).unwrap(),
)
})
.collect(),
)
.expect("verify");
assert!(result);
}
proptest! {
#[test]
fn test_sequence(count in 10u32..500u32) {
tokio::runtime::Runtime::new().unwrap().block_on(async {
let mut leaves: Vec<u32> = (0..count).collect();
let mut rng = thread_rng();
leaves.shuffle(&mut rng);
let leaves_count = rng.gen_range(1..count - 1);
leaves.truncate(leaves_count as usize);
test_sequence_sub_func(count, leaves).await;
});
}
}