diff --git a/src/util/psbt/map/output.rs b/src/util/psbt/map/output.rs index e03049c4..a77a5729 100644 --- a/src/util/psbt/map/output.rs +++ b/src/util/psbt/map/output.rs @@ -13,6 +13,7 @@ // use prelude::*; +use core; use io; @@ -25,7 +26,7 @@ use util::psbt::map::Map; use util::psbt::raw; use util::psbt::Error; -use util::taproot::TapLeafHash; +use util::taproot::{LeafInfo, TapLeafHash}; use util::taproot::{NodeInfo, TaprootBuilder}; @@ -117,6 +118,45 @@ impl TapTree { pub fn into_inner(self) -> TaprootBuilder { 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() + } +} + +/// 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>, +} + +impl<'tree> Iterator for TapTreeIter<'tree> { + type Item = (u8, &'tree Script); + + 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"), + } + } } impl Output {