zebra_state/service/finalized_state/disk_format/upgrade/
add_subtrees.rs

1//! Fully populate the Sapling and Orchard note commitment subtrees for existing blocks in the database.
2
3use std::sync::Arc;
4
5use crossbeam_channel::{Receiver, TryRecvError};
6use hex_literal::hex;
7use itertools::Itertools;
8use semver::Version;
9use tracing::instrument;
10
11use zebra_chain::{
12    block::Height,
13    orchard,
14    parallel::tree::NoteCommitmentTrees,
15    parameters::Network::*,
16    sapling,
17    subtree::{NoteCommitmentSubtree, NoteCommitmentSubtreeIndex},
18};
19
20use crate::service::finalized_state::{
21    disk_format::upgrade::{CancelFormatChange, DiskFormatUpgrade},
22    DiskWriteBatch, ZebraDb,
23};
24
25/// Implements [`DiskFormatUpgrade`] for populating Sapling and Orchard note commitment subtrees.
26pub struct AddSubtrees;
27
28impl DiskFormatUpgrade for AddSubtrees {
29    fn version(&self) -> Version {
30        Version::new(25, 2, 2)
31    }
32
33    fn description(&self) -> &'static str {
34        "add subtrees upgrade"
35    }
36
37    fn prepare(
38        &self,
39        initial_tip_height: Height,
40        upgrade_db: &ZebraDb,
41        cancel_receiver: &Receiver<CancelFormatChange>,
42        older_disk_version: &Version,
43    ) -> Result<(), CancelFormatChange> {
44        let first_version_for_adding_subtrees = Version::new(25, 2, 0);
45        if older_disk_version >= &first_version_for_adding_subtrees {
46            // Clear previous upgrade data, because it was incorrect.
47            reset(initial_tip_height, upgrade_db, cancel_receiver)?;
48        }
49
50        Ok(())
51    }
52
53    /// Runs disk format upgrade for adding Sapling and Orchard note commitment subtrees to database.
54    ///
55    /// Trees are added to the database in reverse height order, so that wallets can sync correctly
56    /// while the upgrade is running.
57    ///
58    /// Returns `Ok` if the upgrade completed, and `Err` if it was cancelled.
59    fn run(
60        &self,
61        initial_tip_height: Height,
62        upgrade_db: &ZebraDb,
63        cancel_receiver: &Receiver<CancelFormatChange>,
64    ) -> Result<(), CancelFormatChange> {
65        // # Consensus
66        //
67        // Zebra stores exactly one note commitment tree for every block with sapling notes.
68        // (It also stores the empty note commitment tree for the genesis block, but we skip that.)
69        //
70        // The consensus rules limit blocks to less than 2^16 sapling and 2^16 orchard outputs. So a
71        // block can't complete multiple level 16 subtrees (or complete an entire subtree by itself).
72        // Currently, with 2MB blocks and v4/v5 sapling and orchard output sizes, the subtree index can
73        // increase by at most 1 every ~20 blocks.
74        //
75        // # Compatibility
76        //
77        // Because wallets search backwards from the chain tip, subtrees need to be added to the
78        // database in reverse height order. (Tip first, genesis last.)
79        //
80        // Otherwise, wallets that sync during the upgrade will be missing some notes.
81
82        // Generate a list of sapling subtree inputs: previous and current trees, and their end heights.
83        let subtrees = upgrade_db
84            .sapling_tree_by_reversed_height_range(..=initial_tip_height)
85            // We need both the tree and its previous tree for each shielded block.
86            .tuple_windows()
87            // Because the iterator is reversed, the larger tree is first.
88            .map(|((end_height, tree), (prev_end_height, prev_tree))| {
89                (prev_end_height, prev_tree, end_height, tree)
90            })
91            // Find new subtrees.
92            .filter(|(_prev_end_height, prev_tree, _end_height, tree)| {
93                tree.contains_new_subtree(prev_tree)
94            });
95
96        for (prev_end_height, prev_tree, end_height, tree) in subtrees {
97            // Return early if the upgrade is cancelled.
98            if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
99                return Err(CancelFormatChange);
100            }
101
102            let subtree =
103                calculate_sapling_subtree(upgrade_db, prev_end_height, prev_tree, end_height, tree);
104            write_sapling_subtree(upgrade_db, subtree);
105        }
106
107        // Generate a list of orchard subtree inputs: previous and current trees, and their end heights.
108        let subtrees = upgrade_db
109            .orchard_tree_by_reversed_height_range(..=initial_tip_height)
110            // We need both the tree and its previous tree for each shielded block.
111            .tuple_windows()
112            // Because the iterator is reversed, the larger tree is first.
113            .map(|((end_height, tree), (prev_end_height, prev_tree))| {
114                (prev_end_height, prev_tree, end_height, tree)
115            })
116            // Find new subtrees.
117            .filter(|(_prev_end_height, prev_tree, _end_height, tree)| {
118                tree.contains_new_subtree(prev_tree)
119            });
120
121        for (prev_end_height, prev_tree, end_height, tree) in subtrees {
122            // Return early if the upgrade is cancelled.
123            if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
124                return Err(CancelFormatChange);
125            }
126
127            let subtree =
128                calculate_orchard_subtree(upgrade_db, prev_end_height, prev_tree, end_height, tree);
129            write_orchard_subtree(upgrade_db, subtree);
130        }
131
132        Ok(())
133    }
134
135    #[allow(clippy::unwrap_in_result)]
136    fn validate(
137        &self,
138        db: &ZebraDb,
139        cancel_receiver: &Receiver<CancelFormatChange>,
140    ) -> Result<Result<(), String>, CancelFormatChange> {
141        // This is redundant in some code paths, but not in others. But it's quick anyway.
142        let quick_result = subtree_format_calculation_pre_checks(db);
143
144        // Check the entire format before returning any errors.
145        let sapling_result = check_sapling_subtrees(db, cancel_receiver)?;
146        let orchard_result = check_orchard_subtrees(db, cancel_receiver)?;
147
148        if quick_result.is_err() || sapling_result.is_err() || orchard_result.is_err() {
149            let err = Err(format!(
150                "missing or invalid subtree(s): \
151             quick: {quick_result:?}, sapling: {sapling_result:?}, orchard: {orchard_result:?}"
152            ));
153            warn!(?err);
154            return Ok(err);
155        }
156
157        Ok(Ok(()))
158    }
159}
160
161/// Reset data from previous upgrades. This data can be complete or incomplete.
162///
163/// Returns `Ok` if the upgrade completed, and `Err` if it was cancelled.
164#[allow(clippy::unwrap_in_result)]
165#[instrument(skip(upgrade_db, cancel_receiver))]
166pub fn reset(
167    _initial_tip_height: Height,
168    upgrade_db: &ZebraDb,
169    cancel_receiver: &Receiver<CancelFormatChange>,
170) -> Result<(), CancelFormatChange> {
171    // Return early if the upgrade is cancelled.
172    if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
173        return Err(CancelFormatChange);
174    }
175
176    // This doesn't delete the maximum index, but the consensus rules make that subtree impossible.
177    // (Adding a note to a full note commitment tree is an error.)
178    //
179    // TODO: convert zs_delete_range() to take std::ops::RangeBounds, and delete the upper bound.
180    let mut batch = DiskWriteBatch::new();
181    batch.delete_range_sapling_subtree(upgrade_db, 0.into(), u16::MAX.into());
182    upgrade_db
183        .write_batch(batch)
184        .expect("deleting old sapling note commitment subtrees is a valid database operation");
185
186    if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
187        return Err(CancelFormatChange);
188    }
189
190    let mut batch = DiskWriteBatch::new();
191    batch.delete_range_orchard_subtree(upgrade_db, 0.into(), u16::MAX.into());
192    upgrade_db
193        .write_batch(batch)
194        .expect("deleting old orchard note commitment subtrees is a valid database operation");
195
196    Ok(())
197}
198
199/// Quickly check that the first calculated subtree is correct.
200///
201/// This allows us to fail the upgrade quickly in tests and during development,
202/// rather than waiting ~20 minutes to see if it failed.
203///
204/// This check runs the first subtree calculation, but it doesn't read the subtree data in the
205/// database. So it can be run before the upgrade is started.
206pub fn subtree_format_calculation_pre_checks(db: &ZebraDb) -> Result<(), String> {
207    // Check the entire format before returning any errors.
208    let sapling_result = quick_check_sapling_subtrees(db);
209    let orchard_result = quick_check_orchard_subtrees(db);
210
211    if sapling_result.is_err() || orchard_result.is_err() {
212        let err = Err(format!(
213            "missing or bad first subtree: sapling: {sapling_result:?}, orchard: {orchard_result:?}"
214        ));
215        warn!(?err);
216        return err;
217    }
218
219    Ok(())
220}
221
222/// A quick test vector that allows us to fail an incorrect upgrade within a few seconds.
223fn first_sapling_mainnet_subtree() -> NoteCommitmentSubtree<sapling::tree::Node> {
224    // This test vector was generated using the command:
225    // ```sh
226    // zcash-cli z_getsubtreesbyindex sapling 0 1
227    // ```
228    NoteCommitmentSubtree {
229        index: 0.into(),
230        root: hex!("754bb593ea42d231a7ddf367640f09bbf59dc00f2c1d2003cc340e0c016b5b13")
231            .as_slice()
232            .try_into()
233            .expect("test vector is valid"),
234        end_height: Height(558822),
235    }
236}
237
238/// A quick test vector that allows us to fail an incorrect upgrade within a few seconds.
239fn first_orchard_mainnet_subtree() -> NoteCommitmentSubtree<orchard::tree::Node> {
240    // This test vector was generated using the command:
241    // ```sh
242    // zcash-cli z_getsubtreesbyindex orchard 0 1
243    // ```
244    NoteCommitmentSubtree {
245        index: 0.into(),
246        root: hex!("d4e323b3ae0cabfb6be4087fec8c66d9a9bbfc354bf1d9588b6620448182063b")
247            .as_slice()
248            .try_into()
249            .expect("test vector is valid"),
250        end_height: Height(1707429),
251    }
252}
253
254/// Quickly check that the first calculated sapling subtree is correct.
255///
256/// This allows us to fail the upgrade quickly in tests and during development,
257/// rather than waiting ~20 minutes to see if it failed.
258///
259/// Returns an error if a note commitment subtree is missing or incorrect.
260fn quick_check_sapling_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
261    // We check the first sapling subtree on mainnet, so skip this check if it isn't available.
262    if db.network() != Mainnet {
263        return Ok(());
264    }
265
266    let Some(NoteCommitmentSubtreeIndex(tip_subtree_index)) =
267        db.sapling_tree_for_tip().subtree_index()
268    else {
269        return Ok(());
270    };
271
272    if tip_subtree_index == 0 && !db.sapling_tree_for_tip().is_complete_subtree() {
273        return Ok(());
274    }
275
276    // Find the first complete subtree: previous and current trees, and their end heights.
277    let first_complete_subtree = db
278        .sapling_tree_by_height_range(..)
279        // We need both the tree and its previous tree for each shielded block.
280        .tuple_windows()
281        .map(|((prev_end_height, prev_tree), (end_height, tree))| {
282            (prev_end_height, prev_tree, end_height, tree)
283        })
284        .find(|(_prev_end_height, prev_tree, _end_height, tree)| {
285            tree.contains_new_subtree(prev_tree)
286        });
287
288    let Some((prev_end_height, prev_tree, end_height, tree)) = first_complete_subtree else {
289        let result = Err("iterator did not find complete subtree, but the tree has it");
290        error!(?result);
291        return result;
292    };
293
294    // Creating this test vector involves a cryptographic check, so only do it once.
295    let expected_subtree = first_sapling_mainnet_subtree();
296
297    let db_subtree = calculate_sapling_subtree(db, prev_end_height, prev_tree, end_height, tree);
298
299    if db_subtree != expected_subtree {
300        let result = Err("first subtree did not match expected test vector");
301        error!(?result, ?db_subtree, ?expected_subtree);
302        return result;
303    }
304
305    Ok(())
306}
307
308/// Quickly check that the first calculated orchard subtree is correct.
309///
310/// This allows us to fail the upgrade quickly in tests and during development,
311/// rather than waiting ~20 minutes to see if it failed.
312///
313/// Returns an error if a note commitment subtree is missing or incorrect.
314fn quick_check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
315    // We check the first orchard subtree on mainnet, so skip this check if it isn't available.
316    if db.network() != Mainnet {
317        return Ok(());
318    }
319
320    let Some(NoteCommitmentSubtreeIndex(tip_subtree_index)) =
321        db.orchard_tree_for_tip().subtree_index()
322    else {
323        return Ok(());
324    };
325
326    if tip_subtree_index == 0 && !db.orchard_tree_for_tip().is_complete_subtree() {
327        return Ok(());
328    }
329
330    // Find the first complete subtree: previous and current trees, and their end heights.
331    let first_complete_subtree = db
332        .orchard_tree_by_height_range(..)
333        // We need both the tree and its previous tree for each shielded block.
334        .tuple_windows()
335        .map(|((prev_end_height, prev_tree), (end_height, tree))| {
336            (prev_end_height, prev_tree, end_height, tree)
337        })
338        .find(|(_prev_end_height, prev_tree, _end_height, tree)| {
339            tree.contains_new_subtree(prev_tree)
340        });
341
342    let Some((prev_end_height, prev_tree, end_height, tree)) = first_complete_subtree else {
343        let result = Err("iterator did not find complete subtree, but the tree has it");
344        error!(?result);
345        return result;
346    };
347
348    // Creating this test vector involves a cryptographic check, so only do it once.
349    let expected_subtree = first_orchard_mainnet_subtree();
350
351    let db_subtree = calculate_orchard_subtree(db, prev_end_height, prev_tree, end_height, tree);
352
353    if db_subtree != expected_subtree {
354        let result = Err("first subtree did not match expected test vector");
355        error!(?result, ?db_subtree, ?expected_subtree);
356        return result;
357    }
358
359    Ok(())
360}
361
362/// Check that Sapling note commitment subtrees were correctly added.
363///
364/// Returns an error if a note commitment subtree is missing or incorrect.
365fn check_sapling_subtrees(
366    db: &ZebraDb,
367    cancel_receiver: &Receiver<CancelFormatChange>,
368) -> Result<Result<(), &'static str>, CancelFormatChange> {
369    let Some(NoteCommitmentSubtreeIndex(mut first_incomplete_subtree_index)) =
370        db.sapling_tree_for_tip().subtree_index()
371    else {
372        return Ok(Ok(()));
373    };
374
375    // If there are no incomplete subtrees in the tree, also expect a subtree for the final index.
376    if db.sapling_tree_for_tip().is_complete_subtree() {
377        first_incomplete_subtree_index += 1;
378    }
379
380    let mut result = Ok(());
381    for index in 0..first_incomplete_subtree_index {
382        // Return early if the format check is cancelled.
383        if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
384            return Err(CancelFormatChange);
385        }
386
387        // Check that there's a continuous range of subtrees from index [0, first_incomplete_subtree_index)
388        let Some(subtree) = db.sapling_subtree_by_index(index) else {
389            result = Err("missing subtree");
390            error!(?result, index);
391            continue;
392        };
393
394        // Check that there was a sapling note at the subtree's end height.
395        let Some(tree) = db.sapling_tree_by_height(&subtree.end_height) else {
396            result = Err("missing note commitment tree at subtree completion height");
397            error!(?result, ?subtree.end_height);
398            continue;
399        };
400
401        // Check the index and root if the sapling note commitment tree at this height is a complete subtree.
402        if let Some((index, node)) = tree.completed_subtree_index_and_root() {
403            if subtree.index != index {
404                result = Err("completed subtree indexes should match");
405                error!(?result);
406            }
407
408            if subtree.root != node {
409                result = Err("completed subtree roots should match");
410                error!(?result);
411            }
412        }
413        // Check that the final note has a greater subtree index if it didn't complete a subtree.
414        else {
415            let prev_height = subtree
416                .end_height
417                .previous()
418                .expect("Note commitment subtrees should not end at the minimal height.");
419
420            let Some(prev_tree) = db.sapling_tree_by_height(&prev_height) else {
421                result = Err("missing note commitment tree below subtree completion height");
422                error!(?result, ?subtree.end_height);
423                continue;
424            };
425
426            let prev_subtree_index = prev_tree.subtree_index();
427            let subtree_index = tree.subtree_index();
428            if subtree_index <= prev_subtree_index {
429                result =
430                    Err("note commitment tree at end height should have incremented subtree index");
431                error!(?result, ?subtree_index, ?prev_subtree_index,);
432            }
433        }
434    }
435
436    let mut subtree_count = 0;
437    for (index, height, tree) in db
438        .sapling_tree_by_height_range(..)
439        // Exclude empty sapling tree and add subtree indexes
440        .filter_map(|(height, tree)| Some((tree.subtree_index()?, height, tree)))
441        // Exclude heights that don't complete a subtree and count completed subtrees
442        .filter_map(|(subtree_index, height, tree)| {
443            if tree.is_complete_subtree() || subtree_index.0 > subtree_count {
444                let subtree_index = subtree_count;
445                subtree_count += 1;
446                Some((subtree_index, height, tree))
447            } else {
448                None
449            }
450        })
451    {
452        // Return early if the format check is cancelled.
453        if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
454            return Err(CancelFormatChange);
455        }
456
457        // Check that there's an entry for every completed sapling subtree root in all sapling trees
458        let Some(subtree) = db.sapling_subtree_by_index(index) else {
459            result = Err("missing subtree");
460            error!(?result, index);
461            continue;
462        };
463
464        // Check that the subtree end height matches that in the sapling trees.
465        if subtree.end_height != height {
466            let is_complete = tree.is_complete_subtree();
467            result = Err("bad sapling subtree end height");
468            error!(?result, ?subtree.end_height, ?height, ?index, ?is_complete, );
469        }
470
471        // Check the root if the sapling note commitment tree at this height is a complete subtree.
472        if let Some((_index, node)) = tree.completed_subtree_index_and_root() {
473            if subtree.root != node {
474                result = Err("completed subtree roots should match");
475                error!(?result);
476            }
477        }
478    }
479
480    if result.is_err() {
481        error!(
482            ?result,
483            ?subtree_count,
484            first_incomplete_subtree_index,
485            "missing or bad sapling subtrees"
486        );
487    }
488
489    Ok(result)
490}
491
492/// Check that Orchard note commitment subtrees were correctly added.
493///
494/// Returns an error if a note commitment subtree is missing or incorrect.
495fn check_orchard_subtrees(
496    db: &ZebraDb,
497    cancel_receiver: &Receiver<CancelFormatChange>,
498) -> Result<Result<(), &'static str>, CancelFormatChange> {
499    let Some(NoteCommitmentSubtreeIndex(mut first_incomplete_subtree_index)) =
500        db.orchard_tree_for_tip().subtree_index()
501    else {
502        return Ok(Ok(()));
503    };
504
505    // If there are no incomplete subtrees in the tree, also expect a subtree for the final index.
506    if db.orchard_tree_for_tip().is_complete_subtree() {
507        first_incomplete_subtree_index += 1;
508    }
509
510    let mut result = Ok(());
511    for index in 0..first_incomplete_subtree_index {
512        // Return early if the format check is cancelled.
513        if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
514            return Err(CancelFormatChange);
515        }
516
517        // Check that there's a continuous range of subtrees from index [0, first_incomplete_subtree_index)
518        let Some(subtree) = db.orchard_subtree_by_index(index) else {
519            result = Err("missing subtree");
520            error!(?result, index);
521            continue;
522        };
523
524        // Check that there was a orchard note at the subtree's end height.
525        let Some(tree) = db.orchard_tree_by_height(&subtree.end_height) else {
526            result = Err("missing note commitment tree at subtree completion height");
527            error!(?result, ?subtree.end_height);
528            continue;
529        };
530
531        // Check the index and root if the orchard note commitment tree at this height is a complete subtree.
532        if let Some((index, node)) = tree.completed_subtree_index_and_root() {
533            if subtree.index != index {
534                result = Err("completed subtree indexes should match");
535                error!(?result);
536            }
537
538            if subtree.root != node {
539                result = Err("completed subtree roots should match");
540                error!(?result);
541            }
542        }
543        // Check that the final note has a greater subtree index if it didn't complete a subtree.
544        else {
545            let prev_height = subtree
546                .end_height
547                .previous()
548                .expect("Note commitment subtrees should not end at the minimal height.");
549
550            let Some(prev_tree) = db.orchard_tree_by_height(&prev_height) else {
551                result = Err("missing note commitment tree below subtree completion height");
552                error!(?result, ?subtree.end_height);
553                continue;
554            };
555
556            let prev_subtree_index = prev_tree.subtree_index();
557            let subtree_index = tree.subtree_index();
558            if subtree_index <= prev_subtree_index {
559                result =
560                    Err("note commitment tree at end height should have incremented subtree index");
561                error!(?result, ?subtree_index, ?prev_subtree_index,);
562            }
563        }
564    }
565
566    let mut subtree_count = 0;
567    for (index, height, tree) in db
568        .orchard_tree_by_height_range(..)
569        // Exclude empty orchard tree and add subtree indexes
570        .filter_map(|(height, tree)| Some((tree.subtree_index()?, height, tree)))
571        // Exclude heights that don't complete a subtree and count completed subtrees
572        .filter_map(|(subtree_index, height, tree)| {
573            if tree.is_complete_subtree() || subtree_index.0 > subtree_count {
574                let subtree_index = subtree_count;
575                subtree_count += 1;
576                Some((subtree_index, height, tree))
577            } else {
578                None
579            }
580        })
581    {
582        // Return early if the format check is cancelled.
583        if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
584            return Err(CancelFormatChange);
585        }
586
587        // Check that there's an entry for every completed orchard subtree root in all orchard trees
588        let Some(subtree) = db.orchard_subtree_by_index(index) else {
589            result = Err("missing subtree");
590            error!(?result, index);
591            continue;
592        };
593
594        // Check that the subtree end height matches that in the orchard trees.
595        if subtree.end_height != height {
596            let is_complete = tree.is_complete_subtree();
597            result = Err("bad orchard subtree end height");
598            error!(?result, ?subtree.end_height, ?height, ?index, ?is_complete, );
599        }
600
601        // Check the root if the orchard note commitment tree at this height is a complete subtree.
602        if let Some((_index, node)) = tree.completed_subtree_index_and_root() {
603            if subtree.root != node {
604                result = Err("completed subtree roots should match");
605                error!(?result);
606            }
607        }
608    }
609
610    if result.is_err() {
611        error!(
612            ?result,
613            ?subtree_count,
614            first_incomplete_subtree_index,
615            "missing or bad orchard subtrees"
616        );
617    }
618
619    Ok(result)
620}
621
622/// Calculates a note commitment subtree for Sapling, reading blocks from `read_db` if needed.
623///
624/// The subtree must be completed by a note commitment in the block at `end_height`.
625/// `tree` is the tree for that block, and `prev_tree` is the tree for the previous block.
626///
627/// `prev_tree` is only used to rebuild the subtree if it was completed without using the last
628/// note commitment in the block at `end_height`.
629///
630/// # Panics
631///
632/// If a subtree is not completed by a note commitment in the block at `end_height`.
633#[must_use = "subtree should be written to the database after it is calculated"]
634#[instrument(skip(read_db, prev_tree, tree))]
635fn calculate_sapling_subtree(
636    read_db: &ZebraDb,
637    prev_end_height: Height,
638    prev_tree: Arc<sapling::tree::NoteCommitmentTree>,
639    end_height: Height,
640    tree: Arc<sapling::tree::NoteCommitmentTree>,
641) -> NoteCommitmentSubtree<sapling::tree::Node> {
642    // If a subtree is completed by a note commitment in the block at `end_height`,
643    // then that subtree can be completed in two different ways:
644    if let Some((index, node)) = tree.completed_subtree_index_and_root() {
645        // If the subtree is completed by the last note commitment in that block,
646        // we already have that subtree root available in the tree.
647        NoteCommitmentSubtree::new(index, end_height, node)
648    } else {
649        // If the subtree is completed without using the last note commitment in the block,
650        // we need to calculate the subtree root, starting with the tree from the previous block.
651
652        // TODO: move the assertion/panic log string formatting into a separate function?
653        let prev_position = prev_tree.position().unwrap_or_else(|| {
654            panic!(
655                "previous block must have a partial subtree:\n\
656                previous subtree:\n\
657                height: {prev_end_height:?}\n\
658                current subtree:\n\
659                height: {end_height:?}"
660            )
661        });
662        let prev_index = prev_tree
663            .subtree_index()
664            .expect("previous block must have a partial subtree");
665        let prev_remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
666
667        let current_position = tree.position().unwrap_or_else(|| {
668            panic!(
669                "current block must have a subtree:\n\
670                previous subtree:\n\
671                height: {prev_end_height:?}\n\
672                index: {prev_index}\n\
673                position: {prev_position}\n\
674                remaining: {prev_remaining_notes}\n\
675                current subtree:\n\
676                height: {end_height:?}"
677            )
678        });
679        let current_index = tree
680            .subtree_index()
681            .expect("current block must have a subtree");
682        let current_remaining_notes = tree.remaining_subtree_leaf_nodes();
683
684        assert_eq!(
685            prev_index.0 + 1,
686            current_index.0,
687            "subtree must have been completed by the current block:\n\
688             previous subtree:\n\
689             height: {prev_end_height:?}\n\
690             index: {prev_index}\n\
691             position: {prev_position}\n\
692             remaining: {prev_remaining_notes}\n\
693             current subtree:\n\
694             height: {end_height:?}\n\
695             index: {current_index}\n\
696             position: {current_position}\n\
697             remaining: {current_remaining_notes}"
698        );
699
700        // Get the missing notes needed to complete the subtree.
701        //
702        // TODO: consider just reading the block's transactions from the database file,
703        //       because we don't use the block header data at all.
704        let block = read_db
705            .block(end_height.into())
706            .expect("height with note commitment tree should have block");
707        let sapling_note_commitments = block
708            .sapling_note_commitments()
709            .take(prev_remaining_notes)
710            .cloned()
711            .collect();
712
713        // This takes less than 1 second per tree, so we don't need to make it cancellable.
714        let (sapling_nct, subtree) = NoteCommitmentTrees::update_sapling_note_commitment_tree(
715            prev_tree,
716            sapling_note_commitments,
717        )
718        .expect("finalized notes should append successfully");
719
720        let (index, node) = subtree.unwrap_or_else(|| {
721            panic!(
722                "already checked that the block completed a subtree:\n\
723                 updated subtree:\n\
724                 index: {:?}\n\
725                 position: {:?}\n\
726                 remaining notes: {}\n\
727                 original previous subtree:\n\
728                 height: {prev_end_height:?}\n\
729                 index: {prev_index}\n\
730                 position: {prev_position}\n\
731                 remaining: {prev_remaining_notes}\n\
732                 original current subtree:\n\
733                 height: {end_height:?}\n\
734                 index: {current_index}\n\
735                 position: {current_position}\n\
736                 remaining: {current_remaining_notes}",
737                sapling_nct.subtree_index(),
738                sapling_nct.position(),
739                sapling_nct.remaining_subtree_leaf_nodes(),
740            )
741        });
742
743        NoteCommitmentSubtree::new(index, end_height, node)
744    }
745}
746
747/// Calculates a note commitment subtree for Orchard, reading blocks from `read_db` if needed.
748///
749/// The subtree must be completed by a note commitment in the block at `end_height`.
750/// `tree` is the tree for that block, and `prev_tree` is the tree for the previous block.
751///
752/// `prev_tree` is only used to rebuild the subtree if it was completed without using the last
753/// note commitment in the block at `end_height`.
754///
755/// # Panics
756///
757/// If a subtree is not completed by a note commitment in the block at `end_height`.
758#[must_use = "subtree should be written to the database after it is calculated"]
759#[instrument(skip(read_db, prev_tree, tree))]
760fn calculate_orchard_subtree(
761    read_db: &ZebraDb,
762    prev_end_height: Height,
763    prev_tree: Arc<orchard::tree::NoteCommitmentTree>,
764    end_height: Height,
765    tree: Arc<orchard::tree::NoteCommitmentTree>,
766) -> NoteCommitmentSubtree<orchard::tree::Node> {
767    // If a subtree is completed by a note commitment in the block at `end_height`,
768    // then that subtree can be completed in two different ways:
769    if let Some((index, node)) = tree.completed_subtree_index_and_root() {
770        // If the subtree is completed by the last note commitment in that block,
771        // we already have that subtree root available in the tree.
772        NoteCommitmentSubtree::new(index, end_height, node)
773    } else {
774        // If the subtree is completed without using the last note commitment in the block,
775        // we need to calculate the subtree root, starting with the tree from the previous block.
776
777        // TODO: move the assertion/panic log string formatting into a separate function?
778        let prev_position = prev_tree.position().unwrap_or_else(|| {
779            panic!(
780                "previous block must have a partial subtree:\n\
781                previous subtree:\n\
782                height: {prev_end_height:?}\n\
783                current subtree:\n\
784                height: {end_height:?}"
785            )
786        });
787        let prev_index = prev_tree
788            .subtree_index()
789            .expect("previous block must have a partial subtree");
790        let prev_remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
791
792        let current_position = tree.position().unwrap_or_else(|| {
793            panic!(
794                "current block must have a subtree:\n\
795                previous subtree:\n\
796                height: {prev_end_height:?}\n\
797                index: {prev_index}\n\
798                position: {prev_position}\n\
799                remaining: {prev_remaining_notes}\n\
800                current subtree:\n\
801                height: {end_height:?}"
802            )
803        });
804        let current_index = tree
805            .subtree_index()
806            .expect("current block must have a subtree");
807        let current_remaining_notes = tree.remaining_subtree_leaf_nodes();
808
809        assert_eq!(
810            prev_index.0 + 1,
811            current_index.0,
812            "subtree must have been completed by the current block:\n\
813             previous subtree:\n\
814             height: {prev_end_height:?}\n\
815             index: {prev_index}\n\
816             position: {prev_position}\n\
817             remaining: {prev_remaining_notes}\n\
818             current subtree:\n\
819             height: {end_height:?}\n\
820             index: {current_index}\n\
821             position: {current_position}\n\
822             remaining: {current_remaining_notes}"
823        );
824
825        // Get the missing notes needed to complete the subtree.
826        //
827        // TODO: consider just reading the block's transactions from the database file,
828        //       because we don't use the block header data at all.
829        let block = read_db
830            .block(end_height.into())
831            .expect("height with note commitment tree should have block");
832        let orchard_note_commitments = block
833            .orchard_note_commitments()
834            .take(prev_remaining_notes)
835            .cloned()
836            .collect();
837
838        // This takes less than 1 second per tree, so we don't need to make it cancellable.
839        let (orchard_nct, subtree) = NoteCommitmentTrees::update_orchard_note_commitment_tree(
840            prev_tree,
841            orchard_note_commitments,
842        )
843        .expect("finalized notes should append successfully");
844
845        let (index, node) = subtree.unwrap_or_else(|| {
846            panic!(
847                "already checked that the block completed a subtree:\n\
848                 updated subtree:\n\
849                 index: {:?}\n\
850                 position: {:?}\n\
851                 remaining notes: {}\n\
852                 original previous subtree:\n\
853                 height: {prev_end_height:?}\n\
854                 index: {prev_index}\n\
855                 position: {prev_position}\n\
856                 remaining: {prev_remaining_notes}\n\
857                 original current subtree:\n\
858                 height: {end_height:?}\n\
859                 index: {current_index}\n\
860                 position: {current_position}\n\
861                 remaining: {current_remaining_notes}",
862                orchard_nct.subtree_index(),
863                orchard_nct.position(),
864                orchard_nct.remaining_subtree_leaf_nodes(),
865            )
866        });
867
868        NoteCommitmentSubtree::new(index, end_height, node)
869    }
870}
871
872/// Writes a Sapling note commitment subtree to `upgrade_db`.
873fn write_sapling_subtree(
874    upgrade_db: &ZebraDb,
875    subtree: NoteCommitmentSubtree<sapling::tree::Node>,
876) {
877    let mut batch = DiskWriteBatch::new();
878
879    batch.insert_sapling_subtree(upgrade_db, &subtree);
880
881    upgrade_db
882        .write_batch(batch)
883        .expect("writing sapling note commitment subtrees should always succeed.");
884
885    if subtree.index.0 % 100 == 0 {
886        info!(end_height = ?subtree.end_height, index = ?subtree.index.0, "calculated and added sapling subtree");
887    }
888    // This log happens about once per second on recent machines with SSD disks.
889    debug!(end_height = ?subtree.end_height, index = ?subtree.index.0, "calculated and added sapling subtree");
890}
891
892/// Writes an Orchard note commitment subtree to `upgrade_db`.
893fn write_orchard_subtree(
894    upgrade_db: &ZebraDb,
895    subtree: NoteCommitmentSubtree<orchard::tree::Node>,
896) {
897    let mut batch = DiskWriteBatch::new();
898
899    batch.insert_orchard_subtree(upgrade_db, &subtree);
900
901    upgrade_db
902        .write_batch(batch)
903        .expect("writing orchard note commitment subtrees should always succeed.");
904
905    if subtree.index.0 % 100 == 0 {
906        info!(end_height = ?subtree.end_height, index = ?subtree.index.0, "calculated and added orchard subtree");
907    }
908    // This log happens about once per second on recent machines with SSD disks.
909    debug!(end_height = ?subtree.end_height, index = ?subtree.index.0, "calculated and added orchard subtree");
910}