miden_crypto/merkle/mmr/full.rs
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417
//! A fully materialized Merkle mountain range (MMR).
//!
//! A MMR is a forest structure, i.e. it is an ordered set of disjoint rooted trees. The trees are
//! ordered by size, from the most to least number of leaves. Every tree is a perfect binary tree,
//! meaning a tree has all its leaves at the same depth, and every inner node has a branch-factor
//! of 2 with both children set.
//!
//! Additionally the structure only supports adding leaves to the right-most tree, the one with the
//! least number of leaves. The structure preserves the invariant that each tree has different
//! depths, i.e. as part of adding a new element to the forest the trees with same depth are
//! merged, creating a new tree with depth d+1, this process is continued until the property is
//! reestablished.
use super::{
super::{InnerNodeInfo, MerklePath},
bit::TrueBitPositionIterator,
leaf_to_corresponding_tree, nodes_in_forest, MmrDelta, MmrError, MmrPeaks, MmrProof, Rpo256,
RpoDigest,
};
use alloc::vec::Vec;
// MMR
// ===============================================================================================
/// A fully materialized Merkle Mountain Range, with every tree in the forest and all their
/// elements.
///
/// Since this is a full representation of the MMR, elements are never removed and the MMR will
/// grow roughly `O(2n)` in number of leaf elements.
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
pub struct Mmr {
/// Refer to the `forest` method documentation for details of the semantics of this value.
pub(super) forest: usize,
/// Contains every element of the forest.
///
/// The trees are in postorder sequential representation. This representation allows for all
/// the elements of every tree in the forest to be stored in the same sequential buffer. It
/// also means new elements can be added to the forest, and merging of trees is very cheap with
/// no need to copy elements.
pub(super) nodes: Vec<RpoDigest>,
}
impl Default for Mmr {
fn default() -> Self {
Self::new()
}
}
impl Mmr {
// CONSTRUCTORS
// ============================================================================================
/// Constructor for an empty `Mmr`.
pub fn new() -> Mmr {
Mmr { forest: 0, nodes: Vec::new() }
}
// ACCESSORS
// ============================================================================================
/// Returns the MMR forest representation.
///
/// The forest value has the following interpretations:
/// - its value is the number of elements in the forest
/// - bit count corresponds to the number of trees in the forest
/// - each true bit position determines the depth of a tree in the forest
pub const fn forest(&self) -> usize {
self.forest
}
// FUNCTIONALITY
// ============================================================================================
/// Given a leaf position, returns the Merkle path to its corresponding peak. If the position
/// is greater-or-equal than the tree size an error is returned.
///
/// Note: The leaf position is the 0-indexed number corresponding to the order the leaves were
/// added, this corresponds to the MMR size _prior_ to adding the element. So the 1st element
/// has position 0, the second position 1, and so on.
pub fn open(&self, pos: usize, target_forest: usize) -> Result<MmrProof, MmrError> {
// find the target tree responsible for the MMR position
let tree_bit =
leaf_to_corresponding_tree(pos, target_forest).ok_or(MmrError::InvalidPosition(pos))?;
// isolate the trees before the target
let forest_before = target_forest & high_bitmask(tree_bit + 1);
let index_offset = nodes_in_forest(forest_before);
// 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);
Ok(MmrProof {
forest: target_forest,
position: pos,
merkle_path: MerklePath::new(path),
})
}
/// Returns the leaf value at position `pos`.
///
/// Note: The leaf position is the 0-indexed number corresponding to the order the leaves were
/// added, this corresponds to the MMR size _prior_ to adding the element. So the 1st element
/// has position 0, the second position 1, and so on.
pub fn get(&self, pos: usize) -> Result<RpoDigest, MmrError> {
// find the target tree responsible for the MMR position
let tree_bit =
leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::InvalidPosition(pos))?;
// isolate the trees before the target
let forest_before = self.forest & high_bitmask(tree_bit + 1);
let index_offset = nodes_in_forest(forest_before);
// 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);
Ok(value)
}
/// Adds a new element to the MMR.
pub fn add(&mut self, el: RpoDigest) {
// Note: every node is also a tree of size 1, adding an element to the forest creates a new
// rooted-tree of size 1. This may temporarily break the invariant that every tree in the
// forest has different sizes, the loop below will eagerly merge trees of same size and
// restore the invariant.
self.nodes.push(el);
let mut left_offset = self.nodes.len().saturating_sub(2);
let mut right = el;
let mut left_tree = 1;
while self.forest & left_tree != 0 {
right = Rpo256::merge(&[self.nodes[left_offset], right]);
self.nodes.push(right);
left_offset = left_offset.saturating_sub(nodes_in_forest(left_tree));
left_tree <<= 1;
}
self.forest += 1;
}
/// Returns an peaks of the MMR for the version specified by `forest`.
pub fn peaks(&self, forest: usize) -> Result<MmrPeaks, MmrError> {
if forest > self.forest {
return Err(MmrError::InvalidPeaks);
}
let peaks: Vec<RpoDigest> = TrueBitPositionIterator::new(forest)
.rev()
.map(|bit| nodes_in_forest(1 << bit))
.scan(0, |offset, el| {
*offset += el;
Some(*offset)
})
.map(|offset| self.nodes[offset - 1])
.collect();
// Safety: the invariant is maintained by the [Mmr]
let peaks = MmrPeaks::new(forest, peaks).unwrap();
Ok(peaks)
}
/// 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, from_forest: usize, to_forest: usize) -> Result<MmrDelta, MmrError> {
if to_forest > self.forest || from_forest > to_forest {
return Err(MmrError::InvalidPeaks);
}
if from_forest == to_forest {
return Ok(MmrDelta { forest: to_forest, data: Vec::new() });
}
let mut result = Vec::new();
// Find the largest tree in this [Mmr] which is new to `from_forest`.
let candidate_trees = to_forest ^ from_forest;
let mut new_high = 1 << candidate_trees.ilog2();
// Collect authentication nodes used for tree merges
// ----------------------------------------------------------------------------------------
// Find the trees from `from_forest` that have been merged into `new_high`.
let mut merges = from_forest & (new_high - 1);
// Find the peaks that are common to `from_forest` and this [Mmr]
let common_trees = from_forest ^ merges;
if merges != 0 {
// Skip the smallest trees unknown to `from_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 `from_forest`.
new_high = 0;
}
// Collect the new [Mmr] peaks
// ----------------------------------------------------------------------------------------
let mut new_peaks = to_forest ^ common_trees ^ new_high;
let old_peaks = to_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: to_forest, data: result })
}
/// An iterator over inner nodes in the MMR. The order of iteration is unspecified.
pub fn inner_nodes(&self) -> MmrNodes {
MmrNodes {
mmr: self,
forest: 0,
last_right: 0,
index: 0,
}
}
// UTILITIES
// ============================================================================================
/// 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,
) -> (RpoDigest, Vec<RpoDigest>) {
// 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);
// 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 indices 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]
};
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();
let value = self.nodes[index_offset + index];
(value, path)
}
}
// CONVERSIONS
// ================================================================================================
impl<T> From<T> for Mmr
where
T: IntoIterator<Item = RpoDigest>,
{
fn from(values: T) -> Self {
let mut mmr = Mmr::new();
for v in values {
mmr.add(v)
}
mmr
}
}
// ITERATOR
// ===============================================================================================
/// Yields inner nodes of the [Mmr].
pub struct MmrNodes<'a> {
/// [Mmr] being yielded, when its `forest` value is matched, the iterations is finished.
mmr: &'a Mmr,
/// Keeps track of the left nodes yielded so far waiting for a right pair, this matches the
/// semantics of the [Mmr]'s forest attribute, since that too works as a buffer of left nodes
/// waiting for a pair to be hashed together.
forest: usize,
/// Keeps track of the last right node yielded, after this value is set, the next iteration
/// will be its parent with its corresponding left node that has been yield already.
last_right: usize,
/// The current index in the `nodes` vector.
index: usize,
}
impl<'a> Iterator for MmrNodes<'a> {
type Item = InnerNodeInfo;
fn next(&mut self) -> Option<Self::Item> {
debug_assert!(self.last_right.count_ones() <= 1, "last_right tracks zero or one element");
// only parent nodes are emitted, remove the single node tree from the forest
let target = self.mmr.forest & (usize::MAX << 1);
if self.forest < target {
if self.last_right == 0 {
// yield the left leaf
debug_assert!(self.last_right == 0, "left must be before right");
self.forest |= 1;
self.index += 1;
// yield the right leaf
debug_assert!((self.forest & 1) == 1, "right must be after left");
self.last_right |= 1;
self.index += 1;
};
debug_assert!(
self.forest & self.last_right != 0,
"parent requires both a left and right",
);
// compute the number of nodes in the right tree, this is the offset to the
// previous left parent
let right_nodes = nodes_in_forest(self.last_right);
// the next parent position is one above the position of the pair
let parent = self.last_right << 1;
// the left node has been paired and the current parent yielded, removed it from the forest
self.forest ^= self.last_right;
if self.forest & parent == 0 {
// this iteration yielded the left parent node
debug_assert!(self.forest & 1 == 0, "next iteration yields a left leaf");
self.last_right = 0;
self.forest ^= parent;
} else {
// the left node of the parent level has been yielded already, this iteration
// was the right parent. Next iteration yields their parent.
self.last_right = parent;
}
// yields a parent
let value = self.mmr.nodes[self.index];
let right = self.mmr.nodes[self.index - 1];
let left = self.mmr.nodes[self.index - 1 - right_nodes];
self.index += 1;
let node = InnerNodeInfo { value, left, right };
Some(node)
} else {
None
}
}
}
// UTILITIES
// ===============================================================================================
/// 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 {
0
} else {
usize::MAX << bit
}
}