1use 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#[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#[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}