zebra_state/service/read/
tree.rs

1//! Reading note commitment trees.
2//!
3//! In the functions in this module:
4//!
5//! The block write task commits blocks to the finalized state before updating
6//! `chain` with a cached copy of the best non-finalized chain from
7//! `NonFinalizedState.chain_set`. Then the block commit task can commit additional blocks to
8//! the finalized state after we've cloned the `chain`.
9//!
10//! This means that some blocks can be in both:
11//! - the cached [`Chain`], and
12//! - the shared finalized [`ZebraDb`] reference.
13
14use std::{collections::BTreeMap, sync::Arc};
15
16use zebra_chain::{
17    orchard, sapling,
18    subtree::{NoteCommitmentSubtreeData, NoteCommitmentSubtreeIndex},
19};
20
21use crate::{
22    service::{finalized_state::ZebraDb, non_finalized_state::Chain},
23    HashOrHeight,
24};
25
26// Doc-only items
27#[allow(unused_imports)]
28use zebra_chain::subtree::NoteCommitmentSubtree;
29
30/// Returns the Sapling
31/// [`NoteCommitmentTree`](sapling::tree::NoteCommitmentTree) specified by a
32/// hash or height, if it exists in the non-finalized `chain` or finalized `db`.
33pub fn sapling_tree<C>(
34    chain: Option<C>,
35    db: &ZebraDb,
36    hash_or_height: HashOrHeight,
37) -> Option<Arc<sapling::tree::NoteCommitmentTree>>
38where
39    C: AsRef<Chain>,
40{
41    // # Correctness
42    //
43    // Since sapling treestates are the same in the finalized and non-finalized
44    // state, we check the most efficient alternative first. (`chain` is always
45    // in memory, but `db` stores blocks on disk, with a memory cache.)
46    chain
47        .and_then(|chain| chain.as_ref().sapling_tree(hash_or_height))
48        .or_else(|| db.sapling_tree_by_hash_or_height(hash_or_height))
49}
50
51/// Returns a list of Sapling [`NoteCommitmentSubtree`]s with indexes in the provided range.
52///
53/// If there is no subtree at the first index in the range, the returned list is empty.
54/// Otherwise, subtrees are continuous up to the finalized tip.
55///
56/// See [`subtrees`] for more details.
57pub fn sapling_subtrees<C>(
58    chain: Option<C>,
59    db: &ZebraDb,
60    range: impl std::ops::RangeBounds<NoteCommitmentSubtreeIndex> + Clone,
61) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<sapling::tree::Node>>
62where
63    C: AsRef<Chain>,
64{
65    subtrees(
66        chain,
67        range,
68        |chain, range| chain.sapling_subtrees_in_range(range),
69        |range| db.sapling_subtree_list_by_index_range(range),
70    )
71}
72
73/// Returns the Orchard
74/// [`NoteCommitmentTree`](orchard::tree::NoteCommitmentTree) specified by a
75/// hash or height, if it exists in the non-finalized `chain` or finalized `db`.
76pub fn orchard_tree<C>(
77    chain: Option<C>,
78    db: &ZebraDb,
79    hash_or_height: HashOrHeight,
80) -> Option<Arc<orchard::tree::NoteCommitmentTree>>
81where
82    C: AsRef<Chain>,
83{
84    // # Correctness
85    //
86    // Since orchard treestates are the same in the finalized and non-finalized
87    // state, we check the most efficient alternative first. (`chain` is always
88    // in memory, but `db` stores blocks on disk, with a memory cache.)
89    chain
90        .and_then(|chain| chain.as_ref().orchard_tree(hash_or_height))
91        .or_else(|| db.orchard_tree_by_hash_or_height(hash_or_height))
92}
93
94/// Returns a list of Orchard [`NoteCommitmentSubtree`]s with indexes in the provided range.
95///
96/// If there is no subtree at the first index in the range, the returned list is empty.
97/// Otherwise, subtrees are continuous up to the finalized tip.
98///
99/// See [`subtrees`] for more details.
100pub fn orchard_subtrees<C>(
101    chain: Option<C>,
102    db: &ZebraDb,
103    range: impl std::ops::RangeBounds<NoteCommitmentSubtreeIndex> + Clone,
104) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<orchard::tree::Node>>
105where
106    C: AsRef<Chain>,
107{
108    subtrees(
109        chain,
110        range,
111        |chain, range| chain.orchard_subtrees_in_range(range),
112        |range| db.orchard_subtree_list_by_index_range(range),
113    )
114}
115
116/// Returns a list of [`NoteCommitmentSubtree`]s in the provided range.
117///
118/// If there is no subtree at the first index in the range, the returned list is empty.
119/// Otherwise, subtrees are continuous up to the finalized tip.
120///
121/// Accepts a `chain` from the non-finalized state, a `range` of subtree indexes to retrieve,
122/// a `read_chain` function for retrieving the `range` of subtrees from `chain`, and
123/// a `read_disk` function for retrieving the `range` from [`ZebraDb`].
124///
125/// Returns a consistent set of subtrees for the supplied chain fork and database.
126/// Avoids reading the database if the subtrees are present in memory.
127///
128/// # Correctness
129///
130/// APIs that return single subtrees can't be used for `read_chain` and `read_disk`, because they
131/// can create an inconsistent list of subtrees after concurrent non-finalized and finalized updates.
132fn subtrees<C, Range, Node, ChainSubtreeFn, DbSubtreeFn>(
133    chain: Option<C>,
134    range: Range,
135    read_chain: ChainSubtreeFn,
136    read_disk: DbSubtreeFn,
137) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<Node>>
138where
139    C: AsRef<Chain>,
140    Node: PartialEq,
141    Range: std::ops::RangeBounds<NoteCommitmentSubtreeIndex> + Clone,
142    ChainSubtreeFn: FnOnce(
143        &Chain,
144        Range,
145    )
146        -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<Node>>,
147    DbSubtreeFn:
148        FnOnce(Range) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<Node>>,
149{
150    use std::ops::Bound::*;
151
152    let start_index = match range.start_bound().cloned() {
153        Included(start_index) => start_index,
154        Excluded(start_index) => (start_index.0 + 1).into(),
155        Unbounded => 0.into(),
156    };
157
158    // # Correctness
159    //
160    // After `chain` was cloned, the StateService can commit additional blocks to the finalized state `db`.
161    // Usually, the subtrees of these blocks are consistent. But if the `chain` is a different fork to `db`,
162    // then the trees can be inconsistent. In that case, if `chain` does not contain a subtree at the first
163    // index in the provided range, we ignore all the trees in `chain` after the first inconsistent tree,
164    // because we know they will be inconsistent as well. (It is cryptographically impossible for tree roots
165    // to be equal once the leaves have diverged.)
166
167    let results = match chain.map(|chain| read_chain(chain.as_ref(), range.clone())) {
168        Some(chain_results) if chain_results.contains_key(&start_index) => return chain_results,
169        Some(chain_results) => {
170            let mut db_results = read_disk(range);
171
172            // Check for inconsistent trees in the fork.
173            for (chain_index, chain_subtree) in chain_results {
174                // If there's no matching index, just update the list of trees.
175                let Some(db_subtree) = db_results.get(&chain_index) else {
176                    db_results.insert(chain_index, chain_subtree);
177                    continue;
178                };
179
180                // We have an outdated chain fork, so skip this subtree and all remaining subtrees.
181                if &chain_subtree != db_subtree {
182                    break;
183                }
184                // Otherwise, the subtree is already in the list, so we don't need to add it.
185            }
186
187            db_results
188        }
189        None => read_disk(range),
190    };
191
192    // Check that we got the start subtree
193    if results.contains_key(&start_index) {
194        results
195    } else {
196        BTreeMap::new()
197    }
198}
199
200/// Get the history tree of the provided chain.
201pub fn history_tree<C>(
202    chain: Option<C>,
203    db: &ZebraDb,
204    hash_or_height: HashOrHeight,
205) -> Option<Arc<zebra_chain::history_tree::HistoryTree>>
206where
207    C: AsRef<Chain>,
208{
209    chain
210        .and_then(|chain| chain.as_ref().history_tree(hash_or_height))
211        .or_else(|| Some(db.history_tree()))
212}