From 2edacb01917e317c75f0f6182a48fe89ca54c322 Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@electriccoin.co>
Date: Sat, 23 Nov 2024 01:27:58 +0000
Subject: [PATCH] incrementalmerkletree: Check `IncrementalWitness` at
 construction

A witness cannot exist for an empty tree.

Part of zcash/incrementalmerkletree#118.
---
 incrementalmerkletree/CHANGELOG.md   |  5 ++++
 incrementalmerkletree/src/witness.rs | 36 +++++++++++++++++-----------
 2 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/incrementalmerkletree/CHANGELOG.md b/incrementalmerkletree/CHANGELOG.md
index a9839f69..a4a8b0cc 100644
--- a/incrementalmerkletree/CHANGELOG.md
+++ b/incrementalmerkletree/CHANGELOG.md
@@ -7,6 +7,11 @@ and this project adheres to Rust's notion of
 
 ## Unreleased
 
+### Changed
+- `incrementalmerkletree::witness`:
+  - `IncrementalWitness::{from_tree, from_parts}` now return `Option<Self>`
+    (returning `None` if a witness cannot be constructed).
+
 ## [0.7.0] - 2024-09-25
 
 ### Changed
diff --git a/incrementalmerkletree/src/witness.rs b/incrementalmerkletree/src/witness.rs
index e8c79af2..47cce853 100644
--- a/incrementalmerkletree/src/witness.rs
+++ b/incrementalmerkletree/src/witness.rs
@@ -25,7 +25,7 @@ use crate::{
 ///
 /// tree.append(TestNode(0));
 /// tree.append(TestNode(1));
-/// let mut witness = IncrementalWitness::from_tree(tree.clone());
+/// let mut witness = IncrementalWitness::from_tree(tree.clone()).expect("tree is not empty");
 /// assert_eq!(witness.witnessed_position(), Position::from(1));
 /// assert_eq!(tree.root(), witness.root());
 ///
@@ -45,30 +45,38 @@ pub struct IncrementalWitness<H, const DEPTH: u8> {
 impl<H, const DEPTH: u8> IncrementalWitness<H, DEPTH> {
     /// Creates an `IncrementalWitness` for the most recent commitment added to the given
     /// [`CommitmentTree`].
-    pub fn from_tree(tree: CommitmentTree<H, DEPTH>) -> Self {
-        IncrementalWitness {
+    ///
+    /// Returns `None` if `tree` is empty (and thus there is no position to witness).
+    pub fn from_tree(tree: CommitmentTree<H, DEPTH>) -> Option<Self> {
+        (!tree.is_empty()).then(|| IncrementalWitness {
             tree,
             filled: vec![],
             cursor_depth: 0,
             cursor: None,
-        }
+        })
     }
 
+    /// Constructs an `IncrementalWitness` from its parts.
+    ///
+    /// Returns `None` if the parts do not form a valid witness, for example if `tree` is
+    /// empty (and thus there is no position to witness).
     pub fn from_parts(
         tree: CommitmentTree<H, DEPTH>,
         filled: Vec<H>,
         cursor: Option<CommitmentTree<H, DEPTH>>,
-    ) -> Self {
-        let mut witness = IncrementalWitness {
-            tree,
-            filled,
-            cursor_depth: 0,
-            cursor,
-        };
+    ) -> Option<Self> {
+        (!tree.is_empty()).then(|| {
+            let mut witness = IncrementalWitness {
+                tree,
+                filled,
+                cursor_depth: 0,
+                cursor,
+            };
 
-        witness.cursor_depth = witness.next_depth();
+            witness.cursor_depth = witness.next_depth();
 
-        witness
+            witness
+        })
     }
 
     pub fn tree(&self) -> &CommitmentTree<H, DEPTH> {
@@ -248,7 +256,7 @@ mod tests {
         for c in 'a'..'h' {
             base_tree.append(c.to_string()).unwrap();
         }
-        let mut witness = IncrementalWitness::from_tree(base_tree);
+        let mut witness = IncrementalWitness::from_tree(base_tree).unwrap();
         for c in 'h'..'z' {
             witness.append(c.to_string()).unwrap();
         }