Skip to content

Commit 7e73ffe

Browse files
committed
Use TreeLike to implement num_tap_leaves
Remove recursive calls and use `TreeLike`'s post order iterator to get the total number of tap leaves for a `concrete::Policy`.
1 parent 051d63e commit 7e73ffe

File tree

1 file changed

+14
-6
lines changed

1 file changed

+14
-6
lines changed

src/policy/concrete.rs

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -681,13 +681,21 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
681681
/// and [`Policy::Threshold`] disjunctions for the `TapTree`.
682682
#[cfg(feature = "compiler")]
683683
fn num_tap_leaves(&self) -> usize {
684-
match self {
685-
Policy::Or(subs) => subs.iter().map(|(_prob, pol)| pol.num_tap_leaves()).sum(),
686-
Policy::Threshold(k, subs) if *k == 1 => {
687-
subs.iter().map(|pol| pol.num_tap_leaves()).sum()
688-
}
689-
_ => 1,
684+
use Policy::*;
685+
686+
let mut nums = vec![];
687+
for data in Arc::new(self).post_order_iter() {
688+
let num_for_child_n = |n| nums[data.child_indices[n]];
689+
690+
let num = match data.node {
691+
Or(subs) => (0..subs.len()).map(num_for_child_n).sum(),
692+
Threshold(k, subs) if *k == 1 => (0..subs.len()).map(num_for_child_n).sum(),
693+
_ => 1,
694+
};
695+
nums.push(num);
690696
}
697+
// Ok to unwrap because we know we processed at least one node.
698+
nums.pop().unwrap()
691699
}
692700

693701
/// Does checks on the number of `TapLeaf`s.

0 commit comments

Comments
 (0)