use std::{iter, marker, ops::Range, u64};
use croaring::Bitmap;
use crate::core::hash::{Hash, ZERO_HASH};
use crate::core::merkle_proof::MerkleProof;
use crate::core::pmmr::{Backend, ReadonlyPMMR};
use crate::core::BlockHeader;
use crate::ser::{PMMRIndexHashable, PMMRable};
pub trait ReadablePMMR {
type Item;
fn get_hash(&self, pos: u64) -> Option<Hash>;
fn get_data(&self, pos: u64) -> Option<Self::Item>;
fn get_from_file(&self, pos: u64) -> Option<Hash>;
fn get_peak_from_file(&self, pos: u64) -> Option<Hash>;
fn get_data_from_file(&self, pos: u64) -> Option<Self::Item>;
fn unpruned_size(&self) -> u64;
fn leaf_pos_iter(&self) -> Box<dyn Iterator<Item = u64> + '_>;
fn leaf_idx_iter(&self, from_idx: u64) -> Box<dyn Iterator<Item = u64> + '_>;
fn n_unpruned_leaves(&self) -> u64;
fn n_unpruned_leaves_to_index(&self, to_index: u64) -> u64;
fn is_empty(&self) -> bool {
self.unpruned_size() == 0
}
fn bag_the_rhs(&self, peak_pos0: u64) -> Option<Hash> {
let size = self.unpruned_size();
let rhs = peaks(size)
.into_iter()
.filter(|&x| x > peak_pos0)
.filter_map(|x| self.get_from_file(x));
let mut res = None;
for peak in rhs.rev() {
res = match res {
None => Some(peak),
Some(rhash) => Some((peak, rhash).hash_with_index(size)),
}
}
res
}
fn peaks(&self) -> Vec<Hash> {
peaks(self.unpruned_size())
.into_iter()
.filter_map(move |pi0| self.get_peak_from_file(pi0))
.collect()
}
fn peak_path(&self, peak_pos0: u64) -> Vec<Hash> {
let rhs = self.bag_the_rhs(peak_pos0);
let mut res = peaks(self.unpruned_size())
.into_iter()
.filter(|&x| x < peak_pos0)
.filter_map(|x| self.get_peak_from_file(x))
.collect::<Vec<_>>();
if let Some(rhs) = rhs {
res.push(rhs);
}
res.reverse();
res
}
fn root(&self) -> Result<Hash, String> {
if self.is_empty() {
return Ok(ZERO_HASH);
}
let mut res = None;
let peaks = self.peaks();
let mmr_size = self.unpruned_size();
for peak in peaks.into_iter().rev() {
res = match res {
None => Some(peak),
Some(rhash) => Some((peak, rhash).hash_with_index(mmr_size)),
}
}
res.ok_or_else(|| "no root, invalid tree".to_owned())
}
fn merkle_proof(&self, pos0: u64) -> Result<MerkleProof, String> {
let size = self.unpruned_size();
debug!("merkle_proof {}, size {}", pos0, size);
if !is_leaf(pos0) {
return Err(format!("not a leaf at pos {}", pos0));
}
self.get_hash(pos0)
.ok_or_else(|| format!("no element at pos {}", pos0))?;
let family_branch = family_branch(pos0, size);
let mut path = family_branch
.iter()
.filter_map(|x| self.get_from_file(x.1))
.collect::<Vec<_>>();
let peak_pos = match family_branch.last() {
Some(&(x, _)) => x,
None => pos0,
};
path.append(&mut self.peak_path(peak_pos));
Ok(MerkleProof {
mmr_size: size,
path,
})
}
}
pub struct PMMR<'a, T, B>
where
T: PMMRable,
B: Backend<T>,
{
pub size: u64,
backend: &'a mut B,
_marker: marker::PhantomData<T>,
}
impl<'a, T, B> PMMR<'a, T, B>
where
T: PMMRable,
B: 'a + Backend<T>,
{
pub fn new(backend: &'a mut B) -> PMMR<'_, T, B> {
PMMR {
backend,
size: 0,
_marker: marker::PhantomData,
}
}
pub fn at(backend: &'a mut B, size: u64) -> PMMR<'_, T, B> {
PMMR {
backend,
size,
_marker: marker::PhantomData,
}
}
pub fn readonly_pmmr(&self) -> ReadonlyPMMR<'_, T, B> {
ReadonlyPMMR::at(&self.backend, self.size)
}
pub fn push(&mut self, leaf: &T) -> Result<u64, String> {
let leaf_pos = self.size;
let mut current_hash = leaf.hash_with_index(leaf_pos);
let mut hashes = vec![current_hash];
let mut pos = leaf_pos;
let (peak_map, height) = peak_map_height(pos);
if height != 0 {
return Err(format!("bad mmr size {}", pos));
}
let mut peak = 1;
while (peak_map & peak) != 0 {
let left_sibling = pos + 1 - 2 * peak;
let left_hash = self
.backend
.get_peak_from_file(left_sibling)
.ok_or("missing left sibling in tree, should not have been pruned")?;
peak *= 2;
pos += 1;
current_hash = (left_hash, current_hash).hash_with_index(pos);
hashes.push(current_hash);
}
self.backend.append(leaf, &hashes)?;
self.size = pos + 1;
Ok(leaf_pos)
}
pub fn push_pruned_subtree(&mut self, hash: Hash, pos0: u64) -> Result<(), String> {
self.backend.append_pruned_subtree(hash, pos0)?;
self.size = pos0 + 1;
let mut pos = pos0;
let mut current_hash = hash;
let (peak_map, _) = peak_map_height(pos);
let mut peak = 1;
while (peak_map & peak) != 0 {
let (parent, sibling) = family(pos);
peak *= 2;
if sibling > pos {
continue;
}
let left_hash = self
.backend
.get_hash(sibling)
.ok_or("missing left sibling in tree, should not have been pruned")?;
pos = parent;
current_hash = (left_hash, current_hash).hash_with_index(parent);
self.backend.append_hash(current_hash)?;
}
self.size = crate::core::pmmr::round_up_to_leaf_pos(pos);
Ok(())
}
pub fn reset_prune_list(&mut self) {
self.backend.reset_prune_list();
}
pub fn remove_from_leaf_set(&mut self, pos0: u64) {
self.backend.remove_from_leaf_set(pos0);
}
pub fn snapshot(&mut self, header: &BlockHeader) -> Result<(), String> {
self.backend.snapshot(header)?;
Ok(())
}
pub fn rewind(&mut self, position: u64, rewind_rm_pos: &Bitmap) -> Result<(), String> {
let leaf_pos = round_up_to_leaf_pos(position);
self.backend.rewind(leaf_pos, rewind_rm_pos)?;
self.size = leaf_pos;
Ok(())
}
pub fn prune(&mut self, pos0: u64) -> Result<bool, String> {
if !is_leaf(pos0) {
return Err(format!("Node at {} is not a leaf, can't prune.", pos0));
}
if self.backend.get_hash(pos0).is_none() {
return Ok(false);
}
self.backend.remove(pos0)?;
Ok(true)
}
pub fn validate(&self) -> Result<(), String> {
for n in 0..self.size {
let height = bintree_postorder_height(n);
if height > 0 {
if let Some(hash) = self.get_hash(n) {
let left_pos = n - (1 << height);
let right_pos = n - 1;
if let Some(left_child_hs) = self.get_from_file(left_pos) {
if let Some(right_child_hs) = self.get_from_file(right_pos) {
if (left_child_hs, right_child_hs).hash_with_index(n) != hash {
return Err(format!(
"Invalid MMR, hash of parent at {} does \
not match children.",
n + 1
));
}
}
}
}
}
}
Ok(())
}
pub fn dump(&self, short: bool) {
let sz = self.unpruned_size();
if sz > 2000 && !short {
return;
}
let start = if short { sz / 8 } else { 0 };
for n in start..(sz / 8 + 1) {
let mut idx = "".to_owned();
let mut hashes = "".to_owned();
for m in (n * 8)..(n + 1) * 8 {
if m >= sz {
break;
}
idx.push_str(&format!("{:>8} ", m));
let ohs = self.get_hash(m);
match ohs {
Some(hs) => hashes.push_str(&format!("{} ", hs)),
None => hashes.push_str(&format!("{:>8} ", "??")),
}
}
debug!("{}", idx);
debug!("{}", hashes);
}
}
pub fn dump_stats(&self) {
debug!("pmmr: unpruned - {}", self.unpruned_size());
self.backend.dump_stats();
}
pub fn dump_from_file(&self, short: bool) {
let sz = self.unpruned_size();
if sz > 2000 && !short {
return;
}
let start = if short { sz / 8 } else { 0 };
for n in start..(sz / 8 + 1) {
let mut idx = "".to_owned();
let mut hashes = "".to_owned();
for m in (n * 8)..(n + 1) * 8 {
if m >= sz {
break;
}
idx.push_str(&format!("{:>8} ", m + 1));
let ohs = self.get_from_file(m);
match ohs {
Some(hs) => hashes.push_str(&format!("{} ", hs)),
None => hashes.push_str(&format!("{:>8} ", " .")),
}
}
debug!("{}", idx);
debug!("{}", hashes);
}
}
}
impl<'a, T, B> ReadablePMMR for PMMR<'a, T, B>
where
T: PMMRable,
B: 'a + Backend<T>,
{
type Item = T::E;
fn get_hash(&self, pos0: u64) -> Option<Hash> {
if pos0 >= self.size {
None
} else if is_leaf(pos0) {
self.backend.get_hash(pos0)
} else {
self.backend.get_from_file(pos0)
}
}
fn get_data(&self, pos0: u64) -> Option<Self::Item> {
if pos0 >= self.size {
None
} else if is_leaf(pos0) {
self.backend.get_data(pos0)
} else {
None
}
}
fn get_from_file(&self, pos0: u64) -> Option<Hash> {
if pos0 >= self.size {
None
} else {
self.backend.get_from_file(pos0)
}
}
fn get_peak_from_file(&self, pos0: u64) -> Option<Hash> {
if pos0 >= self.size {
None
} else {
self.backend.get_peak_from_file(pos0)
}
}
fn get_data_from_file(&self, pos0: u64) -> Option<Self::Item> {
if pos0 >= self.size {
None
} else {
self.backend.get_data_from_file(pos0)
}
}
fn unpruned_size(&self) -> u64 {
self.size
}
fn leaf_pos_iter(&self) -> Box<dyn Iterator<Item = u64> + '_> {
self.backend.leaf_pos_iter()
}
fn leaf_idx_iter(&self, from_idx: u64) -> Box<dyn Iterator<Item = u64> + '_> {
self.backend.leaf_idx_iter(from_idx)
}
fn n_unpruned_leaves(&self) -> u64 {
self.backend.n_unpruned_leaves()
}
fn n_unpruned_leaves_to_index(&self, to_index: u64) -> u64 {
self.backend.n_unpruned_leaves_to_index(to_index)
}
}
const ALL_ONES: u64 = u64::MAX;
pub fn peak_map_height(mut size: u64) -> (u64, u64) {
if size == 0 {
return (0, 0);
}
let mut peak_size = ALL_ONES >> size.leading_zeros();
let mut peak_map = 0;
while peak_size != 0 {
peak_map <<= 1;
if size >= peak_size {
size -= peak_size;
peak_map |= 1;
}
peak_size >>= 1;
}
(peak_map, size)
}
pub fn peak_sizes_height(mut size: u64) -> (Vec<u64>, u64) {
if size == 0 {
return (vec![], 0);
}
let mut peak_size = ALL_ONES >> size.leading_zeros();
let mut peak_sizes = vec![];
while peak_size != 0 {
if size >= peak_size {
peak_sizes.push(peak_size);
size -= peak_size;
}
peak_size >>= 1;
}
(peak_sizes, size)
}
pub fn peaks(size: u64) -> Vec<u64> {
let (peak_sizes, height) = peak_sizes_height(size);
if height == 0 {
peak_sizes
.iter()
.scan(0, |acc, &x| {
*acc += &x;
Some(*acc)
})
.map(|x| x - 1) .collect()
} else {
vec![]
}
}
pub fn n_leaves(size: u64) -> u64 {
let (peak_map, height) = peak_map_height(size);
if height == 0 {
peak_map
} else {
peak_map + 1
}
}
pub fn round_up_to_leaf_pos(pos0: u64) -> u64 {
let (insert_idx, height) = peak_map_height(pos0);
let leaf_idx = if height == 0 {
insert_idx
} else {
insert_idx + 1
};
return insertion_to_pmmr_index(leaf_idx);
}
pub fn insertion_to_pmmr_index(nleaf0: u64) -> u64 {
2 * nleaf0 - nleaf0.count_ones() as u64
}
pub fn pmmr_leaf_to_insertion_index(pos0: u64) -> Option<u64> {
let (insert_idx, height) = peak_map_height(pos0);
if height == 0 {
Some(insert_idx)
} else {
None
}
}
pub fn bintree_postorder_height(pos0: u64) -> u64 {
peak_map_height(pos0).1
}
pub fn is_leaf(pos0: u64) -> bool {
bintree_postorder_height(pos0) == 0
}
pub fn family(pos0: u64) -> (u64, u64) {
let (peak_map, height) = peak_map_height(pos0);
let peak = 1 << height;
if (peak_map & peak) != 0 {
(pos0 + 1, pos0 + 1 - 2 * peak)
} else {
(pos0 + 2 * peak, pos0 + 2 * peak - 1)
}
}
pub fn is_left_sibling(pos0: u64) -> bool {
let (peak_map, height) = peak_map_height(pos0);
let peak = 1 << height;
(peak_map & peak) == 0
}
pub fn family_branch(pos0: u64, size: u64) -> Vec<(u64, u64)> {
let (peak_map, height) = peak_map_height(pos0);
let mut peak = 1 << height;
let mut branch = vec![];
let mut current = pos0;
let mut sibling;
while current + 1 < size {
if (peak_map & peak) != 0 {
current += 1;
sibling = current - 2 * peak;
} else {
current += 2 * peak;
sibling = current - 1;
};
if current >= size {
break;
}
branch.push((current, sibling));
peak <<= 1;
}
branch
}
pub fn bintree_rightmost(pos0: u64) -> u64 {
pos0 - bintree_postorder_height(pos0)
}
pub fn bintree_leftmost(pos0: u64) -> u64 {
let height = bintree_postorder_height(pos0);
pos0 + 2 - (2 << height)
}
pub fn bintree_leaf_pos_iter(pos0: u64) -> Box<dyn Iterator<Item = u64>> {
let leaf_start = pmmr_leaf_to_insertion_index(bintree_leftmost(pos0));
let leaf_end = pmmr_leaf_to_insertion_index(bintree_rightmost(pos0));
let leaf_start = match leaf_start {
Some(l) => l,
None => return Box::new(iter::empty::<u64>()),
};
let leaf_end = match leaf_end {
Some(l) => l,
None => return Box::new(iter::empty::<u64>()),
};
Box::new((leaf_start..=leaf_end).map(|n| insertion_to_pmmr_index(n)))
}
pub fn bintree_pos_iter(pos0: u64) -> impl Iterator<Item = u64> {
let leaf_start = bintree_leftmost(pos0);
(leaf_start..=pos0).into_iter()
}
pub fn bintree_range(pos0: u64) -> Range<u64> {
let height = bintree_postorder_height(pos0);
let leftmost = pos0 + 2 - (2 << height);
leftmost..(pos0 + 1)
}