diff --git a/src/merkle/mmr/partial.rs b/src/merkle/mmr/partial.rs index f6d7da3..629f923 100644 --- a/src/merkle/mmr/partial.rs +++ b/src/merkle/mmr/partial.rs @@ -88,9 +88,11 @@ impl PartialMmr { self.forest } - // Returns a reference to the current peaks in the [PartialMmr] - pub fn peaks(&self) -> &[RpoDigest] { - &self.peaks + // Returns a reference to the current peaks in the [PartialMmr]. + pub fn peaks(&self) -> MmrPeaks { + // expect() is OK here because the constructor ensures that MMR peaks can be constructed + // correctly + MmrPeaks::new(self.forest, self.peaks.clone()).expect("invalid MMR peaks") } /// Given a leaf position, returns the Merkle path to its corresponding peak. If the position diff --git a/src/merkle/mmr/peaks.rs b/src/merkle/mmr/peaks.rs index 04ed141..56e3e47 100644 --- a/src/merkle/mmr/peaks.rs +++ b/src/merkle/mmr/peaks.rs @@ -36,6 +36,14 @@ pub struct MmrPeaks { } impl MmrPeaks { + // CONSTRUCTOR + // -------------------------------------------------------------------------------------------- + + /// Returns new [MmrPeaks] instantiated from the provided vector of peaks and the number of + /// leaves in the underlying MMR. + /// + /// # Errors + /// Returns an error if the number of leaves and the number of peaks are inconsistent. pub fn new(num_leaves: usize, peaks: Vec) -> Result { if num_leaves.count_ones() as usize != peaks.len() { return Err(MmrError::InvalidPeaks); @@ -47,17 +55,23 @@ impl MmrPeaks { // ACCESSORS // -------------------------------------------------------------------------------------------- - /// Returns a count of the MMR's leaves. + /// Returns a count of leaves in the underlying MMR. pub fn num_leaves(&self) -> usize { self.num_leaves } - /// Returns the current peaks of the MMR. + /// Returns the number of peaks of the underlying MMR. + pub fn num_peaks(&self) -> usize { + self.peaks.len() + } + + /// Returns the list of peaks of the underlying MMR. pub fn peaks(&self) -> &[RpoDigest] { &self.peaks } - /// Returns the current num_leaves and peaks of the MMR. + /// Converts this [MmrPeaks] into its components: number of leaves and a vector of peaks of + /// the underlying MMR. pub fn into_parts(self) -> (usize, Vec) { (self.num_leaves, self.peaks) } diff --git a/src/merkle/mmr/tests.rs b/src/merkle/mmr/tests.rs index a92f4c6..de96bba 100644 --- a/src/merkle/mmr/tests.rs +++ b/src/merkle/mmr/tests.rs @@ -755,14 +755,14 @@ fn test_mmr_delta_old_forest() { #[test] fn test_partial_mmr_simple() { let mmr: Mmr = LEAVES.into(); - let acc = mmr.peaks(mmr.forest()).unwrap(); - let mut partial: PartialMmr = acc.clone().into(); + let peaks = mmr.peaks(mmr.forest()).unwrap(); + let mut partial: PartialMmr = peaks.clone().into(); // check initial state of the partial mmr - assert_eq!(partial.peaks(), acc.peaks()); - assert_eq!(partial.forest(), acc.num_leaves()); + assert_eq!(partial.peaks(), peaks); + assert_eq!(partial.forest(), peaks.num_leaves()); assert_eq!(partial.forest(), LEAVES.len()); - assert_eq!(partial.peaks().len(), 3); + assert_eq!(partial.peaks().num_peaks(), 3); assert_eq!(partial.nodes.len(), 0); // check state after adding tracking one element @@ -808,7 +808,7 @@ fn test_partial_mmr_update_single() { partial.apply(delta).unwrap(); assert_eq!(partial.forest(), full.forest()); - assert_eq!(partial.peaks(), full.peaks(full.forest()).unwrap().peaks()); + assert_eq!(partial.peaks(), full.peaks(full.forest()).unwrap()); let proof1 = full.open(i as usize, full.forest()).unwrap(); partial.add(proof1.position, node, &proof1.merkle_path).unwrap();