change(state): Add state requests and support code for the `z_getsubtreesbyindex` RPC (#7408)

* Make NoteCommitmentSubtreeIndex compatible with serde-based RPCs

* Add a stub for z_getsubtreesbyindex

* Define a GetSubtrees RPC response type

* Reject invalid shielded pool names

* Make limit optional

* Define state request and response types for subtrees

* Implement FromDisk for NoteCommitmentSubtreeIndex and add a round-trip test

* Make subtrees compatible with round-trip proptests

* Add finalized state subtree list methods and delete unused methods

* Remove Arc from subtrees in zebra-chain

* Remove Arc from subtrees in zebra-state and use BTreeMap

* Implement subtree list lookups in the non-finalized state and delete unused methods

* Implement consistent concurrent subtree read requests

* Implement ToHex for sapling::Node

* Implement ToHex for orchard::Node

* Implement z_get_subtrees_by_index RPC

* Check for the start_index from the non-finalized state

* Remove an unused mut

* Fix missing doc links

* Fix RPC comments

* Temporarily remove the z_get_subtrees_by_index RPC method

zebra-chain/src/orchard/tree.rs

@@ -20,6 +20,7 @@ use std::{
 use bitvec::prelude::*;
 use bridgetree::{self, NonEmptyFrontier};
 use halo2::pasta::{group::ff::PrimeField, pallas};
+use hex::ToHex;
 use incrementalmerkletree::Hashable;
 use lazy_static::lazy_static;
 use thiserror::Error;
@@ -170,7 +171,7 @@ impl ZcashDeserialize for Root {
 }
 
 /// A node of the Orchard Incremental Note Commitment Tree.
-#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+#[derive(Copy, Clone, Eq, PartialEq)]
 pub struct Node(pallas::Base);
 
 impl Node {
@@ -178,6 +179,16 @@ impl Node {
     pub fn to_repr(&self) -> [u8; 32] {
         self.0.to_repr()
     }
+
+    /// Return the node bytes in big-endian byte-order suitable for printing out byte by byte.
+    ///
+    /// Zebra displays note commitment tree nodes in big-endian byte-order,
+    /// following the u256 convention set by Bitcoin and zcashd.
+    pub fn bytes_in_display_order(&self) -> [u8; 32] {
+        let mut reversed_bytes = self.0.to_repr();
+        reversed_bytes.reverse();
+        reversed_bytes
+    }
 }
 
 impl TryFrom<&[u8]> for Node {
@@ -200,6 +211,40 @@ impl TryFrom<[u8; 32]> for Node {
     }
 }
 
+impl fmt::Display for Node {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.write_str(&self.encode_hex::<String>())
+    }
+}
+
+impl fmt::Debug for Node {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.debug_tuple("orchard::Node")
+            .field(&self.encode_hex::<String>())
+            .finish()
+    }
+}
+
+impl ToHex for &Node {
+    fn encode_hex<T: FromIterator<char>>(&self) -> T {
+        self.bytes_in_display_order().encode_hex()
+    }
+
+    fn encode_hex_upper<T: FromIterator<char>>(&self) -> T {
+        self.bytes_in_display_order().encode_hex_upper()
+    }
+}
+
+impl ToHex for Node {
+    fn encode_hex<T: FromIterator<char>>(&self) -> T {
+        (&self).encode_hex()
+    }
+
+    fn encode_hex_upper<T: FromIterator<char>>(&self) -> T {
+        (&self).encode_hex_upper()
+    }
+}
+
 /// Required to convert [`NoteCommitmentTree`] into [`SerializedTree`].
 ///
 /// Zebra stores Orchard note commitment trees as [`Frontier`][1]s while the

zebra-chain/src/parallel/tree.rs

@@ -16,13 +16,13 @@ pub struct NoteCommitmentTrees {
     pub sapling: Arc<sapling::tree::NoteCommitmentTree>,
 
     /// The sapling note commitment subtree.
-    pub sapling_subtree: Option<Arc<NoteCommitmentSubtree<sapling::tree::Node>>>,
+    pub sapling_subtree: Option<NoteCommitmentSubtree<sapling::tree::Node>>,
 
     /// The orchard note commitment tree.
     pub orchard: Arc<orchard::tree::NoteCommitmentTree>,
 
     /// The orchard note commitment subtree.
-    pub orchard_subtree: Option<Arc<NoteCommitmentSubtree<orchard::tree::Node>>>,
+    pub orchard_subtree: Option<NoteCommitmentSubtree<orchard::tree::Node>>,
 }
 
 /// Note commitment tree errors.

zebra-chain/src/sapling/tree.rs

@@ -19,6 +19,7 @@ use std::{
 
 use bitvec::prelude::*;
 use bridgetree::{self, NonEmptyFrontier};
+use hex::ToHex;
 use incrementalmerkletree::{frontier::Frontier, Hashable};
 
 use lazy_static::lazy_static;
@@ -174,9 +175,49 @@ impl AsRef<[u8; 32]> for Node {
     }
 }
 
+impl fmt::Display for Node {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.write_str(&self.encode_hex::<String>())
+    }
+}
+
 impl fmt::Debug for Node {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        f.debug_tuple("Node").field(&hex::encode(self.0)).finish()
+        f.debug_tuple("sapling::Node")
+            .field(&self.encode_hex::<String>())
+            .finish()
+    }
+}
+
+impl Node {
+    /// Return the node bytes in big-endian byte-order suitable for printing out byte by byte.
+    ///
+    /// Zebra displays note commitment tree nodes in big-endian byte-order,
+    /// following the u256 convention set by Bitcoin and zcashd.
+    pub fn bytes_in_display_order(&self) -> [u8; 32] {
+        let mut reversed_bytes = self.0;
+        reversed_bytes.reverse();
+        reversed_bytes
+    }
+}
+
+impl ToHex for &Node {
+    fn encode_hex<T: FromIterator<char>>(&self) -> T {
+        self.bytes_in_display_order().encode_hex()
+    }
+
+    fn encode_hex_upper<T: FromIterator<char>>(&self) -> T {
+        self.bytes_in_display_order().encode_hex_upper()
+    }
+}
+
+impl ToHex for Node {
+    fn encode_hex<T: FromIterator<char>>(&self) -> T {
+        (&self).encode_hex()
+    }
+
+    fn encode_hex_upper<T: FromIterator<char>>(&self) -> T {
+        (&self).encode_hex_upper()
     }
 }
 

zebra-chain/src/subtree.rs

@@ -1,17 +1,20 @@
 //! Struct representing Sapling/Orchard note commitment subtrees
 
-use std::sync::Arc;
+use serde::{Deserialize, Serialize};
+
+use crate::block::Height;
 
 #[cfg(any(test, feature = "proptest-impl"))]
 use proptest_derive::Arbitrary;
 
-use crate::block::Height;
-
 /// Height at which Zebra tracks subtree roots
 pub const TRACKED_SUBTREE_HEIGHT: u8 = 16;
 
-/// A subtree index
+/// A note commitment subtree index, used to identify a subtree in a shielded pool.
-#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+/// Also used to count subtrees.
+#[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);
 
 impl From<u16> for NoteCommitmentSubtreeIndex {
@@ -20,23 +23,29 @@ impl From<u16> for NoteCommitmentSubtreeIndex {
     }
 }
 
+// TODO:
+// - consider defining sapling::SubtreeRoot and orchard::SubtreeRoot types or type wrappers,
+//   to avoid type confusion between the leaf Node and subtree root types.
+// - rename the `Node` generic to `SubtreeRoot`
+
 /// Subtree root of Sapling or Orchard note commitment tree,
 /// with its associated block height and subtree index.
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
+#[cfg_attr(any(test, feature = "proptest-impl"), derive(Arbitrary))]
 pub struct NoteCommitmentSubtree<Node> {
     /// Index of this subtree
     pub index: NoteCommitmentSubtreeIndex,
-    /// End boundary of this subtree, the block height of its last leaf.
-    pub end: Height,
     /// Root of this subtree.
     pub node: Node,
+    /// End boundary of this subtree, the block height of its last leaf.
+    pub end: Height,
 }
 
 impl<Node> NoteCommitmentSubtree<Node> {
     /// Creates new [`NoteCommitmentSubtree`]
-    pub fn new(index: impl Into<NoteCommitmentSubtreeIndex>, end: Height, node: Node) -> Arc<Self> {
+    pub fn new(index: impl Into<NoteCommitmentSubtreeIndex>, end: Height, node: Node) -> Self {
         let index = index.into();
-        Arc::new(Self { index, end, node })
+        Self { index, end, node }
     }
 
     /// Converts struct to [`NoteCommitmentSubtreeData`].
@@ -47,13 +56,18 @@ impl<Node> NoteCommitmentSubtree<Node> {
 
 /// Subtree root of Sapling or Orchard note commitment tree, with block height, but without the subtree index.
 /// Used for database key-value serialization, where the subtree index is the key, and this struct is the value.
-#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+#[derive(Copy, Clone, Debug, Eq, PartialEq, serde::Serialize)]
 #[cfg_attr(any(test, feature = "proptest-impl"), derive(Arbitrary))]
 pub struct NoteCommitmentSubtreeData<Node> {
-    /// End boundary of this subtree, the block height of its last leaf.
+    /// Merkle root of the 2^16-leaf subtree.
-    pub end: Height,
+    //
-    /// Root of this subtree.
+    // TODO: rename both Rust fields to match the RPC field names
+    #[serde(rename = "root")]
     pub node: Node,
+
+    /// Height of the block containing the note that completed this subtree.
+    #[serde(rename = "end_height")]
+    pub end: Height,
 }
 
 impl<Node> NoteCommitmentSubtreeData<Node> {
@@ -66,7 +80,7 @@ impl<Node> NoteCommitmentSubtreeData<Node> {
     pub fn with_index(
         self,
         index: impl Into<NoteCommitmentSubtreeIndex>,
-    ) -> Arc<NoteCommitmentSubtree<Node>> {
+    ) -> NoteCommitmentSubtree<Node> {
         NoteCommitmentSubtree::new(index, self.end, self.node)
     }
 }

zebra-state/src/request.rs

@@ -15,7 +15,7 @@ use zebra_chain::{
     sapling,
     serialization::SerializationError,
     sprout,
-    subtree::NoteCommitmentSubtree,
+    subtree::{NoteCommitmentSubtree, NoteCommitmentSubtreeIndex},
     transaction::{self, UnminedTx},
     transparent::{self, utxos_from_ordered_utxos},
     value_balance::{ValueBalance, ValueBalanceError},
@@ -236,8 +236,8 @@ impl Treestate {
         sprout: Arc<sprout::tree::NoteCommitmentTree>,
         sapling: Arc<sapling::tree::NoteCommitmentTree>,
         orchard: Arc<orchard::tree::NoteCommitmentTree>,
-        sapling_subtree: Option<Arc<NoteCommitmentSubtree<sapling::tree::Node>>>,
+        sapling_subtree: Option<NoteCommitmentSubtree<sapling::tree::Node>>,
-        orchard_subtree: Option<Arc<NoteCommitmentSubtree<orchard::tree::Node>>>,
+        orchard_subtree: Option<NoteCommitmentSubtree<orchard::tree::Node>>,
         history_tree: Arc<HistoryTree>,
     ) -> Self {
         Self {
@@ -849,6 +849,36 @@ pub enum ReadRequest {
     /// * [`ReadResponse::OrchardTree(None)`](crate::ReadResponse::OrchardTree) otherwise.
     OrchardTree(HashOrHeight),
 
+    /// Returns a list of Sapling note commitment subtrees by their indexes,
+    /// starting at `start_index`, and returning up to `limit` subtrees.
+    ///
+    /// Returns
+    ///
+    /// * [`ReadResponse::SaplingSubtree(BTreeMap<_, NoteCommitmentSubtreeData<_>>))`](crate::ReadResponse::SaplingSubtrees)
+    ///
+    /// If there is no subtree at `start_index`, returns an empty list.
+    SaplingSubtrees {
+        /// The index of the first 2^16-leaf subtree to return.
+        start_index: NoteCommitmentSubtreeIndex,
+        /// The maximum number of subtree values to return.
+        limit: Option<NoteCommitmentSubtreeIndex>,
+    },
+
+    /// Returns a list of Orchard note commitment subtrees by their indexes,
+    /// starting at `start_index`, and returning up to `limit` subtrees.
+    ///
+    /// Returns
+    ///
+    /// * [`ReadResponse::OrchardSubtree(BTreeMap<_, NoteCommitmentSubtreeData<_>>))`](crate::ReadResponse::OrchardSubtrees)
+    ///
+    /// If there is no subtree at `start_index`, returns an empty list.
+    OrchardSubtrees {
+        /// The index of the first 2^16-leaf subtree to return.
+        start_index: NoteCommitmentSubtreeIndex,
+        /// The maximum number of subtree values to return.
+        limit: Option<NoteCommitmentSubtreeIndex>,
+    },
+
     /// Looks up the balance of a set of transparent addresses.
     ///
     /// Returns an [`Amount`](zebra_chain::amount::Amount) with the total
@@ -942,6 +972,8 @@ impl ReadRequest {
             ReadRequest::FindBlockHeaders { .. } => "find_block_headers",
             ReadRequest::SaplingTree { .. } => "sapling_tree",
             ReadRequest::OrchardTree { .. } => "orchard_tree",
+            ReadRequest::SaplingSubtrees { .. } => "sapling_subtrees",
+            ReadRequest::OrchardSubtrees { .. } => "orchard_subtrees",
             ReadRequest::AddressBalance { .. } => "address_balance",
             ReadRequest::TransactionIdsByAddresses { .. } => "transaction_ids_by_addesses",
             ReadRequest::UtxosByAddresses(_) => "utxos_by_addesses",

zebra-state/src/response.rs

@@ -7,6 +7,7 @@ use zebra_chain::{
     block::{self, Block},
     orchard, sapling,
     serialization::DateTime32,
+    subtree::{NoteCommitmentSubtreeData, NoteCommitmentSubtreeIndex},
     transaction::{self, Transaction},
     transparent,
 };
@@ -164,6 +165,18 @@ pub enum ReadResponse {
     /// Response to [`ReadRequest::OrchardTree`] with the specified Orchard note commitment tree.
     OrchardTree(Option<Arc<orchard::tree::NoteCommitmentTree>>),
 
+    /// Response to [`ReadRequest::SaplingSubtrees`] with the specified Sapling note commitment
+    /// subtrees.
+    SaplingSubtrees(
+        BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<sapling::tree::Node>>,
+    ),
+
+    /// Response to [`ReadRequest::OrchardSubtrees`] with the specified Orchard note commitment
+    /// subtrees.
+    OrchardSubtrees(
+        BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<orchard::tree::Node>>,
+    ),
+
     /// Response to [`ReadRequest::AddressBalance`] with the total balance of the addresses.
     AddressBalance(Amount<NonNegative>),
 
@@ -270,6 +283,8 @@ impl TryFrom<ReadResponse> for Response {
             ReadResponse::TransactionIdsForBlock(_)
             | ReadResponse::SaplingTree(_)
             | ReadResponse::OrchardTree(_)
+            | ReadResponse::SaplingSubtrees(_)
+            | ReadResponse::OrchardSubtrees(_)
             | ReadResponse::AddressBalance(_)
             | ReadResponse::AddressesTransactionIds(_)
             | ReadResponse::AddressUtxos(_) => {

zebra-state/src/service.rs

@@ -1502,6 +1502,56 @@ impl Service<ReadRequest> for ReadStateService {
                 .wait_for_panics()
             }
 
+            ReadRequest::SaplingSubtrees { start_index, limit } => {
+                let state = self.clone();
+
+                tokio::task::spawn_blocking(move || {
+                    span.in_scope(move || {
+                        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,
+                                )
+                            },
+                        );
+
+                        // The work is done in the future.
+                        timer.finish(module_path!(), line!(), "ReadRequest::SaplingSubtrees");
+
+                        Ok(ReadResponse::SaplingSubtrees(sapling_subtrees))
+                    })
+                })
+                .wait_for_panics()
+            }
+
+            ReadRequest::OrchardSubtrees { start_index, limit } => {
+                let state = self.clone();
+
+                tokio::task::spawn_blocking(move || {
+                    span.in_scope(move || {
+                        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,
+                                )
+                            },
+                        );
+
+                        // The work is done in the future.
+                        timer.finish(module_path!(), line!(), "ReadRequest::OrchardSubtrees");
+
+                        Ok(ReadResponse::OrchardSubtrees(orchard_subtrees))
+                    })
+                })
+                .wait_for_panics()
+            }
+
             // For the get_address_balance RPC.
             ReadRequest::AddressBalance(addresses) => {
                 let state = self.clone();

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

@@ -80,6 +80,13 @@ impl IntoDisk for orchard::tree::Root {
     }
 }
 
+impl FromDisk for orchard::tree::Root {
+    fn from_bytes(bytes: impl AsRef<[u8]>) -> Self {
+        let array: [u8; 32] = bytes.as_ref().try_into().unwrap();
+        array.try_into().expect("finalized data must be valid")
+    }
+}
+
 impl IntoDisk for NoteCommitmentSubtreeIndex {
     type Bytes = [u8; 2];
 
@@ -88,10 +95,10 @@ impl IntoDisk for NoteCommitmentSubtreeIndex {
     }
 }
 
-impl FromDisk for orchard::tree::Root {
+impl FromDisk for NoteCommitmentSubtreeIndex {
     fn from_bytes(bytes: impl AsRef<[u8]>) -> Self {
-        let array: [u8; 32] = bytes.as_ref().try_into().unwrap();
+        let array: [u8; 2] = bytes.as_ref().try_into().unwrap();
-        array.try_into().expect("finalized data must be valid")
+        Self(u16::from_be_bytes(array))
     }
 }
 

zebra-state/src/service/finalized_state/disk_format/tests/prop.rs

@@ -6,7 +6,7 @@ use zebra_chain::{
     amount::{Amount, NonNegative},
     block::{self, Height},
     orchard, sapling, sprout,
-    subtree::NoteCommitmentSubtreeData,
+    subtree::{NoteCommitmentSubtreeData, NoteCommitmentSubtreeIndex},
     transaction::{self, Transaction},
     transparent,
     value_balance::ValueBalance,
@@ -214,6 +214,15 @@ fn roundtrip_amount() {
     proptest!(|(val in any::<Amount::<NonNegative>>())| assert_value_properties(val));
 }
 
+#[test]
+fn roundtrip_note_commitment_subtree_index() {
+    let _init_guard = zebra_test::init();
+
+    proptest!(|(val in any::<NoteCommitmentSubtreeIndex>())| {
+        assert_value_properties(val)
+    });
+}
+
 // Sprout
 
 #[test]

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

@@ -12,14 +12,17 @@
 //! The [`crate::constants::DATABASE_FORMAT_VERSION`] constant must
 //! be incremented each time the database format (column, serialization, etc) changes.
 
-use std::{collections::HashMap, sync::Arc};
+use std::{
+    collections::{BTreeMap, HashMap},
+    sync::Arc,
+};
 
 use zebra_chain::{
     block::Height,
     orchard,
     parallel::tree::NoteCommitmentTrees,
     sapling, sprout,
-    subtree::{NoteCommitmentSubtree, NoteCommitmentSubtreeData, NoteCommitmentSubtreeIndex},
+    subtree::{NoteCommitmentSubtreeData, NoteCommitmentSubtreeIndex},
     transaction::Transaction,
 };
 
@@ -32,6 +35,10 @@ use crate::{
     BoxError, SemanticallyVerifiedBlock,
 };
 
+// Doc-only items
+#[allow(unused_imports)]
+use zebra_chain::subtree::NoteCommitmentSubtree;
+
 impl ZebraDb {
     // Read shielded methods
 
@@ -173,20 +180,58 @@ impl ZebraDb {
         self.db.zs_range_iter(&sapling_trees, range)
     }
 
-    /// Returns the Sapling note commitment subtree at this 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.
+    ///
+    /// There is no API for retrieving single subtrees by index, because it can accidentally be used
+    /// to create an inconsistent list of subtrees after concurrent non-finalized and finalized
+    /// updates.
     #[allow(clippy::unwrap_in_result)]
-    pub fn sapling_subtree_by_index(
+    pub fn sapling_subtrees_by_index(
         &self,
-        index: impl Into<NoteCommitmentSubtreeIndex> + Copy,
+        start_index: NoteCommitmentSubtreeIndex,
-    ) -> Option<Arc<NoteCommitmentSubtree<sapling::tree::Node>>> {
+        limit: Option<NoteCommitmentSubtreeIndex>,
+    ) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<sapling::tree::Node>> {
         let sapling_subtrees = self
             .db
             .cf_handle("sapling_note_commitment_subtree")
             .unwrap();
 
-        let subtree_data: NoteCommitmentSubtreeData<sapling::tree::Node> =
+        // Calculate the end bound, checking for overflow.
-            self.db.zs_get(&sapling_subtrees, &index.into())?;
+        let exclusive_end_bound: Option<NoteCommitmentSubtreeIndex> = limit
-        Some(subtree_data.with_index(index))
+            .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();
+        }
+
+        // Check that we got the start subtree.
+        if list.get(&start_index).is_some() {
+            list
+        } else {
+            BTreeMap::new()
+        }
     }
 
     // Orchard trees
@@ -203,22 +248,6 @@ impl ZebraDb {
             .expect("Orchard note commitment tree must exist if there is a finalized tip")
     }
 
-    /// Returns the Orchard note commitment subtree at this index
-    #[allow(clippy::unwrap_in_result)]
-    pub fn orchard_subtree_by_index(
-        &self,
-        index: impl Into<NoteCommitmentSubtreeIndex> + Copy,
-    ) -> Option<Arc<NoteCommitmentSubtree<orchard::tree::Node>>> {
-        let orchard_subtrees = self
-            .db
-            .cf_handle("orchard_note_commitment_subtree")
-            .unwrap();
-
-        let subtree_data: NoteCommitmentSubtreeData<orchard::tree::Node> =
-            self.db.zs_get(&orchard_subtrees, &index.into())?;
-        Some(subtree_data.with_index(index))
-    }
-
     /// Returns the Orchard note commitment tree matching the given block height,
     /// or `None` if the height is above the finalized tip.
     #[allow(clippy::unwrap_in_result)]
@@ -260,6 +289,60 @@ impl ZebraDb {
         self.db.zs_range_iter(&orchard_trees, range)
     }
 
+    /// 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.
+    ///
+    /// There is no API for retrieving single subtrees by index, because it can accidentally be used
+    /// to create an inconsistent list of subtrees after concurrent non-finalized and finalized
+    /// updates.
+    #[allow(clippy::unwrap_in_result)]
+    pub fn orchard_subtrees_by_index(
+        &self,
+        start_index: NoteCommitmentSubtreeIndex,
+        limit: Option<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();
+        }
+
+        // Check that we got the start subtree.
+        if list.get(&start_index).is_some() {
+            list
+        } else {
+            BTreeMap::new()
+        }
+    }
+
     /// Returns the shielded note commitment trees of the finalized tip
     /// or the empty trees if the state is empty.
     pub fn note_commitment_trees(&self) -> NoteCommitmentTrees {

zebra-state/src/service/non_finalized_state/chain.rs

@@ -3,7 +3,7 @@
 
 use std::{
     cmp::Ordering,
-    collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque},
+    collections::{BTreeMap, BTreeSet, HashMap, HashSet},
     ops::{Deref, RangeInclusive},
     sync::Arc,
 };
@@ -20,7 +20,7 @@ use zebra_chain::{
     parameters::Network,
     primitives::Groth16Proof,
     sapling, sprout,
-    subtree::{NoteCommitmentSubtree, NoteCommitmentSubtreeIndex},
+    subtree::{NoteCommitmentSubtree, NoteCommitmentSubtreeData, NoteCommitmentSubtreeIndex},
     transaction::Transaction::*,
     transaction::{self, Transaction},
     transparent,
@@ -135,7 +135,8 @@ pub struct Chain {
     /// This extra root is removed when the first non-finalized block is committed.
     pub(crate) sapling_anchors_by_height: BTreeMap<block::Height, sapling::tree::Root>,
     /// A list of Sapling subtrees completed in the non-finalized state
-    pub(crate) sapling_subtrees: VecDeque<Arc<NoteCommitmentSubtree<sapling::tree::Node>>>,
+    pub(crate) sapling_subtrees:
+        BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<sapling::tree::Node>>,
 
     /// The Orchard anchors created by `blocks`.
     ///
@@ -148,7 +149,8 @@ pub struct Chain {
     /// This extra root is removed when the first non-finalized block is committed.
     pub(crate) orchard_anchors_by_height: BTreeMap<block::Height, orchard::tree::Root>,
     /// A list of Orchard subtrees completed in the non-finalized state
-    pub(crate) orchard_subtrees: VecDeque<Arc<NoteCommitmentSubtree<orchard::tree::Node>>>,
+    pub(crate) orchard_subtrees:
+        BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<orchard::tree::Node>>,
 
     // Nullifiers
     //
@@ -351,11 +353,11 @@ impl Chain {
             .expect("The treestate must be present for the root height.");
 
         if treestate.note_commitment_trees.sapling_subtree.is_some() {
-            self.sapling_subtrees.pop_front();
+            self.sapling_subtrees.pop_first();
         }
 
         if treestate.note_commitment_trees.orchard_subtree.is_some() {
-            self.orchard_subtrees.pop_front();
+            self.orchard_subtrees.pop_first();
         }
 
         // Remove the lowest height block from `self.blocks`.
@@ -678,31 +680,45 @@ impl Chain {
             .map(|(_height, tree)| tree.clone())
     }
 
-    /// Returns the Sapling [`NoteCommitmentSubtree`] specified
+    /// Returns the Sapling [`NoteCommitmentSubtree`] that was completed at a block with
-    /// by an index, if it exists in the non-finalized [`Chain`].
+    /// [`HashOrHeight`], if it exists in the non-finalized [`Chain`].
     pub fn sapling_subtree(
         &self,
         hash_or_height: HashOrHeight,
-    ) -> Option<Arc<NoteCommitmentSubtree<sapling::tree::Node>>> {
+    ) -> Option<NoteCommitmentSubtree<sapling::tree::Node>> {
         let height =
             hash_or_height.height_or_else(|hash| self.height_by_hash.get(&hash).cloned())?;
 
         self.sapling_subtrees
             .iter()
-            .find(|subtree| subtree.end == height)
+            .find(|(_index, subtree)| subtree.end == height)
-            .cloned()
+            .map(|(index, subtree)| subtree.with_index(*index))
     }
 
-    /// Returns the Sapling [`NoteCommitmentSubtree`] specified
+    /// Returns a list of Sapling [`NoteCommitmentSubtree`]s at or after `start_index`.
-    /// by a [`HashOrHeight`], if it exists in the non-finalized [`Chain`].
+    /// If `limit` is provided, the list is limited to `limit` entries.
-    pub fn sapling_subtree_by_index(
+    ///
+    /// Unlike the finalized state and `ReadRequest::SaplingSubtrees`, the returned subtrees
+    /// can start after `start_index`. These subtrees are continuous up to the tip.
+    ///
+    /// There is no API for retrieving single subtrees by index, because it can accidentally be
+    /// used to create an inconsistent list of subtrees after concurrent non-finalized and
+    /// finalized updates.
+    pub fn sapling_subtrees_in_range(
         &self,
-        index: NoteCommitmentSubtreeIndex,
+        start_index: NoteCommitmentSubtreeIndex,
-    ) -> Option<Arc<NoteCommitmentSubtree<sapling::tree::Node>>> {
+        limit: Option<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
-            .iter()
+            .range(start_index..)
-            .find(|subtree| subtree.index == index)
+            .take(limit)
-            .cloned()
+            .map(|(index, subtree)| (*index, *subtree))
+            .collect()
     }
 
     /// Adds the Sapling `tree` to the tree and anchor indexes at `height`.
@@ -854,31 +870,45 @@ impl Chain {
             .map(|(_height, tree)| tree.clone())
     }
 
-    /// Returns the Orchard [`NoteCommitmentSubtree`] specified
+    /// Returns the Orchard [`NoteCommitmentSubtree`] that was completed at a block with
-    /// by a [`HashOrHeight`], if it exists in the non-finalized [`Chain`].
+    /// [`HashOrHeight`], if it exists in the non-finalized [`Chain`].
     pub fn orchard_subtree(
         &self,
         hash_or_height: HashOrHeight,
-    ) -> Option<Arc<NoteCommitmentSubtree<orchard::tree::Node>>> {
+    ) -> Option<NoteCommitmentSubtree<orchard::tree::Node>> {
         let height =
             hash_or_height.height_or_else(|hash| self.height_by_hash.get(&hash).cloned())?;
 
         self.orchard_subtrees
             .iter()
-            .find(|subtree| subtree.end == height)
+            .find(|(_index, subtree)| subtree.end == height)
-            .cloned()
+            .map(|(index, subtree)| subtree.with_index(*index))
     }
 
-    /// Returns the Orchard [`NoteCommitmentSubtree`] specified
+    /// Returns a list of Orchard [`NoteCommitmentSubtree`]s at or after `start_index`.
-    /// by an index, if it exists in the non-finalized [`Chain`].
+    /// If `limit` is provided, the list is limited to `limit` entries.
-    pub fn orchard_subtree_by_index(
+    ///
+    /// Unlike the finalized state and `ReadRequest::OrchardSubtrees`, the returned subtrees
+    /// can start after `start_index`. These subtrees are continuous up to the tip.
+    ///
+    /// There is no API for retrieving single subtrees by index, because it can accidentally be
+    /// used to create an inconsistent list of subtrees after concurrent non-finalized and
+    /// finalized updates.
+    pub fn orchard_subtrees_in_range(
         &self,
-        index: NoteCommitmentSubtreeIndex,
+        start_index: NoteCommitmentSubtreeIndex,
-    ) -> Option<Arc<NoteCommitmentSubtree<orchard::tree::Node>>> {
+        limit: Option<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
-            .iter()
+            .range(start_index..)
-            .find(|subtree| subtree.index == index)
+            .take(limit)
-            .cloned()
+            .map(|(index, subtree)| (*index, *subtree))
+            .collect()
     }
 
     /// Adds the Orchard `tree` to the tree and anchor indexes at `height`.
@@ -1354,10 +1384,12 @@ impl Chain {
         self.add_orchard_tree_and_anchor(height, nct.orchard);
 
         if let Some(subtree) = nct.sapling_subtree {
-            self.sapling_subtrees.push_back(subtree)
+            self.sapling_subtrees
+                .insert(subtree.index, subtree.into_data());
         }
         if let Some(subtree) = nct.orchard_subtree {
-            self.orchard_subtrees.push_back(subtree)
+            self.orchard_subtrees
+                .insert(subtree.index, subtree.into_data());
         }
 
         let sapling_root = self.sapling_note_commitment_tree().root();

zebra-state/src/service/read.rs

@@ -39,7 +39,7 @@ pub use find::{
     find_chain_hashes, find_chain_headers, hash_by_height, height_by_hash, next_median_time_past,
     non_finalized_state_contains_block_hash, tip, tip_height,
 };
-pub use tree::{orchard_tree, sapling_tree};
+pub use tree::{orchard_subtrees, orchard_tree, sapling_subtrees, sapling_tree};
 
 #[cfg(feature = "getblocktemplate-rpcs")]
 pub use difficulty::get_block_template_chain_info;

zebra-state/src/service/read/tree.rs

@@ -11,11 +11,11 @@
 //! - the cached [`Chain`], and
 //! - the shared finalized [`ZebraDb`] reference.
 
-use std::sync::Arc;
+use std::{collections::BTreeMap, sync::Arc};
 
 use zebra_chain::{
     orchard, sapling,
-    subtree::{NoteCommitmentSubtree, NoteCommitmentSubtreeIndex},
+    subtree::{NoteCommitmentSubtreeData, NoteCommitmentSubtreeIndex},
 };
 
 use crate::{
@@ -23,6 +23,10 @@ use crate::{
     HashOrHeight,
 };
 
+// Doc-only items
+#[allow(unused_imports)]
+use zebra_chain::subtree::NoteCommitmentSubtree;
+
 /// Returns the Sapling
 /// [`NoteCommitmentTree`](sapling::tree::NoteCommitmentTree) specified by a
 /// hash or height, if it exists in the non-finalized `chain` or finalized `db`.
@@ -44,26 +48,71 @@ where
         .or_else(|| db.sapling_tree_by_hash_or_height(hash_or_height))
 }
 
-/// Returns the Sapling
+/// Returns a list of Sapling [`NoteCommitmentSubtree`]s starting at `start_index`.
-/// [`NoteCommitmentSubtree`] specified by an
+/// If `limit` is provided, the list is limited to `limit` entries.
-/// index, if it exists in the non-finalized `chain` or finalized `db`.
+///
-#[allow(unused)]
+/// If there is no subtree at `start_index` in the non-finalized `chain` or finalized `db`,
-pub fn sapling_subtree<C>(
+/// the returned list is empty. Otherwise, subtrees are continuous and consistent up to the tip.
+///
+/// There is no API for retrieving single subtrees, because it can accidentally be used to create
+/// an inconsistent list of subtrees after concurrent non-finalized and finalized updates.
+pub fn sapling_subtrees<C>(
     chain: Option<C>,
     db: &ZebraDb,
-    index: NoteCommitmentSubtreeIndex,
+    start_index: NoteCommitmentSubtreeIndex,
-) -> Option<Arc<NoteCommitmentSubtree<sapling::tree::Node>>>
+    limit: Option<NoteCommitmentSubtreeIndex>,
+) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<sapling::tree::Node>>
 where
     C: AsRef<Chain>,
 {
     // # Correctness
     //
-    // Since sapling treestates are the same in the finalized and non-finalized
+    // After `chain` was cloned, the StateService can commit additional blocks to the finalized
-    // state, we check the most efficient alternative first. (`chain` is always
+    // state `db`. Usually, the subtrees of these blocks are consistent. But if the `chain` is
-    // in memory, but `db` stores blocks on disk, with a memory cache.)
+    // a different fork to `db`, then the trees can be inconsistent.
-    chain
+    //
-        .and_then(|chain| chain.as_ref().sapling_subtree_by_index(index))
+    // In that case, we ignore all the trees in `chain` after the first inconsistent tree,
-        .or_else(|| db.sapling_subtree_by_index(index))
+    // 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 mut db_list = db.sapling_subtrees_by_index(start_index, limit);
+
+    // 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);
+            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.
+    }
+
+    // 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()
+    }
 }
 
 /// Returns the Orchard
@@ -87,25 +136,71 @@ where
         .or_else(|| db.orchard_tree_by_hash_or_height(hash_or_height))
 }
 
-/// Returns the Orchard [`NoteCommitmentSubtree`] specified by an
+/// Returns a list of Orchard [`NoteCommitmentSubtree`]s starting at `start_index`.
-/// index, if it exists in the non-finalized `chain` or finalized `db`.
+/// If `limit` is provided, the list is limited to `limit` entries.
-#[allow(unused)]
+///
-pub fn orchard_subtree<C>(
+/// If there is no subtree at `start_index` in the non-finalized `chain` or finalized `db`,
+/// the returned list is empty. Otherwise, subtrees are continuous and consistent up to the tip.
+///
+/// There is no API for retrieving single subtrees, because it can accidentally be used to create
+/// an inconsistent list of subtrees.
+pub fn orchard_subtrees<C>(
     chain: Option<C>,
     db: &ZebraDb,
-    index: NoteCommitmentSubtreeIndex,
+    start_index: NoteCommitmentSubtreeIndex,
-) -> Option<Arc<NoteCommitmentSubtree<orchard::tree::Node>>>
+    limit: Option<NoteCommitmentSubtreeIndex>,
+) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<orchard::tree::Node>>
 where
     C: AsRef<Chain>,
 {
     // # Correctness
     //
-    // Since orchard treestates are the same in the finalized and non-finalized
+    // After `chain` was cloned, the StateService can commit additional blocks to the finalized
-    // state, we check the most efficient alternative first. (`chain` is always
+    // state `db`. Usually, the subtrees of these blocks are consistent. But if the `chain` is
-    // in memory, but `db` stores blocks on disk, with a memory cache.)
+    // a different fork to `db`, then the trees can be inconsistent.
-    chain
+    //
-        .and_then(|chain| chain.as_ref().orchard_subtree_by_index(index))
+    // In that case, we ignore all the trees in `chain` after the first inconsistent tree,
-        .or_else(|| db.orchard_subtree_by_index(index))
+    // 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 mut db_list = db.orchard_subtrees_by_index(start_index, limit);
+
+    // 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);
+
+    // 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);
+            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.
+    }
+
+    // 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()
+    }
 }
 
 #[cfg(feature = "getblocktemplate-rpcs")]