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::ops::Deref;
24use std::{slice, vec};
25
26use amplify::num::u7;
27use bc::{
28    ControlBlock, InternalPk, LeafScript, OutputPk, Parity, TapLeafHash, TapMerklePath,
29    TapNodeHash, TapScript,
30};
31use commit_verify::merkle::MerkleBuoy;
32
33use crate::{KeyOrigin, Terminal, XkeyOrigin};
34
35#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Display, Error, From)]
36pub enum InvalidTree {
37    #[from]
38    #[display(doc_comments)]
39    Unfinalized(UnfinalizedTree),
40
41    #[from(FinalizedTree)]
42    #[display("tap tree contains too many script leaves which doesn't fit a single Merkle tree")]
43    MountainRange,
44}
45
46#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Display, Error)]
47#[display("can't add more leaves to an already finalized tap tree")]
48pub struct FinalizedTree;
49
50#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Display, Error)]
51#[display(
52    "unfinalized tap tree containing leaves at level {0} which can't commit into a single Merkle \
53     root"
54)]
55pub struct UnfinalizedTree(pub u7);
56
57#[derive(Clone, Eq, PartialEq, Debug, Default)]
58pub struct TapTreeBuilder {
59    leaves: Vec<LeafInfo>,
60    buoy: MerkleBuoy<u7>,
61    finalized: bool,
62}
63
64impl TapTreeBuilder {
65    pub fn new() -> Self { Self::default() }
66
67    pub fn with_capacity(capacity: usize) -> Self {
68        Self {
69            leaves: Vec::with_capacity(capacity),
70            buoy: zero!(),
71            finalized: false,
72        }
73    }
74
75    pub fn is_finalized(&self) -> bool { self.finalized }
76
77    pub fn push_leaf(&mut self, leaf: LeafInfo) -> Result<bool, FinalizedTree> {
78        if self.finalized {
79            return Err(FinalizedTree);
80        }
81        let depth = leaf.depth;
82        self.leaves.push(leaf);
83        self.buoy.push(depth);
84        if self.buoy.level() == u7::ZERO {
85            self.finalized = true
86        }
87        Ok(self.finalized)
88    }
89
90    pub fn finish(self) -> Result<TapTree, UnfinalizedTree> {
91        if !self.finalized {
92            return Err(UnfinalizedTree(self.buoy.level()));
93        }
94        Ok(TapTree(self.leaves))
95    }
96}
97
98/// Non-empty taproot script tree.
99#[derive(Clone, Eq, PartialEq, Hash, Debug, Default)]
100#[cfg_attr(feature = "serde", derive(Serialize, Deserialize), serde(transparent))]
101pub struct TapTree(Vec<LeafInfo>);
102
103impl Deref for TapTree {
104    type Target = Vec<LeafInfo>;
105    fn deref(&self) -> &Self::Target { &self.0 }
106}
107
108impl IntoIterator for TapTree {
109    type Item = LeafInfo;
110    type IntoIter = vec::IntoIter<LeafInfo>;
111
112    fn into_iter(self) -> Self::IntoIter { self.0.into_iter() }
113}
114
115impl<'a> IntoIterator for &'a TapTree {
116    type Item = &'a LeafInfo;
117    type IntoIter = slice::Iter<'a, LeafInfo>;
118
119    fn into_iter(self) -> Self::IntoIter { self.0.iter() }
120}
121
122impl TapTree {
123    pub fn with_single_leaf(leaf: impl Into<LeafScript>) -> TapTree {
124        Self(vec![LeafInfo {
125            depth: u7::ZERO,
126            script: leaf.into(),
127        }])
128    }
129
130    pub fn from_leaves(leaves: impl IntoIterator<Item = LeafInfo>) -> Result<Self, InvalidTree> {
131        let mut builder = TapTreeBuilder::new();
132        for leaf in leaves {
133            builder.push_leaf(leaf)?;
134        }
135        builder.finish().map_err(InvalidTree::from)
136    }
137
138    pub fn from_builder(builder: TapTreeBuilder) -> Result<Self, UnfinalizedTree> {
139        builder.finish()
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    pub fn into_vec(self) -> Vec<LeafInfo> { self.0 }
151}
152
153#[derive(Clone, Eq, PartialEq, Hash, Debug)]
154#[cfg_attr(feature = "serde", derive(Serialize, Deserialize), serde(rename_all = "camelCase"))]
155pub struct LeafInfo {
156    pub depth: u7,
157    pub script: LeafScript,
158}
159
160impl LeafInfo {
161    pub fn tap_script(depth: u7, script: TapScript) -> Self {
162        LeafInfo {
163            depth,
164            script: LeafScript::from_tap_script(script),
165        }
166    }
167}
168
169#[derive(Getters, Clone, Eq, PartialEq, Debug)]
170#[getter(as_copy)]
171pub struct ControlBlockFactory {
172    internal_pk: InternalPk,
173    output_pk: OutputPk,
174    parity: Parity,
175    merkle_root: TapNodeHash,
176
177    #[getter(skip)]
178    merkle_path: TapMerklePath,
179    #[getter(skip)]
180    remaining_leaves: Vec<LeafInfo>,
181}
182
183impl ControlBlockFactory {
184    #[inline]
185    pub fn with(internal_pk: InternalPk, tap_tree: TapTree) -> Self {
186        let merkle_root = tap_tree.merkle_root();
187        let (output_pk, parity) = internal_pk.to_output_pk(Some(merkle_root));
188        ControlBlockFactory {
189            internal_pk,
190            output_pk,
191            parity,
192            merkle_root,
193            merkle_path: empty!(),
194            remaining_leaves: tap_tree.into_vec(),
195        }
196    }
197
198    #[inline]
199    pub fn into_remaining_leaves(self) -> Vec<LeafInfo> { self.remaining_leaves }
200}
201
202impl Iterator for ControlBlockFactory {
203    type Item = (ControlBlock, LeafScript);
204
205    fn next(&mut self) -> Option<Self::Item> {
206        let leaf = self.remaining_leaves.pop()?;
207        let leaf_script = leaf.script;
208        let control_block = ControlBlock::with(
209            leaf_script.version,
210            self.internal_pk,
211            self.parity,
212            self.merkle_path.clone(),
213        );
214        Some((control_block, leaf_script))
215    }
216}
217
218/// A compact size unsigned integer representing the number of leaf hashes, followed by a list
219/// of leaf hashes, followed by the 4 byte master key fingerprint concatenated with the
220/// derivation path of the public key. The derivation path is represented as 32-bit little
221/// endian unsigned integer indexes concatenated with each other. Public keys are those needed
222/// to spend this output. The leaf hashes are of the leaves which involve this public key. The
223/// internal key does not have leaf hashes, so can be indicated with a hashes len of 0.
224/// Finalizers should remove this field after `PSBT_IN_FINAL_SCRIPTWITNESS` is constructed.
225#[derive(Clone, Eq, PartialEq, Hash, Debug)]
226#[cfg_attr(feature = "serde", derive(Serialize, Deserialize), serde(rename_all = "camelCase"))]
227pub struct TapDerivation {
228    pub leaf_hashes: Vec<TapLeafHash>,
229    pub origin: KeyOrigin,
230}
231
232impl TapDerivation {
233    pub fn with_internal_pk(xpub_origin: XkeyOrigin, terminal: Terminal) -> Self {
234        let origin = KeyOrigin::with(xpub_origin, terminal);
235        TapDerivation {
236            leaf_hashes: empty!(),
237            origin,
238        }
239    }
240}