Merge #815: [Taproot API Project] replace TaprootSpendInfo with new miniscript-specific structure

3d12b42 fuzz: update taptree regression test to also check control blocks (Andrew Poelstra)
c1d4305 tr: add a whole bunch of fixed vector unit tests for TrSpendInfo (Andrew Poelstra)
8bc9400 tr: add conversion from TrSpendInfo to bitcoin::TapTree (Andrew Poelstra)
5fe2d25 tr: replace `bitcoin::TaprootSpendInfo` with `TrSpendInfo`, update psbt (Andrew Poelstra)
795cd4f tr: introduce new `TrSpendInfo` structure which holds a full TapTree (Andrew Poelstra)
27a30e6 psbt: untangle some logic in `update_item_with_descriptor_helper` (Andrew Poelstra)
00495a5 descriptor: add unit test from sanket (Andrew Poelstra)

Pull request description:

  In Miniscript, to compute control blocks, estimate satisfaction costs, or otherwise iterate through all the leaves of a Taptree, we use the `bitcoin::TaprootSpendInfo` structure to maintain a map of all leaves. This map is inappropriate for Miniscript (it may not be appropriate for *anyone* actually..) for a few reasons:

  * It is a map from Tapleaves' encoding as Script to data about the Tapleaves; but in Miniscript the Script encoding isn't primary and isn't even available for non-`ToPublicKey` keys
  * This map structure means that if duplicate leaves exist then only one of the dupes will be accessible.
  * The map structure is also really inefficient; it stores the entire merkle path for each leaf even those these paths significantly overlap, leading to O(n log n) space instead of O(n).
  * Furthermore, we don't need *any* map because we only ever iterate through the entire tree.

  We fix all these issues by introducing a new `TrSpendInfo` struct which stores the entire Merkle tree in a flat representation and can produce an iterator over all the leaves. The iterator item can be used to access the Script, the Miniscript, the leaf version, and the control block for each leaf, while the `TrSpendInfo` structure itself can be used to access the internal and external keys. In other words, this one structure efficiently implements APIs for everything that rust-miniscript needs.

  This completes the Taproot API overhaul project. After this I'll go back to fixing error types, eliminating recursive structures, or overhauling the validation parameters, whichever one seems most tractable from the current state of `master`.

ACKs for top commit:
  sanket1729:
    ACK 3d12b42

Tree-SHA512: 40fa71ebf22b104304ddbd9b3bf9efdd5d509c071a9461c80eb5f669ecdfd4dd7174cb5fcd2af30f32e45aecd370e077a49fd771eb71edf3f6fe4949c9b80d7d

Author: apoelstra

bitcoind-tests/tests/test_desc.rs

@@ -10,6 +10,7 @@ use std::{error, fmt};
 use actual_rand as rand;
 use bitcoin::blockdata::witness::Witness;
 use bitcoin::hashes::{sha256d, Hash};
+use bitcoin::key::TapTweak as _;
 use bitcoin::psbt::Psbt;
 use bitcoin::sighash::SighashCache;
 use bitcoin::taproot::{LeafVersion, TapLeafHash};
@@ -168,16 +169,15 @@ pub fn test_desc_satisfy(
             if let Some(internal_keypair) = internal_keypair {
                 // ---------------------- Tr key spend --------------------
                 let internal_keypair = internal_keypair
-                    .add_xonly_tweak(&secp, &tr.spend_info().tap_tweak().to_scalar())
-                    .expect("Tweaking failed");
+                    .tap_tweak(&secp, tr.spend_info().merkle_root());
                 let sighash_msg = sighash_cache
                     .taproot_key_spend_signature_hash(0, &prevouts, sighash_type)
                     .unwrap();
                 let msg = secp256k1::Message::from_digest(sighash_msg.to_byte_array());
                 let mut aux_rand = [0u8; 32];
                 rand::thread_rng().fill_bytes(&mut aux_rand);
                 let schnorr_sig =
-                    secp.sign_schnorr_with_aux_rand(&msg, &internal_keypair, &aux_rand);
+                    secp.sign_schnorr_with_aux_rand(&msg, &internal_keypair.to_inner(), &aux_rand);
                 psbt.inputs[0].tap_key_sig =
                     Some(taproot::Signature { signature: schnorr_sig, sighash_type });
             } else {

fuzz/fuzz_targets/regression_taptree.rs

@@ -1,3 +1,5 @@
+use std::collections::{HashMap, HashSet};
+
 use descriptor_fuzz::FuzzPk;
 use honggfuzz::fuzz;
 use miniscript::descriptor::Tr;
@@ -27,6 +29,22 @@ fn do_test(data: &[u8]) {
                 new_si.output_key(),
                 "merkle root mismatch (left is old, new is right)",
             );
+
+            // Map every leaf script to a set of all the control blocks
+            let mut new_cbs = HashMap::new();
+            for leaf in new_si.leaves() {
+                new_cbs
+                    .entry(leaf.script())
+                    .or_insert(HashSet::new())
+                    .insert(leaf.control_block().clone());
+            }
+            // ...the old code will only ever yield one of them and it's not easy to predict which one
+            for leaf in new_si.leaves() {
+                let old_cb = old_si
+                    .control_block(&(leaf.script().into(), leaf.leaf_version()))
+                    .unwrap();
+                assert!(new_cbs[leaf.script()].contains(&old_cb));
+            }
         }
     }
 }

src/descriptor/mod.rs

@@ -44,7 +44,10 @@ pub use self::iter::PkIter;
 pub use self::segwitv0::{Wpkh, Wsh, WshInner};
 pub use self::sh::{Sh, ShInner};
 pub use self::sortedmulti::SortedMultiVec;
-pub use self::tr::{TapTree, TapTreeDepthError, TapTreeIter, TapTreeIterItem, Tr};
+pub use self::tr::{
+    TapTree, TapTreeDepthError, TapTreeIter, TapTreeIterItem, Tr, TrSpendInfo, TrSpendInfoIter,
+    TrSpendInfoIterItem,
+};
 
 pub mod checksum;
 mod key;
@@ -1547,17 +1550,20 @@ mod tests {
             "tr({},{{pk({}),{{pk({}),or_d(pk({}),pkh({}))}}}})",
             p1, p2, p3, p4, p5
         ))
-        .unwrap()
-        .to_string();
+        .unwrap();
 
         // p5.to_pubkeyhash() = 516ca378e588a7ed71336147e2a72848b20aca1a
         assert_eq!(
-            descriptor,
+            descriptor.to_string(),
             format!(
                 "tr({},{{pk({}),{{pk({}),or_d(pk({}),pkh({}))}}}})#tvu28c0s",
                 p1, p2, p3, p4, p5
             )
-        )
+        );
+        assert_eq!(
+            descriptor.spend_info().merkle_root().unwrap().to_string(),
+            "e1597abcb76f7cbc0792cf04a9c2d4f39caed1ede0afef772064126f28c69b09"
+        );
     }
 
     #[test]

src/descriptor/tr/mod.rs

@@ -2,12 +2,7 @@
 
 use core::{cmp, fmt, hash};
 
-#[cfg(not(test))] // https://github.com/rust-lang/rust/issues/121684
-use bitcoin::secp256k1;
-use bitcoin::taproot::{
-    LeafVersion, TaprootBuilder, TaprootSpendInfo, TAPROOT_CONTROL_BASE_SIZE,
-    TAPROOT_CONTROL_NODE_SIZE,
-};
+use bitcoin::taproot::{TAPROOT_CONTROL_BASE_SIZE, TAPROOT_CONTROL_NODE_SIZE};
 use bitcoin::{opcodes, Address, Network, ScriptBuf, Weight};
 use sync::Arc;
 
@@ -26,8 +21,10 @@ use crate::{
     Threshold, ToPublicKey, TranslateErr, Translator,
 };
 
+mod spend_info;
 mod taptree;
 
+pub use self::spend_info::{TrSpendInfo, TrSpendInfoIter, TrSpendInfoIterItem};
 pub use self::taptree::{TapTree, TapTreeDepthError, TapTreeIter, TapTreeIterItem};
 
 /// A taproot descriptor
@@ -43,7 +40,7 @@ pub struct Tr<Pk: MiniscriptKey> {
     // The inner `Arc` here is because Rust does not allow us to return a reference
     // to the contents of the `Option` from inside a `MutexGuard`. There is no outer
     // `Arc` because when this structure is cloned, we create a whole new mutex.
-    spend_info: Mutex<Option<Arc<TaprootSpendInfo>>>,
+    spend_info: Mutex<Option<Arc<TrSpendInfo<Pk>>>>,
 }
 
 impl<Pk: MiniscriptKey> Clone for Tr<Pk> {
@@ -118,46 +115,21 @@ impl<Pk: MiniscriptKey> Tr<Pk> {
         }
     }
 
-    /// Compute the [`TaprootSpendInfo`] associated with this descriptor if spend data is `None`.
+    /// Obtain the spending information for this [`Tr`].
     ///
-    /// If spend data is already computed (i.e it is not `None`), this does not recompute it.
+    /// The first time this method is called, it computes the full Taproot Merkle tree of
+    /// all branches as well as the output key which appears on-chain. This is fairly
+    /// expensive since it requires hashing every branch and then doing an elliptic curve
+    /// operation. The result is cached and reused on subsequent calls.
     ///
-    /// [`TaprootSpendInfo`] is only required for spending via the script paths.
-    pub fn spend_info(&self) -> Arc<TaprootSpendInfo>
+    /// This data is needed to compute the Taproot output, so this method is implicitly
+    /// called through [`Self::script_pubkey`], [`Self::address`], etc. It is also needed
+    /// to compute the hash needed to sign the output.
+    pub fn spend_info(&self) -> TrSpendInfo<Pk>
     where
         Pk: ToPublicKey,
     {
-        // If the value is already cache, read it
-        // read only panics if the lock is poisoned (meaning other thread having a lock panicked)
-        let read_lock = self.spend_info.lock().expect("Lock poisoned");
-        if let Some(ref spend_info) = *read_lock {
-            return Arc::clone(spend_info);
-        }
-        drop(read_lock);
-
-        // Get a new secp context
-        // This would be cheap operation after static context support from upstream
-        let secp = secp256k1::Secp256k1::verification_only();
-        // Key spend path with no merkle root
-        let data = if self.tree.is_none() {
-            TaprootSpendInfo::new_key_spend(&secp, self.internal_key.to_x_only_pubkey(), None)
-        } else {
-            let mut builder = TaprootBuilder::new();
-            for leaf in self.leaves() {
-                let script = leaf.miniscript().encode();
-                builder = builder
-                    .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
-            match builder.finalize(&secp, self.internal_key.to_x_only_pubkey()) {
-                Ok(data) => data,
-                Err(_) => unreachable!("We know the builder can be finalized"),
-            }
-        };
-        let spend_info = Arc::new(data);
-        *self.spend_info.lock().expect("Lock poisoned") = Some(Arc::clone(&spend_info));
-        spend_info
+        TrSpendInfo::from_tr(self)
     }
 
     /// Checks whether the descriptor is safe.
@@ -506,7 +478,7 @@ where
             absolute_timelock: None,
         };
         let mut min_wit_len = None;
-        for leaf in desc.leaves() {
+        for leaf in spend_info.leaves() {
             let mut satisfaction = if allow_mall {
                 match leaf.miniscript().build_template(provider) {
                     s @ Satisfaction { stack: Witness::Stack(_), .. } => s,
@@ -523,12 +495,10 @@ where
                 _ => unreachable!(),
             };
 
-            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");
+            let script = ScriptBuf::from(leaf.script());
+            let control_block = leaf.control_block().clone();
 
-            wit.push(Placeholder::TapScript(leaf_script.0));
+            wit.push(Placeholder::TapScript(script));
             wit.push(Placeholder::TapControlBlock(control_block));
 
             let wit_size = witness_size(wit);