mmr: added partial mmr

This commit is contained in:
Augusto F. Hack
2023-10-17 17:38:18 +02:00
parent 78aa714b89
commit bde20f9752
13 changed files with 1130 additions and 136 deletions

View File

@@ -13,12 +13,8 @@
use super::{
super::{InnerNodeInfo, MerklePath, RpoDigest, Vec},
bit::TrueBitPositionIterator,
MmrPeaks, MmrProof, Rpo256,
leaf_to_corresponding_tree, nodes_in_forest, MmrDelta, MmrError, MmrPeaks, MmrProof, Rpo256,
};
use core::fmt::{Display, Formatter};
#[cfg(feature = "std")]
use std::error::Error;
// MMR
// ===============================================================================================
@@ -43,22 +39,6 @@ pub struct Mmr {
pub(super) nodes: Vec<RpoDigest>,
}
#[derive(Debug, PartialEq, Eq, Copy, Clone)]
pub enum MmrError {
InvalidPosition(usize),
}
impl Display for MmrError {
fn fmt(&self, fmt: &mut Formatter<'_>) -> Result<(), core::fmt::Error> {
match self {
MmrError::InvalidPosition(pos) => write!(fmt, "Mmr does not contain position {pos}"),
}
}
}
#[cfg(feature = "std")]
impl Error for MmrError {}
impl Default for Mmr {
fn default() -> Self {
Self::new()
@@ -100,21 +80,16 @@ impl Mmr {
// find the target tree responsible for the MMR position
let tree_bit =
leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::InvalidPosition(pos))?;
let forest_target = 1usize << tree_bit;
// isolate the trees before the target
let forest_before = self.forest & high_bitmask(tree_bit + 1);
let index_offset = nodes_in_forest(forest_before);
// find the root
let index = nodes_in_forest(forest_target) - 1;
// update the value position from global to the target tree
let relative_pos = pos - forest_before;
// collect the path and the final index of the target value
let (_, path) =
self.collect_merkle_path_and_value(tree_bit, relative_pos, index_offset, index);
let (_, path) = self.collect_merkle_path_and_value(tree_bit, relative_pos, index_offset);
Ok(MmrProof {
forest: self.forest,
@@ -132,21 +107,16 @@ impl Mmr {
// find the target tree responsible for the MMR position
let tree_bit =
leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::InvalidPosition(pos))?;
let forest_target = 1usize << tree_bit;
// isolate the trees before the target
let forest_before = self.forest & high_bitmask(tree_bit + 1);
let index_offset = nodes_in_forest(forest_before);
// find the root
let index = nodes_in_forest(forest_target) - 1;
// update the value position from global to the target tree
let relative_pos = pos - forest_before;
// collect the path and the final index of the target value
let (value, _) =
self.collect_merkle_path_and_value(tree_bit, relative_pos, index_offset, index);
let (value, _) = self.collect_merkle_path_and_value(tree_bit, relative_pos, index_offset);
Ok(value)
}
@@ -185,7 +155,82 @@ impl Mmr {
.map(|offset| self.nodes[offset - 1])
.collect();
MmrPeaks { num_leaves: self.forest, peaks }
// Safety: the invariant is maintained by the [Mmr]
MmrPeaks::new(self.forest, peaks).unwrap()
}
/// Compute the required update to `original_forest`.
///
/// The result is a packed sequence of the authentication elements required to update the trees
/// that have been merged together, followed by the new peaks of the [Mmr].
pub fn get_delta(&self, original_forest: usize) -> Result<MmrDelta, MmrError> {
if original_forest > self.forest {
return Err(MmrError::InvalidPeaks);
}
if original_forest == self.forest {
return Ok(MmrDelta { forest: self.forest, data: Vec::new() });
}
let mut result = Vec::new();
// Find the largest tree in this [Mmr] which is new to `original_forest`.
let candidate_trees = self.forest ^ original_forest;
let mut new_high = 1 << candidate_trees.ilog2();
// Collect authentication nodes used for tree merges
// ----------------------------------------------------------------------------------------
// Find the trees from `original_forest` that have been merged into `new_high`.
let mut merges = original_forest & (new_high - 1);
// Find the peaks that are common to `original_forest` and this [Mmr]
let common_trees = original_forest ^ merges;
if merges != 0 {
// Skip the smallest trees unknown to `original_forest`.
let mut target = 1 << merges.trailing_zeros();
// Collect siblings required to computed the merged tree's peak
while target < new_high {
// Computes the offset to the smallest know peak
// - common_trees: peaks unchanged in the current update, target comes after these.
// - merges: peaks that have not been merged so far, target comes after these.
// - target: tree from which to load the sibling. On the first iteration this is a
// value known by the partial mmr, on subsequent iterations this value is to be
// computed from the known peaks and provided authentication nodes.
let known = nodes_in_forest(common_trees | merges | target);
let sibling = nodes_in_forest(target);
result.push(self.nodes[known + sibling - 1]);
// Update the target and account for tree merges
target <<= 1;
while merges & target != 0 {
target <<= 1;
}
// Remove the merges done so far
merges ^= merges & (target - 1);
}
} else {
// The new high tree may not be the result of any merges, if it is smaller than all the
// trees of `original_forest`.
new_high = 0;
}
// Collect the new [Mmr] peaks
// ----------------------------------------------------------------------------------------
let mut new_peaks = self.forest ^ common_trees ^ new_high;
let old_peaks = self.forest ^ new_peaks;
let mut offset = nodes_in_forest(old_peaks);
while new_peaks != 0 {
let target = 1 << new_peaks.ilog2();
offset += nodes_in_forest(target);
result.push(self.nodes[offset - 1]);
new_peaks ^= target;
}
Ok(MmrDelta { forest: self.forest, data: result })
}
/// An iterator over inner nodes in the MMR. The order of iteration is unspecified.
@@ -202,36 +247,52 @@ impl Mmr {
// ============================================================================================
/// Internal function used to collect the Merkle path of a value.
///
/// The arguments are relative to the target tree. To compute the opening of the second leaf
/// for a tree with depth 2 in the forest `0b110`:
///
/// - `tree_bit`: Depth of the target tree, e.g. 2 for the smallest tree.
/// - `relative_pos`: 0-indexed leaf position in the target tree, e.g. 1 for the second leaf.
/// - `index_offset`: Node count prior to the target tree, e.g. 7 for the tree of depth 3.
fn collect_merkle_path_and_value(
&self,
tree_bit: u32,
relative_pos: usize,
index_offset: usize,
mut index: usize,
) -> (RpoDigest, Vec<RpoDigest>) {
// collect the Merkle path
let mut tree_depth = tree_bit as usize;
let mut path = Vec::with_capacity(tree_depth + 1);
while tree_depth > 0 {
let bit = relative_pos & tree_depth;
let right_offset = index - 1;
let left_offset = right_offset - nodes_in_forest(tree_depth);
// see documentation of `leaf_to_corresponding_tree` for details
let tree_depth = (tree_bit + 1) as usize;
let mut path = Vec::with_capacity(tree_depth);
// Elements to the right have a higher position because they were
// added later. Therefore when the bit is true the node's path is
// to the right, and its sibling to the left.
let sibling = if bit != 0 {
// The tree walk below goes from the root to the leaf, compute the root index to start
let mut forest_target = 1usize << tree_bit;
let mut index = nodes_in_forest(forest_target) - 1;
// Loop until the leaf is reached
while forest_target > 1 {
// Update the depth of the tree to correspond to a subtree
forest_target >>= 1;
// compute the indeces of the right and left subtrees based on the post-order
let right_offset = index - 1;
let left_offset = right_offset - nodes_in_forest(forest_target);
let left_or_right = relative_pos & forest_target;
let sibling = if left_or_right != 0 {
// going down the right subtree, the right child becomes the new root
index = right_offset;
// and the left child is the authentication
self.nodes[index_offset + left_offset]
} else {
index = left_offset;
self.nodes[index_offset + right_offset]
};
tree_depth >>= 1;
path.push(sibling);
}
debug_assert!(path.len() == tree_depth - 1);
// the rest of the codebase has the elements going from leaf to root, adjust it here for
// easy of use/consistency sake
path.reverse();
@@ -241,6 +302,9 @@ impl Mmr {
}
}
// CONVERSIONS
// ================================================================================================
impl<T> From<T> for Mmr
where
T: IntoIterator<Item = RpoDigest>,
@@ -335,32 +399,6 @@ impl<'a> Iterator for MmrNodes<'a> {
// UTILITIES
// ===============================================================================================
/// Given a 0-indexed leaf position and the current forest, return the tree number responsible for
/// the position.
///
/// Note:
/// The result is a tree position `p`, it has the following interpretations. $p+1$ is the depth of
/// the tree, which corresponds to the size of a Merkle proof for that tree. $2^p$ is equal to the
/// number of leaves in this particular tree. and $2^(p+1)-1$ corresponds to size of the tree.
pub(crate) const fn leaf_to_corresponding_tree(pos: usize, forest: usize) -> Option<u32> {
if pos >= forest {
None
} else {
// - each bit in the forest is a unique tree and the bit position its power-of-two size
// - each tree owns a consecutive range of positions equal to its size from left-to-right
// - this means the first tree owns from `0` up to the `2^k_0` first positions, where `k_0`
// is the highest true bit position, the second tree from `2^k_0 + 1` up to `2^k_1` where
// `k_1` is the second higest bit, so on.
// - this means the highest bits work as a category marker, and the position is owned by
// the first tree which doesn't share a high bit with the position
let before = forest & pos;
let after = forest ^ before;
let tree = after.ilog2();
Some(tree)
}
}
/// Return a bitmask for the bits including and above the given position.
pub(crate) const fn high_bitmask(bit: u32) -> usize {
if bit > usize::BITS - 1 {
@@ -369,17 +407,3 @@ pub(crate) const fn high_bitmask(bit: u32) -> usize {
usize::MAX << bit
}
}
/// Return the total number of nodes of a given forest
///
/// Panics:
///
/// This will panic if the forest has size greater than `usize::MAX / 2`
pub(crate) const fn nodes_in_forest(forest: usize) -> usize {
// - the size of a perfect binary tree is $2^{k+1}-1$ or $2*2^k-1$
// - the forest represents the sum of $2^k$ so a single multiplication is necessary
// - the number of `-1` is the same as the number of trees, which is the same as the number
// bits set
let tree_count = forest.count_ones() as usize;
forest * 2 - tree_count
}