cleanup(state): Refactor subtree read functions to accept ranges and check memory first

zebra-chain/src/subtree.rs

@@ -14,7 +14,7 @@ pub const TRACKED_SUBTREE_HEIGHT: u8 = 16;
 
 /// A note commitment subtree index, used to identify a subtree in a shielded pool.
 /// Also used to count subtrees.
-#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)]
+#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize, Default)]
 #[cfg_attr(any(test, feature = "proptest-impl"), derive(Arbitrary))]
 #[serde(transparent)]
 pub struct NoteCommitmentSubtreeIndex(pub u16);

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<ReadRequest> 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<ReadRequest> 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..,
+                                    )
+                                }
                             },
                         );
 

zebra-state/src/service/finalized_state/zebra_db/shielded.rs

@@ -245,64 +245,18 @@ 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<NoteCommitmentSubtreeIndex>,
+        range: impl std::ops::RangeBounds<NoteCommitmentSubtreeIndex>,
     ) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<sapling::tree::Node>> {
         let sapling_subtrees = self
             .db
             .cf_handle("sapling_note_commitment_subtree")
             .unwrap();
 
-        // Calculate the end bound, checking for overflow.
-        let exclusive_end_bound: Option<NoteCommitmentSubtreeIndex> = limit
-            .and_then(|limit| start_index.0.checked_add(limit.0))
-            .map(NoteCommitmentSubtreeIndex);
-
-        let list: BTreeMap<
-            NoteCommitmentSubtreeIndex,
-            NoteCommitmentSubtreeData<sapling::tree::Node>,
-        >;
-
-        if let Some(exclusive_end_bound) = exclusive_end_bound {
-            list = self
-                .db
-                .zs_range_iter(&sapling_subtrees, start_index..exclusive_end_bound)
-                .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..)
-                .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).collect()
     }
 
     /// Get the sapling note commitment subtress for the finalized tip.
@@ -415,64 +369,18 @@ 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<NoteCommitmentSubtreeIndex>,
+        range: impl std::ops::RangeBounds<NoteCommitmentSubtreeIndex>,
     ) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<orchard::tree::Node>> {
         let orchard_subtrees = self
             .db
             .cf_handle("orchard_note_commitment_subtree")
             .unwrap();
 
-        // Calculate the end bound, checking for overflow.
-        let exclusive_end_bound: Option<NoteCommitmentSubtreeIndex> = limit
-            .and_then(|limit| start_index.0.checked_add(limit.0))
-            .map(NoteCommitmentSubtreeIndex);
-
-        let list: BTreeMap<
-            NoteCommitmentSubtreeIndex,
-            NoteCommitmentSubtreeData<orchard::tree::Node>,
-        >;
-
-        if let Some(exclusive_end_bound) = exclusive_end_bound {
-            list = self
-                .db
-                .zs_range_iter(&orchard_subtrees, start_index..exclusive_end_bound)
-                .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..)
-                .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).collect()
     }
 
     /// Get the orchard note commitment subtress for the finalized tip.

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<NoteCommitmentSubtreeIndex>,
+        range: impl std::ops::RangeBounds<NoteCommitmentSubtreeIndex>,
     ) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<sapling::tree::Node>> {
-        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<NoteCommitmentSubtreeIndex>,
+        range: impl std::ops::RangeBounds<NoteCommitmentSubtreeIndex>,
     ) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<orchard::tree::Node>> {
-        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()
     }