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}