From 5fcf98669db8a50bc59197f1ce9b914f9ec55486 Mon Sep 17 00:00:00 2001
From: Bobbin Threadbare <bobbinth@protonmail.com>
Date: Mon, 29 Jan 2024 17:58:00 -0800
Subject: [PATCH] feat: add PartialMmr::from_parts() constructor

---
 src/merkle/mmr/partial.rs | 24 ++++++++++++++++++++----
 src/merkle/mmr/peaks.rs   |  6 ++++++
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/src/merkle/mmr/partial.rs b/src/merkle/mmr/partial.rs
index e18b922..2241ebb 100644
--- a/src/merkle/mmr/partial.rs
+++ b/src/merkle/mmr/partial.rs
@@ -10,6 +10,11 @@ use crate::{
     },
 };
 
+// TYPE ALIASES
+// ================================================================================================
+
+type NodeMap = BTreeMap<InOrderIndex, RpoDigest>;
+
 // PARTIAL MERKLE MOUNTAIN RANGE
 // ================================================================================================
 /// Partially materialized Merkle Mountain Range (MMR), used to efficiently store and update the
@@ -55,7 +60,7 @@ pub struct PartialMmr {
     /// 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>,
+    pub(crate) nodes: NodeMap,
 
     /// Flag indicating if the odd element should be tracked.
     ///
@@ -68,16 +73,27 @@ impl PartialMmr {
     // CONSTRUCTORS
     // --------------------------------------------------------------------------------------------
 
-    /// Constructs a [PartialMmr] from the given [MmrPeaks].
+    /// Returns a new [PartialMmr] instantiated from the specified peaks.
     pub fn from_peaks(peaks: MmrPeaks) -> Self {
         let forest = peaks.num_leaves();
-        let peaks = peaks.peaks().to_vec();
+        let peaks = peaks.into();
         let nodes = BTreeMap::new();
         let track_latest = false;
 
         Self { forest, peaks, nodes, track_latest }
     }
 
+    /// Returns a new [PartialMmr] instantiated from the specified components.
+    ///
+    /// This constructor does not check the consistency between peaks and nodes. If the specified
+    /// peaks are nodes are inconsistent, the returned partial MMR may exhibit undefined behavior.
+    pub fn from_parts(peaks: MmrPeaks, nodes: NodeMap, track_latest: bool) -> Self {
+        let forest = peaks.num_leaves();
+        let peaks = peaks.into();
+
+        Self { forest, peaks, nodes, track_latest }
+    }
+
     // PUBLIC ACCESSORS
     // --------------------------------------------------------------------------------------------
 
@@ -516,7 +532,7 @@ impl From<&PartialMmr> for MmrPeaks {
 
 /// An iterator over every inner node of the [PartialMmr].
 pub struct InnerNodeIterator<'a, I: Iterator<Item = (usize, RpoDigest)>> {
-    nodes: &'a BTreeMap<InOrderIndex, RpoDigest>,
+    nodes: &'a NodeMap,
     leaves: I,
     stack: Vec<(InOrderIndex, RpoDigest)>,
     seen_nodes: BTreeSet<InOrderIndex>,
diff --git a/src/merkle/mmr/peaks.rs b/src/merkle/mmr/peaks.rs
index 5822ef9..82bbea4 100644
--- a/src/merkle/mmr/peaks.rs
+++ b/src/merkle/mmr/peaks.rs
@@ -132,3 +132,9 @@ impl MmrPeaks {
         elements
     }
 }
+
+impl From<MmrPeaks> for Vec<RpoDigest> {
+    fn from(peaks: MmrPeaks) -> Self {
+        peaks.peaks
+    }
+}