Skip to content

Commit 82f58ce

Browse files
committed
tr: move TapTreeIter into taptree module
We are going to move TapTree into the module (hence the name), and before we can do this we need to move the iterator.
1 parent 2232979 commit 82f58ce

File tree

3 files changed

+54
-48
lines changed

3 files changed

+54
-48
lines changed

src/descriptor/mod.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ pub use self::bare::{Bare, Pkh};
4242
pub use self::segwitv0::{Wpkh, Wsh, WshInner};
4343
pub use self::sh::{Sh, ShInner};
4444
pub use self::sortedmulti::SortedMultiVec;
45-
pub use self::tr::{TapTree, Tr};
45+
pub use self::tr::{TapTree, TapTreeIter, TapTreeIterItem, Tr};
4646

4747
pub mod checksum;
4848
mod key;

src/descriptor/tr/mod.rs

Lines changed: 3 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ use crate::{
2828

2929
mod taptree;
3030

31-
pub use self::taptree::TapTreeIterItem;
31+
pub use self::taptree::{TapTreeIter, TapTreeIterItem};
3232

3333
/// A Taproot Tree representation.
3434
// Hidden leaves are not yet supported in descriptor spec. Conceptually, it should
@@ -132,7 +132,7 @@ impl<Pk: MiniscriptKey> TapTree<Pk> {
132132

133133
/// Iterates over all miniscripts in DFS walk order compatible with the
134134
/// PSBT requirements (BIP 371).
135-
pub fn iter(&self) -> TapTreeIter<Pk> { TapTreeIter { stack: vec![(0, self)] } }
135+
pub fn iter(&self) -> TapTreeIter<Pk> { TapTreeIter::from_tree(self) }
136136

137137
// Helper function to translate keys
138138
fn translate_helper<T>(&self, t: &mut T) -> Result<TapTree<T::TargetPk>, TranslateErr<T::Error>>
@@ -201,7 +201,7 @@ impl<Pk: MiniscriptKey> Tr<Pk> {
201201
pub fn iter_scripts(&self) -> TapTreeIter<Pk> {
202202
match self.tree {
203203
Some(ref t) => t.iter(),
204-
None => TapTreeIter { stack: vec![] },
204+
None => TapTreeIter::empty(),
205205
}
206206
}
207207

@@ -451,48 +451,6 @@ impl Tr<DefiniteDescriptorKey> {
451451
}
452452
}
453453

454-
/// Iterator for Taproot structures
455-
/// Yields a pair of (depth, miniscript) in a depth first walk
456-
/// For example, this tree:
457-
/// - N0 -
458-
/// / \\
459-
/// N1 N2
460-
/// / \ / \\
461-
/// A B C N3
462-
/// / \\
463-
/// D E
464-
/// would yield (2, A), (2, B), (2,C), (3, D), (3, E).
465-
///
466-
#[derive(Debug, Clone)]
467-
pub struct TapTreeIter<'a, Pk: MiniscriptKey> {
468-
stack: Vec<(u8, &'a TapTree<Pk>)>,
469-
}
470-
471-
impl<Pk: MiniscriptKey> TapTreeIter<'_, Pk> {
472-
/// Helper function to return an empty iterator from Descriptor::tap_tree_iter.
473-
pub(super) fn empty() -> Self { Self { stack: vec![] } }
474-
}
475-
476-
impl<'a, Pk> Iterator for TapTreeIter<'a, Pk>
477-
where
478-
Pk: MiniscriptKey + 'a,
479-
{
480-
type Item = TapTreeIterItem<'a, Pk>;
481-
482-
fn next(&mut self) -> Option<Self::Item> {
483-
while let Some((depth, last)) = self.stack.pop() {
484-
match *last {
485-
TapTree::Tree { ref left, ref right, height: _ } => {
486-
self.stack.push((depth + 1, right));
487-
self.stack.push((depth + 1, left));
488-
}
489-
TapTree::Leaf(ref ms) => return Some(TapTreeIterItem { node: ms, depth }),
490-
}
491-
}
492-
None
493-
}
494-
}
495-
496454
impl<Pk: FromStrKey> core::str::FromStr for Tr<Pk> {
497455
type Err = Error;
498456

src/descriptor/tr/taptree.rs

Lines changed: 50 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,18 +2,66 @@
22

33
use bitcoin::taproot::{LeafVersion, TapLeafHash};
44

5+
use super::TapTree;
56
use crate::miniscript::context::Tap;
7+
use crate::prelude::Vec;
68
use crate::sync::Arc;
79
use crate::{Miniscript, MiniscriptKey, ToPublicKey};
810

11+
/// Iterator over the leaves of a Taptree.
12+
///
13+
/// Yields a pair of (depth, miniscript) in a depth first walk
14+
/// For example, this tree:
15+
/// - N0 -
16+
/// / \\
17+
/// N1 N2
18+
/// / \ / \\
19+
/// A B C N3
20+
/// / \\
21+
/// D E
22+
/// would yield (2, A), (2, B), (2,C), (3, D), (3, E).
23+
///
24+
#[derive(Debug, Clone)]
25+
pub struct TapTreeIter<'a, Pk: MiniscriptKey> {
26+
stack: Vec<(u8, &'a TapTree<Pk>)>,
27+
}
28+
29+
impl<'tr, Pk: MiniscriptKey> TapTreeIter<'tr, Pk> {
30+
/// An empty iterator.
31+
pub fn empty() -> Self { Self { stack: vec![] } }
32+
33+
/// An iterator over a given tree.
34+
pub(super) fn from_tree(tree: &'tr TapTree<Pk>) -> Self { Self { stack: vec![(0, tree)] } }
35+
}
36+
37+
impl<'a, Pk> Iterator for TapTreeIter<'a, Pk>
38+
where
39+
Pk: MiniscriptKey + 'a,
40+
{
41+
type Item = TapTreeIterItem<'a, Pk>;
42+
43+
fn next(&mut self) -> Option<Self::Item> {
44+
while let Some((depth, last)) = self.stack.pop() {
45+
match *last {
46+
TapTree::Tree { ref left, ref right, height: _ } => {
47+
self.stack.push((depth + 1, right));
48+
self.stack.push((depth + 1, left));
49+
}
50+
TapTree::Leaf(ref ms) => return Some(TapTreeIterItem { node: ms, depth }),
51+
}
52+
}
53+
None
54+
}
55+
}
56+
957
/// Iterator over all of the leaves of a Taproot tree.
1058
///
1159
/// If there is no tree (i.e. this is a keyspend-only Taproot descriptor)
1260
/// then the iterator will yield nothing.
1361
#[derive(Clone, PartialEq, Eq, Debug)]
1462
pub struct TapTreeIterItem<'tr, Pk: MiniscriptKey> {
15-
pub(super) node: &'tr Arc<Miniscript<Pk, Tap>>,
16-
pub(super) depth: u8,
63+
node: &'tr Arc<Miniscript<Pk, Tap>>,
64+
depth: u8,
1765
}
1866

1967
impl<'tr, Pk: MiniscriptKey> TapTreeIterItem<'tr, Pk> {

0 commit comments

Comments
 (0)