zebra_chain/parallel/
tree.rs

1//! Parallel note commitment tree update methods.
2
3use std::sync::Arc;
4
5use thiserror::Error;
6
7use crate::{
8    block::Block,
9    orchard, sapling, sprout,
10    subtree::{NoteCommitmentSubtree, NoteCommitmentSubtreeIndex},
11};
12
13/// An argument wrapper struct for note commitment trees.
14///
15/// The default instance represents the trees and subtrees that correspond to the genesis block.
16#[derive(Clone, Debug, Default, Eq, PartialEq)]
17pub struct NoteCommitmentTrees {
18    /// The sprout note commitment tree.
19    pub sprout: Arc<sprout::tree::NoteCommitmentTree>,
20
21    /// The sapling note commitment tree.
22    pub sapling: Arc<sapling::tree::NoteCommitmentTree>,
23
24    /// The sapling note commitment subtree.
25    pub sapling_subtree: Option<NoteCommitmentSubtree<sapling::tree::Node>>,
26
27    /// The orchard note commitment tree.
28    pub orchard: Arc<orchard::tree::NoteCommitmentTree>,
29
30    /// The orchard note commitment subtree.
31    pub orchard_subtree: Option<NoteCommitmentSubtree<orchard::tree::Node>>,
32}
33
34/// Note commitment tree errors.
35#[derive(Error, Copy, Clone, Debug, Eq, PartialEq, Hash)]
36pub enum NoteCommitmentTreeError {
37    /// A sprout tree error
38    #[error("sprout error: {0}")]
39    Sprout(#[from] sprout::tree::NoteCommitmentTreeError),
40
41    /// A sapling tree error
42    #[error("sapling error: {0}")]
43    Sapling(#[from] sapling::tree::NoteCommitmentTreeError),
44
45    /// A orchard tree error
46    #[error("orchard error: {0}")]
47    Orchard(#[from] orchard::tree::NoteCommitmentTreeError),
48}
49
50impl NoteCommitmentTrees {
51    /// Updates the note commitment trees using the transactions in `block`,
52    /// then re-calculates the cached tree roots, using parallel `rayon` threads.
53    ///
54    /// If any of the tree updates cause an error,
55    /// it will be returned at the end of the parallel batches.
56    #[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        // Prepare arguments for parallel threads
67        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    /// Update the sprout note commitment tree.
133    /// This method modifies the tree inside the `Arc`, if the `Arc` only has one reference.
134    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        // Re-calculate and cache the tree root.
145        let _ = sprout_nct.root();
146
147        Ok(sprout)
148    }
149
150    /// Update the sapling note commitment tree.
151    /// This method modifies the tree inside the `Arc`, if the `Arc` only has one reference.
152    #[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        // It is impossible for blocks to contain more than one level 16 sapling root:
166        // > [NU5 onward] nSpendsSapling, nOutputsSapling, and nActionsOrchard MUST all be less than 2^16.
167        // <https://zips.z.cash/protocol/protocol.pdf#txnconsensus>
168        //
169        // Before NU5, this limit holds due to the minimum size of Sapling outputs (948 bytes)
170        // and the maximum size of a block:
171        // > The size of a block MUST be less than or equal to 2000000 bytes.
172        // <https://zips.z.cash/protocol/protocol.pdf#blockheader>
173        // <https://zips.z.cash/protocol/protocol.pdf#txnencoding>
174        let mut subtree_root = None;
175
176        for sapling_note_commitment in sapling_note_commitments {
177            sapling_nct.append(sapling_note_commitment)?;
178
179            // Subtrees end heights come from the blocks they are completed in,
180            // so we check for new subtrees after appending the note.
181            // (If we check before, subtrees at the end of blocks have the wrong heights.)
182            if let Some(index_and_node) = sapling_nct.completed_subtree_index_and_root() {
183                subtree_root = Some(index_and_node);
184            }
185        }
186
187        // Re-calculate and cache the tree root.
188        let _ = sapling_nct.root();
189
190        Ok((sapling, subtree_root))
191    }
192
193    /// Update the orchard note commitment tree.
194    /// This method modifies the tree inside the `Arc`, if the `Arc` only has one reference.
195    #[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        // It is impossible for blocks to contain more than one level 16 orchard root:
209        // > [NU5 onward] nSpendsSapling, nOutputsSapling, and nActionsOrchard MUST all be less than 2^16.
210        // <https://zips.z.cash/protocol/protocol.pdf#txnconsensus>
211        let mut subtree_root = None;
212
213        for orchard_note_commitment in orchard_note_commitments {
214            orchard_nct.append(orchard_note_commitment)?;
215
216            // Subtrees end heights come from the blocks they are completed in,
217            // so we check for new subtrees after appending the note.
218            // (If we check before, subtrees at the end of blocks have the wrong heights.)
219            if let Some(index_and_node) = orchard_nct.completed_subtree_index_and_root() {
220                subtree_root = Some(index_and_node);
221            }
222        }
223
224        // Re-calculate and cache the tree root.
225        let _ = orchard_nct.root();
226
227        Ok((orchard, subtree_root))
228    }
229}