diff --git a/zebra-state/src/service.rs b/zebra-state/src/service.rs index f36cb488ec2..dce5e17e9ba 100644 --- a/zebra-state/src/service.rs +++ b/zebra-state/src/service.rs @@ -45,6 +45,7 @@ use zebra_chain::{ block::{self, CountedHeader, HeightDiff}, diagnostic::{task::WaitForPanics, CodeTimer}, parameters::{Network, NetworkUpgrade}, + subtree::NoteCommitmentSubtreeIndex, }; use crate::{ @@ -1507,14 +1508,29 @@ impl Service for ReadStateService { tokio::task::spawn_blocking(move || { span.in_scope(move || { + let end_index = limit + .and_then(|limit| start_index.0.checked_add(limit.0)) + .map(NoteCommitmentSubtreeIndex); + let sapling_subtrees = state.non_finalized_state_receiver.with_watch_data( |non_finalized_state| { - read::sapling_subtrees( - non_finalized_state.best_chain(), - &state.db, - start_index, - limit, - ) + if let Some(end_index) = end_index { + read::sapling_subtrees( + non_finalized_state.best_chain(), + &state.db, + start_index..end_index, + ) + } else { + // If there is no end bound, just return all the trees. + // If the end bound would overflow, just returns all the trees, because that's what + // `zcashd` does. (It never calculates an end bound, so it just keeps iterating until + // the trees run out.) + read::sapling_subtrees( + non_finalized_state.best_chain(), + &state.db, + start_index.., + ) + } }, ); @@ -1532,14 +1548,29 @@ impl Service for ReadStateService { tokio::task::spawn_blocking(move || { span.in_scope(move || { + let end_index = limit + .and_then(|limit| start_index.0.checked_add(limit.0)) + .map(NoteCommitmentSubtreeIndex); + let orchard_subtrees = state.non_finalized_state_receiver.with_watch_data( |non_finalized_state| { - read::orchard_subtrees( - non_finalized_state.best_chain(), - &state.db, - start_index, - limit, - ) + if let Some(end_index) = end_index { + read::orchard_subtrees( + non_finalized_state.best_chain(), + &state.db, + start_index..end_index, + ) + } else { + // If there is no end bound, just return all the trees. + // If the end bound would overflow, just returns all the trees, because that's what + // `zcashd` does. (It never calculates an end bound, so it just keeps iterating until + // the trees run out.) + read::orchard_subtrees( + non_finalized_state.best_chain(), + &state.db, + start_index.., + ) + } }, ); diff --git a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs index de7515858f1..260d202eac9 100644 --- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs +++ b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs @@ -245,64 +245,20 @@ impl ZebraDb { Some(subtree_data.with_index(index)) } - /// Returns a list of Sapling [`NoteCommitmentSubtree`]s starting at `start_index`. - /// If `limit` is provided, the list is limited to `limit` entries. - /// - /// If there is no subtree at `start_index`, the returned list is empty. - /// Otherwise, subtrees are continuous up to the finalized tip. - /// - /// # Correctness - /// - /// This method is specifically designed for the `z_getsubtreesbyindex` state request. - /// It might not work for other RPCs or state checks. - pub fn sapling_subtree_list_by_index_for_rpc( + /// Returns a list of Sapling [`NoteCommitmentSubtree`]s in the provided range. + #[allow(clippy::unwrap_in_result)] + pub fn sapling_subtree_list_by_index_range( &self, - start_index: NoteCommitmentSubtreeIndex, - limit: Option, + range: impl std::ops::RangeBounds, ) -> BTreeMap> { let sapling_subtrees = self .db .cf_handle("sapling_note_commitment_subtree") .unwrap(); - // Calculate the end bound, checking for overflow. - let exclusive_end_bound: Option = limit - .and_then(|limit| start_index.0.checked_add(limit.0)) - .map(NoteCommitmentSubtreeIndex); - - let list: BTreeMap< - NoteCommitmentSubtreeIndex, - NoteCommitmentSubtreeData, - >; - - if let Some(exclusive_end_bound) = exclusive_end_bound { - list = self - .db - .zs_range_iter(&sapling_subtrees, start_index..exclusive_end_bound, false) - .collect(); - } else { - // If there is no end bound, just return all the trees. - // If the end bound would overflow, just returns all the trees, because that's what - // `zcashd` does. (It never calculates an end bound, so it just keeps iterating until - // the trees run out.) - list = self - .db - .zs_range_iter(&sapling_subtrees, start_index.., false) - .collect(); - } - - // Make sure the amount of retrieved subtrees does not exceed the given limit. - #[cfg(debug_assertions)] - if let Some(limit) = limit { - assert!(list.len() <= limit.0.into()); - } - - // Check that we got the start subtree. - if list.get(&start_index).is_some() { - list - } else { - BTreeMap::new() - } + self.db + .zs_range_iter(&sapling_subtrees, range, false) + .collect() } /// Get the sapling note commitment subtress for the finalized tip. @@ -415,64 +371,20 @@ impl ZebraDb { Some(subtree_data.with_index(index)) } - /// Returns a list of Orchard [`NoteCommitmentSubtree`]s starting at `start_index`. - /// If `limit` is provided, the list is limited to `limit` entries. - /// - /// If there is no subtree at `start_index`, the returned list is empty. - /// Otherwise, subtrees are continuous up to the finalized tip. - /// - /// # Correctness - /// - /// This method is specifically designed for the `z_getsubtreesbyindex` state request. - /// It might not work for other RPCs or state checks. - pub fn orchard_subtree_list_by_index_for_rpc( + /// Returns a list of Orchard [`NoteCommitmentSubtree`]s in the provided range. + #[allow(clippy::unwrap_in_result)] + pub fn orchard_subtree_list_by_index_range( &self, - start_index: NoteCommitmentSubtreeIndex, - limit: Option, + range: impl std::ops::RangeBounds, ) -> BTreeMap> { let orchard_subtrees = self .db .cf_handle("orchard_note_commitment_subtree") .unwrap(); - // Calculate the end bound, checking for overflow. - let exclusive_end_bound: Option = limit - .and_then(|limit| start_index.0.checked_add(limit.0)) - .map(NoteCommitmentSubtreeIndex); - - let list: BTreeMap< - NoteCommitmentSubtreeIndex, - NoteCommitmentSubtreeData, - >; - - if let Some(exclusive_end_bound) = exclusive_end_bound { - list = self - .db - .zs_range_iter(&orchard_subtrees, start_index..exclusive_end_bound, false) - .collect(); - } else { - // If there is no end bound, just return all the trees. - // If the end bound would overflow, just returns all the trees, because that's what - // `zcashd` does. (It never calculates an end bound, so it just keeps iterating until - // the trees run out.) - list = self - .db - .zs_range_iter(&orchard_subtrees, start_index.., false) - .collect(); - } - - // Make sure the amount of retrieved subtrees does not exceed the given limit. - #[cfg(debug_assertions)] - if let Some(limit) = limit { - assert!(list.len() <= limit.0.into()); - } - - // Check that we got the start subtree. - if list.get(&start_index).is_some() { - list - } else { - BTreeMap::new() - } + self.db + .zs_range_iter(&orchard_subtrees, range, false) + .collect() } /// Get the orchard note commitment subtress for the finalized tip. diff --git a/zebra-state/src/service/non_finalized_state/chain.rs b/zebra-state/src/service/non_finalized_state/chain.rs index b415b41d1fa..e4ccf9c308b 100644 --- a/zebra-state/src/service/non_finalized_state/chain.rs +++ b/zebra-state/src/service/non_finalized_state/chain.rs @@ -698,8 +698,7 @@ impl Chain { .map(|(index, subtree)| subtree.with_index(*index)) } - /// Returns a list of Sapling [`NoteCommitmentSubtree`]s at or after `start_index`. - /// If `limit` is provided, the list is limited to `limit` entries. + /// Returns a list of Sapling [`NoteCommitmentSubtree`]s in the provided range. /// /// Unlike the finalized state and `ReadRequest::SaplingSubtrees`, the returned subtrees /// can start after `start_index`. These subtrees are continuous up to the tip. @@ -709,17 +708,10 @@ impl Chain { /// finalized updates. pub fn sapling_subtrees_in_range( &self, - start_index: NoteCommitmentSubtreeIndex, - limit: Option, + range: impl std::ops::RangeBounds, ) -> BTreeMap> { - let limit = limit - .map(|limit| usize::from(limit.0)) - .unwrap_or(usize::MAX); - - // Since we're working in memory, it's ok to iterate through the whole range here. self.sapling_subtrees - .range(start_index..) - .take(limit) + .range(range) .map(|(index, subtree)| (*index, *subtree)) .collect() } @@ -910,8 +902,7 @@ impl Chain { .map(|(index, subtree)| subtree.with_index(*index)) } - /// Returns a list of Orchard [`NoteCommitmentSubtree`]s at or after `start_index`. - /// If `limit` is provided, the list is limited to `limit` entries. + /// Returns a list of Orchard [`NoteCommitmentSubtree`]s in the provided range. /// /// Unlike the finalized state and `ReadRequest::OrchardSubtrees`, the returned subtrees /// can start after `start_index`. These subtrees are continuous up to the tip. @@ -921,17 +912,10 @@ impl Chain { /// finalized updates. pub fn orchard_subtrees_in_range( &self, - start_index: NoteCommitmentSubtreeIndex, - limit: Option, + range: impl std::ops::RangeBounds, ) -> BTreeMap> { - let limit = limit - .map(|limit| usize::from(limit.0)) - .unwrap_or(usize::MAX); - - // Since we're working in memory, it's ok to iterate through the whole range here. self.orchard_subtrees - .range(start_index..) - .take(limit) + .range(range) .map(|(index, subtree)| (*index, *subtree)) .collect() } diff --git a/zebra-state/src/service/read/tests/vectors.rs b/zebra-state/src/service/read/tests/vectors.rs index 2ac96578cc7..3ff14cb4203 100644 --- a/zebra-state/src/service/read/tests/vectors.rs +++ b/zebra-state/src/service/read/tests/vectors.rs @@ -100,6 +100,101 @@ async fn populated_read_state_responds_correctly() -> Result<()> { Ok(()) } +/// Tests if Zebra combines the note commitment subtrees from the finalized and +/// non-finalized states correctly. +#[tokio::test] +async fn test_read_subtrees() -> Result<()> { + use std::ops::Bound::*; + + let dummy_subtree = |(index, height)| { + NoteCommitmentSubtree::new( + u16::try_from(index).expect("should fit in u16"), + Height(height), + sapling::tree::Node::default(), + ) + }; + + let num_db_subtrees = 10; + let num_chain_subtrees = 2; + let index_offset = usize::try_from(num_db_subtrees).expect("constant should fit in usize"); + let db_height_range = 0..num_db_subtrees; + let chain_height_range = num_db_subtrees..(num_db_subtrees + num_chain_subtrees); + + // Prepare the finalized state. + let db = { + let db = ZebraDb::new(&Config::ephemeral(), Mainnet, true); + let db_subtrees = db_height_range.enumerate().map(dummy_subtree); + for db_subtree in db_subtrees { + let mut db_batch = DiskWriteBatch::new(); + db_batch.insert_sapling_subtree(&db, &db_subtree); + db.write(db_batch) + .expect("Writing a batch with a Sapling subtree should succeed."); + } + db + }; + + // Prepare the non-finalized state. + let chain = { + let mut chain = Chain::default(); + let chain_subtrees = chain_height_range + .enumerate() + .map(|(index, height)| dummy_subtree((index_offset + index, height))); + + for chain_subtree in chain_subtrees { + chain.insert_sapling_subtree(chain_subtree); + } + + Arc::new(chain) + }; + + let modify_chain = |chain: &Arc, index: usize, height| { + let mut chain = chain.as_ref().clone(); + chain.insert_sapling_subtree(dummy_subtree((index, height))); + Some(Arc::new(chain)) + }; + + // There should be 10 entries in db and 2 in chain with no overlap + + // Unbounded range should start at 0 + let all_subtrees = sapling_subtrees(Some(chain.clone()), &db, ..); + assert_eq!(all_subtrees.len(), 12, "should have 12 subtrees in state"); + + // Add a subtree to `chain` that overlaps and is not consistent with the db subtrees + let first_chain_index = index_offset - 1; + let end_height = Height(400_000); + let modified_chain = modify_chain(&chain, first_chain_index, end_height.0); + + // The inconsistent entry and any later entries should be omitted + let all_subtrees = sapling_subtrees(modified_chain.clone(), &db, ..); + assert_eq!(all_subtrees.len(), 10, "should have 10 subtrees in state"); + + let first_chain_index = + NoteCommitmentSubtreeIndex(u16::try_from(first_chain_index).expect("should fit in u16")); + + // Entries should be returned without reading from disk if the chain contains the first subtree index in the range + let mut chain_subtrees = sapling_subtrees(modified_chain, &db, first_chain_index..); + assert_eq!(chain_subtrees.len(), 3, "should have 3 subtrees in chain"); + + let (index, subtree) = chain_subtrees + .pop_first() + .expect("chain_subtrees should not be empty"); + assert_eq!(first_chain_index, index, "subtree indexes should match"); + assert_eq!(end_height, subtree.end, "subtree end heights should match"); + + // Check that Zebra retrieves subtrees correctly when using a range with an Excluded start bound + + let start = 0.into(); + let range = (Excluded(start), Unbounded); + let subtrees = sapling_subtrees(Some(chain), &db, range); + assert_eq!(subtrees.len(), 11); + assert!( + !subtrees.contains_key(&start), + "should not contain excluded start bound" + ); + + Ok(()) +} + /// Tests if Zebra combines the Sapling note commitment subtrees from the finalized and /// non-finalized states correctly. #[tokio::test] @@ -114,7 +209,7 @@ async fn test_sapling_subtrees() -> Result<()> { db.write(db_batch) .expect("Writing a batch with a Sapling subtree should succeed."); - // Prepare the non-fianlized state. + // Prepare the non-finalized state. let chain_subtree = NoteCommitmentSubtree::new(1, Height(3), dummy_subtree_root); let mut chain = Chain::default(); chain.insert_sapling_subtree(chain_subtree); @@ -124,40 +219,40 @@ async fn test_sapling_subtrees() -> Result<()> { // the non-finalized state. // Retrieve only the first subtree and check its properties. - let subtrees = sapling_subtrees(chain.clone(), &db, 0.into(), Some(1.into())); + let subtrees = sapling_subtrees(chain.clone(), &db, NoteCommitmentSubtreeIndex(0)..1.into()); let mut subtrees = subtrees.iter(); assert_eq!(subtrees.len(), 1); assert!(subtrees_eq(subtrees.next().unwrap(), &db_subtree)); // Retrieve both subtrees using a limit and check their properties. - let subtrees = sapling_subtrees(chain.clone(), &db, 0.into(), Some(2.into())); + let subtrees = sapling_subtrees(chain.clone(), &db, NoteCommitmentSubtreeIndex(0)..2.into()); let mut subtrees = subtrees.iter(); assert_eq!(subtrees.len(), 2); assert!(subtrees_eq(subtrees.next().unwrap(), &db_subtree)); assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree)); // Retrieve both subtrees without using a limit and check their properties. - let subtrees = sapling_subtrees(chain.clone(), &db, 0.into(), None); + let subtrees = sapling_subtrees(chain.clone(), &db, NoteCommitmentSubtreeIndex(0)..); let mut subtrees = subtrees.iter(); assert_eq!(subtrees.len(), 2); assert!(subtrees_eq(subtrees.next().unwrap(), &db_subtree)); assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree)); // Retrieve only the second subtree and check its properties. - let subtrees = sapling_subtrees(chain.clone(), &db, 1.into(), Some(1.into())); + let subtrees = sapling_subtrees(chain.clone(), &db, NoteCommitmentSubtreeIndex(1)..2.into()); let mut subtrees = subtrees.iter(); assert_eq!(subtrees.len(), 1); assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree)); // Retrieve only the second subtree, using a limit that would allow for more trees if they were // present, and check its properties. - let subtrees = sapling_subtrees(chain.clone(), &db, 1.into(), Some(2.into())); + let subtrees = sapling_subtrees(chain.clone(), &db, NoteCommitmentSubtreeIndex(1)..3.into()); let mut subtrees = subtrees.iter(); assert_eq!(subtrees.len(), 1); assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree)); // Retrieve only the second subtree, without using any limit, and check its properties. - let subtrees = sapling_subtrees(chain, &db, 1.into(), None); + let subtrees = sapling_subtrees(chain, &db, NoteCommitmentSubtreeIndex(1)..); let mut subtrees = subtrees.iter(); assert_eq!(subtrees.len(), 1); assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree)); @@ -179,7 +274,7 @@ async fn test_orchard_subtrees() -> Result<()> { db.write(db_batch) .expect("Writing a batch with an Orchard subtree should succeed."); - // Prepare the non-fianlized state. + // Prepare the non-finalized state. let chain_subtree = NoteCommitmentSubtree::new(1, Height(3), dummy_subtree_root); let mut chain = Chain::default(); chain.insert_orchard_subtree(chain_subtree); @@ -189,40 +284,40 @@ async fn test_orchard_subtrees() -> Result<()> { // the non-finalized state. // Retrieve only the first subtree and check its properties. - let subtrees = orchard_subtrees(chain.clone(), &db, 0.into(), Some(1.into())); + let subtrees = orchard_subtrees(chain.clone(), &db, NoteCommitmentSubtreeIndex(0)..1.into()); let mut subtrees = subtrees.iter(); assert_eq!(subtrees.len(), 1); assert!(subtrees_eq(subtrees.next().unwrap(), &db_subtree)); // Retrieve both subtrees using a limit and check their properties. - let subtrees = orchard_subtrees(chain.clone(), &db, 0.into(), Some(2.into())); + let subtrees = orchard_subtrees(chain.clone(), &db, NoteCommitmentSubtreeIndex(0)..2.into()); let mut subtrees = subtrees.iter(); assert_eq!(subtrees.len(), 2); assert!(subtrees_eq(subtrees.next().unwrap(), &db_subtree)); assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree)); // Retrieve both subtrees without using a limit and check their properties. - let subtrees = orchard_subtrees(chain.clone(), &db, 0.into(), None); + let subtrees = orchard_subtrees(chain.clone(), &db, NoteCommitmentSubtreeIndex(0)..); let mut subtrees = subtrees.iter(); assert_eq!(subtrees.len(), 2); assert!(subtrees_eq(subtrees.next().unwrap(), &db_subtree)); assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree)); // Retrieve only the second subtree and check its properties. - let subtrees = orchard_subtrees(chain.clone(), &db, 1.into(), Some(1.into())); + let subtrees = orchard_subtrees(chain.clone(), &db, NoteCommitmentSubtreeIndex(1)..2.into()); let mut subtrees = subtrees.iter(); assert_eq!(subtrees.len(), 1); assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree)); // Retrieve only the second subtree, using a limit that would allow for more trees if they were // present, and check its properties. - let subtrees = orchard_subtrees(chain.clone(), &db, 1.into(), Some(2.into())); + let subtrees = orchard_subtrees(chain.clone(), &db, NoteCommitmentSubtreeIndex(1)..3.into()); let mut subtrees = subtrees.iter(); assert_eq!(subtrees.len(), 1); assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree)); // Retrieve only the second subtree, without using any limit, and check its properties. - let subtrees = orchard_subtrees(chain, &db, 1.into(), None); + let subtrees = orchard_subtrees(chain, &db, NoteCommitmentSubtreeIndex(1)..); let mut subtrees = subtrees.iter(); assert_eq!(subtrees.len(), 1); assert!(subtrees_eq(subtrees.next().unwrap(), &chain_subtree)); diff --git a/zebra-state/src/service/read/tree.rs b/zebra-state/src/service/read/tree.rs index 98869546952..ec610a32987 100644 --- a/zebra-state/src/service/read/tree.rs +++ b/zebra-state/src/service/read/tree.rs @@ -48,99 +48,26 @@ where .or_else(|| db.sapling_tree_by_hash_or_height(hash_or_height)) } -/// Returns a list of Sapling [`NoteCommitmentSubtree`]s starting at `start_index`. +/// Returns a list of Sapling [`NoteCommitmentSubtree`]s with indexes in the provided range. /// -/// If `limit` is provided, the list is limited to `limit` entries. If there is no subtree at -/// `start_index` in the non-finalized `chain` or finalized `db`, the returned list is empty. +/// If there is no subtree at the first index in the range, the returned list is empty. +/// Otherwise, subtrees are continuous up to the finalized tip. /// -/// # Correctness -/// -/// 1. After `chain` was cloned, the StateService can commit additional blocks to the finalized -/// state `db`. Usually, the subtrees of these blocks are consistent. But if the `chain` is a -/// different fork to `db`, then the trees can be inconsistent. In that case, we ignore all the -/// trees in `chain` after the first inconsistent tree, because we know they will be inconsistent as -/// well. (It is cryptographically impossible for tree roots to be equal once the leaves have -/// diverged.) -/// 2. APIs that return single subtrees can't be used here, because they can create -/// an inconsistent list of subtrees after concurrent non-finalized and finalized updates. +/// See [`subtrees`] for more details. pub fn sapling_subtrees( chain: Option, db: &ZebraDb, - start_index: NoteCommitmentSubtreeIndex, - limit: Option, + range: impl std::ops::RangeBounds + Clone, ) -> BTreeMap> where C: AsRef, { - let mut db_list = db.sapling_subtree_list_by_index_for_rpc(start_index, limit); - - if let Some(limit) = limit { - let subtrees_num = u16::try_from(db_list.len()) - .expect("There can't be more than `u16::MAX` Sapling subtrees."); - - // Return the subtrees if the amount of them reached the given limit. - if subtrees_num == limit.0 { - return db_list; - } - - // If not, make sure the amount is below the limit. - debug_assert!(subtrees_num < limit.0); - } - - // If there's no chain, then we have the complete list. - let Some(chain) = chain else { - return db_list; - }; - - // Unlike the other methods, this returns any trees in the range, - // even if there is no tree for start_index. - let fork_list = chain.as_ref().sapling_subtrees_in_range(start_index, limit); - - // If there's no subtrees in chain, then we have the complete list. - if fork_list.is_empty() { - return db_list; - }; - - // Check for inconsistent trees in the fork. - for (fork_index, fork_subtree) in fork_list { - // If there's no matching index, just update the list of trees. - let Some(db_subtree) = db_list.get(&fork_index) else { - db_list.insert(fork_index, fork_subtree); - - // Stop adding new subtrees once their amount reaches the given limit. - if let Some(limit) = limit { - let subtrees_num = u16::try_from(db_list.len()) - .expect("There can't be more than `u16::MAX` Sapling subtrees."); - - if subtrees_num == limit.0 { - break; - } - } - - continue; - }; - - // We have an outdated chain fork, so skip this subtree and all remaining subtrees. - if &fork_subtree != db_subtree { - break; - } - - // Otherwise, the subtree is already in the list, so we don't need to add it. - } - - // Make sure the amount of retrieved subtrees does not exceed the given limit. - #[cfg(debug_assertions)] - if let Some(limit) = limit { - assert!(db_list.len() <= limit.0.into()); - } - - // Check that we got the start subtree from the non-finalized or finalized state. - // (The non-finalized state doesn't do this check.) - if db_list.get(&start_index).is_some() { - db_list - } else { - BTreeMap::new() - } + subtrees( + chain, + range, + |chain, range| chain.sapling_subtrees_in_range(range), + |range| db.sapling_subtree_list_by_index_range(range), + ) } /// Returns the Orchard @@ -164,96 +91,107 @@ where .or_else(|| db.orchard_tree_by_hash_or_height(hash_or_height)) } -/// Returns a list of Orchard [`NoteCommitmentSubtree`]s starting at `start_index`. +/// Returns a list of Orchard [`NoteCommitmentSubtree`]s with indexes in the provided range. /// -/// If `limit` is provided, the list is limited to `limit` entries. If there is no subtree at -/// `start_index` in the non-finalized `chain` or finalized `db`, the returned list is empty. +/// If there is no subtree at the first index in the range, the returned list is empty. +/// Otherwise, subtrees are continuous up to the finalized tip. /// -/// # Correctness -/// -/// 1. After `chain` was cloned, the StateService can commit additional blocks to the finalized -/// state `db`. Usually, the subtrees of these blocks are consistent. But if the `chain` is a -/// different fork to `db`, then the trees can be inconsistent. In that case, we ignore all the -/// trees in `chain` after the first inconsistent tree, because we know they will be inconsistent as -/// well. (It is cryptographically impossible for tree roots to be equal once the leaves have -/// diverged.) -/// 2. APIs that return single subtrees can't be used here, because they can create -/// an inconsistent list of subtrees after concurrent non-finalized and finalized updates. +/// See [`subtrees`] for more details. pub fn orchard_subtrees( chain: Option, db: &ZebraDb, - start_index: NoteCommitmentSubtreeIndex, - limit: Option, + range: impl std::ops::RangeBounds + Clone, ) -> BTreeMap> where C: AsRef, { - let mut db_list = db.orchard_subtree_list_by_index_for_rpc(start_index, limit); - - if let Some(limit) = limit { - let subtrees_num = u16::try_from(db_list.len()) - .expect("There can't be more than `u16::MAX` Orchard subtrees."); - - // Return the subtrees if the amount of them reached the given limit. - if subtrees_num == limit.0 { - return db_list; - } - - // If not, make sure the amount is below the limit. - debug_assert!(subtrees_num < limit.0); - } - - // If there's no chain, then we have the complete list. - let Some(chain) = chain else { - return db_list; - }; + subtrees( + chain, + range, + |chain, range| chain.orchard_subtrees_in_range(range), + |range| db.orchard_subtree_list_by_index_range(range), + ) +} - // Unlike the other methods, this returns any trees in the range, - // even if there is no tree for start_index. - let fork_list = chain.as_ref().orchard_subtrees_in_range(start_index, limit); +/// Returns a list of [`NoteCommitmentSubtree`]s in the provided range. +/// +/// If there is no subtree at the first index in the range, the returned list is empty. +/// Otherwise, subtrees are continuous up to the finalized tip. +/// +/// Accepts a `chain` from the non-finalized state, a `range` of subtree indexes to retrieve, +/// a `read_chain` function for retrieving the `range` of subtrees from `chain`, and +/// a `read_disk` function for retrieving the `range` from [`ZebraDb`]. +/// +/// Returns a consistent set of subtrees for the supplied chain fork and database. +/// Avoids reading the database if the subtrees are present in memory. +/// +/// # Correctness +/// +/// APIs that return single subtrees can't be used for `read_chain` and `read_disk`, because they +/// can create an inconsistent list of subtrees after concurrent non-finalized and finalized updates. +fn subtrees( + chain: Option, + range: Range, + read_chain: ChainSubtreeFn, + read_disk: DbSubtreeFn, +) -> BTreeMap> +where + C: AsRef, + Node: PartialEq, + Range: std::ops::RangeBounds + Clone, + ChainSubtreeFn: FnOnce( + &Chain, + Range, + ) + -> BTreeMap>, + DbSubtreeFn: + FnOnce(Range) -> BTreeMap>, +{ + use std::ops::Bound::*; - // If there's no subtrees in chain, then we have the complete list. - if fork_list.is_empty() { - return db_list; + let start_index = match range.start_bound().cloned() { + Included(start_index) => start_index, + Excluded(start_index) => (start_index.0 + 1).into(), + Unbounded => 0.into(), }; - // Check for inconsistent trees in the fork. - for (fork_index, fork_subtree) in fork_list { - // If there's no matching index, just update the list of trees. - let Some(db_subtree) = db_list.get(&fork_index) else { - db_list.insert(fork_index, fork_subtree); - - // Stop adding new subtrees once their amount reaches the given limit. - if let Some(limit) = limit { - let subtrees_num = u16::try_from(db_list.len()) - .expect("There can't be more than `u16::MAX` Orchard subtrees."); - - if subtrees_num == limit.0 { + // # Correctness + // + // After `chain` was cloned, the StateService can commit additional blocks to the finalized state `db`. + // Usually, the subtrees of these blocks are consistent. But if the `chain` is a different fork to `db`, + // then the trees can be inconsistent. In that case, if `chain` does not contain a subtree at the first + // index in the provided range, we ignore all the trees in `chain` after the first inconsistent tree, + // because we know they will be inconsistent as well. (It is cryptographically impossible for tree roots + // to be equal once the leaves have diverged.) + + let results = match chain.map(|chain| read_chain(chain.as_ref(), range.clone())) { + Some(chain_results) if chain_results.contains_key(&start_index) => return chain_results, + Some(chain_results) => { + let mut db_results = read_disk(range); + + // Check for inconsistent trees in the fork. + for (chain_index, chain_subtree) in chain_results { + // If there's no matching index, just update the list of trees. + let Some(db_subtree) = db_results.get(&chain_index) else { + db_results.insert(chain_index, chain_subtree); + continue; + }; + + // We have an outdated chain fork, so skip this subtree and all remaining subtrees. + if &chain_subtree != db_subtree { break; } + // Otherwise, the subtree is already in the list, so we don't need to add it. } - continue; - }; - - // We have an outdated chain fork, so skip this subtree and all remaining subtrees. - if &fork_subtree != db_subtree { - break; + db_results } + None => read_disk(range), + }; - // Otherwise, the subtree is already in the list, so we don't need to add it. - } - - // Make sure the amount of retrieved subtrees does not exceed the given limit. - #[cfg(debug_assertions)] - if let Some(limit) = limit { - assert!(db_list.len() <= limit.0.into()); - } - - // Check that we got the start subtree from the non-finalized or finalized state. - // (The non-finalized state doesn't do this check.) - if db_list.get(&start_index).is_some() { - db_list + // Check that we got the start subtree + if results.contains_key(&start_index) { + results } else { BTreeMap::new() }