e27f8ff594 TapTree iterator implementation (Dr Maxim Orlovsky)

Pull request description:

  Implemented after @sanket1729 suggestion in https://github.com/rust-bitcoin/rust-bitcoin/issues/895#issuecomment-1074366108

  Iterates all scripts present in TapTree in DFS order returning `(depth, script)` pairs.

  I propose to have it as an RC fix since this functionality is really lacking and may be required for many wallets working with Taproot PSBT even outside of the scope where I originally needed it (OP_RETURN tweaks for TapTree described in #895)

ACKs for top commit:
  sanket1729:
    utACK e27f8ff594.
  apoelstra:
    ACK e27f8ff594

Tree-SHA512: b398e468a10534561297f22dba47e340391069734a41999edd85d726890752035053690a22014402879ea40b948160f00310f78771443d382c0bbaf0201dfbe5
This commit is contained in:
Andrew Poelstra 2022-03-28 13:42:39 +00:00
commit 388897bf93
No known key found for this signature in database
GPG Key ID: C588D63CE41B97C1
1 changed files with 41 additions and 1 deletions

View File

@ -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::Item> {
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 {