Augusto F. Hack
1 year ago
No known key found for this signature in database
GPG Key ID: 3F3584B7FB1DFB76
13 changed files with 1129 additions and 135 deletions
Split View
Diff Options
-
+1 -0CHANGELOG.md
-
+26 -2src/merkle/merkle_tree.rs
-
+16 -0src/merkle/mmr/delta.rs
-
+35 -0src/merkle/mmr/error.rs
-
+110 -86src/merkle/mmr/full.rs
-
+136 -0src/merkle/mmr/inorder.rs
-
+54 -2src/merkle/mmr/mod.rs
-
+403 -0src/merkle/mmr/partial.rs
-
+27 -6src/merkle/mmr/peaks.rs
-
+1 -1src/merkle/mmr/proof.rs
-
+305 -37src/merkle/mmr/tests.rs
-
+1 -1src/merkle/mod.rs
-
+14 -0src/merkle/path.rs
@ -0,0 +1,16 @@ |
|||
use super::super::{RpoDigest, Vec};
|
|||
|
|||
/// Container for the update data of a [PartialMmr]
|
|||
#[derive(Debug)]
|
|||
pub struct MmrDelta {
|
|||
/// The new version of the [Mmr]
|
|||
pub forest: usize,
|
|||
|
|||
/// Update data.
|
|||
///
|
|||
/// The data is packed as follows:
|
|||
/// 1. All the elements needed to perform authentication path updates. These are the right
|
|||
/// siblings required to perform tree merges on the [PartialMmr].
|
|||
/// 2. The new peaks.
|
|||
pub data: Vec<RpoDigest>,
|
|||
}
|
|||
@ -0,0 +1,35 @@ |
|||
use crate::merkle::MerkleError;
|
|||
use core::fmt::{Display, Formatter};
|
|||
|
|||
#[cfg(feature = "std")]
|
|||
use std::error::Error;
|
|||
|
|||
#[derive(Debug, PartialEq, Eq, Clone)]
|
|||
pub enum MmrError {
|
|||
InvalidPosition(usize),
|
|||
InvalidPeaks,
|
|||
InvalidPeak,
|
|||
InvalidUpdate,
|
|||
UnknownPeak,
|
|||
MerkleError(MerkleError),
|
|||
}
|
|||
|
|||
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}"),
|
|||
MmrError::InvalidPeaks => write!(fmt, "Invalid peaks count"),
|
|||
MmrError::InvalidPeak => {
|
|||
write!(fmt, "Peak values does not match merkle path computed root")
|
|||
}
|
|||
MmrError::InvalidUpdate => write!(fmt, "Invalid mmr update"),
|
|||
MmrError::UnknownPeak => {
|
|||
write!(fmt, "Peak not in Mmr")
|
|||
}
|
|||
MmrError::MerkleError(err) => write!(fmt, "{}", err),
|
|||
}
|
|||
}
|
|||
}
|
|||
|
|||
#[cfg(feature = "std")]
|
|||
impl Error for MmrError {}
|
|||
@ -0,0 +1,136 @@ |
|||
//! Index for nodes of a binary tree based on an in-order tree walk.
|
|||
//!
|
|||
//! In-order walks have the parent node index split its left and right subtrees. All the left
|
|||
//! children have indexes lower than the parent, meanwhile all the right subtree higher indexes.
|
|||
//! This property makes it is easy to compute changes to the index by adding or subtracting the
|
|||
//! leaves count.
|
|||
use core::num::NonZeroUsize;
|
|||
|
|||
/// Index of nodes in a perfectly balanced binary tree based on an in-order tree walk.
|
|||
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
|
|||
pub struct InOrderIndex {
|
|||
idx: usize,
|
|||
}
|
|||
|
|||
impl InOrderIndex {
|
|||
/// Constructor for a new [InOrderIndex].
|
|||
pub fn new(idx: NonZeroUsize) -> InOrderIndex {
|
|||
InOrderIndex { idx: idx.get() }
|
|||
}
|
|||
|
|||
/// Constructs an index from a leaf position.
|
|||
///
|
|||
/// Panics:
|
|||
///
|
|||
/// If `leaf` is higher than or equal to `usize::MAX / 2`.
|
|||
pub fn from_leaf_pos(leaf: usize) -> InOrderIndex {
|
|||
// Convert the position from 0-indexed to 1-indexed, since the bit manipulation in this
|
|||
// implementation only works 1-indexed counting.
|
|||
let pos = leaf + 1;
|
|||
InOrderIndex { idx: pos * 2 - 1 }
|
|||
}
|
|||
|
|||
/// True if the index is pointing at a leaf.
|
|||
///
|
|||
/// Every odd number represents a leaf.
|
|||
pub fn is_leaf(&self) -> bool {
|
|||
self.idx & 1 == 1
|
|||
}
|
|||
|
|||
/// Returns the level of the index.
|
|||
///
|
|||
/// Starts at level zero for leaves and increases by one for each parent.
|
|||
pub fn level(&self) -> u32 {
|
|||
self.idx.trailing_zeros()
|
|||
}
|
|||
|
|||
/// Returns the index of the left child.
|
|||
///
|
|||
/// Panics:
|
|||
///
|
|||
/// If the index corresponds to a leaf.
|
|||
pub fn left_child(&self) -> InOrderIndex {
|
|||
// The left child is itself a parent, with an index that splits its left/right subtrees. To
|
|||
// go from the parent index to its left child, it is only necessary to subtract the count
|
|||
// of elements on the child's right subtree + 1.
|
|||
let els = 1 << (self.level() - 1);
|
|||
InOrderIndex { idx: self.idx - els }
|
|||
}
|
|||
|
|||
/// Returns the index of the right child.
|
|||
///
|
|||
/// Panics:
|
|||
///
|
|||
/// If the index corresponds to a leaf.
|
|||
pub fn right_child(&self) -> InOrderIndex {
|
|||
// To compute the index of the parent of the right subtree it is sufficient to add the size
|
|||
// of its left subtree + 1.
|
|||
let els = 1 << (self.level() - 1);
|
|||
InOrderIndex { idx: self.idx + els }
|
|||
}
|
|||
|
|||
/// Returns the index of the parent node.
|
|||
pub fn parent(&self) -> InOrderIndex {
|
|||
// If the current index corresponds to a node in a left tree, to go up a level it is
|
|||
// required to add the number of nodes of the right sibling, analogously if the node is a
|
|||
// right child, going up requires subtracting the number of nodes in its left subtree.
|
|||
//
|
|||
// Both of the above operations can be performed by bitwise manipulation. Below the mask
|
|||
// sets the number of trailing zeros to be equal the new level of the index, and the bit
|
|||
// marks the parent.
|
|||
let target = self.level() + 1;
|
|||
let bit = 1 << target;
|
|||
let mask = bit - 1;
|
|||
let idx = self.idx ^ (self.idx & mask);
|
|||
InOrderIndex { idx: idx | bit }
|
|||
}
|
|||
|
|||
/// Returns the index of the sibling node.
|
|||
pub fn sibling(&self) -> InOrderIndex {
|
|||
let parent = self.parent();
|
|||
if *self > parent {
|
|||
parent.left_child()
|
|||
} else {
|
|||
parent.right_child()
|
|||
}
|
|||
}
|
|||
}
|
|||
|
|||
#[cfg(test)]
|
|||
mod test {
|
|||
use super::InOrderIndex;
|
|||
use proptest::prelude::*;
|
|||
|
|||
proptest! {
|
|||
#[test]
|
|||
fn proptest_inorder_index_random(count in 1..1000usize) {
|
|||
let left_pos = count * 2;
|
|||
let right_pos = count * 2 + 1;
|
|||
|
|||
let left = InOrderIndex::from_leaf_pos(left_pos);
|
|||
let right = InOrderIndex::from_leaf_pos(right_pos);
|
|||
|
|||
assert!(left.is_leaf());
|
|||
assert!(right.is_leaf());
|
|||
assert_eq!(left.parent(), right.parent());
|
|||
assert_eq!(left.parent().right_child(), right);
|
|||
assert_eq!(left, right.parent().left_child());
|
|||
assert_eq!(left.sibling(), right);
|
|||
assert_eq!(left, right.sibling());
|
|||
}
|
|||
}
|
|||
|
|||
#[test]
|
|||
fn test_inorder_index_basic() {
|
|||
let left = InOrderIndex::from_leaf_pos(0);
|
|||
let right = InOrderIndex::from_leaf_pos(1);
|
|||
|
|||
assert!(left.is_leaf());
|
|||
assert!(right.is_leaf());
|
|||
assert_eq!(left.parent(), right.parent());
|
|||
assert_eq!(left.parent().right_child(), right);
|
|||
assert_eq!(left, right.parent().left_child());
|
|||
assert_eq!(left.sibling(), right);
|
|||
assert_eq!(left, right.sibling());
|
|||
}
|
|||
}
|
|||
@ -0,0 +1,403 @@ |
|||
use crate::{
|
|||
hash::rpo::{Rpo256, RpoDigest},
|
|||
merkle::{
|
|||
mmr::{leaf_to_corresponding_tree, nodes_in_forest},
|
|||
InOrderIndex, MerklePath, MmrError, MmrPeaks,
|
|||
},
|
|||
utils::collections::{BTreeMap, Vec},
|
|||
};
|
|||
|
|||
use super::{MmrDelta, MmrProof};
|
|||
|
|||
/// Partially materialized [Mmr], used to efficiently store and update the authentication paths for
|
|||
/// a subset of the elements in a full [Mmr].
|
|||
///
|
|||
/// This structure store only the authentication path for a value, the value itself is stored
|
|||
/// separately.
|
|||
#[derive(Debug)]
|
|||
pub struct PartialMmr {
|
|||
/// The version of the [Mmr].
|
|||
///
|
|||
/// This value serves the following purposes:
|
|||
///
|
|||
/// - The forest is a counter for the total number of elements in the [Mmr].
|
|||
/// - Since the [Mmr] is an append-only structure, every change to it causes a change to the
|
|||
/// `forest`, so this value has a dual purpose as a version tag.
|
|||
/// - The bits in the forest also corresponds to the count and size of every perfect binary
|
|||
/// tree that composes the [Mmr] structure, which server to compute indexes and perform
|
|||
/// validation.
|
|||
pub(crate) forest: usize,
|
|||
|
|||
/// The [Mmr] peaks.
|
|||
///
|
|||
/// The peaks are used for two reasons:
|
|||
///
|
|||
/// 1. It authenticates the addition of an element to the [PartialMmr], ensuring only valid
|
|||
/// elements are tracked.
|
|||
/// 2. During a [Mmr] update peaks can be merged by hashing the left and right hand sides. The
|
|||
/// peaks are used as the left hand.
|
|||
///
|
|||
/// All the peaks of every tree in the [Mmr] forest. The peaks are always ordered by number of
|
|||
/// leaves, starting from the peak with most children, to the one with least.
|
|||
pub(crate) peaks: Vec<RpoDigest>,
|
|||
|
|||
/// Authentication nodes used to construct merkle paths for a subset of the [Mmr]'s leaves.
|
|||
///
|
|||
/// This does not include the [Mmr]'s peaks nor the tracked nodes, only the elements required
|
|||
/// to construct their authentication paths. This property is used to detect when elements can
|
|||
/// be safely removed from, because they are no longer required to authenticate any element in
|
|||
/// the [PartialMmr].
|
|||
///
|
|||
/// The elements in the [Mmr] are referenced using a in-order tree index. This indexing scheme
|
|||
/// permits for easy computation of the relative nodes (left/right children, sibling, parent),
|
|||
/// which is useful for traversal. The indexing is also stable, meaning that merges to the
|
|||
/// trees in the [Mmr] can be represented without rewrites of the indexes.
|
|||
pub(crate) nodes: BTreeMap<InOrderIndex, RpoDigest>,
|
|||
|
|||
/// Flag indicating if the odd element should be tracked.
|
|||
///
|
|||
/// This flag is necessary because the sibling of the odd doesn't exist yet, so it can not be
|
|||
/// added into `nodes` to signal the value is being tracked.
|
|||
pub(crate) track_latest: bool,
|
|||
}
|
|||
|
|||
impl PartialMmr {
|
|||
// CONSTRUCTORS
|
|||
// --------------------------------------------------------------------------------------------
|
|||
|
|||
/// Constructs a [PartialMmr] from the given [MmrPeaks].
|
|||
pub fn from_peaks(accumulator: MmrPeaks) -> Self {
|
|||
let forest = accumulator.num_leaves();
|
|||
let peaks = accumulator.peaks().to_vec();
|
|||
let nodes = BTreeMap::new();
|
|||
let track_latest = false;
|
|||
|
|||
Self { forest, peaks, nodes, track_latest }
|
|||
}
|
|||
|
|||
// ACCESSORS
|
|||
// --------------------------------------------------------------------------------------------
|
|||
|
|||
// Gets the current `forest`.
|
|||
//
|
|||
// This value corresponds to the version of the [PartialMmr] and the number of leaves in it.
|
|||
pub fn forest(&self) -> usize {
|
|||
self.forest
|
|||
}
|
|||
|
|||
// Returns a reference to the current peaks in the [PartialMmr]
|
|||
pub fn peaks(&self) -> &[RpoDigest] {
|
|||
&self.peaks
|
|||
}
|
|||
|
|||
/// 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. If the requested value is not
|
|||
/// tracked returns `None`.
|
|||
///
|
|||
/// 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) -> Result<Option<MmrProof>, MmrError> {
|
|||
let tree_bit =
|
|||
leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::InvalidPosition(pos))?;
|
|||
let depth = tree_bit as usize;
|
|||
|
|||
let mut nodes = Vec::with_capacity(depth);
|
|||
let mut idx = InOrderIndex::from_leaf_pos(pos);
|
|||
|
|||
while let Some(node) = self.nodes.get(&idx.sibling()) {
|
|||
nodes.push(*node);
|
|||
idx = idx.parent();
|
|||
}
|
|||
|
|||
// If there are nodes then the path must be complete, otherwise it is a bug
|
|||
debug_assert!(nodes.is_empty() || nodes.len() == depth);
|
|||
|
|||
if nodes.len() != depth {
|
|||
// The requested `pos` is not being tracked.
|
|||
Ok(None)
|
|||
} else {
|
|||
Ok(Some(MmrProof {
|
|||
forest: self.forest,
|
|||
position: pos,
|
|||
merkle_path: MerklePath::new(nodes),
|
|||
}))
|
|||
}
|
|||
}
|
|||
|
|||
// MODIFIERS
|
|||
// --------------------------------------------------------------------------------------------
|
|||
|
|||
/// Add the authentication path represented by [MerklePath] if it is valid.
|
|||
///
|
|||
/// The `index` refers to the global position of the leaf in the [Mmr], these are 0-indexed
|
|||
/// values assigned in a strictly monotonic fashion as elements are inserted into the [Mmr],
|
|||
/// this value corresponds to the values used in the [Mmr] structure.
|
|||
///
|
|||
/// The `node` corresponds to the value at `index`, and `path` is the authentication path for
|
|||
/// that element up to its corresponding Mmr peak. The `node` is only used to compute the root
|
|||
/// from the authentication path to valid the data, only the authentication data is saved in
|
|||
/// the structure. If the value is required it should be stored out-of-band.
|
|||
pub fn add(
|
|||
&mut self,
|
|||
index: usize,
|
|||
node: RpoDigest,
|
|||
path: &MerklePath,
|
|||
) -> Result<(), MmrError> {
|
|||
// Checks there is a tree with same depth as the authentication path, if not the path is
|
|||
// invalid.
|
|||
let tree = 1 << path.depth();
|
|||
if tree & self.forest == 0 {
|
|||
return Err(MmrError::UnknownPeak);
|
|||
};
|
|||
|
|||
if index + 1 == self.forest
|
|||
&& path.depth() == 0
|
|||
&& self.peaks.last().map_or(false, |v| *v == node)
|
|||
{
|
|||
self.track_latest = true;
|
|||
return Ok(());
|
|||
}
|
|||
|
|||
// ignore the trees smaller than the target (these elements are position after the current
|
|||
// target and don't affect the target index)
|
|||
let target_forest = self.forest ^ (self.forest & (tree - 1));
|
|||
let peak_pos = (target_forest.count_ones() - 1) as usize;
|
|||
|
|||
// translate from mmr index to merkle path
|
|||
let path_idx = index - (target_forest ^ tree);
|
|||
|
|||
// Compute the root of the authentication path, and check it matches the current version of
|
|||
// the PartialMmr.
|
|||
let computed = path.compute_root(path_idx as u64, node).map_err(MmrError::MerkleError)?;
|
|||
if self.peaks[peak_pos] != computed {
|
|||
return Err(MmrError::InvalidPeak);
|
|||
}
|
|||
|
|||
let mut idx = InOrderIndex::from_leaf_pos(index);
|
|||
for node in path.nodes() {
|
|||
self.nodes.insert(idx.sibling(), *node);
|
|||
idx = idx.parent();
|
|||
}
|
|||
|
|||
Ok(())
|
|||
}
|
|||
|
|||
/// Remove a leaf of the [PartialMmr] and the unused nodes from the authentication path.
|
|||
///
|
|||
/// Note: `leaf_pos` corresponds to the position the [Mmr] and not on an individual tree.
|
|||
pub fn remove(&mut self, leaf_pos: usize) {
|
|||
let mut idx = InOrderIndex::from_leaf_pos(leaf_pos);
|
|||
|
|||
self.nodes.remove(&idx.sibling());
|
|||
|
|||
// `idx` represent the element that can be computed by the authentication path, because
|
|||
// these elements can be computed they are not saved for the authentication of the current
|
|||
// target. In other words, if the idx is present it was added for the authentication of
|
|||
// another element, and no more elements should be removed otherwise it would remove that
|
|||
// element's authentication data.
|
|||
while !self.nodes.contains_key(&idx) {
|
|||
idx = idx.parent();
|
|||
self.nodes.remove(&idx.sibling());
|
|||
}
|
|||
}
|
|||
|
|||
/// Applies updates to the [PartialMmr].
|
|||
pub fn apply(&mut self, delta: MmrDelta) -> Result<(), MmrError> {
|
|||
if delta.forest < self.forest {
|
|||
return Err(MmrError::InvalidPeaks);
|
|||
}
|
|||
|
|||
if delta.forest == self.forest {
|
|||
if !delta.data.is_empty() {
|
|||
return Err(MmrError::InvalidUpdate);
|
|||
}
|
|||
|
|||
return Ok(());
|
|||
}
|
|||
|
|||
// find the tree merges
|
|||
let changes = self.forest ^ delta.forest;
|
|||
let largest = 1 << changes.ilog2();
|
|||
let merges = self.forest & (largest - 1);
|
|||
|
|||
debug_assert!(
|
|||
!self.track_latest || (merges & 1) == 1,
|
|||
"if there is an odd element, a merge is required"
|
|||
);
|
|||
|
|||
// count the number elements needed to produce largest from the current state
|
|||
let (merge_count, new_peaks) = if merges != 0 {
|
|||
let depth = largest.trailing_zeros();
|
|||
let skipped = merges.trailing_zeros();
|
|||
let computed = merges.count_ones() - 1;
|
|||
let merge_count = depth - skipped - computed;
|
|||
|
|||
let new_peaks = delta.forest & (largest - 1);
|
|||
|
|||
(merge_count, new_peaks)
|
|||
} else {
|
|||
(0, changes)
|
|||
};
|
|||
|
|||
// verify the delta size
|
|||
if (delta.data.len() as u32) != merge_count + new_peaks.count_ones() {
|
|||
return Err(MmrError::InvalidUpdate);
|
|||
}
|
|||
|
|||
// keeps track of how many data elements from the update have been consumed
|
|||
let mut update_count = 0;
|
|||
|
|||
if merges != 0 {
|
|||
// starts at the smallest peak and follows the merged peaks
|
|||
let mut peak_idx = forest_to_root_index(self.forest);
|
|||
|
|||
// match order of the update data while applying it
|
|||
self.peaks.reverse();
|
|||
|
|||
// set to true when the data is needed for authentication paths updates
|
|||
let mut track = self.track_latest;
|
|||
self.track_latest = false;
|
|||
|
|||
let mut peak_count = 0;
|
|||
let mut target = 1 << merges.trailing_zeros();
|
|||
let mut new = delta.data[0];
|
|||
update_count += 1;
|
|||
|
|||
while target < largest {
|
|||
// check if either the left or right subtrees have saved for authentication paths.
|
|||
// If so, turn tracking on to update those paths.
|
|||
if target != 1 && !track {
|
|||
let left_child = peak_idx.left_child();
|
|||
let right_child = peak_idx.right_child();
|
|||
track = self.nodes.contains_key(&left_child)
|
|||
| self.nodes.contains_key(&right_child);
|
|||
}
|
|||
|
|||
// update data only contains the nodes from the right subtrees, left nodes are
|
|||
// either previously known peaks or computed values
|
|||
let (left, right) = if target & merges != 0 {
|
|||
let peak = self.peaks[peak_count];
|
|||
peak_count += 1;
|
|||
(peak, new)
|
|||
} else {
|
|||
let update = delta.data[update_count];
|
|||
update_count += 1;
|
|||
(new, update)
|
|||
};
|
|||
|
|||
if track {
|
|||
self.nodes.insert(peak_idx.sibling(), right);
|
|||
}
|
|||
|
|||
peak_idx = peak_idx.parent();
|
|||
new = Rpo256::merge(&[left, right]);
|
|||
target <<= 1;
|
|||
}
|
|||
|
|||
debug_assert!(peak_count == (merges.count_ones() as usize));
|
|||
|
|||
// restore the peaks order
|
|||
self.peaks.reverse();
|
|||
// remove the merged peaks
|
|||
self.peaks.truncate(self.peaks.len() - peak_count);
|
|||
// add the newly computed peak, the result of the merges
|
|||
self.peaks.push(new);
|
|||
}
|
|||
|
|||
// The rest of the update data is composed of peaks. None of these elements can contain
|
|||
// tracked elements because the peaks were unknown, and it is not possible to add elements
|
|||
// for tacking without authenticating it to a peak.
|
|||
self.peaks.extend_from_slice(&delta.data[update_count..]);
|
|||
self.forest = delta.forest;
|
|||
|
|||
debug_assert!(self.peaks.len() == (self.forest.count_ones() as usize));
|
|||
|
|||
Ok(())
|
|||
}
|
|||
}
|
|||
|
|||
// CONVERSIONS
|
|||
// ================================================================================================
|
|||
|
|||
impl From<MmrPeaks> for PartialMmr {
|
|||
fn from(peaks: MmrPeaks) -> Self {
|
|||
Self::from_peaks(peaks)
|
|||
}
|
|||
}
|
|||
|
|||
impl From<PartialMmr> for MmrPeaks {
|
|||
fn from(partial_mmr: PartialMmr) -> Self {
|
|||
// Safety: the [PartialMmr] maintains the constraints the number of true bits in the forest
|
|||
// matches the number of peaks, as required by the [MmrPeaks]
|
|||
MmrPeaks::new(partial_mmr.forest, partial_mmr.peaks).unwrap()
|
|||
}
|
|||
}
|
|||
|
|||
impl From<&MmrPeaks> for PartialMmr {
|
|||
fn from(peaks: &MmrPeaks) -> Self {
|
|||
Self::from_peaks(peaks.clone())
|
|||
}
|
|||
}
|
|||
|
|||
impl From<&PartialMmr> for MmrPeaks {
|
|||
fn from(partial_mmr: &PartialMmr) -> Self {
|
|||
// Safety: the [PartialMmr] maintains the constraints the number of true bits in the forest
|
|||
// matches the number of peaks, as required by the [MmrPeaks]
|
|||
MmrPeaks::new(partial_mmr.forest, partial_mmr.peaks.clone()).unwrap()
|
|||
}
|
|||
}
|
|||
|
|||
// UTILS
|
|||
// ================================================================================================
|
|||
|
|||
/// Given the description of a `forest`, returns the index of the root element of the smallest tree
|
|||
/// in it.
|
|||
pub fn forest_to_root_index(forest: usize) -> InOrderIndex {
|
|||
// Count total size of all trees in the forest.
|
|||
let nodes = nodes_in_forest(forest);
|
|||
|
|||
// Add the count for the parent nodes that separate each tree. These are allocated but
|
|||
// currently empty, and correspond to the nodes that will be used once the trees are merged.
|
|||
let open_trees = (forest.count_ones() - 1) as usize;
|
|||
|
|||
// Remove the count of the right subtree of the target tree, target tree root index comes
|
|||
// before the subtree for the in-order tree walk.
|
|||
let right_subtree_count = ((1u32 << forest.trailing_zeros()) - 1) as usize;
|
|||
|
|||
let idx = nodes + open_trees - right_subtree_count;
|
|||
|
|||
InOrderIndex::new(idx.try_into().unwrap())
|
|||
}
|
|||
|
|||
#[cfg(test)]
|
|||
mod test {
|
|||
use super::forest_to_root_index;
|
|||
use crate::merkle::InOrderIndex;
|
|||
|
|||
#[test]
|
|||
fn test_forest_to_root_index() {
|
|||
fn idx(pos: usize) -> InOrderIndex {
|
|||
InOrderIndex::new(pos.try_into().unwrap())
|
|||
}
|
|||
|
|||
// When there is a single tree in the forest, the index is equivalent to the number of
|
|||
// leaves in that tree, which is `2^n`.
|
|||
assert_eq!(forest_to_root_index(0b0001), idx(1));
|
|||
assert_eq!(forest_to_root_index(0b0010), idx(2));
|
|||
assert_eq!(forest_to_root_index(0b0100), idx(4));
|
|||
assert_eq!(forest_to_root_index(0b1000), idx(8));
|
|||
|
|||
assert_eq!(forest_to_root_index(0b0011), idx(5));
|
|||
assert_eq!(forest_to_root_index(0b0101), idx(9));
|
|||
assert_eq!(forest_to_root_index(0b1001), idx(17));
|
|||
assert_eq!(forest_to_root_index(0b0111), idx(13));
|
|||
assert_eq!(forest_to_root_index(0b1011), idx(21));
|
|||
assert_eq!(forest_to_root_index(0b1111), idx(29));
|
|||
|
|||
assert_eq!(forest_to_root_index(0b0110), idx(10));
|
|||
assert_eq!(forest_to_root_index(0b1010), idx(18));
|
|||
assert_eq!(forest_to_root_index(0b1100), idx(20));
|
|||
assert_eq!(forest_to_root_index(0b1110), idx(26));
|
|||
}
|
|||
}
|
|||