diff --git a/src/util/psbt/map/output.rs b/src/util/psbt/map/output.rs index c49eb760a3..855524d5de 100644 --- a/src/util/psbt/map/output.rs +++ b/src/util/psbt/map/output.rs @@ -26,7 +26,7 @@ use util::psbt::map::Map; use util::psbt::raw; use util::psbt::Error; -use util::taproot::{LeafInfo, TapLeafHash}; +use util::taproot::{ScriptLeaf, TapLeafHash}; use util::taproot::{NodeInfo, TaprootBuilder}; @@ -155,43 +155,33 @@ impl TapTree { self.0 } - /// Returns iterator for a taproot script tree, operating in DFS order over leaf depth and - /// leaf script pairs. - pub fn iter(&self) -> TapTreeIter { - self.into_iter() + /// Returns [`TapTreeIter`] iterator for a taproot script tree, operating in DFS order over + /// tree [`ScriptLeaf`]s. + pub fn script_leaves(&self) -> TapTreeIter { + match (self.0.branch().len(), self.0.branch().last()) { + (1, Some(Some(root))) => { + TapTreeIter { + leaf_iter: root.leaves.iter() + } + } + // This should be unreachable as we Taptree is already finalized + _ => unreachable!("non-finalized tree builder inside TapTree"), + } } } /// Iterator for a taproot script tree, operating in DFS order over leaf depth and /// leaf script pairs. pub struct TapTreeIter<'tree> { - leaf_iter: core::slice::Iter<'tree, LeafInfo>, + leaf_iter: core::slice::Iter<'tree, ScriptLeaf>, } impl<'tree> Iterator for TapTreeIter<'tree> { - type Item = (u8, &'tree Script); + type Item = &'tree ScriptLeaf; + #[inline] fn next(&mut self) -> Option { - self.leaf_iter.next().map(|leaf_info| { - (leaf_info.merkle_branch.as_inner().len() as u8, &leaf_info.script) - }) - } -} - -impl<'tree> IntoIterator for &'tree TapTree { - type Item = (u8, &'tree Script); - type IntoIter = TapTreeIter<'tree>; - - fn into_iter(self) -> Self::IntoIter { - match (self.0.branch().len(), self.0.branch().last()) { - (1, Some(Some(root))) => { - TapTreeIter { - leaf_iter: root.leaves.iter() - } - } - // This should be unreachable as we Taptree is already finalized - _ => unreachable!("non-finalized tree builder inside TapTree"), - } + self.leaf_iter.next() } } diff --git a/src/util/psbt/serialize.rs b/src/util/psbt/serialize.rs index 438693f902..3c946ec773 100644 --- a/src/util/psbt/serialize.rs +++ b/src/util/psbt/serialize.rs @@ -327,9 +327,9 @@ impl Serialize for TapTree { // // TaprootMerkleBranch can only have len atmost 128(TAPROOT_CONTROL_MAX_NODE_COUNT). // safe to cast from usize to u8 - buf.push(leaf_info.merkle_branch.as_inner().len() as u8); - buf.push(leaf_info.ver.to_consensus()); - leaf_info.script.consensus_encode(&mut buf).expect("Vecs dont err"); + buf.push(leaf_info.merkle_branch().as_inner().len() as u8); + buf.push(leaf_info.leaf_version().to_consensus()); + leaf_info.script().consensus_encode(&mut buf).expect("Vecs dont err"); } buf } diff --git a/src/util/taproot.rs b/src/util/taproot.rs index 325feaa65a..512d6a9f83 100644 --- a/src/util/taproot.rs +++ b/src/util/taproot.rs @@ -560,7 +560,7 @@ pub struct NodeInfo { /// Merkle hash for this node. pub(crate) hash: sha256::Hash, /// Information about leaves inside this node. - pub(crate) leaves: Vec, + pub(crate) leaves: Vec, /// Tracks information on hidden nodes below this node. pub(crate) has_hidden_nodes: bool, } @@ -577,9 +577,9 @@ impl NodeInfo { /// Creates a new leaf [`NodeInfo`] with given [`Script`] and [`LeafVersion`]. pub fn new_leaf_with_ver(script: Script, ver: LeafVersion) -> Self { - let leaf = LeafInfo::new(script, ver); + let leaf = ScriptLeaf::new(script, ver); Self { - hash: leaf.hash(), + hash: sha256::Hash::from_inner(leaf.leaf_hash().into_inner()), leaves: vec![leaf], has_hidden_nodes: false, } @@ -608,17 +608,17 @@ impl NodeInfo { /// Store information about taproot leaf node. #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] -pub(crate) struct LeafInfo { +pub struct ScriptLeaf { /// The underlying script. - pub(crate) script: Script, + script: Script, /// The leaf version. - pub(crate) ver: LeafVersion, + ver: LeafVersion, /// The merkle proof (hashing partners) to get this node. - pub(crate) merkle_branch: TaprootMerkleBranch, + merkle_branch: TaprootMerkleBranch, } -impl LeafInfo { - /// Creates an new [`LeafInfo`] from `script` and `ver` and no merkle branch. +impl ScriptLeaf { + /// Creates an new [`ScriptLeaf`] from `script` and `ver` and no merkle branch. fn new(script: Script, ver: LeafVersion) -> Self { Self { script: script, @@ -627,10 +627,37 @@ impl LeafInfo { } } - /// Computes a leaf hash for this [`LeafInfo`]. - fn hash(&self) -> sha256::Hash { - let leaf_hash = TapLeafHash::from_script(&self.script, self.ver); - sha256::Hash::from_inner(leaf_hash.into_inner()) + /// Returns the depth of this script leaf in the tap tree. + #[inline] + pub fn depth(&self) -> u8 { + // The depth is guaranteed to be < 127 by the TaprootBuilder type. + // TODO: Following MSRV bump implement via `try_into().expect("")`. + self.merkle_branch.0.len() as u8 + } + + /// Computes a leaf hash for this [`ScriptLeaf`]. + #[inline] + pub fn leaf_hash(&self) -> TapLeafHash { + TapLeafHash::from_script(&self.script, self.ver) + } + + /// Returns reference to the leaf script. + #[inline] + pub fn script(&self) -> &Script { + &self.script + } + + /// Returns leaf version of the script. + #[inline] + pub fn leaf_version(&self) -> LeafVersion { + self.ver + } + + /// Returns reference to the merkle proof (hashing partners) to get this + /// node in form of [`TaprootMerkleBranch`]. + #[inline] + pub fn merkle_branch(&self) -> &TaprootMerkleBranch { + &self.merkle_branch } }