From b43902f0e6619c2c3b2c47077c42e0c8f6534937 Mon Sep 17 00:00:00 2001
From: Kris Nuttycombe <kris@nutty.land>
Date: Wed, 24 May 2023 11:27:41 -0600
Subject: [PATCH] Add the ability to compute a `MerklePath` from a Frontier

This adds a thin wrapper around `NonEmptyWitness::Frontier` that uses
the Frontier's statically known depth to be able to compute a
`MerklePath` with the same statically known depth.
---
 incrementalmerkletree/src/frontier.rs | 46 +++++++++++++++++++++++++--
 1 file changed, 43 insertions(+), 3 deletions(-)

diff --git a/incrementalmerkletree/src/frontier.rs b/incrementalmerkletree/src/frontier.rs
index a93e4bf..170a4f7 100644
--- a/incrementalmerkletree/src/frontier.rs
+++ b/incrementalmerkletree/src/frontier.rs
@@ -1,7 +1,7 @@
 use std::convert::TryFrom;
 use std::mem::size_of;
 
-use crate::{Address, Hashable, Level, Position, Source};
+use crate::{Address, Hashable, Level, MerklePath, Position, Source};
 
 #[cfg(feature = "legacy-api")]
 use {std::collections::VecDeque, std::iter::repeat};
@@ -248,6 +248,27 @@ impl<H: Hashable + Clone, const DEPTH: u8> Frontier<H, DEPTH> {
                 frontier.root(Some(DEPTH.into()))
             })
     }
+
+    /// Constructs a Merkle path that is suitable as a witness for the leaf at the tip of this
+    /// frontier by using empty roots for the right-hand ommers. This is generally only useful
+    /// for testing, so is not exposed in the public API.
+    ///
+    /// Returns `Ok(Some(MerklePath))` if successful, `Ok(None)` if the frontier is empty,
+    /// or an error containing the address of the failure.
+    pub fn witness<F>(&self, complement_nodes: F) -> Result<Option<MerklePath<H, DEPTH>>, Address>
+    where
+        F: Fn(Address) -> Option<H>,
+    {
+        self.frontier
+            .as_ref()
+            .map(|f| {
+                f.witness(DEPTH, complement_nodes).map(|path_elems| {
+                    MerklePath::from_parts(path_elems, f.position())
+                        .expect("Path length should be equal to frontier depth.")
+                })
+            })
+            .transpose()
+    }
 }
 
 #[cfg(feature = "legacy-api")]
@@ -604,9 +625,9 @@ mod tests {
     }
 
     #[test]
-    fn frontier_witness() {
+    fn nonempty_frontier_witness() {
         let mut frontier = NonEmptyFrontier::<String>::new("a".to_string());
-        for c in 'b'..'h' {
+        for c in 'b'..='g' {
             frontier.append(c.to_string());
         }
         let bridge_value_at = |addr: Address| match <u8>::from(addr.level()) {
@@ -623,6 +644,25 @@ mod tests {
         );
     }
 
+    #[test]
+    fn frontier_witness() {
+        let mut frontier = Frontier::<String, 4>::empty();
+        for c in 'a'..='g' {
+            frontier.append(c.to_string());
+        }
+
+        assert_eq!(
+            frontier
+                .witness(|addr| Some(String::empty_root(addr.level())))
+                .map(|maybe_p| maybe_p.map(|p| p.path_elems().to_vec())),
+            Ok(Some(
+                ["_", "ef", "abcd", "________"]
+                    .map(|v| v.to_string())
+                    .to_vec()
+            )),
+        );
+    }
+
     #[test]
     #[cfg(feature = "legacy-api")]
     fn test_commitment_tree_complete() {