From e3f173e5210b57a4c12ba184330eb634d863e56d Mon Sep 17 00:00:00 2001
From: Dr Maxim Orlovsky <orlovsky@pandoracore.com>
Date: Thu, 31 Mar 2022 15:12:05 +0200
Subject: [PATCH] Require taproot tree depth argument always to be u8

---
 src/util/psbt/serialize.rs |  2 +-
 src/util/taproot.rs        | 24 ++++++++++++------------
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/src/util/psbt/serialize.rs b/src/util/psbt/serialize.rs
index bb1443a5..a25126ac 100644
--- a/src/util/psbt/serialize.rs
+++ b/src/util/psbt/serialize.rs
@@ -352,7 +352,7 @@ impl Deserialize for TapTree {
 
             let leaf_version = LeafVersion::from_consensus(*version)
                 .map_err(|_| encode::Error::ParseFailed("Leaf Version Error"))?;
-            builder = builder.add_leaf_with_ver(usize::from(*depth), script, leaf_version)
+            builder = builder.add_leaf_with_ver(*depth, script, leaf_version)
                 .map_err(|_| encode::Error::ParseFailed("Tree not in DFS order"))?;
         }
         if builder.is_complete() {
diff --git a/src/util/taproot.rs b/src/util/taproot.rs
index 69d381b7..ff8ae5c9 100644
--- a/src/util/taproot.rs
+++ b/src/util/taproot.rs
@@ -423,7 +423,7 @@ impl TaprootBuilder {
     /// are not provided in DFS walk order. The depth of the root node is 0.
     pub fn add_leaf_with_ver(
         self,
-        depth: usize,
+        depth: u8,
         script: Script,
         ver: LeafVersion,
     ) -> Result<Self, TaprootBuilderError> {
@@ -435,13 +435,13 @@ impl TaprootBuilder {
     /// leaves are not provided in DFS walk order. The depth of the root node is 0.
     ///
     /// See [`TaprootBuilder::add_leaf_with_ver`] for adding a leaf with specific version.
-    pub fn add_leaf(self, depth: usize, script: Script) -> Result<Self, TaprootBuilderError> {
+    pub fn add_leaf(self, depth: u8, script: Script) -> Result<Self, TaprootBuilderError> {
         self.add_leaf_with_ver(depth, script, LeafVersion::TapScript)
     }
 
     /// Adds a hidden/omitted node at `depth` to the builder. Errors if the leaves are not provided
     /// in DFS walk order. The depth of the root node is 0.
-    pub fn add_hidden(self, depth: usize, hash: sha256::Hash) -> Result<Self, TaprootBuilderError> {
+    pub fn add_hidden(self, depth: u8, hash: sha256::Hash) -> Result<Self, TaprootBuilderError> {
         let node = NodeInfo::new_hidden(hash);
         self.insert(node, depth)
     }
@@ -473,20 +473,20 @@ impl TaprootBuilder {
     }
 
     /// Inserts a leaf at `depth`.
-    fn insert(mut self, mut node: NodeInfo, mut depth: usize) -> Result<Self, TaprootBuilderError> {
+    fn insert(mut self, mut node: NodeInfo, mut depth: u8) -> Result<Self, TaprootBuilderError> {
         // early error on invalid depth. Though this will be checked later
         // while constructing TaprootMerkelBranch
-        if depth > TAPROOT_CONTROL_MAX_NODE_COUNT {
-            return Err(TaprootBuilderError::InvalidMerkleTreeDepth(depth));
+        if depth as usize > TAPROOT_CONTROL_MAX_NODE_COUNT {
+            return Err(TaprootBuilderError::InvalidMerkleTreeDepth(depth as usize));
         }
         // We cannot insert a leaf at a lower depth while a deeper branch is unfinished. Doing
         // so would mean the add_leaf/add_hidden invocations do not correspond to a DFS traversal of a
         // binary tree.
-        if depth + 1 < self.branch.len() {
+        if depth as usize + 1 < self.branch.len() {
             return Err(TaprootBuilderError::NodeNotInDfsOrder);
         }
 
-        while self.branch.len() == depth + 1 {
+        while self.branch.len() == depth as usize + 1 {
             let child = match self.branch.pop() {
                 None => unreachable!("Len of branch checked to be >= 1"),
                 Some(Some(child)) => child,
@@ -507,14 +507,14 @@ impl TaprootBuilder {
             depth -= 1;
         }
 
-        if self.branch.len() < depth + 1 {
+        if self.branch.len() < depth as usize + 1 {
             // add enough nodes so that we can insert node at depth `depth`
-            let num_extra_nodes = depth + 1 - self.branch.len();
+            let num_extra_nodes = depth as usize + 1 - self.branch.len();
             self.branch
                 .extend((0..num_extra_nodes).into_iter().map(|_| None));
         }
         // Push the last node to the branch
-        self.branch[depth] = Some(node);
+        self.branch[depth as usize] = Some(node);
         Ok(self)
     }
 }
@@ -1270,7 +1270,7 @@ mod test {
             v: &serde_json::Value,
             mut builder: TaprootBuilder,
             leaves: &mut Vec<(Script, LeafVersion)>,
-            depth: usize,
+            depth: u8,
         ) -> TaprootBuilder {
             if v.is_null() {
                 // nothing to push