Skip to content

Commit 9280db2

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 3fc604a commit 9280db2

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
@@ -668,13 +668,21 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
668668
/// and [`Policy::Threshold`] disjunctions for the `TapTree`.
669669
#[cfg(feature = "compiler")]
670670
fn num_tap_leaves(&self) -> usize {
671-
match self {
672-
Policy::Or(subs) => subs.iter().map(|(_prob, pol)| pol.num_tap_leaves()).sum(),
673-
Policy::Threshold(k, subs) if *k == 1 => {
674-
subs.iter().map(|pol| pol.num_tap_leaves()).sum()
675-
}
676-
_ => 1,
671+
use Policy::*;
672+
673+
let mut nums = vec![];
674+
for data in Arc::new(self).post_order_iter() {
675+
let num_for_child_n = |n| nums[data.child_indices[n]];
676+
677+
let num = match data.node {
678+
Or(subs) => (0..subs.len()).map(num_for_child_n).sum(),
679+
Threshold(k, subs) if *k == 1 => (0..subs.len()).map(num_for_child_n).sum(),
680+
_ => 1,
681+
};
682+
nums.push(num);
677683
}
684+
// Ok to unwrap because we know we processed at least one node.
685+
nums.pop().unwrap()
678686
}
679687

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

0 commit comments

Comments
 (0)