derive/
taptree.rs

1// Modern, minimalistic & standard-compliant cold wallet library.
2//
3// SPDX-License-Identifier: Apache-2.0
4//
5// Written in 2020-2024 by
6//     Dr Maxim Orlovsky <orlovsky@lnp-bp.org>
7//
8// Copyright (C) 2020-2024 LNP/BP Standards Association. All rights reserved.
9// Copyright (C) 2020-2024 Dr Maxim Orlovsky. All rights reserved.
10//
11// Licensed under the Apache License, Version 2.0 (the "License");
12// you may not use this file except in compliance with the License.
13// You may obtain a copy of the License at
14//
15//     http://www.apache.org/licenses/LICENSE-2.0
16//
17// Unless required by applicable law or agreed to in writing, software
18// distributed under the License is distributed on an "AS IS" BASIS,
19// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
20// See the License for the specific language governing permissions and
21// limitations under the License.
22
23use std::fmt::{self, Display, Formatter, Write};
24use std::ops::Deref;
25use std::{slice, vec};
26
27use amplify::num::u7;
28use bc::{
29    ControlBlock, InternalPk, LeafScript, OutputPk, Parity, TapLeafHash, TapMerklePath,
30    TapNodeHash, TapScript,
31};
32use commit_verify::merkle::MerkleBuoy;
33
34use crate::{KeyOrigin, Terminal, XkeyOrigin};
35
36#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Display, Error, From)]
37pub enum InvalidTree {
38    #[from]
39    #[display(doc_comments)]
40    Unfinalized(UnfinalizedTree),
41
42    #[from(FinalizedTree)]
43    #[display("tap tree contains too many script leaves which doesn't fit a single Merkle tree")]
44    MountainRange,
45}
46
47#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Display, Error)]
48#[display("can't add more leaves to an already finalized tap tree")]
49pub struct FinalizedTree;
50
51#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Display, Error)]
52#[display(
53    "unfinalized tap tree containing leaves at level {0} which can't commit into a single Merkle \
54     root"
55)]
56pub struct UnfinalizedTree(pub u7);
57
58#[derive(Clone, Eq, PartialEq, Debug, Default)]
59pub struct TapTreeBuilder<L = LeafScript> {
60    leaves: Vec<LeafInfo<L>>,
61    buoy: MerkleBuoy<u7>,
62    finalized: bool,
63}
64
65impl<L> TapTreeBuilder<L> {
66    pub fn new() -> Self {
67        Self {
68            leaves: none!(),
69            buoy: default!(),
70            finalized: false,
71        }
72    }
73
74    pub fn with_capacity(capacity: usize) -> Self {
75        Self {
76            leaves: Vec::with_capacity(capacity),
77            buoy: zero!(),
78            finalized: false,
79        }
80    }
81
82    pub fn is_finalized(&self) -> bool { self.finalized }
83
84    pub fn with_leaf(mut self, leaf: LeafInfo<L>) -> Result<Self, FinalizedTree> {
85        self.push_leaf(leaf)?;
86        Ok(self)
87    }
88
89    pub fn push_leaf(&mut self, leaf: LeafInfo<L>) -> Result<bool, FinalizedTree> {
90        if self.finalized {
91            return Err(FinalizedTree);
92        }
93        let depth = leaf.depth;
94        self.leaves.push(leaf);
95        self.buoy.push(depth);
96        if self.buoy.level() == u7::ZERO {
97            self.finalized = true
98        }
99        Ok(self.finalized)
100    }
101
102    pub fn finish(self) -> Result<TapTree<L>, UnfinalizedTree> {
103        if !self.finalized {
104            return Err(UnfinalizedTree(self.buoy.level()));
105        }
106        Ok(TapTree(self.leaves))
107    }
108}
109
110/// Non-empty taproot script tree.
111#[derive(Clone, Eq, PartialEq, Hash, Debug, Default)]
112#[cfg_attr(feature = "serde", derive(Serialize, Deserialize), serde(transparent))]
113pub struct TapTree<L = LeafScript>(Vec<LeafInfo<L>>);
114
115impl<L> Deref for TapTree<L> {
116    type Target = Vec<LeafInfo<L>>;
117    fn deref(&self) -> &Self::Target { &self.0 }
118}
119
120impl<L> IntoIterator for TapTree<L> {
121    type Item = LeafInfo<L>;
122    type IntoIter = vec::IntoIter<LeafInfo<L>>;
123
124    fn into_iter(self) -> Self::IntoIter { self.0.into_iter() }
125}
126
127impl<'a, L> IntoIterator for &'a TapTree<L> {
128    type Item = &'a LeafInfo<L>;
129    type IntoIter = slice::Iter<'a, LeafInfo<L>>;
130
131    fn into_iter(self) -> Self::IntoIter { self.0.iter() }
132}
133
134impl TapTree {
135    pub fn with_single_leaf(leaf: impl Into<LeafScript>) -> TapTree {
136        Self(vec![LeafInfo {
137            depth: u7::ZERO,
138            script: leaf.into(),
139        }])
140    }
141
142    pub fn merkle_root(&self) -> TapNodeHash {
143        if self.0.len() == 1 {
144            TapLeafHash::with_leaf_script(&self.0[0].script).into()
145        } else {
146            todo!("#10 implement TapTree::merkle_root for trees with more than one leaf")
147        }
148    }
149}
150
151impl<L> TapTree<L> {
152    pub fn from_leaves(leaves: impl IntoIterator<Item = LeafInfo<L>>) -> Result<Self, InvalidTree> {
153        let mut builder = TapTreeBuilder::<L>::new();
154        for leaf in leaves {
155            builder.push_leaf(leaf)?;
156        }
157        builder.finish().map_err(InvalidTree::from)
158    }
159
160    pub fn from_builder(builder: TapTreeBuilder<L>) -> Result<Self, UnfinalizedTree> {
161        builder.finish()
162    }
163
164    pub fn into_vec(self) -> Vec<LeafInfo<L>> { self.0 }
165
166    pub fn map<M>(self, f: impl Fn(L) -> M) -> TapTree<M> {
167        TapTree(
168            self.into_iter()
169                .map(|leaf| LeafInfo {
170                    depth: leaf.depth,
171                    script: f(leaf.script),
172                })
173                .collect(),
174        )
175    }
176}
177
178impl<L: Display> Display for TapTree<L> {
179    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
180        let mut buoy = MerkleBuoy::<u7>::default();
181
182        let mut depth = u7::ZERO;
183        for leaf in &self.0 {
184            for _ in depth.into_u8()..leaf.depth.into_u8() {
185                f.write_char('{')?;
186            }
187            buoy.push(leaf.depth);
188            if depth == leaf.depth {
189                f.write_char(',')?;
190            }
191            depth = leaf.depth;
192            for _ in buoy.level().into_u8()..depth.into_u8() {
193                f.write_char('}')?;
194            }
195            debug_assert_ne!(buoy.level(), u7::ZERO);
196        }
197        debug_assert_eq!(buoy.level(), u7::ZERO);
198        Ok(())
199    }
200}
201
202#[derive(Clone, Eq, PartialEq, Hash, Debug)]
203#[cfg_attr(feature = "serde", derive(Serialize, Deserialize), serde(rename_all = "camelCase"))]
204pub struct LeafInfo<L = LeafScript> {
205    pub depth: u7,
206    pub script: L,
207}
208
209impl LeafInfo<LeafScript> {
210    pub fn tap_script(depth: u7, script: TapScript) -> Self {
211        LeafInfo {
212            depth,
213            script: LeafScript::from_tap_script(script),
214        }
215    }
216}
217
218#[derive(Getters, Clone, Eq, PartialEq, Debug)]
219#[getter(as_copy)]
220pub struct ControlBlockFactory {
221    internal_pk: InternalPk,
222    output_pk: OutputPk,
223    parity: Parity,
224    merkle_root: TapNodeHash,
225
226    #[getter(skip)]
227    merkle_path: TapMerklePath,
228    #[getter(skip)]
229    remaining_leaves: Vec<LeafInfo>,
230}
231
232impl ControlBlockFactory {
233    #[inline]
234    pub fn with(internal_pk: InternalPk, tap_tree: TapTree) -> Self {
235        let merkle_root = tap_tree.merkle_root();
236        let (output_pk, parity) = internal_pk.to_output_pk(Some(merkle_root));
237        ControlBlockFactory {
238            internal_pk,
239            output_pk,
240            parity,
241            merkle_root,
242            merkle_path: empty!(),
243            remaining_leaves: tap_tree.into_vec(),
244        }
245    }
246
247    #[inline]
248    pub fn into_remaining_leaves(self) -> Vec<LeafInfo> { self.remaining_leaves }
249}
250
251impl Iterator for ControlBlockFactory {
252    type Item = (ControlBlock, LeafScript);
253
254    fn next(&mut self) -> Option<Self::Item> {
255        let leaf = self.remaining_leaves.pop()?;
256        let leaf_script = leaf.script;
257        let control_block = ControlBlock::with(
258            leaf_script.version,
259            self.internal_pk,
260            self.parity,
261            self.merkle_path.clone(),
262        );
263        Some((control_block, leaf_script))
264    }
265}
266
267/// A compact size unsigned integer representing the number of leaf hashes, followed by a list
268/// of leaf hashes, followed by the 4 byte master key fingerprint concatenated with the
269/// derivation path of the public key. The derivation path is represented as 32-bit little
270/// endian unsigned integer indexes concatenated with each other. Public keys are those needed
271/// to spend this output. The leaf hashes are of the leaves which involve this public key. The
272/// internal key does not have leaf hashes, so can be indicated with a hashes len of 0.
273/// Finalizers should remove this field after `PSBT_IN_FINAL_SCRIPTWITNESS` is constructed.
274#[derive(Clone, Eq, PartialEq, Hash, Debug)]
275#[cfg_attr(feature = "serde", derive(Serialize, Deserialize), serde(rename_all = "camelCase"))]
276pub struct TapDerivation {
277    pub leaf_hashes: Vec<TapLeafHash>,
278    pub origin: KeyOrigin,
279}
280
281impl TapDerivation {
282    pub fn with_internal_pk(xpub_origin: XkeyOrigin, terminal: Terminal) -> Self {
283        let origin = KeyOrigin::with(xpub_origin, terminal);
284        TapDerivation {
285            leaf_hashes: empty!(),
286            origin,
287        }
288    }
289}
derive/taptree.rs

derive/
taptree.rs
