tr: encapsulate object which is yielded by TapTreeIter

All this does is wrap the existing (depth, script) return value in a
structure that has a pile of accessors on it. The accessors include
a couple of expensive computations which are currently done manually.

A later commit/PR will introduce a `SpendInfo` structure which will
cache these computations and amortize much of the hashing needed to
produce control blocks. But these new accessors simplify some things.

Author: apoelstra

bitcoind-tests/tests/test_desc.rs

@@ -186,9 +186,9 @@ pub fn test_desc_satisfy(
             // ------------------ script spend -------------
             let x_only_keypairs_reqd: Vec<(secp256k1::Keypair, TapLeafHash)> = tr
                 .iter_scripts()
-                .flat_map(|(_depth, ms)| {
-                    let leaf_hash = TapLeafHash::from_script(&ms.encode(), LeafVersion::TapScript);
-                    ms.iter_pk().filter_map(move |pk| {
+                .flat_map(|leaf| {
+                    let leaf_hash = TapLeafHash::from_script(&leaf.compute_script(), LeafVersion::TapScript);
+                    leaf.miniscript().iter_pk().filter_map(move |pk| {
                         let i = x_only_pks.iter().position(|&x| x.to_public_key() == pk);
                         i.map(|idx| (xonly_keypairs[idx], leaf_hash))
                     })

examples/taproot.rs

@@ -67,15 +67,17 @@ fn main() {
 
         // Iterate through scripts
         let mut iter = p.iter_scripts();
+        let mut next = iter.next().unwrap();
         assert_eq!(
-            iter.next().unwrap(),
+            (next.depth(), next.miniscript().as_ref()),
             (
                 1u8,
                 &Miniscript::<String, Tap>::from_str("and_v(vc:pk_k(In),older(9))").unwrap()
             )
         );
+        next = iter.next().unwrap();
         assert_eq!(
-            iter.next().unwrap(),
+            (next.depth(), next.miniscript().as_ref()),
             (1u8, &Miniscript::<String, Tap>::from_str("and_v(v:pk(hA),pk(S))").unwrap())
         );
         assert_eq!(iter.next(), None);

src/descriptor/mod.rs

@@ -988,8 +988,9 @@ impl<Pk: FromStrKey> FromStr for Descriptor<Pk> {
             // FIXME preserve weird/broken behavior from 12.x.
             // See https://github.com/rust-bitcoin/rust-miniscript/issues/734
             ret.sanity_check()?;
-            for (_, ms) in inner.iter_scripts() {
-                ms.ext_check(&crate::miniscript::analyzable::ExtParams::sane())?;
+            for item in inner.iter_scripts() {
+                item.miniscript()
+                    .ext_check(&crate::miniscript::analyzable::ExtParams::sane())?;
             }
         }
         Ok(ret)

src/descriptor/tr/mod.rs

@@ -26,6 +26,10 @@ use crate::{
     Threshold, ToPublicKey, TranslateErr, Translator,
 };
 
+mod taptree;
+
+pub use self::taptree::TapTreeIterItem;
+
 /// A Taproot Tree representation.
 // Hidden leaves are not yet supported in descriptor spec. Conceptually, it should
 // be simple to integrate those here, but it is best to wait on core for the exact syntax.
@@ -226,10 +230,10 @@ impl<Pk: MiniscriptKey> Tr<Pk> {
             TaprootSpendInfo::new_key_spend(&secp, self.internal_key.to_x_only_pubkey(), None)
         } else {
             let mut builder = TaprootBuilder::new();
-            for (depth, ms) in self.iter_scripts() {
-                let script = ms.encode();
+            for leaf in self.iter_scripts() {
+                let script = leaf.miniscript().encode();
                 builder = builder
-                    .add_leaf(depth, script)
+                    .add_leaf(leaf.depth(), script)
                     .expect("Computing spend data on a valid Tree should always succeed");
             }
             // Assert builder cannot error here because we have a well formed descriptor
@@ -245,8 +249,8 @@ impl<Pk: MiniscriptKey> Tr<Pk> {
 
     /// Checks whether the descriptor is safe.
     pub fn sanity_check(&self) -> Result<(), Error> {
-        for (_depth, ms) in self.iter_scripts() {
-            ms.sanity_check()?;
+        for leaf in self.iter_scripts() {
+            leaf.miniscript().sanity_check()?;
         }
         Ok(())
     }
@@ -276,11 +280,11 @@ impl<Pk: MiniscriptKey> Tr<Pk> {
 
         let wu = tree
             .iter()
-            .filter_map(|(depth, ms)| {
-                let script_size = ms.script_size();
-                let max_sat_elems = ms.max_satisfaction_witness_elements().ok()?;
-                let max_sat_size = ms.max_satisfaction_size().ok()?;
-                let control_block_size = control_block_len(depth);
+            .filter_map(|leaf| {
+                let script_size = leaf.miniscript().script_size();
+                let max_sat_elems = leaf.miniscript().max_satisfaction_witness_elements().ok()?;
+                let max_sat_size = leaf.miniscript().max_satisfaction_size().ok()?;
+                let control_block_size = control_block_len(leaf.depth());
 
                 // stack varint difference (+1 for ctrl block, witness script already included)
                 let stack_varint_diff = varint_len(max_sat_elems + 1) - varint_len(0);
@@ -326,11 +330,11 @@ impl<Pk: MiniscriptKey> Tr<Pk> {
         };
 
         tree.iter()
-            .filter_map(|(depth, ms)| {
-                let script_size = ms.script_size();
-                let max_sat_elems = ms.max_satisfaction_witness_elements().ok()?;
-                let max_sat_size = ms.max_satisfaction_size().ok()?;
-                let control_block_size = control_block_len(depth);
+            .filter_map(|leaf| {
+                let script_size = leaf.miniscript().script_size();
+                let max_sat_elems = leaf.miniscript().max_satisfaction_witness_elements().ok()?;
+                let max_sat_size = leaf.miniscript().max_satisfaction_size().ok()?;
+                let control_block_size = control_block_len(leaf.depth());
                 Some(
                     // scriptSig len byte
                     4 +
@@ -473,7 +477,7 @@ impl<'a, Pk> Iterator for TapTreeIter<'a, Pk>
 where
     Pk: MiniscriptKey + 'a,
 {
-    type Item = (u8, &'a Miniscript<Pk, Tap>);
+    type Item = TapTreeIterItem<'a, Pk>;
 
     fn next(&mut self) -> Option<Self::Item> {
         while let Some((depth, last)) = self.stack.pop() {
@@ -482,7 +486,7 @@ where
                     self.stack.push((depth + 1, right));
                     self.stack.push((depth + 1, left));
                 }
-                TapTree::Leaf(ref ms) => return Some((depth, ms)),
+                TapTree::Leaf(ref ms) => return Some(TapTreeIterItem { node: ms, depth }),
             }
         }
         None
@@ -629,7 +633,7 @@ impl<Pk: MiniscriptKey> ForEachKey<Pk> for Tr<Pk> {
     fn for_each_key<'a, F: FnMut(&'a Pk) -> bool>(&'a self, mut pred: F) -> bool {
         let script_keys_res = self
             .iter_scripts()
-            .all(|(_d, ms)| ms.for_each_key(&mut pred));
+            .all(|leaf| leaf.miniscript().for_each_key(&mut pred));
         script_keys_res && pred(&self.internal_key)
     }
 }
@@ -673,14 +677,14 @@ where
             absolute_timelock: None,
         };
         let mut min_wit_len = None;
-        for (_depth, ms) in desc.iter_scripts() {
+        for leaf in desc.iter_scripts() {
             let mut satisfaction = if allow_mall {
-                match ms.build_template(provider) {
+                match leaf.miniscript().build_template(provider) {
                     s @ Satisfaction { stack: Witness::Stack(_), .. } => s,
                     _ => continue, // No witness for this script in tr descriptor, look for next one
                 }
             } else {
-                match ms.build_template_mall(provider) {
+                match leaf.miniscript().build_template_mall(provider) {
                     s @ Satisfaction { stack: Witness::Stack(_), .. } => s,
                     _ => continue, // No witness for this script in tr descriptor, look for next one
                 }
@@ -690,7 +694,7 @@ where
                 _ => unreachable!(),
             };
 
-            let leaf_script = (ms.encode(), LeafVersion::TapScript);
+            let leaf_script = (leaf.compute_script(), LeafVersion::TapScript);
             let control_block = spend_info
                 .control_block(&leaf_script)
                 .expect("Control block must exist in script map for every known leaf");

src/descriptor/tr/taptree.rs

@@ -0,0 +1,58 @@
+// SPDX-License-Identifier: CC0-1.0
+
+use bitcoin::taproot::{LeafVersion, TapLeafHash};
+
+use crate::miniscript::context::Tap;
+use crate::sync::Arc;
+use crate::{Miniscript, MiniscriptKey, ToPublicKey};
+
+/// Iterator over all of the leaves of a Taproot tree.
+///
+/// If there is no tree (i.e. this is a keyspend-only Taproot descriptor)
+/// then the iterator will yield nothing.
+#[derive(Clone, PartialEq, Eq, Debug)]
+pub struct TapTreeIterItem<'tr, Pk: MiniscriptKey> {
+    pub(super) node: &'tr Arc<Miniscript<Pk, Tap>>,
+    pub(super) depth: u8,
+}
+
+impl<'tr, Pk: MiniscriptKey> TapTreeIterItem<'tr, Pk> {
+    /// The Tapscript in the leaf.
+    ///
+    /// To obtain a [`bitcoin::Script`] from this node, call [`Miniscript::encode`]
+    /// on the returned value.
+    #[inline]
+    pub fn miniscript(&self) -> &'tr Arc<Miniscript<Pk, Tap>> { self.node }
+
+    /// The depth of this leaf.
+    ///
+    /// This is useful for reconstructing the shape of the tree.
+    #[inline]
+    pub fn depth(&self) -> u8 { self.depth }
+
+    /// The Tapleaf version of this leaf.
+    ///
+    /// This function returns a constant value, since there is only one version in use
+    /// on the Bitcoin network; however, it may be useful to use this method in case
+    /// you wish to be forward-compatible with future versions supported by this
+    /// library.
+    #[inline]
+    pub fn leaf_version(&self) -> LeafVersion { LeafVersion::TapScript }
+}
+
+impl<Pk: ToPublicKey> TapTreeIterItem<'_, Pk> {
+    /// Computes the Bitcoin Script of the leaf.
+    ///
+    /// This function is potentially expensive.
+    #[inline]
+    pub fn compute_script(&self) -> bitcoin::ScriptBuf { self.node.encode() }
+
+    /// Computes the [`TapLeafHash`] of the leaf.
+    ///
+    /// This function is potentially expensive, since it serializes the full Bitcoin
+    /// Script of the leaf and hashes this data.
+    #[inline]
+    pub fn compute_tap_leaf_hash(&self) -> TapLeafHash {
+        TapLeafHash::from_script(&self.compute_script(), self.leaf_version())
+    }
+}

src/psbt/mod.rs

@@ -1112,14 +1112,12 @@ fn update_item_with_descriptor_helper<F: PsbtFields>(
 
             let mut builder = taproot::TaprootBuilder::new();
 
-            for ((_depth_der, ms_derived), (depth, ms)) in
-                tr_derived.iter_scripts().zip(tr_xpk.iter_scripts())
-            {
-                debug_assert_eq!(_depth_der, depth);
-                let leaf_script = (ms_derived.encode(), LeafVersion::TapScript);
+            for (leaf_derived, leaf) in tr_derived.iter_scripts().zip(tr_xpk.iter_scripts()) {
+                debug_assert_eq!(leaf_derived.depth(), leaf.depth());
+                let leaf_script = (leaf_derived.compute_script(), leaf_derived.leaf_version());
                 let tapleaf_hash = TapLeafHash::from_script(&leaf_script.0, leaf_script.1);
                 builder = builder
-                    .add_leaf(depth, leaf_script.0.clone())
+                    .add_leaf(leaf_derived.depth(), leaf_script.0.clone())
                     .expect("Computing spend data on a valid tree should always succeed");
                 if let Some(tap_scripts) = item.tap_scripts() {
                     let control_block = spend_info
@@ -1128,7 +1126,11 @@ fn update_item_with_descriptor_helper<F: PsbtFields>(
                     tap_scripts.insert(control_block, leaf_script);
                 }
 
-                for (pk_pkh_derived, pk_pkh_xpk) in ms_derived.iter_pk().zip(ms.iter_pk()) {
+                for (pk_pkh_derived, pk_pkh_xpk) in leaf_derived
+                    .miniscript()
+                    .iter_pk()
+                    .zip(leaf.miniscript().iter_pk())
+                {
                     let (xonly, xpk) = (pk_pkh_derived.to_x_only_pubkey(), pk_pkh_xpk);
 
                     let xpk_full_derivation_path = xpk