From 837061f2f9b3e3cdbad9f3362e5644895c14eb77 Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Thu, 12 Oct 2023 19:01:14 -0400
Subject: [PATCH 01/13] refactors sapling_subtrees to use ranges

---
 zebra-state/src/service.rs                    | 28 +++++--
 .../finalized_state/zebra_db/shielded.rs      | 45 +----------
 .../src/service/non_finalized_state/chain.rs  | 14 +---
 zebra-state/src/service/read/tree.rs          | 78 ++++---------------
 4 files changed, 43 insertions(+), 122 deletions(-)

diff --git a/zebra-state/src/service.rs b/zebra-state/src/service.rs
index f36cb488ec2..e30042c4a6a 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<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..,
+                                    )
+                                }
                             },
                         );
 
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 58d9e43a6c1..d45b7c9be58 100644
--- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
@@ -248,8 +248,7 @@ 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.
+    /// Returns a list of Sapling [`NoteCommitmentSubtree`]s in the provided range.
     ///
     /// If there is no subtree at `start_index`, the returned list is empty.
     /// Otherwise, subtrees are continuous up to the finalized tip.
@@ -261,52 +260,14 @@ impl ZebraDb {
     #[allow(clippy::unwrap_in_result)]
     pub fn sapling_subtree_list_by_index_for_rpc(
         &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.
diff --git a/zebra-state/src/service/non_finalized_state/chain.rs b/zebra-state/src/service/non_finalized_state/chain.rs
index b415b41d1fa..d2e3819573c 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<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()
     }
diff --git a/zebra-state/src/service/read/tree.rs b/zebra-state/src/service/read/tree.rs
index 98869546952..10cb22f3323 100644
--- a/zebra-state/src/service/read/tree.rs
+++ b/zebra-state/src/service/read/tree.rs
@@ -66,78 +66,30 @@ where
 pub fn sapling_subtrees<C>(
     chain: Option<C>,
     db: &ZebraDb,
-    start_index: NoteCommitmentSubtreeIndex,
-    limit: Option<NoteCommitmentSubtreeIndex>,
+    range: impl std::ops::RangeBounds<NoteCommitmentSubtreeIndex> + Clone,
 ) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<sapling::tree::Node>>
 where
     C: AsRef<Chain>,
 {
-    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;
-    };
+    use std::ops::Bound::*;
 
-    // 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;
+    let start_index = match range.start_bound().cloned() {
+        Included(start_index) | Excluded(start_index) => start_index,
+        Unbounded => panic!("range provided to sapling_subtrees() must have start bound"),
     };
 
-    // 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;
+    let results = match chain.map(|chain| chain.as_ref().sapling_subtrees_in_range(range.clone())) {
+        Some(results) if results.contains_key(&start_index) => results,
+        Some(mut results) => {
+            results.append(&mut db.sapling_subtree_list_by_index_for_rpc(range));
+            results
         }
+        None => db.sapling_subtree_list_by_index_for_rpc(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()
     }

From a7edaa97029821780598129c16eef0e773357248 Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Thu, 12 Oct 2023 19:08:57 -0400
Subject: [PATCH 02/13] refactor orchard_subtrees() to use ranges

---
 zebra-state/src/service.rs                    | 27 +++++--
 .../finalized_state/zebra_db/shielded.rs      | 45 +----------
 .../src/service/non_finalized_state/chain.rs  | 14 +---
 zebra-state/src/service/read/tree.rs          | 78 ++++---------------
 4 files changed, 42 insertions(+), 122 deletions(-)

diff --git a/zebra-state/src/service.rs b/zebra-state/src/service.rs
index e30042c4a6a..dce5e17e9ba 100644
--- a/zebra-state/src/service.rs
+++ b/zebra-state/src/service.rs
@@ -1548,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..,
+                                    )
+                                }
                             },
                         );
 
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 d45b7c9be58..26552d20296 100644
--- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
@@ -382,8 +382,7 @@ 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.
+    /// Returns a list of Orchard [`NoteCommitmentSubtree`]s in the provided range.
     ///
     /// If there is no subtree at `start_index`, the returned list is empty.
     /// Otherwise, subtrees are continuous up to the finalized tip.
@@ -395,52 +394,14 @@ impl ZebraDb {
     #[allow(clippy::unwrap_in_result)]
     pub fn orchard_subtree_list_by_index_for_rpc(
         &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.
diff --git a/zebra-state/src/service/non_finalized_state/chain.rs b/zebra-state/src/service/non_finalized_state/chain.rs
index d2e3819573c..e4ccf9c308b 100644
--- a/zebra-state/src/service/non_finalized_state/chain.rs
+++ b/zebra-state/src/service/non_finalized_state/chain.rs
@@ -902,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.
@@ -913,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()
     }
diff --git a/zebra-state/src/service/read/tree.rs b/zebra-state/src/service/read/tree.rs
index 10cb22f3323..d8015dd231b 100644
--- a/zebra-state/src/service/read/tree.rs
+++ b/zebra-state/src/service/read/tree.rs
@@ -134,78 +134,30 @@ where
 pub fn orchard_subtrees<C>(
     chain: Option<C>,
     db: &ZebraDb,
-    start_index: NoteCommitmentSubtreeIndex,
-    limit: Option<NoteCommitmentSubtreeIndex>,
+    range: impl std::ops::RangeBounds<NoteCommitmentSubtreeIndex> + Clone,
 ) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<orchard::tree::Node>>
 where
     C: AsRef<Chain>,
 {
-    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;
-    };
-
-    // 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);
+    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) | Excluded(start_index) => start_index,
+        Unbounded => panic!("range provided to orchard_subtrees() must have start bound"),
     };
 
-    // 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 {
-                    break;
-                }
-            }
-
-            continue;
-        };
-
-        // We have an outdated chain fork, so skip this subtree and all remaining subtrees.
-        if &fork_subtree != db_subtree {
-            break;
+    let results = match chain.map(|chain| chain.as_ref().orchard_subtrees_in_range(range.clone())) {
+        Some(results) if results.contains_key(&start_index) => results,
+        Some(mut results) => {
+            results.append(&mut db.orchard_subtree_list_by_index_for_rpc(range));
+            results
         }
+        None => db.orchard_subtree_list_by_index_for_rpc(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()
     }

From f1aaf2b5108a91d29066fe38d0e382cb306fcda7 Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Thu, 12 Oct 2023 20:55:53 -0400
Subject: [PATCH 03/13] Restores check for inconsistent subtree roots

---
 zebra-state/src/service/read/tree.rs | 48 +++++++++++++++++++++++-----
 1 file changed, 40 insertions(+), 8 deletions(-)

diff --git a/zebra-state/src/service/read/tree.rs b/zebra-state/src/service/read/tree.rs
index d8015dd231b..d9985eb4185 100644
--- a/zebra-state/src/service/read/tree.rs
+++ b/zebra-state/src/service/read/tree.rs
@@ -79,10 +79,26 @@ where
     };
 
     let results = match chain.map(|chain| chain.as_ref().sapling_subtrees_in_range(range.clone())) {
-        Some(results) if results.contains_key(&start_index) => results,
-        Some(mut results) => {
-            results.append(&mut db.sapling_subtree_list_by_index_for_rpc(range));
-            results
+        Some(chain_results) if chain_results.contains_key(&start_index) => return chain_results,
+        Some(chain_results) => {
+            let mut db_results = db.sapling_subtree_list_by_index_for_rpc(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.
+            }
+
+            db_results
         }
         None => db.sapling_subtree_list_by_index_for_rpc(range),
     };
@@ -147,10 +163,26 @@ where
     };
 
     let results = match chain.map(|chain| chain.as_ref().orchard_subtrees_in_range(range.clone())) {
-        Some(results) if results.contains_key(&start_index) => results,
-        Some(mut results) => {
-            results.append(&mut db.orchard_subtree_list_by_index_for_rpc(range));
-            results
+        Some(chain_results) if chain_results.contains_key(&start_index) => return chain_results,
+        Some(chain_results) => {
+            let mut db_results = db.orchard_subtree_list_by_index_for_rpc(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.
+            }
+
+            db_results
         }
         None => db.orchard_subtree_list_by_index_for_rpc(range),
     };

From 5b55813901227e1613a87a9c0ea9f1a07454896a Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Thu, 12 Oct 2023 23:30:09 -0400
Subject: [PATCH 04/13] updates vectors tests

---
 zebra-state/src/service/read/tests/vectors.rs | 24 +++++++++----------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/zebra-state/src/service/read/tests/vectors.rs b/zebra-state/src/service/read/tests/vectors.rs
index 2ac96578cc7..cfe9274b2a7 100644
--- a/zebra-state/src/service/read/tests/vectors.rs
+++ b/zebra-state/src/service/read/tests/vectors.rs
@@ -124,40 +124,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));
@@ -189,40 +189,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));

From ffc4f7557415cc2a6ede67031d5cf3a6ec75382d Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Tue, 17 Oct 2023 21:52:10 -0400
Subject: [PATCH 05/13] Applies some suggestions from code review and adds
 `read::tree::subtrees` function

---
 zebra-state/src/service/read/tree.rs | 77 ++++++++++++----------------
 1 file changed, 34 insertions(+), 43 deletions(-)

diff --git a/zebra-state/src/service/read/tree.rs b/zebra-state/src/service/read/tree.rs
index d9985eb4185..f1355351af2 100644
--- a/zebra-state/src/service/read/tree.rs
+++ b/zebra-state/src/service/read/tree.rs
@@ -71,44 +71,12 @@ pub fn sapling_subtrees<C>(
 where
     C: AsRef<Chain>,
 {
-    use std::ops::Bound::*;
-
-    let start_index = match range.start_bound().cloned() {
-        Included(start_index) | Excluded(start_index) => start_index,
-        Unbounded => panic!("range provided to sapling_subtrees() must have start bound"),
-    };
-
-    let results = match chain.map(|chain| chain.as_ref().sapling_subtrees_in_range(range.clone())) {
-        Some(chain_results) if chain_results.contains_key(&start_index) => return chain_results,
-        Some(chain_results) => {
-            let mut db_results = db.sapling_subtree_list_by_index_for_rpc(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.
-            }
-
-            db_results
-        }
-        None => db.sapling_subtree_list_by_index_for_rpc(range),
-    };
-
-    // Check that we got the start subtree
-    if results.contains_key(&start_index) {
-        results
-    } else {
-        BTreeMap::new()
-    }
+    subtrees(
+        chain,
+        range,
+        |chain, range| chain.sapling_subtrees_in_range(range),
+        |range| db.sapling_subtree_list_by_index_for_rpc(range),
+    )
 }
 
 /// Returns the Orchard
@@ -154,18 +122,41 @@ pub fn orchard_subtrees<C>(
 ) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<orchard::tree::Node>>
 where
     C: AsRef<Chain>,
+{
+    subtrees(
+        chain,
+        range,
+        |chain, range| chain.orchard_subtrees_in_range(range),
+        |range| db.orchard_subtree_list_by_index_for_rpc(range),
+    )
+}
+
+/// Returns a list of [`NoteCommitmentSubtree`]s starting at `start_index`.
+fn subtrees<C, R, N, F1, F2>(
+    chain: Option<C>,
+    range: R,
+    read_chain: F1,
+    read_disk: F2,
+) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<N>>
+where
+    C: AsRef<Chain>,
+    N: PartialEq,
+    R: std::ops::RangeBounds<NoteCommitmentSubtreeIndex> + Clone,
+    F1: FnOnce(&Chain, R) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<N>>,
+    F2: FnOnce(R) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<N>>,
 {
     use std::ops::Bound::*;
 
     let start_index = match range.start_bound().cloned() {
-        Included(start_index) | Excluded(start_index) => start_index,
-        Unbounded => panic!("range provided to orchard_subtrees() must have start bound"),
+        Included(start_index) => start_index,
+        Excluded(start_index) => (start_index.0 + 1).into(),
+        Unbounded => 0.into(),
     };
 
-    let results = match chain.map(|chain| chain.as_ref().orchard_subtrees_in_range(range.clone())) {
+    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 = db.orchard_subtree_list_by_index_for_rpc(range);
+            let mut db_results = read_disk(range);
 
             // Check for inconsistent trees in the fork.
             for (chain_index, chain_subtree) in chain_results {
@@ -184,7 +175,7 @@ where
 
             db_results
         }
-        None => db.orchard_subtree_list_by_index_for_rpc(range),
+        None => read_disk(range),
     };
 
     // Check that we got the start subtree

From 0d7b13b089c7b43069956293d665db6d79534e54 Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Wed, 18 Oct 2023 16:23:06 -0400
Subject: [PATCH 06/13] updates correctness comment & db read method names

---
 .../finalized_state/zebra_db/shielded.rs      | 20 +------
 zebra-state/src/service/read/tree.rs          | 53 +++++++++----------
 2 files changed, 26 insertions(+), 47 deletions(-)

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 26552d20296..27d445fae15 100644
--- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
@@ -249,16 +249,8 @@ impl ZebraDb {
     }
 
     /// Returns a list of Sapling [`NoteCommitmentSubtree`]s in the provided range.
-    ///
-    /// 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.
     #[allow(clippy::unwrap_in_result)]
-    pub fn sapling_subtree_list_by_index_for_rpc(
+    pub fn sapling_subtree_list_by_index_range(
         &self,
         range: impl std::ops::RangeBounds<NoteCommitmentSubtreeIndex>,
     ) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<sapling::tree::Node>> {
@@ -383,16 +375,8 @@ impl ZebraDb {
     }
 
     /// Returns a list of Orchard [`NoteCommitmentSubtree`]s in the provided range.
-    ///
-    /// 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.
     #[allow(clippy::unwrap_in_result)]
-    pub fn orchard_subtree_list_by_index_for_rpc(
+    pub fn orchard_subtree_list_by_index_range(
         &self,
         range: impl std::ops::RangeBounds<NoteCommitmentSubtreeIndex>,
     ) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<orchard::tree::Node>> {
diff --git a/zebra-state/src/service/read/tree.rs b/zebra-state/src/service/read/tree.rs
index f1355351af2..efc9684e551 100644
--- a/zebra-state/src/service/read/tree.rs
+++ b/zebra-state/src/service/read/tree.rs
@@ -50,19 +50,8 @@ where
 
 /// 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` in the non-finalized `chain` or finalized `db`, the returned list is empty.
-///
-/// # 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.
+/// 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.
 pub fn sapling_subtrees<C>(
     chain: Option<C>,
     db: &ZebraDb,
@@ -75,7 +64,7 @@ where
         chain,
         range,
         |chain, range| chain.sapling_subtrees_in_range(range),
-        |range| db.sapling_subtree_list_by_index_for_rpc(range),
+        |range| db.sapling_subtree_list_by_index_range(range),
     )
 }
 
@@ -100,21 +89,10 @@ 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.
-///
-/// # 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.
+/// 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.
 pub fn orchard_subtrees<C>(
     chain: Option<C>,
     db: &ZebraDb,
@@ -127,11 +105,19 @@ where
         chain,
         range,
         |chain, range| chain.orchard_subtrees_in_range(range),
-        |range| db.orchard_subtree_list_by_index_for_rpc(range),
+        |range| db.orchard_subtree_list_by_index_range(range),
     )
 }
 
 /// Returns a list of [`NoteCommitmentSubtree`]s starting at `start_index`.
+///
+/// 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
+///
+/// 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.
 fn subtrees<C, R, N, F1, F2>(
     chain: Option<C>,
     range: R,
@@ -153,6 +139,15 @@ where
         Unbounded => 0.into(),
     };
 
+    // # 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) => {

From 8fe86fe9ccc1e581e2d6147779a95cd0c1e46757 Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Wed, 18 Oct 2023 20:04:35 -0400
Subject: [PATCH 07/13] adds test_subtrees

---
 zebra-state/src/service/read/tests/vectors.rs | 69 ++++++++++++++++++-
 1 file changed, 67 insertions(+), 2 deletions(-)

diff --git a/zebra-state/src/service/read/tests/vectors.rs b/zebra-state/src/service/read/tests/vectors.rs
index cfe9274b2a7..8d3b606b7f9 100644
--- a/zebra-state/src/service/read/tests/vectors.rs
+++ b/zebra-state/src/service/read/tests/vectors.rs
@@ -100,6 +100,71 @@ 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_subtrees() -> Result<()> {
+    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<Chain>, 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
+    // The inconsistent entry and any later entries should be omitted
+    let modified_chain = modify_chain(&chain, index_offset - 1, 400_000);
+    let all_subtrees = sapling_subtrees(modified_chain, &db, ..);
+    assert_eq!(all_subtrees.len(), 10, "should have 10 subtrees in state");
+
+    Ok(())
+}
+
 /// Tests if Zebra combines the Sapling note commitment subtrees from the finalized and
 /// non-finalized states correctly.
 #[tokio::test]
@@ -114,7 +179,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);
@@ -179,7 +244,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);

From 4ae71e57d6e4eaa4497f38309d58afbe91a55902 Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Wed, 18 Oct 2023 20:26:23 -0400
Subject: [PATCH 08/13] tests in-memory reads

---
 zebra-state/src/service/read/tests/vectors.rs | 23 ++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/zebra-state/src/service/read/tests/vectors.rs b/zebra-state/src/service/read/tests/vectors.rs
index 8d3b606b7f9..0c59782f311 100644
--- a/zebra-state/src/service/read/tests/vectors.rs
+++ b/zebra-state/src/service/read/tests/vectors.rs
@@ -103,7 +103,7 @@ async fn populated_read_state_responds_correctly() -> Result<()> {
 /// Tests if Zebra combines the note commitment subtrees from the finalized and
 /// non-finalized states correctly.
 #[tokio::test]
-async fn test_subtrees() -> Result<()> {
+async fn test_read_subtrees() -> Result<()> {
     let dummy_subtree = |(index, height)| {
         NoteCommitmentSubtree::new(
             u16::try_from(index).expect("should fit in u16"),
@@ -152,16 +152,33 @@ async fn test_subtrees() -> Result<()> {
     };
 
     // 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 modified_chain = modify_chain(&chain, index_offset - 1, 400_000);
-    let all_subtrees = sapling_subtrees(modified_chain, &db, ..);
+    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");
+
     Ok(())
 }
 

From 7ec6f217cea89db8649327cd7c2cb15a07ae618e Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Wed, 18 Oct 2023 20:44:19 -0400
Subject: [PATCH 09/13] test that subtree read fns work right with excluded
 start bound

---
 zebra-chain/src/subtree.rs                    |  2 +-
 zebra-state/src/service/read/tests/vectors.rs | 27 +++++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/zebra-chain/src/subtree.rs b/zebra-chain/src/subtree.rs
index 52e3ef3d7d0..b4a74343379 100644
--- a/zebra-chain/src/subtree.rs
+++ b/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);
diff --git a/zebra-state/src/service/read/tests/vectors.rs b/zebra-state/src/service/read/tests/vectors.rs
index 0c59782f311..7236b70ab66 100644
--- a/zebra-state/src/service/read/tests/vectors.rs
+++ b/zebra-state/src/service/read/tests/vectors.rs
@@ -104,6 +104,8 @@ async fn populated_read_state_responds_correctly() -> Result<()> {
 /// non-finalized states correctly.
 #[tokio::test]
 async fn test_read_subtrees() -> Result<()> {
+    use std::ops::Bound::{self, *};
+
     let dummy_subtree = |(index, height)| {
         NoteCommitmentSubtree::new(
             u16::try_from(index).expect("should fit in u16"),
@@ -179,6 +181,31 @@ async fn test_read_subtrees() -> Result<()> {
     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
+
+    #[derive(Clone, Default, PartialEq, Eq, Hash)]
+    struct RangeExclusiveStart<Idx> {
+        pub start: Idx,
+        pub end: Option<Idx>,
+    }
+
+    impl<T> std::ops::RangeBounds<T> for RangeExclusiveStart<T> {
+        fn start_bound(&self) -> Bound<&T> {
+            Excluded(&self.start)
+        }
+        fn end_bound(&self) -> Bound<&T> {
+            self.end.as_ref().map(Excluded).unwrap_or(Unbounded)
+        }
+    }
+
+    let range = RangeExclusiveStart::default();
+    let subtrees = sapling_subtrees(Some(chain), &db, range.clone());
+    assert_eq!(subtrees.len(), 11);
+    assert!(
+        !subtrees.contains_key(&range.start),
+        "should not contain excluded start bound"
+    );
+
     Ok(())
 }
 

From 9737d22a7c63365dd721641c9958472b99c95b4a Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Thu, 19 Oct 2023 16:45:41 -0400
Subject: [PATCH 10/13] applies suggestions from code review

---
 zebra-state/src/service/read/tests/vectors.rs | 24 ++++---------------
 1 file changed, 5 insertions(+), 19 deletions(-)

diff --git a/zebra-state/src/service/read/tests/vectors.rs b/zebra-state/src/service/read/tests/vectors.rs
index 7236b70ab66..3ff14cb4203 100644
--- a/zebra-state/src/service/read/tests/vectors.rs
+++ b/zebra-state/src/service/read/tests/vectors.rs
@@ -104,7 +104,7 @@ async fn populated_read_state_responds_correctly() -> Result<()> {
 /// non-finalized states correctly.
 #[tokio::test]
 async fn test_read_subtrees() -> Result<()> {
-    use std::ops::Bound::{self, *};
+    use std::ops::Bound::*;
 
     let dummy_subtree = |(index, height)| {
         NoteCommitmentSubtree::new(
@@ -183,26 +183,12 @@ async fn test_read_subtrees() -> Result<()> {
 
     // Check that Zebra retrieves subtrees correctly when using a range with an Excluded start bound
 
-    #[derive(Clone, Default, PartialEq, Eq, Hash)]
-    struct RangeExclusiveStart<Idx> {
-        pub start: Idx,
-        pub end: Option<Idx>,
-    }
-
-    impl<T> std::ops::RangeBounds<T> for RangeExclusiveStart<T> {
-        fn start_bound(&self) -> Bound<&T> {
-            Excluded(&self.start)
-        }
-        fn end_bound(&self) -> Bound<&T> {
-            self.end.as_ref().map(Excluded).unwrap_or(Unbounded)
-        }
-    }
-
-    let range = RangeExclusiveStart::default();
-    let subtrees = sapling_subtrees(Some(chain), &db, range.clone());
+    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(&range.start),
+        !subtrees.contains_key(&start),
         "should not contain excluded start bound"
     );
 

From 668462d586cfe5a409c56b189478c374296e6b2d Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Thu, 19 Oct 2023 20:01:58 -0400
Subject: [PATCH 11/13] Applies suggestions from code review.

---
 zebra-chain/src/subtree.rs           |  2 +-
 zebra-state/src/service/read/tree.rs | 41 ++++++++++++++++++++--------
 2 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/zebra-chain/src/subtree.rs b/zebra-chain/src/subtree.rs
index b4a74343379..52e3ef3d7d0 100644
--- a/zebra-chain/src/subtree.rs
+++ b/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, Default)]
+#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)]
 #[cfg_attr(any(test, feature = "proptest-impl"), derive(Arbitrary))]
 #[serde(transparent)]
 pub struct NoteCommitmentSubtreeIndex(pub u16);
diff --git a/zebra-state/src/service/read/tree.rs b/zebra-state/src/service/read/tree.rs
index efc9684e551..a866a0194a5 100644
--- a/zebra-state/src/service/read/tree.rs
+++ b/zebra-state/src/service/read/tree.rs
@@ -52,6 +52,8 @@ where
 ///
 /// 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.
+///
+/// See [`subtrees`] for more details.
 pub fn sapling_subtrees<C>(
     chain: Option<C>,
     db: &ZebraDb,
@@ -93,6 +95,8 @@ where
 ///
 /// 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.
+///
+/// See [`subtrees`] for more details.
 pub fn orchard_subtrees<C>(
     chain: Option<C>,
     db: &ZebraDb,
@@ -109,27 +113,40 @@ where
     )
 }
 
-/// Returns a list of [`NoteCommitmentSubtree`]s starting at `start_index`.
+/// 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 without calling `read_disk` if `read_chain` returns a subtree at the first subtree index
+/// in the provided `range`. Otherwise, if `read_chain` and `read_disk` return different subtrees for
+/// the same subtree index, ignores all the trees in `chain` after the first inconsistent tree.
+///
 /// # Correctness
 ///
-/// 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.
-fn subtrees<C, R, N, F1, F2>(
+/// 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<C, Range, Node, ChainSubtreeFn, DbSubtreeFn>(
     chain: Option<C>,
-    range: R,
-    read_chain: F1,
-    read_disk: F2,
-) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<N>>
+    range: Range,
+    read_chain: ChainSubtreeFn,
+    read_disk: DbSubtreeFn,
+) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<Node>>
 where
     C: AsRef<Chain>,
-    N: PartialEq,
-    R: std::ops::RangeBounds<NoteCommitmentSubtreeIndex> + Clone,
-    F1: FnOnce(&Chain, R) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<N>>,
-    F2: FnOnce(R) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<N>>,
+    Node: PartialEq,
+    Range: std::ops::RangeBounds<NoteCommitmentSubtreeIndex> + Clone,
+    ChainSubtreeFn: FnOnce(
+        &Chain,
+        Range,
+    )
+        -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<Node>>,
+    DbSubtreeFn:
+        FnOnce(Range) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<Node>>,
 {
     use std::ops::Bound::*;
 

From 3416c11264c6c7f4d0cfb3a47429b73f7dde3600 Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Thu, 19 Oct 2023 20:55:18 -0400
Subject: [PATCH 12/13] Updates docs applying suggestion from code review

---
 zebra-state/src/service/read/tree.rs | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/zebra-state/src/service/read/tree.rs b/zebra-state/src/service/read/tree.rs
index a866a0194a5..ec610a32987 100644
--- a/zebra-state/src/service/read/tree.rs
+++ b/zebra-state/src/service/read/tree.rs
@@ -48,7 +48,7 @@ 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 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.
@@ -122,9 +122,8 @@ where
 /// a `read_chain` function for retrieving the `range` of subtrees from `chain`, and
 /// a `read_disk` function for retrieving the `range` from [`ZebraDb`].
 ///
-/// Returns without calling `read_disk` if `read_chain` returns a subtree at the first subtree index
-/// in the provided `range`. Otherwise, if `read_chain` and `read_disk` return different subtrees for
-/// the same subtree index, ignores all the trees in `chain` after the first inconsistent tree.
+/// 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
 ///

From d3963698ee6e2804d6102dde6789b36eab95adbc Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Thu, 19 Oct 2023 22:09:12 -0400
Subject: [PATCH 13/13] adds new arg to zs_range_iter calls

---
 .../src/service/finalized_state/zebra_db/shielded.rs      | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

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 4d10bb24b4c..260d202eac9 100644
--- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
@@ -256,7 +256,9 @@ impl ZebraDb {
             .cf_handle("sapling_note_commitment_subtree")
             .unwrap();
 
-        self.db.zs_range_iter(&sapling_subtrees, range).collect()
+        self.db
+            .zs_range_iter(&sapling_subtrees, range, false)
+            .collect()
     }
 
     /// Get the sapling note commitment subtress for the finalized tip.
@@ -380,7 +382,9 @@ impl ZebraDb {
             .cf_handle("orchard_note_commitment_subtree")
             .unwrap();
 
-        self.db.zs_range_iter(&orchard_subtrees, range).collect()
+        self.db
+            .zs_range_iter(&orchard_subtrees, range, false)
+            .collect()
     }
 
     /// Get the orchard note commitment subtress for the finalized tip.