Skip to content

Commit 795cd4f

Browse files
committed
tr: introduce new TrSpendInfo structure which holds a full TapTree
This commit introduces a new data structure but **does not** use it. The next commit will do this. I have separated them so that this one, which introduces a bunch of algorithmic code, can be reviewed separately from the API-breaking one. When computing a `Tr` output, we need to encode all its tapleaves into Script, put these into a Merkle tree and tweak the internal key with the root of this tree. When spending from one of the branches of this output, we need the Merkle path to that output. We currently do this by using the `TaprootSpendInfo` structure from rust-bitcoin. This is not a very good fit for rust-miniscript, because it constructs a map from Tapscripts to their control blocks. This is slow and memory-wasteful to construct, and while it makes random access fairly fast, it makes sequential access pretty slow. In Miniscript we almost always want sequential access, because all of our algorithms are some form of either "try every possibility and choose the optimum" or "aggregate every possibility". It also means that if there are multiple leaves with the same script, only one copy will ever be accessible. (If they are at different depths, the low-depth one will be yielded, but if they are at the same depth it's effectively random which one will get priority.) Having multiple copies of the same script is a pointless malleability vector, but this behavior is still surprising and annoying to have to think about. To replace `bitcoin::TaprootSpendInfo` we create a new structure `TrSpendInfo`. This structure doesn't maintain any maps: it stores a full Merkleized taptree in such a way that it can efficiently yield all of the leaves' control blocks in-order. It is likely that at some point we will want to upport this, or some variant of it, into rust-bitcoin, since for typical usecases it's much faster to use and construct.
1 parent 27a30e6 commit 795cd4f

File tree

4 files changed

+337
-3
lines changed

4 files changed

+337
-3
lines changed

src/descriptor/mod.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,10 @@ pub use self::bare::{Bare, Pkh};
4242
pub use self::segwitv0::{Wpkh, Wsh, WshInner};
4343
pub use self::sh::{Sh, ShInner};
4444
pub use self::sortedmulti::SortedMultiVec;
45-
pub use self::tr::{TapTree, TapTreeDepthError, TapTreeIter, TapTreeIterItem, Tr};
45+
pub use self::tr::{
46+
TapTree, TapTreeDepthError, TapTreeIter, TapTreeIterItem, Tr, TrSpendInfo, TrSpendInfoIter,
47+
TrSpendInfoIterItem,
48+
};
4649

4750
pub mod checksum;
4851
mod key;

src/descriptor/tr/mod.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,10 @@ use crate::{
2626
Threshold, ToPublicKey, TranslateErr, Translator,
2727
};
2828

29+
mod spend_info;
2930
mod taptree;
3031

32+
pub use self::spend_info::{TrSpendInfo, TrSpendInfoIter, TrSpendInfoIterItem};
3133
pub use self::taptree::{TapTree, TapTreeDepthError, TapTreeIter, TapTreeIterItem};
3234

3335
/// A taproot descriptor

src/descriptor/tr/spend_info.rs

Lines changed: 325 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,325 @@
1+
// SPDX-License-Identifier: CC0-1.0
2+
3+
//! Taproot Spending Information
4+
//!
5+
//! Provides a structure which can be used to obtain control blocks and other information
6+
//! needed for Taproot spends.
7+
//!
8+
9+
use bitcoin::key::{Parity, TapTweak as _, TweakedPublicKey, UntweakedPublicKey};
10+
use bitcoin::secp256k1::Secp256k1;
11+
use bitcoin::taproot::{ControlBlock, LeafVersion, TapLeafHash, TapNodeHash, TaprootMerkleBranch};
12+
use bitcoin::{Script, ScriptBuf};
13+
14+
use crate::miniscript::context::Tap;
15+
use crate::prelude::Vec;
16+
use crate::sync::Arc;
17+
use crate::{Miniscript, MiniscriptKey, ToPublicKey};
18+
19+
/// Utility structure which maintains a stack of bits (at most 128) using a u128.
20+
///
21+
/// Will panic if the user attempts to push more than 128 bits; we assume in this
22+
/// module that we are starting with a validated [`super::TapTree`] and therefore
23+
/// that this can't happen.
24+
#[derive(Default)]
25+
struct BitStack128 {
26+
inner: u128,
27+
height: u8,
28+
}
29+
30+
impl BitStack128 {
31+
fn push(&mut self, bit: bool) {
32+
if bit {
33+
self.inner |= 1u128 << self.height;
34+
} else {
35+
self.inner &= !(1u128 << self.height);
36+
}
37+
self.height += 1;
38+
}
39+
40+
fn pop(&mut self) -> Option<bool> {
41+
if self.height > 0 {
42+
self.height -= 1;
43+
Some(self.inner & (1u128 << self.height) != 0)
44+
} else {
45+
None
46+
}
47+
}
48+
}
49+
50+
/// A structure which can be used to obtain control blocks and other information
51+
/// needed for Taproot spends.
52+
///
53+
/// Conceptually, this object is a copy of the Taproot tree with each leave annotated
54+
/// with extra information that can be used to compute its control block.
55+
pub struct TrSpendInfo<Pk: MiniscriptKey> {
56+
internal_key: UntweakedPublicKey,
57+
output_key: TweakedPublicKey,
58+
output_key_parity: Parity,
59+
/// The nodes of the tree, in pre-order, i.e. left-to-right depth-first order.
60+
nodes: Vec<TrSpendInfoNode<Pk>>,
61+
}
62+
63+
impl<Pk: ToPublicKey> TrSpendInfo<Pk> {
64+
fn nodes_from_tap_tree(tree: &super::TapTree<Pk>) -> Vec<TrSpendInfoNode<Pk>> {
65+
let mut nodes = vec![];
66+
let mut parent_stack = Vec::with_capacity(128); // FIXME use ArrayVec here
67+
for leaf in tree.leaves() {
68+
let depth = usize::from(leaf.depth());
69+
let script = leaf.miniscript().encode();
70+
71+
let leaf_hash = TapLeafHash::from_script(&script, leaf.leaf_version());
72+
let mut current_hash = TapNodeHash::from(leaf_hash);
73+
74+
// 1. If this node increases our depth, add parents.
75+
while parent_stack.len() < depth {
76+
// When we encounter a leaf we put all of its parent nodes into the
77+
// result. We set the "sibling hash" to a dummy value (specifically,
78+
// `current_hash`, because it's convenient and the right type).
79+
parent_stack.push((false, nodes.len()));
80+
nodes.push(TrSpendInfoNode { sibling_hash: current_hash, leaf_data: None });
81+
}
82+
// If parent_stack.len() < depth then we pushed things onto the stack in
83+
// the previous step so that we now have equality. Meanwhile, it is
84+
// impossible for parent_stack.len() > depth because we pop things off
85+
// the stack in step 3 below.
86+
assert_eq!(depth, parent_stack.len());
87+
88+
// 2. Add the node.
89+
//
90+
// Again, we don't know the sibling hash yet so we use the current hash.
91+
// But this time the current hash isn't an arbitrary dummy value -- in the
92+
// next step we will have an invariant that incomplete nodes' "sibling hashes"
93+
// are set to the nodes' own hashes.
94+
//
95+
// We will use this hash to compute the parent's hash then replace it with
96+
// the actual sibling hash. We do this for every node EXCEPT the root node,
97+
// whose "sibling hash" will then wind up being equal to the Merkle root
98+
// of the whole tree.
99+
nodes.push(TrSpendInfoNode {
100+
sibling_hash: current_hash,
101+
leaf_data: Some(LeafData {
102+
script,
103+
miniscript: Arc::clone(leaf.miniscript()),
104+
leaf_hash,
105+
}),
106+
});
107+
108+
// 3. Recursively complete nodes as long as we are on a right branch.
109+
//
110+
// As described above, for each parent node, we compute its hash and store it
111+
// in `sibling_hash`. At that point we're done with the childrens' hashes so
112+
// we finally replace those with their sibling hashes.
113+
let mut cur_index = nodes.len() - 1;
114+
while let Some((done_left_child, parent_idx)) = parent_stack.pop() {
115+
if done_left_child {
116+
let lchild_hash = nodes[parent_idx + 1].sibling_hash;
117+
// Set current node's "sibling hash" to its own hash.
118+
let new_merkle_root = TapNodeHash::from_node_hashes(lchild_hash, current_hash);
119+
nodes[parent_idx].sibling_hash = new_merkle_root;
120+
// Set the children's sibling hashes to each others' hashes.
121+
nodes[parent_idx + 1].sibling_hash = current_hash;
122+
nodes[cur_index].sibling_hash = lchild_hash;
123+
// Recurse.
124+
current_hash = new_merkle_root;
125+
cur_index = parent_idx;
126+
} else {
127+
// Once we hit a left branch we can't do anything until we see the next leaf.
128+
parent_stack.push((true, parent_idx));
129+
break;
130+
}
131+
}
132+
}
133+
debug_assert_eq!(parent_stack.len(), 0);
134+
debug_assert_ne!(nodes.len(), 0);
135+
136+
nodes
137+
}
138+
139+
/// Constructs a [`TrSpendInfo`] for a [`super::Tr`].
140+
pub fn from_tr(tr: &super::Tr<Pk>) -> Self {
141+
let internal_key = tr.internal_key().to_x_only_pubkey();
142+
143+
let nodes = match tr.tap_tree() {
144+
Some(tree) => Self::nodes_from_tap_tree(tree),
145+
None => vec![],
146+
};
147+
148+
let secp = Secp256k1::verification_only();
149+
let (output_key, output_key_parity) =
150+
internal_key.tap_tweak(&secp, nodes.first().map(|node| node.sibling_hash));
151+
152+
TrSpendInfo { internal_key, output_key, output_key_parity, nodes }
153+
}
154+
155+
/// If this [`TrSpendInfo`] has an associated Taproot tree, return its Merkle root.
156+
pub fn merkle_root(&self) -> Option<TapNodeHash> {
157+
// As described in `nodes_from_tap_tree`, the "sibling hash" of the root node
158+
// is actually the Merkle root of the whole tree.
159+
self.nodes.first().map(|node| node.sibling_hash)
160+
}
161+
162+
/// The internal key of the Taproot output.
163+
///
164+
/// This returns the x-only public key which appears on-chain. For the abstroct
165+
/// public key, use the `internal_key` method on the original [`super::Tr`] used to
166+
/// create this object.
167+
pub fn internal_key(&self) -> UntweakedPublicKey { self.internal_key }
168+
169+
// I don't really like these names, but they're used in rust-bitcoin so we'll stick
170+
// with them and just doc-alias them to better names so they show up in search results.
171+
/// The external key of the Taproot output.
172+
#[doc(alias = "external_key")]
173+
pub fn output_key(&self) -> TweakedPublicKey { self.output_key }
174+
175+
/// The parity of the external key of the Taproot output.
176+
#[doc(alias = "external_key_parity")]
177+
pub fn output_key_parity(&self) -> Parity { self.output_key_parity }
178+
179+
/// An iterator over the leaves of the Taptree.
180+
///
181+
/// This yields the same leaves in the same order as [`super::Tr::leaves`] on the original
182+
/// [`super::Tr`]. However, in addition to yielding the leaves and their depths, it also
183+
/// yields their scripts, leafhashes, and control blocks.
184+
pub fn leaves(&self) -> TrSpendInfoIter<Pk> {
185+
TrSpendInfoIter {
186+
spend_info: self,
187+
index: 0,
188+
merkle_stack: Vec::with_capacity(128),
189+
done_left_stack: BitStack128::default(),
190+
}
191+
}
192+
}
193+
194+
/// An internal node of the spend
195+
#[derive(Debug)]
196+
struct TrSpendInfoNode<Pk: MiniscriptKey> {
197+
sibling_hash: TapNodeHash,
198+
leaf_data: Option<LeafData<Pk>>,
199+
}
200+
201+
#[derive(Debug)]
202+
struct LeafData<Pk: MiniscriptKey> {
203+
script: ScriptBuf,
204+
miniscript: Arc<Miniscript<Pk, Tap>>,
205+
leaf_hash: TapLeafHash,
206+
}
207+
208+
/// An iterator over the leaves of a Taproot tree. Produced by [`TrSpendInfo::leaves`].
209+
///
210+
/// This is conceptually similar to [`super::TapTreeIter`], which can be obtained by
211+
/// calling [`super::TapTree::leaves`]. That iterator goes over the leaves of the tree,
212+
/// yielding the Miniscripts of the leaves and their depth.
213+
///
214+
/// This iterator goes over the leaves in the same order, yielding the data that actually
215+
/// goes on chain: their scripts, control blocks, etc.
216+
pub struct TrSpendInfoIter<'sp, Pk: MiniscriptKey> {
217+
spend_info: &'sp TrSpendInfo<Pk>,
218+
index: usize,
219+
merkle_stack: Vec<TapNodeHash>,
220+
done_left_stack: BitStack128,
221+
}
222+
223+
impl<'sp, Pk: MiniscriptKey> Iterator for TrSpendInfoIter<'sp, Pk> {
224+
type Item = TrSpendInfoIterItem<'sp, Pk>;
225+
226+
fn next(&mut self) -> Option<Self::Item> {
227+
while self.index < self.spend_info.nodes.len() {
228+
let current_node = &self.spend_info.nodes[self.index];
229+
if self.index > 0 {
230+
self.merkle_stack.push(current_node.sibling_hash);
231+
}
232+
self.index += 1;
233+
234+
if let Some(ref leaf) = current_node.leaf_data {
235+
// leaf
236+
let mut merkle_stack = self.merkle_stack.clone();
237+
merkle_stack.reverse();
238+
self.merkle_stack.pop();
239+
240+
loop {
241+
match self.done_left_stack.pop() {
242+
None => break, // this leaf is the root node
243+
Some(false) => {
244+
self.done_left_stack.push(true);
245+
break;
246+
}
247+
Some(true) => {
248+
self.merkle_stack.pop();
249+
}
250+
}
251+
}
252+
253+
return Some(TrSpendInfoIterItem {
254+
script: &leaf.script,
255+
miniscript: &leaf.miniscript,
256+
leaf_hash: leaf.leaf_hash,
257+
control_block: ControlBlock {
258+
leaf_version: LeafVersion::TapScript,
259+
output_key_parity: self.spend_info.output_key_parity,
260+
internal_key: self.spend_info.internal_key,
261+
merkle_branch: TaprootMerkleBranch::try_from(merkle_stack)
262+
.expect("merkle stack guaranteed to be within allowable length"),
263+
},
264+
});
265+
} else {
266+
// internal node
267+
self.done_left_stack.push(false);
268+
}
269+
}
270+
None
271+
}
272+
}
273+
274+
/// Item yielded from a [`TrSpendInfoIter`].
275+
#[derive(Clone, PartialEq, Eq, Debug)]
276+
pub struct TrSpendInfoIterItem<'tr, Pk: MiniscriptKey> {
277+
script: &'tr Script,
278+
miniscript: &'tr Arc<Miniscript<Pk, Tap>>,
279+
leaf_hash: TapLeafHash,
280+
control_block: ControlBlock,
281+
}
282+
283+
impl<'sp, Pk: MiniscriptKey> TrSpendInfoIterItem<'sp, Pk> {
284+
/// The Tapscript of this leaf.
285+
#[inline]
286+
pub fn script(&self) -> &'sp Script { self.script }
287+
288+
/// The Tapscript of this leaf, in Miniscript form.
289+
#[inline]
290+
pub fn miniscript(&self) -> &'sp Arc<Miniscript<Pk, Tap>> { self.miniscript }
291+
292+
/// The depth of the leaf in the tree.
293+
///
294+
/// This value is returned as `u8` since it is guaranteed to be <= 128 by the Taproot
295+
/// consensus rules.
296+
#[inline]
297+
pub fn depth(&self) -> u8 {
298+
self.control_block.merkle_branch.len() as u8 // cast ok, length limited to 128
299+
}
300+
301+
/// The Tapleaf version of this leaf.
302+
///
303+
/// This function returns a constant value, since there is only one version in use
304+
/// on the Bitcoin network; however, it may be useful to use this method in case
305+
/// you wish to be forward-compatible with future versions supported by this
306+
/// library.
307+
#[inline]
308+
pub fn leaf_version(&self) -> LeafVersion { self.control_block.leaf_version }
309+
310+
/// The hash of this leaf.
311+
///
312+
/// This hash, prefixed with the leaf's [`Self::leaf_version`], is what is directly
313+
/// committed in the Taproot tree.
314+
#[inline]
315+
pub fn leaf_hash(&self) -> TapLeafHash { self.leaf_hash }
316+
317+
/// The control block of this leaf.
318+
///
319+
/// Unlike the other data obtainable from [`TrSpendInfoIterItem`], this one is computed
320+
/// dynamically during iteration and therefore will not outlive the iterator item. If
321+
/// you need access to multiple control blocks at once, will likely need to clone and
322+
/// store them separately.
323+
#[inline]
324+
pub fn control_block(&self) -> &ControlBlock { &self.control_block }
325+
}

src/descriptor/tr/taptree.rs

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -212,14 +212,18 @@ impl<'tr, Pk: MiniscriptKey> TapTreeIterItem<'tr, Pk> {
212212
impl<Pk: ToPublicKey> TapTreeIterItem<'_, Pk> {
213213
/// Computes the Bitcoin Script of the leaf.
214214
///
215-
/// This function is potentially expensive.
215+
/// This function is potentially expensive. If you are calling this method on
216+
/// all (or many) of the leaves of the tree, you may instead want to call
217+
/// [`super::Tr::spend_info`] and use the [`super::TrSpendInfo::leaves`] iterator instead.
216218
#[inline]
217219
pub fn compute_script(&self) -> bitcoin::ScriptBuf { self.node.encode() }
218220

219221
/// Computes the [`TapLeafHash`] of the leaf.
220222
///
221223
/// This function is potentially expensive, since it serializes the full Bitcoin
222-
/// Script of the leaf and hashes this data.
224+
/// Script of the leaf and hashes this data. If you are calling this method on
225+
/// all (or many) of the leaves of the tree, you may instead want to call
226+
/// [`super::Tr::spend_info`] and use the [`super::TrSpendInfo::leaves`] iterator instead.
223227
#[inline]
224228
pub fn compute_tap_leaf_hash(&self) -> TapLeafHash {
225229
TapLeafHash::from_script(&self.compute_script(), self.leaf_version())

0 commit comments

Comments
 (0)