change(state): Add note subtree index handling to zebra-state, but don't write them to the finalized state yet (#7334)

* zebra-chain changes from the subtree-boundaries branch

```sh
git checkout -b subtree-boundaries-zebra-chain main
git checkout origin/subtree-boundaries zebra-chain
git commit
```

* Temporarily populate new subtree fields with None - for revert

This temporary commit needs to be reverted in the next PR.

* Applies suggestions from code review

* removes from_repr_unchecked methods

* simplifies loop

* adds subtrees to zebra-state

* uses split_at, from_repr, & updates state-db-upgrades.md

* Update book/src/dev/state-db-upgrades.md

Co-authored-by: teor <teor@riseup.net>

* renames partial_subtree to subtree_data

* tests that subtree serialization format

* adds raw data format serialization round-trip test

* decrements minor version and skips inserting subtrees in db

---------

Co-authored-by: teor <teor@riseup.net>

book/src/dev/state-db-upgrades.md

@@ -92,10 +92,12 @@ We use the following rocksdb column families:
 | `sapling_nullifiers`               | `sapling::Nullifier`   | `()`                          | Create  |
 | `sapling_anchors`                  | `sapling::tree::Root`  | `()`                          | Create  |
 | `sapling_note_commitment_tree`     | `block::Height`        | `sapling::NoteCommitmentTree` | Create  |
+| `sapling_note_commitment_subtree`  | `block::Height`        | `NoteCommitmentSubtreeData`   | Create  |
 | *Orchard*                          |                        |                               |         |
 | `orchard_nullifiers`               | `orchard::Nullifier`   | `()`                          | Create  |
 | `orchard_anchors`                  | `orchard::tree::Root`  | `()`                          | Create  |
 | `orchard_note_commitment_tree`     | `block::Height`        | `orchard::NoteCommitmentTree` | Create  |
+| `orchard_note_commitment_subtree`  | `block::Height`        | `NoteCommitmentSubtreeData`   | Create  |
 | *Chain*                            |                        |                               |         |
 | `history_tree`                     | `block::Height`        | `NonEmptyHistoryTree`         | Delete  |
 | `tip_chain_value_pool`             | `()`                   | `ValueBalance`                | Update  |
@@ -118,6 +120,8 @@ Block and Transaction Data:
   used instead of a `BTreeSet<OutputLocation>` value, to improve database performance
 - `AddressTransaction`: `AddressLocation \|\| TransactionLocation`
   used instead of a `BTreeSet<TransactionLocation>` value, to improve database performance
+- `NoteCommitmentSubtreeIndex`: 16 bits, big-endian, unsigned
+- `NoteCommitmentSubtreeData<{sapling, orchard}::tree::Node>`: `Height \|\| {sapling, orchard}::tree::Node`
 
 We use big-endian encoding for keys, to allow database index prefix searches.
 
@@ -334,6 +338,9 @@ So they should not be used for consensus-critical checks.
   as a "Merkle tree frontier" which is basically a (logarithmic) subset of
   the Merkle tree nodes as required to insert new items.
 
+- The `{sapling, orchard}_note_commitment_subtree` stores the completion height and
+  root for every completed level 16 note commitment subtree, for the specific pool.
+
 - `history_tree` stores the ZIP-221 history tree state at the tip of the finalized
   state. There is always a single entry for it. The tree is stored as the set of "peaks"
   of the "Merkle mountain range" tree structure, which is what is required to

zebra-chain/src/orchard/arbitrary.rs

@@ -125,14 +125,37 @@ impl Arbitrary for Flags {
     type Strategy = BoxedStrategy<Self>;
 }
 
+fn pallas_base_strat() -> BoxedStrategy<pallas::Base> {
+    (vec(any::<u8>(), 64))
+        .prop_map(|bytes| {
+            let bytes = bytes.try_into().expect("vec is the correct length");
+            pallas::Base::from_uniform_bytes(&bytes)
+        })
+        .boxed()
+}
+
 impl Arbitrary for tree::Root {
     type Parameters = ();
 
     fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
-        (vec(any::<u8>(), 64))
+        pallas_base_strat()
-            .prop_map(|bytes| {
+            .prop_map(|base| {
-                let bytes = bytes.try_into().expect("vec is the correct length");
+                Self::try_from(base.to_repr())
-                Self::try_from(pallas::Base::from_uniform_bytes(&bytes).to_repr())
+                    .expect("a valid generated Orchard note commitment tree root")
+            })
+            .boxed()
+    }
+
+    type Strategy = BoxedStrategy<Self>;
+}
+
+impl Arbitrary for tree::Node {
+    type Parameters = ();
+
+    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
+        pallas_base_strat()
+            .prop_map(|base| {
+                Self::try_from(base.to_repr())
                     .expect("a valid generated Orchard note commitment tree root")
             })
             .boxed()

zebra-chain/src/orchard/tree.rs

@@ -184,11 +184,19 @@ impl TryFrom<&[u8]> for Node {
     type Error = &'static str;
 
     fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
-        Option::<pallas::Base>::from(pallas::Base::from_repr(
+        <[u8; 32]>::try_from(bytes)
-            bytes.try_into().map_err(|_| "wrong byte slice len")?,
+            .map_err(|_| "wrong byte slice len")?
-        ))
+            .try_into()
-        .map(Node)
+    }
-        .ok_or("invalid Pallas field element")
+}
+
+impl TryFrom<[u8; 32]> for Node {
+    type Error = &'static str;
+
+    fn try_from(bytes: [u8; 32]) -> Result<Self, Self::Error> {
+        Option::<pallas::Base>::from(pallas::Base::from_repr(bytes))
+            .map(Node)
+            .ok_or("invalid Pallas field element")
     }
 }
 

zebra-chain/src/sapling/arbitrary.rs

@@ -119,16 +119,30 @@ fn spendauth_verification_key_bytes() -> impl Strategy<Value = ValidatingKey> {
     })
 }
 
+fn jubjub_base_strat() -> BoxedStrategy<jubjub::Base> {
+    (vec(any::<u8>(), 64))
+        .prop_map(|bytes| {
+            let bytes = bytes.try_into().expect("vec is the correct length");
+            jubjub::Base::from_bytes_wide(&bytes)
+        })
+        .boxed()
+}
+
 impl Arbitrary for tree::Root {
     type Parameters = ();
 
     fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
-        (vec(any::<u8>(), 64))
+        jubjub_base_strat().prop_map(tree::Root).boxed()
-            .prop_map(|bytes| {
+    }
-                let bytes = bytes.try_into().expect("vec is the correct length");
+
-                tree::Root(jubjub::Base::from_bytes_wide(&bytes))
+    type Strategy = BoxedStrategy<Self>;
-            })
+}
-            .boxed()
+
+impl Arbitrary for tree::Node {
+    type Parameters = ();
+
+    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
+        jubjub_base_strat().prop_map(tree::Node::from).boxed()
     }
 
     type Strategy = BoxedStrategy<Self>;

zebra-chain/src/subtree.rs

@@ -2,6 +2,9 @@
 
 use std::sync::Arc;
 
+#[cfg(any(test, feature = "proptest-impl"))]
+use proptest_derive::Arbitrary;
+
 use crate::block::Height;
 
 /// Height at which Zebra tracks subtree roots
@@ -35,11 +38,17 @@ impl<Node> NoteCommitmentSubtree<Node> {
         let index = index.into();
         Arc::new(Self { index, end, node })
     }
+
+    /// Converts struct to [`NoteCommitmentSubtreeData`].
+    pub fn into_data(self) -> NoteCommitmentSubtreeData<Node> {
+        NoteCommitmentSubtreeData::new(self.end, self.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)]
+#[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.
     pub end: Height,

zebra-state/src/request.rs

@@ -15,6 +15,7 @@ use zebra_chain::{
     sapling,
     serialization::SerializationError,
     sprout,
+    subtree::NoteCommitmentSubtree,
     transaction::{self, UnminedTx},
     transparent::{self, utxos_from_ordered_utxos},
     value_balance::{ValueBalance, ValueBalanceError},
@@ -235,15 +236,17 @@ 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>>>,
+        orchard_subtree: Option<Arc<NoteCommitmentSubtree<orchard::tree::Node>>>,
         history_tree: Arc<HistoryTree>,
     ) -> Self {
         Self {
             note_commitment_trees: NoteCommitmentTrees {
                 sprout,
                 sapling,
-                sapling_subtree: None,
+                sapling_subtree,
                 orchard,
-                orchard_subtree: None,
+                orchard_subtree,
             },
             history_tree,
         }

zebra-state/src/service/finalized_state/disk_db.rs

@@ -517,10 +517,12 @@ impl DiskDb {
         "sapling_nullifiers",
         "sapling_anchors",
         "sapling_note_commitment_tree",
+        "sapling_note_commitment_subtree",
         // Orchard
         "orchard_nullifiers",
         "orchard_anchors",
         "orchard_note_commitment_tree",
+        "orchard_note_commitment_subtree",
         // Chain
         "history_tree",
         "tip_chain_value_pool",

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

@@ -7,10 +7,16 @@
 
 use bincode::Options;
 
-use zebra_chain::{orchard, sapling, sprout};
+use zebra_chain::{
+    block::Height,
+    orchard, sapling, sprout,
+    subtree::{NoteCommitmentSubtreeData, NoteCommitmentSubtreeIndex},
+};
 
 use crate::service::finalized_state::disk_format::{FromDisk, IntoDisk};
 
+use super::block::HEIGHT_DISK_BYTES;
+
 impl IntoDisk for sprout::Nullifier {
     type Bytes = [u8; 32];
 
@@ -74,6 +80,14 @@ impl IntoDisk for orchard::tree::Root {
     }
 }
 
+impl IntoDisk for NoteCommitmentSubtreeIndex {
+    type Bytes = [u8; 2];
+
+    fn as_bytes(&self) -> Self::Bytes {
+        self.0.to_be_bytes()
+    }
+}
+
 impl FromDisk for orchard::tree::Root {
     fn from_bytes(bytes: impl AsRef<[u8]>) -> Self {
         let array: [u8; 32] = bytes.as_ref().try_into().unwrap();
@@ -140,3 +154,49 @@ impl FromDisk for orchard::tree::NoteCommitmentTree {
             .expect("deserialization format should match the serialization format used by IntoDisk")
     }
 }
+
+impl IntoDisk for sapling::tree::Node {
+    type Bytes = Vec<u8>;
+
+    fn as_bytes(&self) -> Self::Bytes {
+        self.as_ref().to_vec()
+    }
+}
+
+impl IntoDisk for orchard::tree::Node {
+    type Bytes = Vec<u8>;
+
+    fn as_bytes(&self) -> Self::Bytes {
+        self.to_repr().to_vec()
+    }
+}
+
+impl<Node: IntoDisk<Bytes = Vec<u8>>> IntoDisk for NoteCommitmentSubtreeData<Node> {
+    type Bytes = Vec<u8>;
+
+    fn as_bytes(&self) -> Self::Bytes {
+        [self.end.as_bytes().to_vec(), self.node.as_bytes()].concat()
+    }
+}
+
+impl FromDisk for sapling::tree::Node {
+    fn from_bytes(bytes: impl AsRef<[u8]>) -> Self {
+        Self::try_from(bytes.as_ref()).expect("trusted data should deserialize successfully")
+    }
+}
+
+impl FromDisk for orchard::tree::Node {
+    fn from_bytes(bytes: impl AsRef<[u8]>) -> Self {
+        Self::try_from(bytes.as_ref()).expect("trusted data should deserialize successfully")
+    }
+}
+
+impl<Node: FromDisk> FromDisk for NoteCommitmentSubtreeData<Node> {
+    fn from_bytes(disk_bytes: impl AsRef<[u8]>) -> Self {
+        let (height_bytes, node_bytes) = disk_bytes.as_ref().split_at(HEIGHT_DISK_BYTES);
+        Self::new(
+            Height::from_bytes(height_bytes),
+            Node::from_bytes(node_bytes),
+        )
+    }
+}

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

@@ -6,6 +6,7 @@ use zebra_chain::{
     amount::{Amount, NonNegative},
     block::{self, Height},
     orchard, sapling, sprout,
+    subtree::NoteCommitmentSubtreeData,
     transaction::{self, Transaction},
     transparent,
     value_balance::ValueBalance,
@@ -361,6 +362,16 @@ fn roundtrip_sapling_tree_root() {
     proptest!(|(val in any::<sapling::tree::Root>())| assert_value_properties(val));
 }
 
+#[test]
+fn roundtrip_sapling_subtree_data() {
+    let _init_guard = zebra_test::init();
+
+    proptest!(|(mut val in any::<NoteCommitmentSubtreeData<sapling::tree::Node>>())| {
+        val.end = val.end.clamp(Height(0), MAX_ON_DISK_HEIGHT);
+        assert_value_properties(val)
+    });
+}
+
 // TODO: test note commitment tree round-trip, after implementing proptest::Arbitrary
 
 // Orchard
@@ -436,6 +447,16 @@ fn roundtrip_orchard_tree_root() {
     proptest!(|(val in any::<orchard::tree::Root>())| assert_value_properties(val));
 }
 
+#[test]
+fn roundtrip_orchard_subtree_data() {
+    let _init_guard = zebra_test::init();
+
+    proptest!(|(mut val in any::<NoteCommitmentSubtreeData<orchard::tree::Node>>())| {
+        val.end = val.end.clamp(Height(0), MAX_ON_DISK_HEIGHT);
+        assert_value_properties(val)
+    });
+}
+
 // TODO: test note commitment tree round-trip, after implementing proptest::Arbitrary
 
 // Chain

zebra-state/src/service/finalized_state/disk_format/tests/snapshots/column_family_names.snap

@@ -1,6 +1,6 @@
 ---
 source: zebra-state/src/service/finalized_state/disk_format/tests/snapshot.rs
-assertion_line: 72
+assertion_line: 81
 expression: cf_names
 ---
 [
@@ -12,9 +12,11 @@ expression: cf_names
   "height_by_hash",
   "history_tree",
   "orchard_anchors",
+  "orchard_note_commitment_subtree",
   "orchard_note_commitment_tree",
   "orchard_nullifiers",
   "sapling_anchors",
+  "sapling_note_commitment_subtree",
   "sapling_note_commitment_tree",
   "sapling_nullifiers",
   "sprout_anchors",

zebra-state/src/service/finalized_state/disk_format/tests/snapshots/empty_column_families@mainnet_0.snap

@@ -1,14 +1,15 @@
 ---
 source: zebra-state/src/service/finalized_state/disk_format/tests/snapshot.rs
-assertion_line: 154
 expression: empty_column_families
 ---
 [
   "balance_by_transparent_addr: no entries",
   "history_tree: no entries",
   "orchard_anchors: no entries",
+  "orchard_note_commitment_subtree: no entries",
   "orchard_nullifiers: no entries",
   "sapling_anchors: no entries",
+  "sapling_note_commitment_subtree: no entries",
   "sapling_nullifiers: no entries",
   "sprout_anchors: no entries",
   "sprout_nullifiers: no entries",

zebra-state/src/service/finalized_state/disk_format/tests/snapshots/empty_column_families@mainnet_1.snap

@@ -4,7 +4,9 @@ expression: empty_column_families
 ---
 [
   "history_tree: no entries",
+  "orchard_note_commitment_subtree: no entries",
   "orchard_nullifiers: no entries",
+  "sapling_note_commitment_subtree: no entries",
   "sapling_nullifiers: no entries",
   "sprout_nullifiers: no entries",
 ]

zebra-state/src/service/finalized_state/disk_format/tests/snapshots/empty_column_families@mainnet_2.snap

@@ -4,7 +4,9 @@ expression: empty_column_families
 ---
 [
   "history_tree: no entries",
+  "orchard_note_commitment_subtree: no entries",
   "orchard_nullifiers: no entries",
+  "sapling_note_commitment_subtree: no entries",
   "sapling_nullifiers: no entries",
   "sprout_nullifiers: no entries",
 ]

zebra-state/src/service/finalized_state/disk_format/tests/snapshots/empty_column_families@no_blocks.snap

@@ -1,6 +1,6 @@
 ---
 source: zebra-state/src/service/finalized_state/disk_format/tests/snapshot.rs
-assertion_line: 154
+assertion_line: 166
 expression: empty_column_families
 ---
 [
@@ -11,9 +11,11 @@ expression: empty_column_families
   "height_by_hash: no entries",
   "history_tree: no entries",
   "orchard_anchors: no entries",
+  "orchard_note_commitment_subtree: no entries",
   "orchard_note_commitment_tree: no entries",
   "orchard_nullifiers: no entries",
   "sapling_anchors: no entries",
+  "sapling_note_commitment_subtree: no entries",
   "sapling_note_commitment_tree: no entries",
   "sapling_nullifiers: no entries",
   "sprout_anchors: no entries",

zebra-state/src/service/finalized_state/disk_format/tests/snapshots/empty_column_families@testnet_0.snap

@@ -1,14 +1,15 @@
 ---
 source: zebra-state/src/service/finalized_state/disk_format/tests/snapshot.rs
-assertion_line: 154
 expression: empty_column_families
 ---
 [
   "balance_by_transparent_addr: no entries",
   "history_tree: no entries",
   "orchard_anchors: no entries",
+  "orchard_note_commitment_subtree: no entries",
   "orchard_nullifiers: no entries",
   "sapling_anchors: no entries",
+  "sapling_note_commitment_subtree: no entries",
   "sapling_nullifiers: no entries",
   "sprout_anchors: no entries",
   "sprout_nullifiers: no entries",

zebra-state/src/service/finalized_state/disk_format/tests/snapshots/empty_column_families@testnet_1.snap

@@ -4,7 +4,9 @@ expression: empty_column_families
 ---
 [
   "history_tree: no entries",
+  "orchard_note_commitment_subtree: no entries",
   "orchard_nullifiers: no entries",
+  "sapling_note_commitment_subtree: no entries",
   "sapling_nullifiers: no entries",
   "sprout_nullifiers: no entries",
 ]

zebra-state/src/service/finalized_state/disk_format/tests/snapshots/empty_column_families@testnet_2.snap

@@ -4,7 +4,9 @@ expression: empty_column_families
 ---
 [
   "history_tree: no entries",
+  "orchard_note_commitment_subtree: no entries",
   "orchard_nullifiers: no entries",
+  "sapling_note_commitment_subtree: no entries",
   "sapling_nullifiers: no entries",
   "sprout_nullifiers: no entries",
 ]

zebra-state/src/service/finalized_state/tests/vectors.rs

@@ -10,12 +10,13 @@ use rand::random;
 use halo2::pasta::{group::ff::PrimeField, pallas};
 
 use zebra_chain::{
+    block::Height,
     orchard::{
-        tree::legacy::LegacyNoteCommitmentTree as LegacyOrchardNoteCommitmentTree,
+        self, tree::legacy::LegacyNoteCommitmentTree as LegacyOrchardNoteCommitmentTree,
         tree::NoteCommitmentTree as OrchardNoteCommitmentTree,
     },
     sapling::{
-        tree::legacy::LegacyNoteCommitmentTree as LegacySaplingNoteCommitmentTree,
+        self, tree::legacy::LegacyNoteCommitmentTree as LegacySaplingNoteCommitmentTree,
         tree::NoteCommitmentTree as SaplingNoteCommitmentTree,
     },
     sprout::{
@@ -23,6 +24,7 @@ use zebra_chain::{
         tree::NoteCommitmentTree as SproutNoteCommitmentTree,
         NoteCommitment as SproutNoteCommitment,
     },
+    subtree::NoteCommitmentSubtreeData,
 };
 
 use crate::service::finalized_state::disk_format::{FromDisk, IntoDisk};
@@ -172,8 +174,14 @@ fn sapling_note_commitment_tree_serialization() {
     // The purpose of this test is to make sure the serialization format does
     // not change by accident.
     let expected_serialized_tree_hex = "0102007c3ea01a6e3a3d90cf59cd789e467044b5cd78eb2c84cc6816f960746d0e036c0162324ff2c329e99193a74d28a585a3c167a93bf41a255135529c913bd9b1e66601ddaa1ab86de5c153993414f34ba97e9674c459dfadde112b89eeeafa0e5a204c";
+    let expected_serialized_subtree: &str =
+        "0186a0ddaa1ab86de5c153993414f34ba97e9674c459dfadde112b89eeeafa0e5a204c";
 
-    sapling_checks(incremental_tree, expected_serialized_tree_hex);
+    sapling_checks(
+        incremental_tree,
+        expected_serialized_tree_hex,
+        expected_serialized_subtree,
+    );
 }
 
 /// Check that the sapling tree database serialization format has not changed for one commitment.
@@ -205,8 +213,14 @@ fn sapling_note_commitment_tree_serialization_one() {
     // The purpose of this test is to make sure the serialization format does
     // not change by accident.
     let expected_serialized_tree_hex = "010000225747f3b5d5dab4e5a424f81f85c904ff43286e0f3fd07ef0b8c6a627b1145800012c60c7de033d7539d123fb275011edfe08d57431676981d162c816372063bc71";
+    let expected_serialized_subtree: &str =
+        "0186a02c60c7de033d7539d123fb275011edfe08d57431676981d162c816372063bc71";
 
-    sapling_checks(incremental_tree, expected_serialized_tree_hex);
+    sapling_checks(
+        incremental_tree,
+        expected_serialized_tree_hex,
+        expected_serialized_subtree,
+    );
 }
 
 /// Check that the sapling tree database serialization format has not changed when the number of
@@ -251,8 +265,14 @@ fn sapling_note_commitment_tree_serialization_pow2() {
     // The purpose of this test is to make sure the serialization format does
     // not change by accident.
     let expected_serialized_tree_hex = "010701f43e3aac61e5a753062d4d0508c26ceaf5e4c0c58ba3c956e104b5d2cf67c41c3a3661bc12b72646c94bc6c92796e81953985ee62d80a9ec3645a9a95740ac15025991131c5c25911b35fcea2a8343e2dfd7a4d5b45493390e0cb184394d91c349002df68503da9247dfde6585cb8c9fa94897cf21735f8fc1b32116ef474de05c01d23765f3d90dfd97817ed6d995bd253d85967f77b9f1eaef6ecbcb0ef6796812";
+    let expected_serialized_subtree =
+        "0186a0d23765f3d90dfd97817ed6d995bd253d85967f77b9f1eaef6ecbcb0ef6796812";
 
-    sapling_checks(incremental_tree, expected_serialized_tree_hex);
+    sapling_checks(
+        incremental_tree,
+        expected_serialized_tree_hex,
+        expected_serialized_subtree,
+    );
 }
 
 /// Check that the orchard tree database serialization format has not changed.
@@ -298,8 +318,14 @@ fn orchard_note_commitment_tree_serialization() {
     // The purpose of this test is to make sure the serialization format does
     // not change by accident.
     let expected_serialized_tree_hex = "010200ee9488053a30c596b43014105d3477e6f578c89240d1d1ee1743b77bb6adc40a01a34b69a4e4d9ccf954d46e5da1004d361a5497f511aeb4d481d23c0be177813301a0be6dab19bc2c65d8299258c16e14d48ec4d4959568c6412aa85763c222a702";
+    let expected_serialized_subtree =
+        "0186a0a0be6dab19bc2c65d8299258c16e14d48ec4d4959568c6412aa85763c222a702";
 
-    orchard_checks(incremental_tree, expected_serialized_tree_hex);
+    orchard_checks(
+        incremental_tree,
+        expected_serialized_tree_hex,
+        expected_serialized_subtree,
+    );
 }
 
 /// Check that the orchard tree database serialization format has not changed for one commitment.
@@ -333,8 +359,14 @@ fn orchard_note_commitment_tree_serialization_one() {
     // The purpose of this test is to make sure the serialization format does
     // not change by accident.
     let expected_serialized_tree_hex = "01000068135cf49933229099a44ec99a75e1e1cb4640f9b5bdec6b3223856fea16390a000178afd4da59c541e9c2f317f9aff654f1fb38d14dc99431cbbfa93601c7068117";
+    let expected_serialized_subtree =
+        "0186a078afd4da59c541e9c2f317f9aff654f1fb38d14dc99431cbbfa93601c7068117";
 
-    orchard_checks(incremental_tree, expected_serialized_tree_hex);
+    orchard_checks(
+        incremental_tree,
+        expected_serialized_tree_hex,
+        expected_serialized_subtree,
+    );
 }
 
 /// Check that the orchard tree database serialization format has not changed when the number of
@@ -379,8 +411,14 @@ fn orchard_note_commitment_tree_serialization_pow2() {
     // The purpose of this test is to make sure the serialization format does
     // not change by accident.
     let expected_serialized_tree_hex = "01010178315008fb2998b430a5731d6726207dc0f0ec81ea64af5cf612956901e72f0eee9488053a30c596b43014105d3477e6f578c89240d1d1ee1743b77bb6adc40a0001d3d525931005e45f5a29bc82524e871e5ee1b6d77839deb741a6e50cd99fdf1a";
+    let expected_serialized_subtree =
+        "0186a0d3d525931005e45f5a29bc82524e871e5ee1b6d77839deb741a6e50cd99fdf1a";
 
-    orchard_checks(incremental_tree, expected_serialized_tree_hex);
+    orchard_checks(
+        incremental_tree,
+        expected_serialized_tree_hex,
+        expected_serialized_subtree,
+    );
 }
 
 fn sprout_checks(incremental_tree: SproutNoteCommitmentTree, expected_serialized_tree_hex: &str) {
@@ -433,8 +471,12 @@ fn sprout_checks(incremental_tree: SproutNoteCommitmentTree, expected_serialized
     assert_eq!(re_serialized_legacy_tree, re_serialized_tree);
 }
 
-fn sapling_checks(incremental_tree: SaplingNoteCommitmentTree, expected_serialized_tree_hex: &str) {
+fn sapling_checks(
-    let serialized_tree = incremental_tree.as_bytes();
+    incremental_tree: SaplingNoteCommitmentTree,
+    expected_serialized_tree_hex: &str,
+    expected_serialized_subtree: &str,
+) {
+    let serialized_tree: Vec<u8> = incremental_tree.as_bytes();
 
     assert_eq!(hex::encode(&serialized_tree), expected_serialized_tree_hex);
 
@@ -481,9 +523,35 @@ fn sapling_checks(incremental_tree: SaplingNoteCommitmentTree, expected_serializ
 
     assert_eq!(serialized_tree, re_serialized_tree);
     assert_eq!(re_serialized_legacy_tree, re_serialized_tree);
+
+    // Check subtree format
+
+    let subtree = NoteCommitmentSubtreeData::new(
+        Height(100000),
+        sapling::tree::Node::from_bytes(incremental_tree.hash()),
+    );
+
+    let serialized_subtree = subtree.as_bytes();
+
+    assert_eq!(
+        hex::encode(&serialized_subtree),
+        expected_serialized_subtree
+    );
+
+    let deserialized_subtree =
+        NoteCommitmentSubtreeData::<sapling::tree::Node>::from_bytes(&serialized_subtree);
+
+    assert_eq!(
+        subtree, deserialized_subtree,
+        "(de)serialization should not modify subtree value"
+    );
 }
 
-fn orchard_checks(incremental_tree: OrchardNoteCommitmentTree, expected_serialized_tree_hex: &str) {
+fn orchard_checks(
+    incremental_tree: OrchardNoteCommitmentTree,
+    expected_serialized_tree_hex: &str,
+    expected_serialized_subtree: &str,
+) {
     let serialized_tree = incremental_tree.as_bytes();
 
     assert_eq!(hex::encode(&serialized_tree), expected_serialized_tree_hex);
@@ -531,4 +599,26 @@ fn orchard_checks(incremental_tree: OrchardNoteCommitmentTree, expected_serializ
 
     assert_eq!(serialized_tree, re_serialized_tree);
     assert_eq!(re_serialized_legacy_tree, re_serialized_tree);
+
+    // Check subtree format
+
+    let subtree = NoteCommitmentSubtreeData::new(
+        Height(100000),
+        orchard::tree::Node::from_bytes(incremental_tree.hash()),
+    );
+
+    let serialized_subtree = subtree.as_bytes();
+
+    assert_eq!(
+        hex::encode(&serialized_subtree),
+        expected_serialized_subtree
+    );
+
+    let deserialized_subtree =
+        NoteCommitmentSubtreeData::<orchard::tree::Node>::from_bytes(&serialized_subtree);
+
+    assert_eq!(
+        subtree, deserialized_subtree,
+        "(de)serialization should not modify subtree value"
+    );
 }

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

@@ -15,7 +15,11 @@
 use std::{collections::HashMap, sync::Arc};
 
 use zebra_chain::{
-    block::Height, orchard, parallel::tree::NoteCommitmentTrees, sapling, sprout,
+    block::Height,
+    orchard,
+    parallel::tree::NoteCommitmentTrees,
+    sapling, sprout,
+    subtree::{NoteCommitmentSubtree, NoteCommitmentSubtreeData, NoteCommitmentSubtreeIndex},
     transaction::Transaction,
 };
 
@@ -178,6 +182,22 @@ impl ZebraDb {
         Some(Arc::new(tree))
     }
 
+    /// Returns the Sapling note commitment subtree at this index
+    #[allow(clippy::unwrap_in_result)]
+    pub fn sapling_subtree_by_index(
+        &self,
+        index: impl Into<NoteCommitmentSubtreeIndex> + Copy,
+    ) -> Option<Arc<NoteCommitmentSubtree<sapling::tree::Node>>> {
+        let sapling_subtrees = self
+            .db
+            .cf_handle("sapling_note_commitment_subtree")
+            .unwrap();
+
+        let subtree_data: NoteCommitmentSubtreeData<sapling::tree::Node> =
+            self.db.zs_get(&sapling_subtrees, &index.into())?;
+        Some(subtree_data.with_index(index))
+    }
+
     /// Returns the Orchard note commitment tree of the finalized tip
     /// or the empty tree if the state is empty.
     pub fn orchard_tree(&self) -> Arc<orchard::tree::NoteCommitmentTree> {
@@ -190,6 +210,22 @@ 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)]
@@ -312,6 +348,9 @@ impl DiskWriteBatch {
         let sapling_tree_cf = db.cf_handle("sapling_note_commitment_tree").unwrap();
         let orchard_tree_cf = db.cf_handle("orchard_note_commitment_tree").unwrap();
 
+        let _sapling_subtree_cf = db.cf_handle("sapling_note_commitment_subtree").unwrap();
+        let _orchard_subtree_cf = db.cf_handle("orchard_note_commitment_subtree").unwrap();
+
         let height = finalized.verified.height;
         let trees = finalized.treestate.note_commitment_trees.clone();
 
@@ -357,6 +396,16 @@ impl DiskWriteBatch {
             self.zs_insert(&orchard_tree_cf, height, trees.orchard);
         }
 
+        // TODO: Increment DATABASE_FORMAT_MINOR_VERSION and uncomment these insertions
+
+        // if let Some(subtree) = trees.sapling_subtree {
+        //     self.zs_insert(&sapling_subtree_cf, subtree.index, subtree.into_data());
+        // }
+
+        // if let Some(subtree) = trees.orchard_subtree {
+        //     self.zs_insert(&orchard_subtree_cf, subtree.index, subtree.into_data());
+        // }
+
         self.prepare_history_batch(db, finalized)
     }
 

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

@@ -3,7 +3,7 @@
 
 use std::{
     cmp::Ordering,
-    collections::{BTreeMap, BTreeSet, HashMap, HashSet},
+    collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque},
     ops::{Deref, RangeInclusive},
     sync::Arc,
 };
@@ -20,6 +20,7 @@ use zebra_chain::{
     parameters::Network,
     primitives::Groth16Proof,
     sapling, sprout,
+    subtree::{NoteCommitmentSubtree, NoteCommitmentSubtreeIndex},
     transaction::Transaction::*,
     transaction::{self, Transaction},
     transparent,
@@ -133,6 +134,8 @@ pub struct Chain {
     /// When a chain is forked from the finalized tip, also contains the finalized tip root.
     /// 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>>>,
 
     /// The Orchard anchors created by `blocks`.
     ///
@@ -144,6 +147,8 @@ pub struct Chain {
     /// When a chain is forked from the finalized tip, also contains the finalized tip root.
     /// 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>>>,
 
     // Nullifiers
     //
@@ -221,9 +226,11 @@ impl Chain {
             sapling_anchors: MultiSet::new(),
             sapling_anchors_by_height: Default::default(),
             sapling_trees_by_height: Default::default(),
+            sapling_subtrees: Default::default(),
             orchard_anchors: MultiSet::new(),
             orchard_anchors_by_height: Default::default(),
             orchard_trees_by_height: Default::default(),
+            orchard_subtrees: Default::default(),
             sprout_nullifiers: Default::default(),
             sapling_nullifiers: Default::default(),
             orchard_nullifiers: Default::default(),
@@ -343,6 +350,14 @@ impl Chain {
             .treestate(block_height.into())
             .expect("The treestate must be present for the root height.");
 
+        if treestate.note_commitment_trees.sapling_subtree.is_some() {
+            self.sapling_subtrees.pop_front();
+        }
+
+        if treestate.note_commitment_trees.orchard_subtree.is_some() {
+            self.orchard_subtrees.pop_front();
+        }
+
         // Remove the lowest height block from `self.blocks`.
         let block = self
             .blocks
@@ -663,6 +678,33 @@ impl Chain {
             .map(|(_height, tree)| tree.clone())
     }
 
+    /// Returns the Sapling [`NoteCommitmentSubtree`] specified
+    /// by an index, if it exists in the non-finalized [`Chain`].
+    pub fn sapling_subtree(
+        &self,
+        hash_or_height: HashOrHeight,
+    ) -> Option<Arc<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)
+            .cloned()
+    }
+
+    /// Returns the Sapling [`NoteCommitmentSubtree`](sapling::tree::NoteCommitmentSubtree) specified
+    /// by a [`HashOrHeight`], if it exists in the non-finalized [`Chain`].
+    pub fn sapling_subtree_by_index(
+        &self,
+        index: NoteCommitmentSubtreeIndex,
+    ) -> Option<Arc<NoteCommitmentSubtree<sapling::tree::Node>>> {
+        self.sapling_subtrees
+            .iter()
+            .find(|subtree| subtree.index == index)
+            .cloned()
+    }
+
     /// Adds the Sapling `tree` to the tree and anchor indexes at `height`.
     ///
     /// `height` can be either:
@@ -812,6 +854,33 @@ impl Chain {
             .map(|(_height, tree)| tree.clone())
     }
 
+    /// Returns the Orchard [`NoteCommitmentSubtree`](orchard::tree::NoteCommitmentSubtree) specified
+    /// by a [`HashOrHeight`], if it exists in the non-finalized [`Chain`].
+    pub fn orchard_subtree(
+        &self,
+        hash_or_height: HashOrHeight,
+    ) -> Option<Arc<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)
+            .cloned()
+    }
+
+    /// Returns the Orchard [`NoteCommitmentSubtree`] specified
+    /// by an index, if it exists in the non-finalized [`Chain`].
+    pub fn orchard_subtree_by_index(
+        &self,
+        index: NoteCommitmentSubtreeIndex,
+    ) -> Option<Arc<NoteCommitmentSubtree<orchard::tree::Node>>> {
+        self.orchard_subtrees
+            .iter()
+            .find(|subtree| subtree.index == index)
+            .cloned()
+    }
+
     /// Adds the Orchard `tree` to the tree and anchor indexes at `height`.
     ///
     /// `height` can be either:
@@ -1004,11 +1073,15 @@ impl Chain {
         let sapling_tree = self.sapling_tree(hash_or_height)?;
         let orchard_tree = self.orchard_tree(hash_or_height)?;
         let history_tree = self.history_tree(hash_or_height)?;
+        let sapling_subtree = self.sapling_subtree(hash_or_height);
+        let orchard_subtree = self.orchard_subtree(hash_or_height);
 
         Some(Treestate::new(
             sprout_tree,
             sapling_tree,
             orchard_tree,
+            sapling_subtree,
+            orchard_subtree,
             history_tree,
         ))
     }
@@ -1280,6 +1353,13 @@ impl Chain {
         self.add_sapling_tree_and_anchor(height, nct.sapling);
         self.add_orchard_tree_and_anchor(height, nct.orchard);
 
+        if let Some(subtree) = nct.sapling_subtree {
+            self.sapling_subtrees.push_back(subtree)
+        }
+        if let Some(subtree) = nct.orchard_subtree {
+            self.orchard_subtrees.push_back(subtree)
+        }
+
         let sapling_root = self.sapling_note_commitment_tree().root();
         let orchard_root = self.orchard_note_commitment_tree().root();
 

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

@@ -13,7 +13,10 @@
 
 use std::sync::Arc;
 
-use zebra_chain::{orchard, sapling};
+use zebra_chain::{
+    orchard, sapling,
+    subtree::{NoteCommitmentSubtree, NoteCommitmentSubtreeIndex},
+};
 
 use crate::{
     service::{finalized_state::ZebraDb, non_finalized_state::Chain},
@@ -41,6 +44,28 @@ where
         .or_else(|| db.sapling_tree_by_hash_or_height(hash_or_height))
 }
 
+/// Returns the Sapling
+/// [`NoteCommitmentSubtree`](NoteCommitmentSubtree) specified by an
+/// index, if it exists in the non-finalized `chain` or finalized `db`.
+#[allow(unused)]
+pub fn sapling_subtree<C>(
+    chain: Option<C>,
+    db: &ZebraDb,
+    index: NoteCommitmentSubtreeIndex,
+) -> Option<Arc<NoteCommitmentSubtree<sapling::tree::Node>>>
+where
+    C: AsRef<Chain>,
+{
+    // # Correctness
+    //
+    // Since sapling treestates are the same in the finalized and non-finalized
+    // state, we check the most efficient alternative first. (`chain` is always
+    // in memory, but `db` stores blocks on disk, with a memory cache.)
+    chain
+        .and_then(|chain| chain.as_ref().sapling_subtree_by_index(index))
+        .or_else(|| db.sapling_subtree_by_index(index))
+}
+
 /// Returns the Orchard
 /// [`NoteCommitmentTree`](orchard::tree::NoteCommitmentTree) specified by a
 /// hash or height, if it exists in the non-finalized `chain` or finalized `db`.
@@ -62,6 +87,28 @@ where
         .or_else(|| db.orchard_tree_by_hash_or_height(hash_or_height))
 }
 
+/// Returns the Orchard
+/// [`NoteCommitmentSubtree`](NoteCommitmentSubtree) specified by an
+/// index, if it exists in the non-finalized `chain` or finalized `db`.
+#[allow(unused)]
+pub fn orchard_subtree<C>(
+    chain: Option<C>,
+    db: &ZebraDb,
+    index: NoteCommitmentSubtreeIndex,
+) -> Option<Arc<NoteCommitmentSubtree<orchard::tree::Node>>>
+where
+    C: AsRef<Chain>,
+{
+    // # Correctness
+    //
+    // Since orchard treestates are the same in the finalized and non-finalized
+    // state, we check the most efficient alternative first. (`chain` is always
+    // in memory, but `db` stores blocks on disk, with a memory cache.)
+    chain
+        .and_then(|chain| chain.as_ref().orchard_subtree_by_index(index))
+        .or_else(|| db.orchard_subtree_by_index(index))
+}
+
 #[cfg(feature = "getblocktemplate-rpcs")]
 /// Get the history tree of the provided chain.
 pub fn history_tree<C>(