zebra_chain/parallel/
tree.rs1use std::sync::Arc;
4
5use thiserror::Error;
6
7use crate::{
8 block::Block,
9 orchard, sapling, sprout,
10 subtree::{NoteCommitmentSubtree, NoteCommitmentSubtreeIndex},
11};
12
13#[derive(Clone, Debug, Default, Eq, PartialEq)]
17pub struct NoteCommitmentTrees {
18 pub sprout: Arc<sprout::tree::NoteCommitmentTree>,
20
21 pub sapling: Arc<sapling::tree::NoteCommitmentTree>,
23
24 pub sapling_subtree: Option<NoteCommitmentSubtree<sapling::tree::Node>>,
26
27 pub orchard: Arc<orchard::tree::NoteCommitmentTree>,
29
30 pub orchard_subtree: Option<NoteCommitmentSubtree<orchard::tree::Node>>,
32}
33
34#[derive(Error, Copy, Clone, Debug, Eq, PartialEq, Hash)]
36pub enum NoteCommitmentTreeError {
37 #[error("sprout error: {0}")]
39 Sprout(#[from] sprout::tree::NoteCommitmentTreeError),
40
41 #[error("sapling error: {0}")]
43 Sapling(#[from] sapling::tree::NoteCommitmentTreeError),
44
45 #[error("orchard error: {0}")]
47 Orchard(#[from] orchard::tree::NoteCommitmentTreeError),
48}
49
50impl NoteCommitmentTrees {
51 #[allow(clippy::unwrap_in_result)]
57 pub fn update_trees_parallel(
58 &mut self,
59 block: &Arc<Block>,
60 ) -> Result<(), NoteCommitmentTreeError> {
61 let block = block.clone();
62 let height = block
63 .coinbase_height()
64 .expect("height was already validated");
65
66 let NoteCommitmentTrees {
68 sprout,
69 sapling,
70 orchard,
71 ..
72 } = self.clone();
73
74 let sprout_note_commitments: Vec<_> = block.sprout_note_commitments().cloned().collect();
75 let sapling_note_commitments: Vec<_> = block.sapling_note_commitments().cloned().collect();
76 let orchard_note_commitments: Vec<_> = block.orchard_note_commitments().cloned().collect();
77
78 let mut sprout_result = None;
79 let mut sapling_result = None;
80 let mut orchard_result = None;
81
82 rayon::in_place_scope_fifo(|scope| {
83 if !sprout_note_commitments.is_empty() {
84 scope.spawn_fifo(|_scope| {
85 sprout_result = Some(Self::update_sprout_note_commitment_tree(
86 sprout,
87 sprout_note_commitments,
88 ));
89 });
90 }
91
92 if !sapling_note_commitments.is_empty() {
93 scope.spawn_fifo(|_scope| {
94 sapling_result = Some(Self::update_sapling_note_commitment_tree(
95 sapling,
96 sapling_note_commitments,
97 ));
98 });
99 }
100
101 if !orchard_note_commitments.is_empty() {
102 scope.spawn_fifo(|_scope| {
103 orchard_result = Some(Self::update_orchard_note_commitment_tree(
104 orchard,
105 orchard_note_commitments,
106 ));
107 });
108 }
109 });
110
111 if let Some(sprout_result) = sprout_result {
112 self.sprout = sprout_result?;
113 }
114
115 if let Some(sapling_result) = sapling_result {
116 let (sapling, subtree_root) = sapling_result?;
117 self.sapling = sapling;
118 self.sapling_subtree =
119 subtree_root.map(|(idx, node)| NoteCommitmentSubtree::new(idx, height, node));
120 };
121
122 if let Some(orchard_result) = orchard_result {
123 let (orchard, subtree_root) = orchard_result?;
124 self.orchard = orchard;
125 self.orchard_subtree =
126 subtree_root.map(|(idx, node)| NoteCommitmentSubtree::new(idx, height, node));
127 };
128
129 Ok(())
130 }
131
132 fn update_sprout_note_commitment_tree(
135 mut sprout: Arc<sprout::tree::NoteCommitmentTree>,
136 sprout_note_commitments: Vec<sprout::NoteCommitment>,
137 ) -> Result<Arc<sprout::tree::NoteCommitmentTree>, NoteCommitmentTreeError> {
138 let sprout_nct = Arc::make_mut(&mut sprout);
139
140 for sprout_note_commitment in sprout_note_commitments {
141 sprout_nct.append(sprout_note_commitment)?;
142 }
143
144 let _ = sprout_nct.root();
146
147 Ok(sprout)
148 }
149
150 #[allow(clippy::unwrap_in_result)]
153 pub fn update_sapling_note_commitment_tree(
154 mut sapling: Arc<sapling::tree::NoteCommitmentTree>,
155 sapling_note_commitments: Vec<sapling::tree::NoteCommitmentUpdate>,
156 ) -> Result<
157 (
158 Arc<sapling::tree::NoteCommitmentTree>,
159 Option<(NoteCommitmentSubtreeIndex, sapling::tree::Node)>,
160 ),
161 NoteCommitmentTreeError,
162 > {
163 let sapling_nct = Arc::make_mut(&mut sapling);
164
165 let mut subtree_root = None;
175
176 for sapling_note_commitment in sapling_note_commitments {
177 sapling_nct.append(sapling_note_commitment)?;
178
179 if let Some(index_and_node) = sapling_nct.completed_subtree_index_and_root() {
183 subtree_root = Some(index_and_node);
184 }
185 }
186
187 let _ = sapling_nct.root();
189
190 Ok((sapling, subtree_root))
191 }
192
193 #[allow(clippy::unwrap_in_result)]
196 pub fn update_orchard_note_commitment_tree(
197 mut orchard: Arc<orchard::tree::NoteCommitmentTree>,
198 orchard_note_commitments: Vec<orchard::tree::NoteCommitmentUpdate>,
199 ) -> Result<
200 (
201 Arc<orchard::tree::NoteCommitmentTree>,
202 Option<(NoteCommitmentSubtreeIndex, orchard::tree::Node)>,
203 ),
204 NoteCommitmentTreeError,
205 > {
206 let orchard_nct = Arc::make_mut(&mut orchard);
207
208 let mut subtree_root = None;
212
213 for orchard_note_commitment in orchard_note_commitments {
214 orchard_nct.append(orchard_note_commitment)?;
215
216 if let Some(index_and_node) = orchard_nct.completed_subtree_index_and_root() {
220 subtree_root = Some(index_and_node);
221 }
222 }
223
224 let _ = orchard_nct.root();
226
227 Ok((orchard, subtree_root))
228 }
229}