feat(shielded): Store Sapling & Orchard note commitment trees in finalized and non-finalized state (#3818)

* Query Sapling & Orchard trees by height in the finalized state

* Add Sapling & Orchard trees to the non-finalized state

* Add a TODO about concurrent read-only access to Sprout tree

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

* Update the database format version

* Keep only the most recent Sprout tree in the database

* Check that the database returns empty trees for the genesis block

* Assert that the database returns the highest trees

* Document how to update insta snapshots

* Add note commitment tree insta snapshot tests

* Add comments about cached tree roots in snapshots

* Add snapshot data for sapling and orchard trees

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: teor <teor@riseup.net>

zebra-state/src/constants.rs

@@ -18,7 +18,7 @@ pub use zebra_chain::transparent::MIN_TRANSPARENT_COINBASE_MATURITY;
 pub const MAX_BLOCK_REORG_HEIGHT: u32 = MIN_TRANSPARENT_COINBASE_MATURITY - 1;
 
 /// The database format version, incremented each time the database format changes.
-pub const DATABASE_FORMAT_VERSION: u32 = 12;
+pub const DATABASE_FORMAT_VERSION: u32 = 13;
 
 /// The maximum number of blocks to check for NU5 transactions,
 /// before we assume we are on a pre-NU5 legacy chain.

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

@@ -9,8 +9,11 @@
 //!
 //! # Fixing Test Failures
 //!
-//! If this test fails, run `cargo insta review` to update the test snapshots,
+//! If this test fails, run:
-//! then commit the `test_*.snap` files using git.
+//! ```sh
+//! cargo insta test --review --delete-unreferenced-snapshots
+//! ```
+//! to update the test snapshots, then commit the `test_*.snap` files using git.
 //!
 //! # Snapshot Format
 //!
@@ -134,6 +137,9 @@ fn snapshot_raw_rocksdb_column_family_data(db: &DiskDb, original_cf_names: &[Str
             // distinguish column family names from empty column families
             empty_column_families.push(format!("{}: no entries", cf_name));
         } else {
+            // The note commitment tree snapshots will change if the trees do not have cached roots.
+            // But we expect them to always have cached roots,
+            // because those roots are used to populate the anchor column families.
             insta::assert_ron_snapshot!(format!("{}_raw_data", cf_name), cf_data);
         }
 

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

@@ -1,10 +1,14 @@
 ---
 source: zebra-state/src/service/finalized_state/disk_format/tests/snapshot.rs
-assertion_line: 125
+assertion_line: 137
 expression: cf_data
 
 ---
 [
+  KV(
+    k: "00000000",
+    v: "0000",
+  ),
   KV(
     k: "00000001",
     v: "0001ae2935f1dfd8a24aed7c70df7de3a668eb7a49b1319880dde2bbd9031ae5d82f",

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

@@ -1,10 +1,18 @@
 ---
 source: zebra-state/src/service/finalized_state/disk_format/tests/snapshot.rs
-assertion_line: 125
+assertion_line: 137
 expression: cf_data
 
 ---
 [
+  KV(
+    k: "00000000",
+    v: "0000",
+  ),
+  KV(
+    k: "00000001",
+    v: "0001ae2935f1dfd8a24aed7c70df7de3a668eb7a49b1319880dde2bbd9031ae5d82f",
+  ),
   KV(
     k: "00000002",
     v: "0001ae2935f1dfd8a24aed7c70df7de3a668eb7a49b1319880dde2bbd9031ae5d82f",

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

@@ -1,10 +1,14 @@
 ---
 source: zebra-state/src/service/finalized_state/disk_format/tests/snapshot.rs
-assertion_line: 125
+assertion_line: 137
 expression: cf_data
 
 ---
 [
+  KV(
+    k: "00000000",
+    v: "0000",
+  ),
   KV(
     k: "00000001",
     v: "0001ae2935f1dfd8a24aed7c70df7de3a668eb7a49b1319880dde2bbd9031ae5d82f",

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

@@ -1,10 +1,18 @@
 ---
 source: zebra-state/src/service/finalized_state/disk_format/tests/snapshot.rs
-assertion_line: 125
+assertion_line: 137
 expression: cf_data
 
 ---
 [
+  KV(
+    k: "00000000",
+    v: "0000",
+  ),
+  KV(
+    k: "00000001",
+    v: "0001ae2935f1dfd8a24aed7c70df7de3a668eb7a49b1319880dde2bbd9031ae5d82f",
+  ),
   KV(
     k: "00000002",
     v: "0001ae2935f1dfd8a24aed7c70df7de3a668eb7a49b1319880dde2bbd9031ae5d82f",

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

@@ -1,10 +1,14 @@
 ---
 source: zebra-state/src/service/finalized_state/disk_format/tests/snapshot.rs
-assertion_line: 125
+assertion_line: 137
 expression: cf_data
 
 ---
 [
+  KV(
+    k: "00000000",
+    v: "0000",
+  ),
   KV(
     k: "00000001",
     v: "0001fbc2f4300c01f0b7820d00e3347c8da4ee614674376cbc45359daa54f9b5493e",

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

@@ -1,10 +1,18 @@
 ---
 source: zebra-state/src/service/finalized_state/disk_format/tests/snapshot.rs
-assertion_line: 125
+assertion_line: 137
 expression: cf_data
 
 ---
 [
+  KV(
+    k: "00000000",
+    v: "0000",
+  ),
+  KV(
+    k: "00000001",
+    v: "0001fbc2f4300c01f0b7820d00e3347c8da4ee614674376cbc45359daa54f9b5493e",
+  ),
   KV(
     k: "00000002",
     v: "0001fbc2f4300c01f0b7820d00e3347c8da4ee614674376cbc45359daa54f9b5493e",

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

@@ -1,10 +1,14 @@
 ---
 source: zebra-state/src/service/finalized_state/disk_format/tests/snapshot.rs
-assertion_line: 125
+assertion_line: 137
 expression: cf_data
 
 ---
 [
+  KV(
+    k: "00000000",
+    v: "0000",
+  ),
   KV(
     k: "00000001",
     v: "0001fbc2f4300c01f0b7820d00e3347c8da4ee614674376cbc45359daa54f9b5493e",

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

@@ -1,10 +1,18 @@
 ---
 source: zebra-state/src/service/finalized_state/disk_format/tests/snapshot.rs
-assertion_line: 125
+assertion_line: 137
 expression: cf_data
 
 ---
 [
+  KV(
+    k: "00000000",
+    v: "0000",
+  ),
+  KV(
+    k: "00000001",
+    v: "0001fbc2f4300c01f0b7820d00e3347c8da4ee614674376cbc45359daa54f9b5493e",
+  ),
   KV(
     k: "00000002",
     v: "0001fbc2f4300c01f0b7820d00e3347c8da4ee614674376cbc45359daa54f9b5493e",

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

@@ -1,8 +1,7 @@
 ---
 source: zebra-state/src/service/finalized_state/disk_format/tests/snapshot.rs
-assertion_line: 125
+assertion_line: 137
 expression: cf_data
-
 ---
 [
   KV(

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

@@ -1,8 +1,7 @@
 ---
 source: zebra-state/src/service/finalized_state/disk_format/tests/snapshot.rs
-assertion_line: 125
+assertion_line: 137
 expression: cf_data
-
 ---
 [
   KV(

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

@@ -1,8 +1,7 @@
 ---
 source: zebra-state/src/service/finalized_state/disk_format/tests/snapshot.rs
-assertion_line: 125
+assertion_line: 137
 expression: cf_data
-
 ---
 [
   KV(

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

@@ -1,8 +1,7 @@
 ---
 source: zebra-state/src/service/finalized_state/disk_format/tests/snapshot.rs
-assertion_line: 125
+assertion_line: 137
 expression: cf_data
-
 ---
 [
   KV(

zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshot.rs

@@ -24,12 +24,16 @@
 //!
 //! # Fixing Test Failures
 //!
-//! If this test fails, run `cargo insta review` to update the test snapshots,
+//! If this test fails, run:
-//! then commit the `test_*.snap` files using git.
+//! ```sh
+//! cargo insta test --review --delete-unreferenced-snapshots
+//! ```
+//! to update the test snapshots, then commit the `test_*.snap` files using git.
 //!
 //! # TODO
 //!
-//! Test shielded data, and data activated in Overwinter and later network upgrades.
+//! Test the rest of the shielded data,
+//! and data activated in Overwinter and later network upgrades.
 
 use std::sync::Arc;
 
@@ -37,7 +41,9 @@ use serde::{Deserialize, Serialize};
 
 use zebra_chain::{
     block::{self, Block, Height},
+    orchard,
     parameters::Network::{self, *},
+    sapling,
     serialization::{ZcashDeserializeInto, ZcashSerialize},
     transaction::Transaction,
 };
@@ -196,9 +202,24 @@ fn snapshot_block_and_transaction_data(state: &FinalizedState) {
     insta::assert_ron_snapshot!("tip", tip.map(Tip::from));
 
     if let Some((max_height, tip_block_hash)) = tip {
+        // Check that the database returns empty note commitment trees for the
+        // genesis block.
+        let sapling_tree = state
+            .sapling_note_commitment_tree_by_height(&block::Height::MIN)
+            .expect("the genesis block in the database has a Sapling tree");
+        let orchard_tree = state
+            .orchard_note_commitment_tree_by_height(&block::Height::MIN)
+            .expect("the genesis block in the database has an Orchard tree");
+
+        assert_eq!(sapling_tree, sapling::tree::NoteCommitmentTree::default());
+        assert_eq!(orchard_tree, orchard::tree::NoteCommitmentTree::default());
+
         let mut stored_block_hashes = Vec::new();
         let mut stored_blocks = Vec::new();
 
+        let mut stored_sapling_trees = Vec::new();
+        let mut stored_orchard_trees = Vec::new();
+
         let mut stored_transaction_hashes = Vec::new();
         let mut stored_transactions = Vec::new();
 
@@ -216,6 +237,15 @@ fn snapshot_block_and_transaction_data(state: &FinalizedState) {
                 .block(query_height.into())
                 .expect("heights up to tip have blocks");
 
+            let sapling_tree_by_height = state
+                .sapling_note_commitment_tree_by_height(&query_height)
+                .expect("heights up to tip have Sapling trees");
+            let orchard_tree_by_height = state
+                .orchard_note_commitment_tree_by_height(&query_height)
+                .expect("heights up to tip have Orchard trees");
+            let sapling_tree_at_tip = state.sapling_note_commitment_tree();
+            let orchard_tree_at_tip = state.db.orchard_note_commitment_tree();
+
             // We don't need to snapshot the heights,
             // because they are fully determined by the tip and block hashes.
             //
@@ -224,6 +254,9 @@ fn snapshot_block_and_transaction_data(state: &FinalizedState) {
 
             if query_height == max_height {
                 assert_eq!(stored_block_hash, tip_block_hash);
+
+                assert_eq!(sapling_tree_at_tip, sapling_tree_by_height);
+                assert_eq!(orchard_tree_at_tip, orchard_tree_by_height);
             }
 
             assert_eq!(
@@ -236,6 +269,9 @@ fn snapshot_block_and_transaction_data(state: &FinalizedState) {
             stored_block_hashes.push((stored_height, BlockHash(stored_block_hash.to_string())));
             stored_blocks.push(BlockData::new(stored_height, &stored_block));
 
+            stored_sapling_trees.push((stored_height, sapling_tree_by_height));
+            stored_orchard_trees.push((stored_height, orchard_tree_by_height));
+
             // Check block transaction hashes and transactions.
             for tx_index in 0..stored_block.transactions.len() {
                 let transaction = &stored_block.transactions[tx_index];
@@ -256,6 +292,7 @@ fn snapshot_block_and_transaction_data(state: &FinalizedState) {
             stored_block_hashes
         );
         assert!(is_sorted(&stored_blocks), "unsorted: {:?}", stored_blocks);
+
         assert!(
             is_sorted(&stored_transaction_hashes),
             "unsorted: {:?}",
@@ -267,10 +304,18 @@ fn snapshot_block_and_transaction_data(state: &FinalizedState) {
             stored_transactions
         );
 
-        // The blocks, transactions, and their hashes are in height/index order, and we want to snapshot that order.
+        // The blocks, trees, transactions, and their hashes are in height/index order,
+        // and we want to snapshot that order.
         // So we don't sort the vectors before snapshotting.
         insta::assert_ron_snapshot!("block_hashes", stored_block_hashes);
         insta::assert_ron_snapshot!("blocks", stored_blocks);
+
+        // These snapshots will change if the trees do not have cached roots.
+        // But we expect them to always have cached roots,
+        // because those roots are used to populate the anchor column families.
+        insta::assert_ron_snapshot!("sapling_trees", stored_sapling_trees);
+        insta::assert_ron_snapshot!("orchard_trees", stored_orchard_trees);
+
         insta::assert_ron_snapshot!("transaction_hashes", stored_transaction_hashes);
         insta::assert_ron_snapshot!("transactions", stored_transactions);
     }

zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshots/orchard_trees@mainnet_0.snap

@@ -0,0 +1,14 @@
+---
+source: zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshot.rs
+expression: stored_orchard_trees
+---
+[
+  (Height(0), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: Some(Root(Base(
+      bytes: (174, 41, 53, 241, 223, 216, 162, 74, 237, 124, 112, 223, 125, 227, 166, 104, 235, 122, 73, 177, 49, 152, 128, 221, 226, 187, 217, 3, 26, 229, 216, 47),
+    ))),
+  )),
+]

zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshots/orchard_trees@mainnet_1.snap

@@ -0,0 +1,20 @@
+---
+source: zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshot.rs
+expression: stored_orchard_trees
+---
+[
+  (Height(0), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: None,
+  )),
+  (Height(1), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: Some(Root(Base(
+      bytes: (174, 41, 53, 241, 223, 216, 162, 74, 237, 124, 112, 223, 125, 227, 166, 104, 235, 122, 73, 177, 49, 152, 128, 221, 226, 187, 217, 3, 26, 229, 216, 47),
+    ))),
+  )),
+]

zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshots/orchard_trees@mainnet_2.snap

@@ -0,0 +1,28 @@
+---
+source: zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshot.rs
+expression: stored_orchard_trees
+---
+[
+  (Height(0), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: None,
+  )),
+  (Height(1), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: Some(Root(Base(
+      bytes: (174, 41, 53, 241, 223, 216, 162, 74, 237, 124, 112, 223, 125, 227, 166, 104, 235, 122, 73, 177, 49, 152, 128, 221, 226, 187, 217, 3, 26, 229, 216, 47),
+    ))),
+  )),
+  (Height(2), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: Some(Root(Base(
+      bytes: (174, 41, 53, 241, 223, 216, 162, 74, 237, 124, 112, 223, 125, 227, 166, 104, 235, 122, 73, 177, 49, 152, 128, 221, 226, 187, 217, 3, 26, 229, 216, 47),
+    ))),
+  )),
+]

zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshots/orchard_trees@testnet_0.snap

@@ -0,0 +1,14 @@
+---
+source: zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshot.rs
+expression: stored_orchard_trees
+---
+[
+  (Height(0), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: Some(Root(Base(
+      bytes: (174, 41, 53, 241, 223, 216, 162, 74, 237, 124, 112, 223, 125, 227, 166, 104, 235, 122, 73, 177, 49, 152, 128, 221, 226, 187, 217, 3, 26, 229, 216, 47),
+    ))),
+  )),
+]

zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshots/orchard_trees@testnet_1.snap

@@ -0,0 +1,20 @@
+---
+source: zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshot.rs
+expression: stored_orchard_trees
+---
+[
+  (Height(0), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: None,
+  )),
+  (Height(1), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: Some(Root(Base(
+      bytes: (174, 41, 53, 241, 223, 216, 162, 74, 237, 124, 112, 223, 125, 227, 166, 104, 235, 122, 73, 177, 49, 152, 128, 221, 226, 187, 217, 3, 26, 229, 216, 47),
+    ))),
+  )),
+]

zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshots/orchard_trees@testnet_2.snap

@@ -0,0 +1,28 @@
+---
+source: zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshot.rs
+expression: stored_orchard_trees
+---
+[
+  (Height(0), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: None,
+  )),
+  (Height(1), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: Some(Root(Base(
+      bytes: (174, 41, 53, 241, 223, 216, 162, 74, 237, 124, 112, 223, 125, 227, 166, 104, 235, 122, 73, 177, 49, 152, 128, 221, 226, 187, 217, 3, 26, 229, 216, 47),
+    ))),
+  )),
+  (Height(2), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: Some(Root(Base(
+      bytes: (174, 41, 53, 241, 223, 216, 162, 74, 237, 124, 112, 223, 125, 227, 166, 104, 235, 122, 73, 177, 49, 152, 128, 221, 226, 187, 217, 3, 26, 229, 216, 47),
+    ))),
+  )),
+]

zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshots/sapling_trees@mainnet_0.snap

@@ -0,0 +1,14 @@
+---
+source: zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshot.rs
+expression: stored_sapling_trees
+---
+[
+  (Height(0), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: Some(Root(Fq(
+      bytes: (251, 194, 244, 48, 12, 1, 240, 183, 130, 13, 0, 227, 52, 124, 141, 164, 238, 97, 70, 116, 55, 108, 188, 69, 53, 157, 170, 84, 249, 181, 73, 62),
+    ))),
+  )),
+]

zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshots/sapling_trees@mainnet_1.snap

@@ -0,0 +1,20 @@
+---
+source: zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshot.rs
+expression: stored_sapling_trees
+---
+[
+  (Height(0), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: None,
+  )),
+  (Height(1), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: Some(Root(Fq(
+      bytes: (251, 194, 244, 48, 12, 1, 240, 183, 130, 13, 0, 227, 52, 124, 141, 164, 238, 97, 70, 116, 55, 108, 188, 69, 53, 157, 170, 84, 249, 181, 73, 62),
+    ))),
+  )),
+]

zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshots/sapling_trees@mainnet_2.snap

@@ -0,0 +1,28 @@
+---
+source: zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshot.rs
+expression: stored_sapling_trees
+---
+[
+  (Height(0), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: None,
+  )),
+  (Height(1), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: Some(Root(Fq(
+      bytes: (251, 194, 244, 48, 12, 1, 240, 183, 130, 13, 0, 227, 52, 124, 141, 164, 238, 97, 70, 116, 55, 108, 188, 69, 53, 157, 170, 84, 249, 181, 73, 62),
+    ))),
+  )),
+  (Height(2), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: Some(Root(Fq(
+      bytes: (251, 194, 244, 48, 12, 1, 240, 183, 130, 13, 0, 227, 52, 124, 141, 164, 238, 97, 70, 116, 55, 108, 188, 69, 53, 157, 170, 84, 249, 181, 73, 62),
+    ))),
+  )),
+]

zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshots/sapling_trees@testnet_0.snap

@@ -0,0 +1,14 @@
+---
+source: zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshot.rs
+expression: stored_sapling_trees
+---
+[
+  (Height(0), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: Some(Root(Fq(
+      bytes: (251, 194, 244, 48, 12, 1, 240, 183, 130, 13, 0, 227, 52, 124, 141, 164, 238, 97, 70, 116, 55, 108, 188, 69, 53, 157, 170, 84, 249, 181, 73, 62),
+    ))),
+  )),
+]

zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshots/sapling_trees@testnet_1.snap

@@ -0,0 +1,20 @@
+---
+source: zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshot.rs
+expression: stored_sapling_trees
+---
+[
+  (Height(0), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: None,
+  )),
+  (Height(1), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: Some(Root(Fq(
+      bytes: (251, 194, 244, 48, 12, 1, 240, 183, 130, 13, 0, 227, 52, 124, 141, 164, 238, 97, 70, 116, 55, 108, 188, 69, 53, 157, 170, 84, 249, 181, 73, 62),
+    ))),
+  )),
+]

zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshots/sapling_trees@testnet_2.snap

@@ -0,0 +1,28 @@
+---
+source: zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshot.rs
+expression: stored_sapling_trees
+---
+[
+  (Height(0), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: None,
+  )),
+  (Height(1), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: Some(Root(Fq(
+      bytes: (251, 194, 244, 48, 12, 1, 240, 183, 130, 13, 0, 227, 52, 124, 141, 164, 238, 97, 70, 116, 55, 108, 188, 69, 53, 157, 170, 84, 249, 181, 73, 62),
+    ))),
+  )),
+  (Height(2), NoteCommitmentTree(
+    inner: Frontier(
+      frontier: None,
+    ),
+    cached_root: Some(Root(Fq(
+      bytes: (251, 194, 244, 48, 12, 1, 240, 183, 130, 13, 0, 227, 52, 124, 141, 164, 238, 97, 70, 116, 55, 108, 188, 69, 53, 157, 170, 84, 249, 181, 73, 62),
+    ))),
+  )),
+]

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

@@ -13,7 +13,7 @@
 //! be incremented each time the database format (column, serialization, etc) changes.
 
 use zebra_chain::{
-    history_tree::HistoryTree, orchard, parameters::Network, sapling, sprout,
+    block::Height, history_tree::HistoryTree, orchard, parameters::Network, sapling, sprout,
     transaction::Transaction,
 };
 
@@ -117,6 +117,17 @@ impl ZebraDb {
             .expect("Sapling note commitment tree must exist if there is a finalized tip")
     }
 
+    /// Returns the Sapling note commitment tree matching the given block height.
+    #[allow(dead_code)]
+    pub fn sapling_note_commitment_tree_by_height(
+        &self,
+        height: &Height,
+    ) -> Option<sapling::tree::NoteCommitmentTree> {
+        let sapling_trees = self.db.cf_handle("sapling_note_commitment_tree").unwrap();
+
+        self.db.zs_get(sapling_trees, height)
+    }
+
     /// Returns the Orchard note commitment tree of the finalized tip
     /// or the empty tree if the state is empty.
     pub fn orchard_note_commitment_tree(&self) -> orchard::tree::NoteCommitmentTree {
@@ -133,6 +144,17 @@ impl ZebraDb {
             .expect("Orchard note commitment tree must exist if there is a finalized tip")
     }
 
+    /// Returns the Orchard note commitment tree matching the given block height.
+    #[allow(dead_code)]
+    pub fn orchard_note_commitment_tree_by_height(
+        &self,
+        height: &Height,
+    ) -> Option<orchard::tree::NoteCommitmentTree> {
+        let orchard_trees = self.db.cf_handle("orchard_note_commitment_tree").unwrap();
+
+        self.db.zs_get(orchard_trees, height)
+    }
+
     /// 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 {
@@ -244,14 +266,16 @@ impl DiskWriteBatch {
         self.zs_insert(sapling_anchors, sapling_root, ());
         self.zs_insert(orchard_anchors, orchard_root, ());
 
-        // Update the trees in state
+        // Delete the previously stored Sprout note commitment tree.
         let current_tip_height = *height - 1;
         if let Some(h) = current_tip_height {
             self.zs_delete(sprout_note_commitment_tree_cf, h);
-            self.zs_delete(sapling_note_commitment_tree_cf, h);
-            self.zs_delete(orchard_note_commitment_tree_cf, h);
         }
 
+        // TODO: if we ever need concurrent read-only access to the sprout tree,
+        // store it by `()`, not height. Otherwise, the ReadStateService could
+        // access a height that was just deleted by a concurrent StateService
+        // write. This requires a database version update.
         self.zs_insert(
             sprout_note_commitment_tree_cf,
             height,

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

@@ -49,12 +49,20 @@ pub struct Chain {
     /// The Sprout note commitment tree of the tip of this `Chain`,
     /// including all finalized notes, and the non-finalized notes in this chain.
     pub(super) sprout_note_commitment_tree: sprout::tree::NoteCommitmentTree,
+    /// The Sprout note commitment tree for each anchor.
+    /// This is required for interstitial states.
+    pub(crate) sprout_trees_by_anchor:
+        HashMap<sprout::tree::Root, sprout::tree::NoteCommitmentTree>,
     /// The Sapling note commitment tree of the tip of this `Chain`,
     /// including all finalized notes, and the non-finalized notes in this chain.
     pub(super) sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree,
+    /// The Sapling note commitment tree for each height.
+    pub(crate) sapling_trees_by_height: BTreeMap<block::Height, sapling::tree::NoteCommitmentTree>,
     /// The Orchard note commitment tree of the tip of this `Chain`,
     /// including all finalized notes, and the non-finalized notes in this chain.
     pub(super) orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree,
+    /// The Orchard note commitment tree for each height.
+    pub(crate) orchard_trees_by_height: BTreeMap<block::Height, orchard::tree::NoteCommitmentTree>,
     /// The ZIP-221 history tree of the tip of this `Chain`,
     /// including all finalized blocks, and the non-finalized `blocks` in this chain.
     pub(crate) history_tree: HistoryTree,
@@ -63,10 +71,6 @@ pub struct Chain {
     pub(crate) sprout_anchors: MultiSet<sprout::tree::Root>,
     /// The Sprout anchors created by each block in `blocks`.
     pub(crate) sprout_anchors_by_height: BTreeMap<block::Height, sprout::tree::Root>,
-    /// The Sprout note commitment tree for each anchor.
-    /// This is required for interstitial states.
-    pub(crate) sprout_trees_by_anchor:
-        HashMap<sprout::tree::Root, sprout::tree::NoteCommitmentTree>,
     /// The Sapling anchors created by `blocks`.
     pub(crate) sapling_anchors: MultiSet<sapling::tree::Root>,
     /// The Sapling anchors created by each block in `blocks`.
@@ -124,8 +128,10 @@ impl Chain {
             sprout_trees_by_anchor: Default::default(),
             sapling_anchors: MultiSet::new(),
             sapling_anchors_by_height: Default::default(),
+            sapling_trees_by_height: Default::default(),
             orchard_anchors: MultiSet::new(),
             orchard_anchors_by_height: Default::default(),
+            orchard_trees_by_height: Default::default(),
             sprout_nullifiers: Default::default(),
             sapling_nullifiers: Default::default(),
             orchard_nullifiers: Default::default(),
@@ -162,7 +168,9 @@ impl Chain {
             self.sprout_note_commitment_tree.root() == other.sprout_note_commitment_tree.root() &&
             self.sprout_trees_by_anchor == other.sprout_trees_by_anchor &&
             self.sapling_note_commitment_tree.root() == other.sapling_note_commitment_tree.root() &&
+            self.sapling_trees_by_height== other.sapling_trees_by_height &&
             self.orchard_note_commitment_tree.root() == other.orchard_note_commitment_tree.root() &&
+            self.orchard_trees_by_height== other.orchard_trees_by_height &&
 
             // history tree
             self.history_tree.as_ref().map(NonEmptyHistoryTree::hash) == other.history_tree.as_ref().map(NonEmptyHistoryTree::hash) &&
@@ -416,12 +424,14 @@ impl Chain {
             created_utxos: self.created_utxos.clone(),
             spent_utxos: self.spent_utxos.clone(),
             sprout_note_commitment_tree,
+            sprout_trees_by_anchor: self.sprout_trees_by_anchor.clone(),
             sapling_note_commitment_tree,
+            sapling_trees_by_height: self.sapling_trees_by_height.clone(),
             orchard_note_commitment_tree,
+            orchard_trees_by_height: self.orchard_trees_by_height.clone(),
             sprout_anchors: self.sprout_anchors.clone(),
             sapling_anchors: self.sapling_anchors.clone(),
             orchard_anchors: self.orchard_anchors.clone(),
-            sprout_trees_by_anchor: self.sprout_trees_by_anchor.clone(),
             sprout_anchors_by_height: self.sprout_anchors_by_height.clone(),
             sapling_anchors_by_height: self.sapling_anchors_by_height.clone(),
             orchard_anchors_by_height: self.orchard_anchors_by_height.clone(),
@@ -553,6 +563,12 @@ impl UpdateWith<ContextuallyValidBlock> for Chain {
             self.update_chain_tip_with(orchard_shielded_data)?;
         }
 
+        // Update the note commitment trees indexed by height.
+        self.sapling_trees_by_height
+            .insert(height, self.sapling_note_commitment_tree.clone());
+        self.orchard_trees_by_height
+            .insert(height, self.orchard_note_commitment_tree.clone());
+
         // Having updated all the note commitment trees and nullifier sets in
         // this block, the roots of the note commitment trees as of the last
         // transaction are the treestates of this block.
@@ -688,6 +704,9 @@ impl UpdateWith<ContextuallyValidBlock> for Chain {
             self.sapling_anchors.remove(&anchor),
             "Sapling anchor must be present if block was added to chain"
         );
+        self.sapling_trees_by_height
+            .remove(&height)
+            .expect("Sapling note commitment tree must be present if block was added to chain");
 
         let anchor = self
             .orchard_anchors_by_height
@@ -697,6 +716,9 @@ impl UpdateWith<ContextuallyValidBlock> for Chain {
             self.orchard_anchors.remove(&anchor),
             "Orchard anchor must be present if block was added to chain"
         );
+        self.orchard_trees_by_height
+            .remove(&height)
+            .expect("Orchard note commitment tree must be present if block was added to chain");
 
         // revert the chain value pool balances, if needed
         self.revert_chain_with(chain_value_pool_change, position);

zebra-state/src/service/non_finalized_state/tests/prop.rs

@@ -600,8 +600,11 @@ fn different_blocks_different_chains() -> Result<()> {
 
                 // note commitment trees
                 chain1.sprout_note_commitment_tree = chain2.sprout_note_commitment_tree.clone();
+                chain1.sprout_trees_by_anchor = chain2.sprout_trees_by_anchor.clone();
                 chain1.sapling_note_commitment_tree = chain2.sapling_note_commitment_tree.clone();
+                chain1.sapling_trees_by_height = chain2.sapling_trees_by_height.clone();
                 chain1.orchard_note_commitment_tree = chain2.orchard_note_commitment_tree.clone();
+                chain1.orchard_trees_by_height = chain2.orchard_trees_by_height.clone();
 
                 // history tree
                 chain1.history_tree = chain2.history_tree.clone();
@@ -609,7 +612,6 @@ fn different_blocks_different_chains() -> Result<()> {
                 // anchors
                 chain1.sprout_anchors = chain2.sprout_anchors.clone();
                 chain1.sprout_anchors_by_height = chain2.sprout_anchors_by_height.clone();
-                chain1.sprout_trees_by_anchor = chain2.sprout_trees_by_anchor.clone();
                 chain1.sapling_anchors = chain2.sapling_anchors.clone();
                 chain1.sapling_anchors_by_height = chain2.sapling_anchors_by_height.clone();
                 chain1.orchard_anchors = chain2.orchard_anchors.clone();