Skip to content

Commit

Permalink
feat: improve MMR api (#324)
Browse files Browse the repository at this point in the history
  • Loading branch information
@bobbinth
bobbinth authored Aug 18, 2024
1 parent d92fae7 commit a12e62f
Show file tree
Hide file tree
Showing 4 changed files with 117 additions and 92 deletions.
5 changes: 5 additions & 0 deletions CHANGELOG.md
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,8 @@
## 0.11.0 (TBD)

- [BREAKING]: renamed `Mmr::open()` into `Mmr::open_at()` and `Mmr::peaks()` into `Mmr::peaks_at()` (#234).
- Added `Mmr::open()` and `Mmr::peaks()` which rely on `Mmr::open_at()` and `Mmr::peaks()` respectively (#234).

## 0.10.0 (2024-08-06)

* Added more `RpoDigest` and `RpxDigest` conversions (#311).
Expand Down
41 changes: 33 additions & 8 deletions src/merkle/mmr/full.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -72,19 +72,36 @@ impl Mmr {
// 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.
/// Returns an [MmrProof] for the leaf at the specified position.
///
/// 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> {
///
/// # Errors
/// Returns an error if the specified leaf position is out of bounds for this MMR.
pub fn open(&self, pos: usize) -> Result<MmrProof, MmrError> {
self.open_at(pos, self.forest)
}

/// Returns an [MmrProof] for the leaf at the specified position using the state of the MMR
/// at the specified `forest`.
///
/// 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.
///
/// # Errors
/// Returns an error if:
/// - The specified leaf position is out of bounds for this MMR.
/// - The specified `forest` value is not valid for this MMR.
pub fn open_at(&self, pos: usize, 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))?;
leaf_to_corresponding_tree(pos, forest).ok_or(MmrError::InvalidPosition(pos))?;

// isolate the trees before the target
let forest_before = target_forest & high_bitmask(tree_bit + 1);
let forest_before = forest & high_bitmask(tree_bit + 1);
let index_offset = nodes_in_forest(forest_before);

// update the value position from global to the target tree
Expand All @@ -94,7 +111,7 @@ impl Mmr {
let (_, path) = self.collect_merkle_path_and_value(tree_bit, relative_pos, index_offset);

Ok(MmrProof {
forest: target_forest,
forest,
position: pos,
merkle_path: MerklePath::new(path),
})
Expand Down Expand Up @@ -145,8 +162,16 @@ impl Mmr {
self.forest += 1;
}

/// Returns an peaks of the MMR for the version specified by `forest`.
pub fn peaks(&self, forest: usize) -> Result<MmrPeaks, MmrError> {
/// Returns the current peaks of the MMR.
pub fn peaks(&self) -> MmrPeaks {
self.peaks_at(self.forest).expect("failed to get peaks at current forest")
}

/// Returns the peaks of the MMR at the state specified by `forest`.
///
/// # Errors
/// Returns an error if the specified `forest` value is not valid for this MMR.
pub fn peaks_at(&self, forest: usize) -> Result<MmrPeaks, MmrError> {
if forest > self.forest {
return Err(MmrError::InvalidPeaks);
}
Expand Down
44 changes: 21 additions & 23 deletions src/merkle/mmr/partial.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -688,18 +688,18 @@ mod tests {
// build an MMR with 10 nodes (2 peaks) and a partial MMR based on it
let mut mmr = Mmr::default();
(0..10).for_each(|i| mmr.add(int_to_node(i)));
let mut partial_mmr: PartialMmr = mmr.peaks(mmr.forest()).unwrap().into();
let mut partial_mmr: PartialMmr = mmr.peaks().into();

// add authentication path for position 1 and 8
{
let node = mmr.get(1).unwrap();
let proof = mmr.open(1, mmr.forest()).unwrap();
let proof = mmr.open(1).unwrap();
partial_mmr.track(1, node, &proof.merkle_path).unwrap();
}

{
let node = mmr.get(8).unwrap();
let proof = mmr.open(8, mmr.forest()).unwrap();
let proof = mmr.open(8).unwrap();
partial_mmr.track(8, node, &proof.merkle_path).unwrap();
}

Expand All @@ -712,7 +712,7 @@ mod tests {
validate_apply_delta(&mmr, &mut partial_mmr);
{
let node = mmr.get(12).unwrap();
let proof = mmr.open(12, mmr.forest()).unwrap();
let proof = mmr.open(12).unwrap();
partial_mmr.track(12, node, &proof.merkle_path).unwrap();
assert!(partial_mmr.track_latest);
}
Expand All @@ -737,7 +737,7 @@ mod tests {
let nodes_delta = partial.apply(delta).unwrap();

// new peaks were computed correctly
assert_eq!(mmr.peaks(mmr.forest()).unwrap(), partial.peaks());
assert_eq!(mmr.peaks(), partial.peaks());

let mut expected_nodes = nodes_before;
for (key, value) in nodes_delta {
Expand All @@ -753,7 +753,7 @@ mod tests {
let index_value: u64 = index.into();
let pos = index_value / 2;
let proof1 = partial.open(pos as usize).unwrap().unwrap();
let proof2 = mmr.open(pos as usize, mmr.forest()).unwrap();
let proof2 = mmr.open(pos as usize).unwrap();
assert_eq!(proof1, proof2);
}
}
Expand All @@ -762,16 +762,16 @@ mod tests {
fn test_partial_mmr_inner_nodes_iterator() {
// build the MMR
let mmr: Mmr = LEAVES.into();
let first_peak = mmr.peaks(mmr.forest).unwrap().peaks()[0];
let first_peak = mmr.peaks().peaks()[0];

// -- test single tree ----------------------------

// get path and node for position 1
let node1 = mmr.get(1).unwrap();
let proof1 = mmr.open(1, mmr.forest()).unwrap();
let proof1 = mmr.open(1).unwrap();

// create partial MMR and add authentication path to node at position 1
let mut partial_mmr: PartialMmr = mmr.peaks(mmr.forest()).unwrap().into();
let mut partial_mmr: PartialMmr = mmr.peaks().into();
partial_mmr.track(1, node1, &proof1.merkle_path).unwrap();

// empty iterator should have no nodes
Expand All @@ -789,13 +789,13 @@ mod tests {
// -- test no duplicates --------------------------

// build the partial MMR
let mut partial_mmr: PartialMmr = mmr.peaks(mmr.forest()).unwrap().into();
let mut partial_mmr: PartialMmr = mmr.peaks().into();

let node0 = mmr.get(0).unwrap();
let proof0 = mmr.open(0, mmr.forest()).unwrap();
let proof0 = mmr.open(0).unwrap();

let node2 = mmr.get(2).unwrap();
let proof2 = mmr.open(2, mmr.forest()).unwrap();
let proof2 = mmr.open(2).unwrap();

partial_mmr.track(0, node0, &proof0.merkle_path).unwrap();
partial_mmr.track(1, node1, &proof1.merkle_path).unwrap();
Expand Down Expand Up @@ -826,10 +826,10 @@ mod tests {
// -- test multiple trees -------------------------

// build the partial MMR
let mut partial_mmr: PartialMmr = mmr.peaks(mmr.forest()).unwrap().into();
let mut partial_mmr: PartialMmr = mmr.peaks().into();

let node5 = mmr.get(5).unwrap();
let proof5 = mmr.open(5, mmr.forest()).unwrap();
let proof5 = mmr.open(5).unwrap();

partial_mmr.track(1, node1, &proof1.merkle_path).unwrap();
partial_mmr.track(5, node5, &proof5.merkle_path).unwrap();
Expand All @@ -841,7 +841,7 @@ mod tests {
let index1 = NodeIndex::new(2, 1).unwrap();
let index5 = NodeIndex::new(1, 1).unwrap();

let second_peak = mmr.peaks(mmr.forest).unwrap().peaks()[1];
let second_peak = mmr.peaks().peaks()[1];

let path1 = store.get_path(first_peak, index1).unwrap().path;
let path5 = store.get_path(second_peak, index5).unwrap().path;
Expand All @@ -860,8 +860,7 @@ mod tests {
mmr.add(el);
partial_mmr.add(el, false);

let mmr_peaks = mmr.peaks(mmr.forest()).unwrap();
assert_eq!(mmr_peaks, partial_mmr.peaks());
assert_eq!(mmr.peaks(), partial_mmr.peaks());
assert_eq!(mmr.forest(), partial_mmr.forest());
}
}
Expand All @@ -877,12 +876,11 @@ mod tests {
mmr.add(el);
partial_mmr.add(el, true);

let mmr_peaks = mmr.peaks(mmr.forest()).unwrap();
assert_eq!(mmr_peaks, partial_mmr.peaks());
assert_eq!(mmr.peaks(), partial_mmr.peaks());
assert_eq!(mmr.forest(), partial_mmr.forest());

for pos in 0..i {
let mmr_proof = mmr.open(pos as usize, mmr.forest()).unwrap();
let mmr_proof = mmr.open(pos as usize).unwrap();
let partialmmr_proof = partial_mmr.open(pos as usize).unwrap().unwrap();
assert_eq!(mmr_proof, partialmmr_proof);
}
Expand All @@ -894,8 +892,8 @@ mod tests {
let mut mmr = Mmr::from((0..7).map(int_to_node));

// derive a partial Mmr from it which tracks authentication path to leaf 5
let mut partial_mmr = PartialMmr::from_peaks(mmr.peaks(mmr.forest()).unwrap());
let path_to_5 = mmr.open(5, mmr.forest()).unwrap().merkle_path;
let mut partial_mmr = PartialMmr::from_peaks(mmr.peaks());
let path_to_5 = mmr.open(5).unwrap().merkle_path;
let leaf_at_5 = mmr.get(5).unwrap();
partial_mmr.track(5, leaf_at_5, &path_to_5).unwrap();

Expand All @@ -905,6 +903,6 @@ mod tests {
partial_mmr.add(leaf_at_7, false);

// the openings should be the same
assert_eq!(mmr.open(5, mmr.forest()).unwrap(), partial_mmr.open(5).unwrap().unwrap());
assert_eq!(mmr.open(5).unwrap(), partial_mmr.open(5).unwrap().unwrap());
}
}
Loading

0 comments on commit a12e62f

Please sign in to comment.