From 873a72ff981c428ef7f7b9467ccf8952ff11ddcc Mon Sep 17 00:00:00 2001
From: Kris Nuttycombe <kris@nutty.land>
Date: Thu, 6 Apr 2023 19:48:10 -0600
Subject: [PATCH] Return `MerklePath` from `ShardTree::witness`

---
 shardtree/src/lib.rs | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/shardtree/src/lib.rs b/shardtree/src/lib.rs
index 5dd267d..8d5f8bd 100644
--- a/shardtree/src/lib.rs
+++ b/shardtree/src/lib.rs
@@ -7,7 +7,7 @@ use either::Either;
 use std::collections::{BTreeMap, BTreeSet};
 use std::rc::Rc;
 
-use incrementalmerkletree::{Address, Hashable, Level, Position, Retention};
+use incrementalmerkletree::{Address, Hashable, Level, MerklePath, Position, Retention};
 
 bitflags! {
     pub struct RetentionFlags: u8 {
@@ -2078,7 +2078,7 @@ impl<
         &self,
         position: Position,
         checkpoint_depth: usize,
-    ) -> Result<Vec<H>, QueryError> {
+    ) -> Result<MerklePath<H, DEPTH>, QueryError> {
         let max_leaf_position = self
             .max_leaf_position(checkpoint_depth)
             .and_then(|v| v.ok_or_else(|| QueryError::TreeIncomplete(vec![Self::root_addr()])))?;
@@ -2105,7 +2105,7 @@ impl<
                 cur_addr = cur_addr.parent();
             }
 
-            Ok(witness)
+            Ok(MerklePath::from_parts(witness, position).unwrap())
         }
     }
 
@@ -2652,7 +2652,9 @@ mod tests {
         }
 
         fn witness(&self, position: Position, checkpoint_depth: usize) -> Option<Vec<H>> {
-            ShardTree::witness(self, position, checkpoint_depth).ok()
+            ShardTree::witness(self, position, checkpoint_depth)
+                .ok()
+                .map(|p| p.path_elems().to_vec())
         }
 
         fn remove_mark(&mut self, position: Position) -> bool {