Skip to content

Commit 1612c37

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 c1836ab commit 1612c37

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
@@ -680,13 +680,21 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
680680
/// and [`Policy::Threshold`] disjunctions for the `TapTree`.
681681
#[cfg(feature = "compiler")]
682682
fn num_tap_leaves(&self) -> usize {
683-
match self {
684-
Policy::Or(subs) => subs.iter().map(|(_prob, pol)| pol.num_tap_leaves()).sum(),
685-
Policy::Threshold(k, subs) if *k == 1 => {
686-
subs.iter().map(|pol| pol.num_tap_leaves()).sum()
687-
}
688-
_ => 1,
683+
use Policy::*;
684+
685+
let mut nums = vec![];
686+
for data in Arc::new(self).post_order_iter() {
687+
let num_for_child_n = |n| nums[data.child_indices[n]];
688+
689+
let num = match data.node {
690+
Or(subs) => (0..subs.len()).map(num_for_child_n).sum(),
691+
Threshold(k, subs) if *k == 1 => (0..subs.len()).map(num_for_child_n).sum(),
692+
_ => 1,
693+
};
694+
nums.push(num);
689695
}
696+
// Ok to unwrap because we know we processed at least one node.
697+
nums.pop().unwrap()
690698
}
691699

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

0 commit comments

Comments
 (0)