From ec684a4abf86bd589f9822f3b7ce6df27e6bd73a Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 31 Aug 2023 13:25:57 +1000
Subject: [PATCH 01/33] Copy the add_subtrees upgrade from the original branch

---
 .../finalized_state/disk_format/upgrade.rs    |   2 +
 .../disk_format/upgrade/add_subtrees.rs       | 168 ++++++++++++++++++
 2 files changed, 170 insertions(+)
 create mode 100644 zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs

diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
index a693ac174c8..958b75fbd55 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
@@ -25,6 +25,8 @@ use crate::{
     Config,
 };
 
+mod add_subtrees;
+
 /// The kind of database format change we're performing.
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub enum DbFormatChange {
diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
new file mode 100644
index 00000000000..65bfda4ad4d
--- /dev/null
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -0,0 +1,168 @@
+//! Fully populate the Sapling and Orchard note commitment subtrees for existing blocks in the database.
+
+use std::sync::{
+    atomic::{self, AtomicBool},
+    Arc,
+};
+
+use zebra_chain::{
+    block::Height, orchard::tree::NoteCommitmentTree as OrchardNoteCommitmentTree,
+    sapling::tree::NoteCommitmentTree as SaplingNoteCommitmentTree, subtree::NoteCommitmentSubtree,
+};
+
+use crate::service::finalized_state::{DiskWriteBatch, ZebraDb};
+
+/// Runs disk format upgrade for adding Sapling and Orchard note commitment subtrees to database.
+pub fn _run(
+    initial_tip_height: Height,
+    upgrade_db: &ZebraDb,
+    should_cancel_format_change: Arc<AtomicBool>,
+) {
+    let mut subtree_count = 0;
+    let mut prev_tree: Option<_> = None;
+    for (height, tree) in upgrade_db.sapling_tree_by_height_range(..=initial_tip_height) {
+        if should_cancel_format_change.load(atomic::Ordering::Relaxed) {
+            break;
+        }
+
+        let Some(frontier) = tree.frontier() else {
+            prev_tree = Some(tree);
+            continue;
+        };
+
+        // Blocks cannot complete multiple level 16 subtrees,
+        // the subtree index can increase by a maximum of 1 every ~20 blocks.
+        let subtree_address = SaplingNoteCommitmentTree::subtree_address(frontier);
+        if subtree_address.index() <= subtree_count {
+            prev_tree = Some(tree);
+            continue;
+        }
+
+        let (index, node) = if SaplingNoteCommitmentTree::is_complete_subtree(frontier) {
+            tree.completed_subtree_index_and_root()
+                .expect("already checked is_complete_subtree()")
+        } else {
+            let mut sapling_nct = Arc::try_unwrap(
+                prev_tree
+                    .take()
+                    .expect("should have some previous sapling frontier"),
+            )
+            .unwrap_or_else(|shared_tree| (*shared_tree).clone());
+
+            let block = upgrade_db
+                .block(height.into())
+                .expect("height with note commitment tree should have block");
+
+            let sapling_note_commitments: Vec<_> = block
+                .transactions
+                .iter()
+                .flat_map(|tx| tx.sapling_note_commitments())
+                .cloned()
+                .collect();
+
+            for sapling_note_commitment in sapling_note_commitments {
+                sapling_nct
+                    .append(sapling_note_commitment)
+                    .expect("finalized notes should append successfully");
+
+                if sapling_nct
+                    .frontier()
+                    .map_or(false, SaplingNoteCommitmentTree::is_complete_subtree)
+                {
+                    break;
+                }
+            }
+
+            sapling_nct
+                .completed_subtree_index_and_root()
+                .expect("already checked is_complete_subtree()")
+        };
+
+        let subtree = NoteCommitmentSubtree::new(index, height, node);
+
+        let mut batch = DiskWriteBatch::new();
+
+        batch.insert_sapling_subtree(upgrade_db, subtree);
+
+        upgrade_db
+            .write_batch(batch)
+            .expect("writing sapling note commitment subtrees should always succeed.");
+
+        subtree_count += 1;
+        prev_tree = Some(tree);
+    }
+
+    let mut subtree_count = 0;
+    let mut prev_tree: Option<_> = None;
+    for (height, tree) in upgrade_db.orchard_tree_by_height_range(..=initial_tip_height) {
+        if should_cancel_format_change.load(atomic::Ordering::Relaxed) {
+            break;
+        }
+
+        let Some(frontier) = tree.frontier() else {
+            prev_tree = Some(tree);
+            continue;
+        };
+
+        // Blocks cannot complete multiple level 16 subtrees,
+        // the subtree index can increase by a maximum of 1 every ~20 blocks.
+        let subtree_address = OrchardNoteCommitmentTree::subtree_address(frontier);
+        if subtree_address.index() <= subtree_count {
+            prev_tree = Some(tree);
+            continue;
+        }
+
+        let (index, node) = if OrchardNoteCommitmentTree::is_complete_subtree(frontier) {
+            tree.completed_subtree_index_and_root()
+                .expect("already checked is_complete_subtree()")
+        } else {
+            let mut orchard_nct = Arc::try_unwrap(
+                prev_tree
+                    .take()
+                    .expect("should have some previous orchard frontier"),
+            )
+            .unwrap_or_else(|shared_tree| (*shared_tree).clone());
+
+            let block = upgrade_db
+                .block(height.into())
+                .expect("height with note commitment tree should have block");
+
+            let orchard_note_commitments: Vec<_> = block
+                .transactions
+                .iter()
+                .flat_map(|tx| tx.orchard_note_commitments())
+                .cloned()
+                .collect();
+
+            for orchard_note_commitment in orchard_note_commitments {
+                orchard_nct
+                    .append(orchard_note_commitment)
+                    .expect("finalized notes should append successfully");
+
+                if orchard_nct
+                    .frontier()
+                    .map_or(false, OrchardNoteCommitmentTree::is_complete_subtree)
+                {
+                    break;
+                }
+            }
+
+            orchard_nct
+                .completed_subtree_index_and_root()
+                .expect("already checked is_complete_subtree()")
+        };
+
+        let subtree = NoteCommitmentSubtree::new(index, height, node);
+
+        let mut batch = DiskWriteBatch::new();
+
+        batch.insert_orchard_subtree(upgrade_db, subtree);
+
+        upgrade_db
+            .write_batch(batch)
+            .expect("writing orchard note commitment subtrees should always succeed.");
+
+        subtree_count += 1;
+        prev_tree = Some(tree);
+    }
+}

From 38c90d520625a4c22fbe2378d1102d1fa8b87c20 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 31 Aug 2023 13:29:44 +1000
Subject: [PATCH 02/33] Copy the database write changes in shielded.rs from the
 original branch

---
 .../finalized_state/zebra_db/shielded.rs      | 48 +++++++++++++++----
 1 file changed, 38 insertions(+), 10 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
index 57f9ae0ba2a..485203db432 100644
--- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
@@ -354,8 +354,8 @@ 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 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();
@@ -402,19 +402,32 @@ 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.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());
-        // }
+        if let Some(subtree) = trees.orchard_subtree {
+            self.zs_insert(&orchard_subtree_cf, subtree.index, subtree.into_data());
+        }
 
         self.prepare_history_batch(db, finalized)
     }
 
+    // Sapling tree methods
+
+    /// Inserts the Sapling note commitment subtree.
+    pub fn insert_sapling_subtree(
+        &mut self,
+        zebra_db: &ZebraDb,
+        subtree: Arc<NoteCommitmentSubtree<sapling::tree::Node>>,
+    ) {
+        let sapling_subtree_cf = zebra_db
+            .db
+            .cf_handle("sapling_note_commitment_subtree")
+            .unwrap();
+        self.zs_insert(&sapling_subtree_cf, subtree.index, subtree.into_data());
+    }
+
     /// Deletes the Sapling note commitment tree at the given [`Height`].
     pub fn delete_sapling_tree(&mut self, zebra_db: &ZebraDb, height: &Height) {
         let sapling_tree_cf = zebra_db
@@ -436,6 +449,21 @@ impl DiskWriteBatch {
         self.zs_delete_range(&sapling_tree_cf, from, to);
     }
 
+    // Orchard tree methods
+
+    /// Inserts the Orchard note commitment subtree.
+    pub fn insert_orchard_subtree(
+        &mut self,
+        zebra_db: &ZebraDb,
+        subtree: Arc<NoteCommitmentSubtree<orchard::tree::Node>>,
+    ) {
+        let orchard_subtree_cf = zebra_db
+            .db
+            .cf_handle("orchard_note_commitment_subtree")
+            .unwrap();
+        self.zs_insert(&orchard_subtree_cf, subtree.index, subtree.into_data());
+    }
+
     /// Deletes the Orchard note commitment tree at the given [`Height`].
     pub fn delete_orchard_tree(&mut self, zebra_db: &ZebraDb, height: &Height) {
         let orchard_tree_cf = zebra_db

From dcb5d473e3251b2b026eb508f066f1a7353eb30f Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 31 Aug 2023 13:36:32 +1000
Subject: [PATCH 03/33]  Copy the tree API changes from the original branch

---
 zebra-chain/src/orchard/tree.rs | 5 +++++
 zebra-chain/src/sapling/tree.rs | 5 +++++
 2 files changed, 10 insertions(+)

diff --git a/zebra-chain/src/orchard/tree.rs b/zebra-chain/src/orchard/tree.rs
index c73a3556be5..f0973ea4042 100644
--- a/zebra-chain/src/orchard/tree.rs
+++ b/zebra-chain/src/orchard/tree.rs
@@ -347,6 +347,11 @@ impl NoteCommitmentTree {
         }
     }
 
+    /// Returns frontier of non-empty tree, or None.
+    pub fn frontier(&self) -> Option<&NonEmptyFrontier<Node>> {
+        self.inner.value()
+    }
+
     /// Returns true if the most recently appended leaf completes the subtree
     pub fn is_complete_subtree(tree: &NonEmptyFrontier<Node>) -> bool {
         tree.position()
diff --git a/zebra-chain/src/sapling/tree.rs b/zebra-chain/src/sapling/tree.rs
index a3ea78b4278..07f7a02e665 100644
--- a/zebra-chain/src/sapling/tree.rs
+++ b/zebra-chain/src/sapling/tree.rs
@@ -332,6 +332,11 @@ impl NoteCommitmentTree {
         }
     }
 
+    /// Returns frontier of non-empty tree, or None.
+    pub fn frontier(&self) -> Option<&NonEmptyFrontier<Node>> {
+        self.inner.value()
+    }
+
     /// Returns true if the most recently appended leaf completes the subtree
     pub fn is_complete_subtree(tree: &NonEmptyFrontier<Node>) -> bool {
         tree.position()

From d01096ce6165bf66a8d9da7ee8d5b5f1b119bc1e Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 31 Aug 2023 13:50:08 +1000
Subject: [PATCH 04/33] Simplify subtree APIs to avoid exposing frontiers

---
 zebra-chain/src/orchard/tree.rs               | 32 +++++++----
 zebra-chain/src/sapling/tree.rs               | 30 +++++++----
 .../disk_format/upgrade/add_subtrees.rs       | 53 +++++++++----------
 3 files changed, 66 insertions(+), 49 deletions(-)

diff --git a/zebra-chain/src/orchard/tree.rs b/zebra-chain/src/orchard/tree.rs
index f0973ea4042..2ad0e9450e3 100644
--- a/zebra-chain/src/orchard/tree.rs
+++ b/zebra-chain/src/orchard/tree.rs
@@ -347,33 +347,43 @@ impl NoteCommitmentTree {
         }
     }
 
-    /// Returns frontier of non-empty tree, or None.
-    pub fn frontier(&self) -> Option<&NonEmptyFrontier<Node>> {
+    /// Returns frontier of non-empty tree, or `None` if the tree is empty.
+    fn frontier(&self) -> Option<&NonEmptyFrontier<Node>> {
         self.inner.value()
     }
 
     /// Returns true if the most recently appended leaf completes the subtree
-    pub fn is_complete_subtree(tree: &NonEmptyFrontier<Node>) -> bool {
+    pub fn is_complete_subtree(&self) -> bool {
+        let Some(tree) = self.frontier() else {
+            // An empty tree can't be a complete subtree.
+            return false;
+        };
+
         tree.position()
             .is_complete_subtree(TRACKED_SUBTREE_HEIGHT.into())
     }
 
-    /// Returns subtree address at [`TRACKED_SUBTREE_HEIGHT`]
-    pub fn subtree_address(tree: &NonEmptyFrontier<Node>) -> incrementalmerkletree::Address {
-        incrementalmerkletree::Address::above_position(
+    /// Returns the subtree address at [`TRACKED_SUBTREE_HEIGHT`].
+    /// Returns `None` if the tree is empty.
+    pub fn subtree_address(&self) -> Option<incrementalmerkletree::Address> {
+        let tree = self.frontier()?;
+
+        Some(incrementalmerkletree::Address::above_position(
             TRACKED_SUBTREE_HEIGHT.into(),
             tree.position(),
-        )
+        ))
     }
 
     /// Returns subtree index and root if the most recently appended leaf completes the subtree
     #[allow(clippy::unwrap_in_result)]
     pub fn completed_subtree_index_and_root(&self) -> Option<(u16, Node)> {
-        let value = self.inner.value()?;
-        Self::is_complete_subtree(value).then_some(())?;
-        let address = Self::subtree_address(value);
+        if !self.is_complete_subtree() {
+            return None;
+        }
+
+        let address = self.subtree_address()?;
         let index = address.index().try_into().expect("should fit in u16");
-        let root = value.root(Some(TRACKED_SUBTREE_HEIGHT.into()));
+        let root = self.frontier()?.root(Some(TRACKED_SUBTREE_HEIGHT.into()));
 
         Some((index, root))
     }
diff --git a/zebra-chain/src/sapling/tree.rs b/zebra-chain/src/sapling/tree.rs
index 07f7a02e665..0170e6d385a 100644
--- a/zebra-chain/src/sapling/tree.rs
+++ b/zebra-chain/src/sapling/tree.rs
@@ -333,32 +333,42 @@ impl NoteCommitmentTree {
     }
 
     /// Returns frontier of non-empty tree, or None.
-    pub fn frontier(&self) -> Option<&NonEmptyFrontier<Node>> {
+    fn frontier(&self) -> Option<&NonEmptyFrontier<Node>> {
         self.inner.value()
     }
 
     /// Returns true if the most recently appended leaf completes the subtree
-    pub fn is_complete_subtree(tree: &NonEmptyFrontier<Node>) -> bool {
+    pub fn is_complete_subtree(&self) -> bool {
+        let Some(tree) = self.frontier() else {
+            // An empty tree can't be a complete subtree.
+            return false;
+        };
+
         tree.position()
             .is_complete_subtree(TRACKED_SUBTREE_HEIGHT.into())
     }
 
-    /// Returns subtree address at [`TRACKED_SUBTREE_HEIGHT`]
-    pub fn subtree_address(tree: &NonEmptyFrontier<Node>) -> incrementalmerkletree::Address {
-        incrementalmerkletree::Address::above_position(
+    /// Returns the subtree address at [`TRACKED_SUBTREE_HEIGHT`].
+    /// Returns `None` if the tree is empty.
+    pub fn subtree_address(&self) -> Option<incrementalmerkletree::Address> {
+        let tree = self.frontier()?;
+
+        Some(incrementalmerkletree::Address::above_position(
             TRACKED_SUBTREE_HEIGHT.into(),
             tree.position(),
-        )
+        ))
     }
 
     /// Returns subtree index and root if the most recently appended leaf completes the subtree
     #[allow(clippy::unwrap_in_result)]
     pub fn completed_subtree_index_and_root(&self) -> Option<(u16, Node)> {
-        let value = self.inner.value()?;
-        Self::is_complete_subtree(value).then_some(())?;
-        let address = Self::subtree_address(value);
+        if !self.is_complete_subtree() {
+            return None;
+        }
+
+        let address = self.subtree_address()?;
         let index = address.index().try_into().expect("should fit in u16");
-        let root = value.root(Some(TRACKED_SUBTREE_HEIGHT.into()));
+        let root = self.frontier()?.root(Some(TRACKED_SUBTREE_HEIGHT.into()));
 
         Some((index, root))
     }
diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
index 65bfda4ad4d..e3d6fb089fd 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -5,10 +5,7 @@ use std::sync::{
     Arc,
 };
 
-use zebra_chain::{
-    block::Height, orchard::tree::NoteCommitmentTree as OrchardNoteCommitmentTree,
-    sapling::tree::NoteCommitmentTree as SaplingNoteCommitmentTree, subtree::NoteCommitmentSubtree,
-};
+use zebra_chain::{block::Height, subtree::NoteCommitmentSubtree};
 
 use crate::service::finalized_state::{DiskWriteBatch, ZebraDb};
 
@@ -25,20 +22,19 @@ pub fn _run(
             break;
         }
 
-        let Some(frontier) = tree.frontier() else {
+        // Blocks cannot complete multiple level 16 subtrees,
+        // the subtree index can increase by a maximum of 1 every ~20 blocks.
+        let Some(subtree_address) = tree.subtree_address() else {
             prev_tree = Some(tree);
             continue;
         };
 
-        // Blocks cannot complete multiple level 16 subtrees,
-        // the subtree index can increase by a maximum of 1 every ~20 blocks.
-        let subtree_address = SaplingNoteCommitmentTree::subtree_address(frontier);
         if subtree_address.index() <= subtree_count {
             prev_tree = Some(tree);
             continue;
         }
 
-        let (index, node) = if SaplingNoteCommitmentTree::is_complete_subtree(frontier) {
+        let (index, node) = if tree.is_complete_subtree() {
             tree.completed_subtree_index_and_root()
                 .expect("already checked is_complete_subtree()")
         } else {
@@ -65,17 +61,18 @@ pub fn _run(
                     .append(sapling_note_commitment)
                     .expect("finalized notes should append successfully");
 
-                if sapling_nct
-                    .frontier()
-                    .map_or(false, SaplingNoteCommitmentTree::is_complete_subtree)
-                {
+                // The loop always breaks on this condition,
+                // because we checked the block has enough commitments,
+                // and that the final commitment in the block doesn't complete a subtree.
+                if sapling_nct.is_complete_subtree() {
                     break;
                 }
             }
 
-            sapling_nct
-                .completed_subtree_index_and_root()
-                .expect("already checked is_complete_subtree()")
+            sapling_nct.completed_subtree_index_and_root().expect(
+                "already checked is_complete_subtree(),\
+                     and that the block must complete a subtree",
+            )
         };
 
         let subtree = NoteCommitmentSubtree::new(index, height, node);
@@ -99,20 +96,19 @@ pub fn _run(
             break;
         }
 
-        let Some(frontier) = tree.frontier() else {
+        // Blocks cannot complete multiple level 16 subtrees,
+        // the subtree index can increase by a maximum of 1 every ~20 blocks.
+        let Some(subtree_address) = tree.subtree_address() else {
             prev_tree = Some(tree);
             continue;
         };
 
-        // Blocks cannot complete multiple level 16 subtrees,
-        // the subtree index can increase by a maximum of 1 every ~20 blocks.
-        let subtree_address = OrchardNoteCommitmentTree::subtree_address(frontier);
         if subtree_address.index() <= subtree_count {
             prev_tree = Some(tree);
             continue;
         }
 
-        let (index, node) = if OrchardNoteCommitmentTree::is_complete_subtree(frontier) {
+        let (index, node) = if tree.is_complete_subtree() {
             tree.completed_subtree_index_and_root()
                 .expect("already checked is_complete_subtree()")
         } else {
@@ -139,17 +135,18 @@ pub fn _run(
                     .append(orchard_note_commitment)
                     .expect("finalized notes should append successfully");
 
-                if orchard_nct
-                    .frontier()
-                    .map_or(false, OrchardNoteCommitmentTree::is_complete_subtree)
-                {
+                // The loop always breaks on this condition,
+                // because we checked the block has enough commitments,
+                // and that the final commitment in the block doesn't complete a subtree.
+                if orchard_nct.is_complete_subtree() {
                     break;
                 }
             }
 
-            orchard_nct
-                .completed_subtree_index_and_root()
-                .expect("already checked is_complete_subtree()")
+            orchard_nct.completed_subtree_index_and_root().expect(
+                "already checked is_complete_subtree(),\
+                     and that the block must complete a subtree",
+            )
         };
 
         let subtree = NoteCommitmentSubtree::new(index, height, node);

From cf05e626e9c234775d893f28fe3f942d17c25aa7 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 31 Aug 2023 13:51:45 +1000
Subject: [PATCH 05/33] Fix a dead code warning by re-using existing methods

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

diff --git a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
index 485203db432..ce3b5373a77 100644
--- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
@@ -354,9 +354,6 @@ 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();
 
@@ -403,11 +400,11 @@ impl DiskWriteBatch {
         }
 
         if let Some(subtree) = trees.sapling_subtree {
-            self.zs_insert(&sapling_subtree_cf, subtree.index, subtree.into_data());
+            self.insert_sapling_subtree(zebra_db, subtree);
         }
 
         if let Some(subtree) = trees.orchard_subtree {
-            self.zs_insert(&orchard_subtree_cf, subtree.index, subtree.into_data());
+            self.insert_orchard_subtree(zebra_db, subtree);
         }
 
         self.prepare_history_batch(db, finalized)

From da5851f76b0a58be4f1cd4543b10757871d6ae7d Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 31 Aug 2023 14:01:55 +1000
Subject: [PATCH 06/33] Use mpsc::Receiver<CancelFormatChange> in the subtree
 upgrade

---
 .../disk_format/upgrade/add_subtrees.rs       | 33 ++++++++++++-------
 1 file changed, 22 insertions(+), 11 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
index e3d6fb089fd..e7298504e8b 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -1,25 +1,25 @@
 //! Fully populate the Sapling and Orchard note commitment subtrees for existing blocks in the database.
 
-use std::sync::{
-    atomic::{self, AtomicBool},
-    Arc,
-};
+use std::sync::{mpsc, Arc};
 
 use zebra_chain::{block::Height, subtree::NoteCommitmentSubtree};
 
-use crate::service::finalized_state::{DiskWriteBatch, ZebraDb};
+use crate::service::finalized_state::{
+    disk_format::upgrade::CancelFormatChange, DiskWriteBatch, ZebraDb,
+};
 
 /// Runs disk format upgrade for adding Sapling and Orchard note commitment subtrees to database.
-pub fn _run(
+pub fn run(
     initial_tip_height: Height,
     upgrade_db: &ZebraDb,
-    should_cancel_format_change: Arc<AtomicBool>,
+    cancel_receiver: mpsc::Receiver<CancelFormatChange>,
 ) {
     let mut subtree_count = 0;
     let mut prev_tree: Option<_> = None;
     for (height, tree) in upgrade_db.sapling_tree_by_height_range(..=initial_tip_height) {
-        if should_cancel_format_change.load(atomic::Ordering::Relaxed) {
-            break;
+        // Return early if there is a cancel signal.
+        if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
+            return;
         }
 
         // Blocks cannot complete multiple level 16 subtrees,
@@ -57,6 +57,11 @@ pub fn _run(
                 .collect();
 
             for sapling_note_commitment in sapling_note_commitments {
+                // Return early if there is a cancel signal.
+                if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
+                    return;
+                }
+
                 sapling_nct
                     .append(sapling_note_commitment)
                     .expect("finalized notes should append successfully");
@@ -92,8 +97,9 @@ pub fn _run(
     let mut subtree_count = 0;
     let mut prev_tree: Option<_> = None;
     for (height, tree) in upgrade_db.orchard_tree_by_height_range(..=initial_tip_height) {
-        if should_cancel_format_change.load(atomic::Ordering::Relaxed) {
-            break;
+        // Return early if there is a cancel signal.
+        if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
+            return;
         }
 
         // Blocks cannot complete multiple level 16 subtrees,
@@ -131,6 +137,11 @@ pub fn _run(
                 .collect();
 
             for orchard_note_commitment in orchard_note_commitments {
+                // Return early if there is a cancel signal.
+                if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
+                    return;
+                }
+
                 orchard_nct
                     .append(orchard_note_commitment)
                     .expect("finalized notes should append successfully");

From fbadf4c4302fe69a9da869ee26eb35b0adbd6456 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 31 Aug 2023 14:02:19 +1000
Subject: [PATCH 07/33] Run the subtree upgrade on startup

---
 .../finalized_state/disk_format/upgrade.rs    | 27 ++++++++++++++++---
 1 file changed, 24 insertions(+), 3 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
index 958b75fbd55..42efb427792 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
@@ -275,7 +275,7 @@ impl DbFormatChange {
             return;
         };
 
-        // Start of a database upgrade task.
+        // Note commitment tree de-duplication database upgrade task.
 
         let version_for_pruning_trees =
             Version::parse("25.1.1").expect("Hardcoded version string should be valid.");
@@ -337,7 +337,7 @@ impl DbFormatChange {
             }
 
             // Before marking the state as upgraded, check that the upgrade completed successfully.
-            Self::check_for_duplicate_trees(db);
+            Self::check_for_duplicate_trees(db.clone());
 
             // Mark the database as upgraded. Zebra won't repeat the upgrade anymore once the
             // database is marked, so the upgrade MUST be complete at this point.
@@ -348,7 +348,28 @@ impl DbFormatChange {
             Self::mark_as_upgraded_to(&version_for_pruning_trees, &config, network);
         }
 
-        // End of a database upgrade task.
+        // Note commitment subtree creation database upgrade task.
+
+        let version_for_adding_subtrees =
+            Version::parse("25.2.0").expect("Hardcoded version string should be valid.");
+
+        // Check if we need to add note commitment subtrees to the database.
+        if older_disk_version < version_for_adding_subtrees {
+            add_subtrees::run(initial_tip_height, &db, cancel_receiver);
+
+            // Before marking the state as upgraded, check that the upgrade completed successfully.
+            //
+            // TODO: do this check in all the same places as check_for_duplicate_trees()
+            //Self::check_for_continuous_subtrees(db);
+
+            // Mark the database as upgraded. Zebra won't repeat the upgrade anymore once the
+            // database is marked, so the upgrade MUST be complete at this point.
+            info!(
+                ?newer_running_version,
+                "Zebra automatically upgraded the database format to:"
+            );
+            Self::mark_as_upgraded_to(&version_for_adding_subtrees, &config, network);
+        }
 
         // # New Upgrades Usually Go Here
         //

From 6950ad6c6c1cc66caca1d19833196f4912cff239 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 31 Aug 2023 13:22:52 +1000
Subject: [PATCH 08/33] Bump the database format version to 25.2.0

---
 zebra-state/src/constants.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/zebra-state/src/constants.rs b/zebra-state/src/constants.rs
index 6c454b6d884..8abf3750d90 100644
--- a/zebra-state/src/constants.rs
+++ b/zebra-state/src/constants.rs
@@ -48,11 +48,11 @@ pub(crate) const DATABASE_FORMAT_VERSION: u64 = 25;
 /// - adding new column families,
 /// - changing the format of a column family in a compatible way, or
 /// - breaking changes with compatibility code in all supported Zebra versions.
-pub(crate) const DATABASE_FORMAT_MINOR_VERSION: u64 = 1;
+pub(crate) const DATABASE_FORMAT_MINOR_VERSION: u64 = 2;
 
 /// The database format patch version, incremented each time the on-disk database format has a
 /// significant format compatibility fix.
-pub(crate) const DATABASE_FORMAT_PATCH_VERSION: u64 = 1;
+pub(crate) const DATABASE_FORMAT_PATCH_VERSION: u64 = 0;
 
 /// The name of the file containing the minor and patch database versions.
 ///

From ac7fb850a376c2c2aee332886379de8b73bf0d3c Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 31 Aug 2023 14:34:17 +1000
Subject: [PATCH 09/33] Fix a confusing 'upgrade complete' log

---
 .../service/finalized_state/disk_format/upgrade.rs  | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
index 42efb427792..0c5da2d1ca6 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
@@ -341,10 +341,6 @@ impl DbFormatChange {
 
             // Mark the database as upgraded. Zebra won't repeat the upgrade anymore once the
             // database is marked, so the upgrade MUST be complete at this point.
-            info!(
-                ?newer_running_version,
-                "Zebra automatically upgraded the database format to:"
-            );
             Self::mark_as_upgraded_to(&version_for_pruning_trees, &config, network);
         }
 
@@ -364,13 +360,14 @@ impl DbFormatChange {
 
             // Mark the database as upgraded. Zebra won't repeat the upgrade anymore once the
             // database is marked, so the upgrade MUST be complete at this point.
-            info!(
-                ?newer_running_version,
-                "Zebra automatically upgraded the database format to:"
-            );
             Self::mark_as_upgraded_to(&version_for_adding_subtrees, &config, network);
         }
 
+        info!(
+            ?newer_running_version,
+            "Zebra automatically upgraded the database format to:"
+        );
+
         // # New Upgrades Usually Go Here
         //
         // New code goes above this comment!

From 08624143f6b4415981cf2baf7eae8ebac40f96d4 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 31 Aug 2023 14:40:00 +1000
Subject: [PATCH 10/33] Clarify some comments and error messages

---
 .../disk_format/upgrade/add_subtrees.rs       | 20 +++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
index e7298504e8b..431859403be 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -22,13 +22,14 @@ pub fn run(
             return;
         }
 
-        // Blocks cannot complete multiple level 16 subtrees,
-        // the subtree index can increase by a maximum of 1 every ~20 blocks.
+        // Empty note commitment trees can't contain subtrees.
         let Some(subtree_address) = tree.subtree_address() else {
             prev_tree = Some(tree);
             continue;
         };
 
+        // Blocks cannot complete multiple level 16 subtrees,
+        // the subtree index can increase by a maximum of 1 every ~20 blocks.
         if subtree_address.index() <= subtree_count {
             prev_tree = Some(tree);
             continue;
@@ -75,8 +76,9 @@ pub fn run(
             }
 
             sapling_nct.completed_subtree_index_and_root().expect(
-                "already checked is_complete_subtree(),\
-                     and that the block must complete a subtree",
+                "block should have completed a subtree before its final note commitment: \
+                 already checked is_complete_subtree(),\
+                 and that the block must complete a subtree",
             )
         };
 
@@ -102,13 +104,14 @@ pub fn run(
             return;
         }
 
-        // Blocks cannot complete multiple level 16 subtrees,
-        // the subtree index can increase by a maximum of 1 every ~20 blocks.
+        // Empty note commitment trees can't contain subtrees.
         let Some(subtree_address) = tree.subtree_address() else {
             prev_tree = Some(tree);
             continue;
         };
 
+        // Blocks cannot complete multiple level 16 subtrees,
+        // the subtree index can increase by a maximum of 1 every ~20 blocks.
         if subtree_address.index() <= subtree_count {
             prev_tree = Some(tree);
             continue;
@@ -155,8 +158,9 @@ pub fn run(
             }
 
             orchard_nct.completed_subtree_index_and_root().expect(
-                "already checked is_complete_subtree(),\
-                     and that the block must complete a subtree",
+                "block should have completed a subtree before its final note commitment: \
+                 already checked is_complete_subtree(),\
+                 and that the block must complete a subtree",
             )
         };
 

From 8c244a0e2e88fae0e7b4d13597145c7042a3c0bc Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 31 Aug 2023 14:59:30 +1000
Subject: [PATCH 11/33] Simplify prev_tree unwrap to avoid an (impossible?)
 concurrency bug

---
 .../disk_format/upgrade/add_subtrees.rs       | 20 ++++++++-----------
 1 file changed, 8 insertions(+), 12 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
index 431859403be..c8aeae1099b 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -39,12 +39,10 @@ pub fn run(
             tree.completed_subtree_index_and_root()
                 .expect("already checked is_complete_subtree()")
         } else {
-            let mut sapling_nct = Arc::try_unwrap(
-                prev_tree
-                    .take()
-                    .expect("should have some previous sapling frontier"),
-            )
-            .unwrap_or_else(|shared_tree| (*shared_tree).clone());
+            let mut prev_tree = prev_tree
+                .take()
+                .expect("should have some previous sapling frontier");
+            let sapling_nct = Arc::make_mut(&mut prev_tree);
 
             let block = upgrade_db
                 .block(height.into())
@@ -121,12 +119,10 @@ pub fn run(
             tree.completed_subtree_index_and_root()
                 .expect("already checked is_complete_subtree()")
         } else {
-            let mut orchard_nct = Arc::try_unwrap(
-                prev_tree
-                    .take()
-                    .expect("should have some previous orchard frontier"),
-            )
-            .unwrap_or_else(|shared_tree| (*shared_tree).clone());
+            let mut prev_tree = prev_tree
+                .take()
+                .expect("should have some previous orchard frontier");
+            let orchard_nct = Arc::make_mut(&mut prev_tree);
 
             let block = upgrade_db
                 .block(height.into())

From cdb1e62a616f7b28e26f348911505de7eb932e4c Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 31 Aug 2023 15:22:08 +1000
Subject: [PATCH 12/33] Use separate subtree writing functions

---
 zebra-chain/src/orchard/tree.rs               |  4 +-
 zebra-chain/src/parallel/tree.rs              | 10 ++-
 zebra-chain/src/sapling/tree.rs               |  4 +-
 zebra-chain/src/subtree.rs                    | 10 ++-
 .../disk_format/upgrade/add_subtrees.rs       | 82 ++++++++++++-------
 5 files changed, 72 insertions(+), 38 deletions(-)

diff --git a/zebra-chain/src/orchard/tree.rs b/zebra-chain/src/orchard/tree.rs
index 2ad0e9450e3..2b035865d78 100644
--- a/zebra-chain/src/orchard/tree.rs
+++ b/zebra-chain/src/orchard/tree.rs
@@ -31,7 +31,7 @@ use crate::{
     serialization::{
         serde_helpers, ReadZcashExt, SerializationError, ZcashDeserialize, ZcashSerialize,
     },
-    subtree::TRACKED_SUBTREE_HEIGHT,
+    subtree::{NoteCommitmentSubtreeIndex, TRACKED_SUBTREE_HEIGHT},
 };
 
 pub mod legacy;
@@ -376,7 +376,7 @@ impl NoteCommitmentTree {
 
     /// Returns subtree index and root if the most recently appended leaf completes the subtree
     #[allow(clippy::unwrap_in_result)]
-    pub fn completed_subtree_index_and_root(&self) -> Option<(u16, Node)> {
+    pub fn completed_subtree_index_and_root(&self) -> Option<(NoteCommitmentSubtreeIndex, Node)> {
         if !self.is_complete_subtree() {
             return None;
         }
diff --git a/zebra-chain/src/parallel/tree.rs b/zebra-chain/src/parallel/tree.rs
index b5c922444f2..eb4190425bb 100644
--- a/zebra-chain/src/parallel/tree.rs
+++ b/zebra-chain/src/parallel/tree.rs
@@ -4,7 +4,11 @@ use std::sync::Arc;
 
 use thiserror::Error;
 
-use crate::{block::Block, orchard, sapling, sprout, subtree::NoteCommitmentSubtree};
+use crate::{
+    block::Block,
+    orchard, sapling, sprout,
+    subtree::{NoteCommitmentSubtree, NoteCommitmentSubtreeIndex},
+};
 
 /// An argument wrapper struct for note commitment trees.
 #[derive(Clone, Debug)]
@@ -163,7 +167,7 @@ impl NoteCommitmentTrees {
     ) -> Result<
         (
             Arc<sapling::tree::NoteCommitmentTree>,
-            Option<(u16, sapling::tree::Node)>,
+            Option<(NoteCommitmentSubtreeIndex, sapling::tree::Node)>,
         ),
         NoteCommitmentTreeError,
     > {
@@ -202,7 +206,7 @@ impl NoteCommitmentTrees {
     ) -> Result<
         (
             Arc<orchard::tree::NoteCommitmentTree>,
-            Option<(u16, orchard::tree::Node)>,
+            Option<(NoteCommitmentSubtreeIndex, orchard::tree::Node)>,
         ),
         NoteCommitmentTreeError,
     > {
diff --git a/zebra-chain/src/sapling/tree.rs b/zebra-chain/src/sapling/tree.rs
index 0170e6d385a..abd07519194 100644
--- a/zebra-chain/src/sapling/tree.rs
+++ b/zebra-chain/src/sapling/tree.rs
@@ -32,7 +32,7 @@ use crate::{
     serialization::{
         serde_helpers, ReadZcashExt, SerializationError, ZcashDeserialize, ZcashSerialize,
     },
-    subtree::TRACKED_SUBTREE_HEIGHT,
+    subtree::{NoteCommitmentSubtreeIndex, TRACKED_SUBTREE_HEIGHT},
 };
 
 pub mod legacy;
@@ -361,7 +361,7 @@ impl NoteCommitmentTree {
 
     /// Returns subtree index and root if the most recently appended leaf completes the subtree
     #[allow(clippy::unwrap_in_result)]
-    pub fn completed_subtree_index_and_root(&self) -> Option<(u16, Node)> {
+    pub fn completed_subtree_index_and_root(&self) -> Option<(NoteCommitmentSubtreeIndex, Node)> {
         if !self.is_complete_subtree() {
             return None;
         }
diff --git a/zebra-chain/src/subtree.rs b/zebra-chain/src/subtree.rs
index 24a62cfa9b0..7c8a9b456b0 100644
--- a/zebra-chain/src/subtree.rs
+++ b/zebra-chain/src/subtree.rs
@@ -1,6 +1,6 @@
 //! Struct representing Sapling/Orchard note commitment subtrees
 
-use std::sync::Arc;
+use std::{num::TryFromIntError, sync::Arc};
 
 #[cfg(any(test, feature = "proptest-impl"))]
 use proptest_derive::Arbitrary;
@@ -20,6 +20,14 @@ impl From<u16> for NoteCommitmentSubtreeIndex {
     }
 }
 
+impl TryFrom<u64> for NoteCommitmentSubtreeIndex {
+    type Error = TryFromIntError;
+
+    fn try_from(value: u64) -> Result<Self, Self::Error> {
+        u16::try_from(value).map(Self)
+    }
+}
+
 /// Subtree root of Sapling or Orchard note commitment tree,
 /// with its associated block height and subtree index.
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
index c8aeae1099b..5109f45e19e 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -2,7 +2,11 @@
 
 use std::sync::{mpsc, Arc};
 
-use zebra_chain::{block::Height, subtree::NoteCommitmentSubtree};
+use zebra_chain::{
+    block::Height,
+    orchard, sapling,
+    subtree::{NoteCommitmentSubtree, NoteCommitmentSubtreeIndex},
+};
 
 use crate::service::finalized_state::{
     disk_format::upgrade::CancelFormatChange, DiskWriteBatch, ZebraDb,
@@ -35,9 +39,8 @@ pub fn run(
             continue;
         }
 
-        let (index, node) = if tree.is_complete_subtree() {
-            tree.completed_subtree_index_and_root()
-                .expect("already checked is_complete_subtree()")
+        if let Some((index, node)) = tree.completed_subtree_index_and_root() {
+            write_sapling_subtree(upgrade_db, index, height, node);
         } else {
             let mut prev_tree = prev_tree
                 .take()
@@ -73,22 +76,14 @@ pub fn run(
                 }
             }
 
-            sapling_nct.completed_subtree_index_and_root().expect(
+            let (index, node) = sapling_nct.completed_subtree_index_and_root().expect(
                 "block should have completed a subtree before its final note commitment: \
                  already checked is_complete_subtree(),\
                  and that the block must complete a subtree",
-            )
-        };
-
-        let subtree = NoteCommitmentSubtree::new(index, height, node);
+            );
 
-        let mut batch = DiskWriteBatch::new();
-
-        batch.insert_sapling_subtree(upgrade_db, subtree);
-
-        upgrade_db
-            .write_batch(batch)
-            .expect("writing sapling note commitment subtrees should always succeed.");
+            write_sapling_subtree(upgrade_db, index, height, node);
+        };
 
         subtree_count += 1;
         prev_tree = Some(tree);
@@ -115,9 +110,8 @@ pub fn run(
             continue;
         }
 
-        let (index, node) = if tree.is_complete_subtree() {
-            tree.completed_subtree_index_and_root()
-                .expect("already checked is_complete_subtree()")
+        if let Some((index, node)) = tree.completed_subtree_index_and_root() {
+            write_orchard_subtree(upgrade_db, index, height, node);
         } else {
             let mut prev_tree = prev_tree
                 .take()
@@ -153,24 +147,52 @@ pub fn run(
                 }
             }
 
-            orchard_nct.completed_subtree_index_and_root().expect(
+            let (index, node) = orchard_nct.completed_subtree_index_and_root().expect(
                 "block should have completed a subtree before its final note commitment: \
                  already checked is_complete_subtree(),\
                  and that the block must complete a subtree",
-            )
+            );
+
+            write_orchard_subtree(upgrade_db, index, height, node);
         };
 
-        let subtree = NoteCommitmentSubtree::new(index, height, node);
+        subtree_count += 1;
+        prev_tree = Some(tree);
+    }
+}
 
-        let mut batch = DiskWriteBatch::new();
+/// Writes a Sapling note commitment subtree to `upgrade_db`.
+fn write_sapling_subtree(
+    upgrade_db: &ZebraDb,
+    index: NoteCommitmentSubtreeIndex,
+    height: Height,
+    node: sapling::tree::Node,
+) {
+    let subtree = NoteCommitmentSubtree::new(index, height, node);
 
-        batch.insert_orchard_subtree(upgrade_db, subtree);
+    let mut batch = DiskWriteBatch::new();
 
-        upgrade_db
-            .write_batch(batch)
-            .expect("writing orchard note commitment subtrees should always succeed.");
+    batch.insert_sapling_subtree(upgrade_db, subtree);
 
-        subtree_count += 1;
-        prev_tree = Some(tree);
-    }
+    upgrade_db
+        .write_batch(batch)
+        .expect("writing sapling note commitment subtrees should always succeed.");
+}
+
+/// Writes a Orchard note commitment subtree to `upgrade_db`.
+fn write_orchard_subtree(
+    upgrade_db: &ZebraDb,
+    index: NoteCommitmentSubtreeIndex,
+    height: Height,
+    node: orchard::tree::Node,
+) {
+    let subtree = NoteCommitmentSubtree::new(index, height, node);
+
+    let mut batch = DiskWriteBatch::new();
+
+    batch.insert_orchard_subtree(upgrade_db, subtree);
+
+    upgrade_db
+        .write_batch(batch)
+        .expect("writing orchard note commitment subtrees should always succeed.");
 }

From 63728dad26fd2e2adc385d8bfaee76c4f5ab7a13 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 31 Aug 2023 15:37:41 +1000
Subject: [PATCH 13/33] Use common note commitment list code

---
 zebra-chain/src/block.rs                      | 30 ++++++++++++++-----
 zebra-chain/src/parallel/tree.rs              | 21 ++-----------
 .../disk_format/upgrade/add_subtrees.rs       | 22 +++-----------
 3 files changed, 30 insertions(+), 43 deletions(-)

diff --git a/zebra-chain/src/block.rs b/zebra-chain/src/block.rs
index d472e930c7f..8b1e17ba35f 100644
--- a/zebra-chain/src/block.rs
+++ b/zebra-chain/src/block.rs
@@ -2,6 +2,8 @@
 
 use std::{collections::HashMap, fmt, ops::Neg, sync::Arc};
 
+use halo2::pasta::pallas;
+
 use crate::{
     amount::NegativeAllowed,
     block::merkle::AuthDataRoot,
@@ -152,16 +154,30 @@ impl Block {
 
     /// Access the [`orchard::Nullifier`]s from all transactions in this block.
     pub fn orchard_nullifiers(&self) -> impl Iterator<Item = &orchard::Nullifier> {
-        // Work around a compiler panic (ICE) with flat_map():
-        // https://github.com/rust-lang/rust/issues/105044
-        #[allow(clippy::needless_collect)]
-        let nullifiers: Vec<_> = self
-            .transactions
+        self.transactions
             .iter()
             .flat_map(|transaction| transaction.orchard_nullifiers())
-            .collect();
+    }
+
+    /// Access the [`sprout::NoteCommitment`]s from all transactions in this block.
+    pub fn sprout_note_commitments(&self) -> impl Iterator<Item = &sprout::NoteCommitment> {
+        self.transactions
+            .iter()
+            .flat_map(|transaction| transaction.sprout_note_commitments())
+    }
 
-        nullifiers.into_iter()
+    /// Access the [sapling note commitments](jubjub::Fq) from all transactions in this block.
+    pub fn sapling_note_commitments(&self) -> impl Iterator<Item = &jubjub::Fq> {
+        self.transactions
+            .iter()
+            .flat_map(|transaction| transaction.sapling_note_commitments())
+    }
+
+    /// Access the [orchard note commitments](pallas::Base) from all transactions in this block.
+    pub fn orchard_note_commitments(&self) -> impl Iterator<Item = &pallas::Base> {
+        self.transactions
+            .iter()
+            .flat_map(|transaction| transaction.orchard_note_commitments())
     }
 
     /// Count how many Sapling transactions exist in a block,
diff --git a/zebra-chain/src/parallel/tree.rs b/zebra-chain/src/parallel/tree.rs
index eb4190425bb..23d8f96d8cb 100644
--- a/zebra-chain/src/parallel/tree.rs
+++ b/zebra-chain/src/parallel/tree.rs
@@ -69,24 +69,9 @@ impl NoteCommitmentTrees {
             ..
         } = self.clone();
 
-        let sprout_note_commitments: Vec<_> = block
-            .transactions
-            .iter()
-            .flat_map(|tx| tx.sprout_note_commitments())
-            .cloned()
-            .collect();
-        let sapling_note_commitments: Vec<_> = block
-            .transactions
-            .iter()
-            .flat_map(|tx| tx.sapling_note_commitments())
-            .cloned()
-            .collect();
-        let orchard_note_commitments: Vec<_> = block
-            .transactions
-            .iter()
-            .flat_map(|tx| tx.orchard_note_commitments())
-            .cloned()
-            .collect();
+        let sprout_note_commitments: Vec<_> = block.sprout_note_commitments().cloned().collect();
+        let sapling_note_commitments: Vec<_> = block.sapling_note_commitments().cloned().collect();
+        let orchard_note_commitments: Vec<_> = block.orchard_note_commitments().cloned().collect();
 
         let mut sprout_result = None;
         let mut sapling_result = None;
diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
index 5109f45e19e..e959338461a 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -51,21 +51,14 @@ pub fn run(
                 .block(height.into())
                 .expect("height with note commitment tree should have block");
 
-            let sapling_note_commitments: Vec<_> = block
-                .transactions
-                .iter()
-                .flat_map(|tx| tx.sapling_note_commitments())
-                .cloned()
-                .collect();
-
-            for sapling_note_commitment in sapling_note_commitments {
+            for sapling_note_commitment in block.sapling_note_commitments() {
                 // Return early if there is a cancel signal.
                 if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
                     return;
                 }
 
                 sapling_nct
-                    .append(sapling_note_commitment)
+                    .append(*sapling_note_commitment)
                     .expect("finalized notes should append successfully");
 
                 // The loop always breaks on this condition,
@@ -122,21 +115,14 @@ pub fn run(
                 .block(height.into())
                 .expect("height with note commitment tree should have block");
 
-            let orchard_note_commitments: Vec<_> = block
-                .transactions
-                .iter()
-                .flat_map(|tx| tx.orchard_note_commitments())
-                .cloned()
-                .collect();
-
-            for orchard_note_commitment in orchard_note_commitments {
+            for orchard_note_commitment in block.orchard_note_commitments() {
                 // Return early if there is a cancel signal.
                 if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
                     return;
                 }
 
                 orchard_nct
-                    .append(orchard_note_commitment)
+                    .append(*orchard_note_commitment)
                     .expect("finalized notes should append successfully");
 
                 // The loop always breaks on this condition,

From 8e2b2841147165780a77a820b645d4c5e08f1844 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 31 Aug 2023 15:40:54 +1000
Subject: [PATCH 14/33] Fix subtree completion condition and add asserts

---
 .../disk_format/upgrade/add_subtrees.rs       | 28 ++++++++++++++++---
 1 file changed, 24 insertions(+), 4 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
index e959338461a..69417e88d11 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -34,12 +34,19 @@ pub fn run(
 
         // Blocks cannot complete multiple level 16 subtrees,
         // the subtree index can increase by a maximum of 1 every ~20 blocks.
-        if subtree_address.index() <= subtree_count {
+        // If this block does complete a subtree, the subtree is either completed by a note before
+        // the final note (so the final note is in the next subtree), or by the final note
+        // (so the final note is the end of this subtree).
+        if subtree_address.index() <= subtree_count && !tree.is_complete_subtree() {
             prev_tree = Some(tree);
             continue;
         }
 
         if let Some((index, node)) = tree.completed_subtree_index_and_root() {
+            assert_eq!(
+                index, subtree_count,
+                "trees are inserted in order with no gaps"
+            );
             write_sapling_subtree(upgrade_db, index, height, node);
         } else {
             let mut prev_tree = prev_tree
@@ -75,6 +82,10 @@ pub fn run(
                  and that the block must complete a subtree",
             );
 
+            assert_eq!(
+                index, subtree_count,
+                "trees are inserted in order with no gaps"
+            );
             write_sapling_subtree(upgrade_db, index, height, node);
         };
 
@@ -96,14 +107,19 @@ pub fn run(
             continue;
         };
 
-        // Blocks cannot complete multiple level 16 subtrees,
-        // the subtree index can increase by a maximum of 1 every ~20 blocks.
-        if subtree_address.index() <= subtree_count {
+        // Blocks cannot complete multiple level 16 subtrees.
+        // If a block does complete a subtree, it is either inside the block, or at the end.
+        // (See the detailed comment for Sapling.)
+        if subtree_address.index() <= subtree_count && !tree.is_complete_subtree() {
             prev_tree = Some(tree);
             continue;
         }
 
         if let Some((index, node)) = tree.completed_subtree_index_and_root() {
+            assert_eq!(
+                index, subtree_count,
+                "trees are inserted in order with no gaps"
+            );
             write_orchard_subtree(upgrade_db, index, height, node);
         } else {
             let mut prev_tree = prev_tree
@@ -139,6 +155,10 @@ pub fn run(
                  and that the block must complete a subtree",
             );
 
+            assert_eq!(
+                index, subtree_count,
+                "trees are inserted in order with no gaps"
+            );
             write_orchard_subtree(upgrade_db, index, height, node);
         };
 

From b110ddee6ba359a2766637cc3f91c72df7f14394 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 31 Aug 2023 15:55:56 +1000
Subject: [PATCH 15/33] Simplify subtree API and avoid exposing Address

---
 zebra-chain/src/orchard/tree.rs                | 18 +++++++++++-------
 zebra-chain/src/sapling/tree.rs                | 18 +++++++++++-------
 zebra-chain/src/subtree.rs                     |  8 ++++++++
 .../disk_format/upgrade/add_subtrees.rs        | 16 ++++++++--------
 4 files changed, 38 insertions(+), 22 deletions(-)

diff --git a/zebra-chain/src/orchard/tree.rs b/zebra-chain/src/orchard/tree.rs
index 2b035865d78..58b6980e4c3 100644
--- a/zebra-chain/src/orchard/tree.rs
+++ b/zebra-chain/src/orchard/tree.rs
@@ -363,26 +363,30 @@ impl NoteCommitmentTree {
             .is_complete_subtree(TRACKED_SUBTREE_HEIGHT.into())
     }
 
-    /// Returns the subtree address at [`TRACKED_SUBTREE_HEIGHT`].
+    /// Returns the subtree index at [`TRACKED_SUBTREE_HEIGHT`].
     /// Returns `None` if the tree is empty.
-    pub fn subtree_address(&self) -> Option<incrementalmerkletree::Address> {
+    #[allow(clippy::unwrap_in_result)]
+    pub fn subtree_index(&self) -> Option<NoteCommitmentSubtreeIndex> {
         let tree = self.frontier()?;
 
-        Some(incrementalmerkletree::Address::above_position(
+        let index = incrementalmerkletree::Address::above_position(
             TRACKED_SUBTREE_HEIGHT.into(),
             tree.position(),
-        ))
+        )
+        .index()
+        .try_into()
+        .expect("fits in u16");
+
+        Some(index)
     }
 
     /// Returns subtree index and root if the most recently appended leaf completes the subtree
-    #[allow(clippy::unwrap_in_result)]
     pub fn completed_subtree_index_and_root(&self) -> Option<(NoteCommitmentSubtreeIndex, Node)> {
         if !self.is_complete_subtree() {
             return None;
         }
 
-        let address = self.subtree_address()?;
-        let index = address.index().try_into().expect("should fit in u16");
+        let index = self.subtree_index()?;
         let root = self.frontier()?.root(Some(TRACKED_SUBTREE_HEIGHT.into()));
 
         Some((index, root))
diff --git a/zebra-chain/src/sapling/tree.rs b/zebra-chain/src/sapling/tree.rs
index abd07519194..27fd5e3c680 100644
--- a/zebra-chain/src/sapling/tree.rs
+++ b/zebra-chain/src/sapling/tree.rs
@@ -348,26 +348,30 @@ impl NoteCommitmentTree {
             .is_complete_subtree(TRACKED_SUBTREE_HEIGHT.into())
     }
 
-    /// Returns the subtree address at [`TRACKED_SUBTREE_HEIGHT`].
+    /// Returns the subtree index at [`TRACKED_SUBTREE_HEIGHT`].
     /// Returns `None` if the tree is empty.
-    pub fn subtree_address(&self) -> Option<incrementalmerkletree::Address> {
+    #[allow(clippy::unwrap_in_result)]
+    pub fn subtree_index(&self) -> Option<NoteCommitmentSubtreeIndex> {
         let tree = self.frontier()?;
 
-        Some(incrementalmerkletree::Address::above_position(
+        let index = incrementalmerkletree::Address::above_position(
             TRACKED_SUBTREE_HEIGHT.into(),
             tree.position(),
-        ))
+        )
+        .index()
+        .try_into()
+        .expect("fits in u16");
+
+        Some(index)
     }
 
     /// Returns subtree index and root if the most recently appended leaf completes the subtree
-    #[allow(clippy::unwrap_in_result)]
     pub fn completed_subtree_index_and_root(&self) -> Option<(NoteCommitmentSubtreeIndex, Node)> {
         if !self.is_complete_subtree() {
             return None;
         }
 
-        let address = self.subtree_address()?;
-        let index = address.index().try_into().expect("should fit in u16");
+        let index = self.subtree_index()?;
         let root = self.frontier()?.root(Some(TRACKED_SUBTREE_HEIGHT.into()));
 
         Some((index, root))
diff --git a/zebra-chain/src/subtree.rs b/zebra-chain/src/subtree.rs
index 7c8a9b456b0..3d34a3b46c4 100644
--- a/zebra-chain/src/subtree.rs
+++ b/zebra-chain/src/subtree.rs
@@ -28,6 +28,14 @@ impl TryFrom<u64> for NoteCommitmentSubtreeIndex {
     }
 }
 
+// If we want to automatically convert NoteCommitmentSubtreeIndex to the generic integer literal
+// type, we can only implement conversion into u64. (Or u16, but not both.)
+impl From<NoteCommitmentSubtreeIndex> for u64 {
+    fn from(value: NoteCommitmentSubtreeIndex) -> Self {
+        value.0.into()
+    }
+}
+
 /// Subtree root of Sapling or Orchard note commitment tree,
 /// with its associated block height and subtree index.
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
index 69417e88d11..26ec608a038 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -27,7 +27,7 @@ pub fn run(
         }
 
         // Empty note commitment trees can't contain subtrees.
-        let Some(subtree_address) = tree.subtree_address() else {
+        let Some(end_of_block_subtree_index) = tree.subtree_index() else {
             prev_tree = Some(tree);
             continue;
         };
@@ -37,14 +37,14 @@ pub fn run(
         // If this block does complete a subtree, the subtree is either completed by a note before
         // the final note (so the final note is in the next subtree), or by the final note
         // (so the final note is the end of this subtree).
-        if subtree_address.index() <= subtree_count && !tree.is_complete_subtree() {
+        if end_of_block_subtree_index.0 <= subtree_count && !tree.is_complete_subtree() {
             prev_tree = Some(tree);
             continue;
         }
 
         if let Some((index, node)) = tree.completed_subtree_index_and_root() {
             assert_eq!(
-                index, subtree_count,
+                index.0, subtree_count,
                 "trees are inserted in order with no gaps"
             );
             write_sapling_subtree(upgrade_db, index, height, node);
@@ -83,7 +83,7 @@ pub fn run(
             );
 
             assert_eq!(
-                index, subtree_count,
+                index.0, subtree_count,
                 "trees are inserted in order with no gaps"
             );
             write_sapling_subtree(upgrade_db, index, height, node);
@@ -102,7 +102,7 @@ pub fn run(
         }
 
         // Empty note commitment trees can't contain subtrees.
-        let Some(subtree_address) = tree.subtree_address() else {
+        let Some(end_of_block_subtree_index) = tree.subtree_index() else {
             prev_tree = Some(tree);
             continue;
         };
@@ -110,14 +110,14 @@ pub fn run(
         // Blocks cannot complete multiple level 16 subtrees.
         // If a block does complete a subtree, it is either inside the block, or at the end.
         // (See the detailed comment for Sapling.)
-        if subtree_address.index() <= subtree_count && !tree.is_complete_subtree() {
+        if end_of_block_subtree_index.0 <= subtree_count && !tree.is_complete_subtree() {
             prev_tree = Some(tree);
             continue;
         }
 
         if let Some((index, node)) = tree.completed_subtree_index_and_root() {
             assert_eq!(
-                index, subtree_count,
+                index.0, subtree_count,
                 "trees are inserted in order with no gaps"
             );
             write_orchard_subtree(upgrade_db, index, height, node);
@@ -156,7 +156,7 @@ pub fn run(
             );
 
             assert_eq!(
-                index, subtree_count,
+                index.0, subtree_count,
                 "trees are inserted in order with no gaps"
             );
             write_orchard_subtree(upgrade_db, index, height, node);

From 96183d97f8996ebd89e3c41e02ae4a1fa117007a Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 31 Aug 2023 16:08:34 +1000
Subject: [PATCH 16/33] Fix API compatibility when Arcs are removed

---
 .../finalized_state/disk_format/upgrade/add_subtrees.rs   | 4 ++--
 .../src/service/finalized_state/zebra_db/shielded.rs      | 8 ++++----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
index 26ec608a038..928bd8ddabb 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -178,7 +178,7 @@ fn write_sapling_subtree(
 
     let mut batch = DiskWriteBatch::new();
 
-    batch.insert_sapling_subtree(upgrade_db, subtree);
+    batch.insert_sapling_subtree(upgrade_db, &subtree);
 
     upgrade_db
         .write_batch(batch)
@@ -196,7 +196,7 @@ fn write_orchard_subtree(
 
     let mut batch = DiskWriteBatch::new();
 
-    batch.insert_orchard_subtree(upgrade_db, subtree);
+    batch.insert_orchard_subtree(upgrade_db, &subtree);
 
     upgrade_db
         .write_batch(batch)
diff --git a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
index ce3b5373a77..1f874f4a418 100644
--- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
@@ -400,11 +400,11 @@ impl DiskWriteBatch {
         }
 
         if let Some(subtree) = trees.sapling_subtree {
-            self.insert_sapling_subtree(zebra_db, subtree);
+            self.insert_sapling_subtree(zebra_db, &subtree);
         }
 
         if let Some(subtree) = trees.orchard_subtree {
-            self.insert_orchard_subtree(zebra_db, subtree);
+            self.insert_orchard_subtree(zebra_db, &subtree);
         }
 
         self.prepare_history_batch(db, finalized)
@@ -416,7 +416,7 @@ impl DiskWriteBatch {
     pub fn insert_sapling_subtree(
         &mut self,
         zebra_db: &ZebraDb,
-        subtree: Arc<NoteCommitmentSubtree<sapling::tree::Node>>,
+        subtree: &NoteCommitmentSubtree<sapling::tree::Node>,
     ) {
         let sapling_subtree_cf = zebra_db
             .db
@@ -452,7 +452,7 @@ impl DiskWriteBatch {
     pub fn insert_orchard_subtree(
         &mut self,
         zebra_db: &ZebraDb,
-        subtree: Arc<NoteCommitmentSubtree<orchard::tree::Node>>,
+        subtree: &NoteCommitmentSubtree<orchard::tree::Node>,
     ) {
         let orchard_subtree_cf = zebra_db
             .db

From a024e07e7a37c37535a74c9763e340ce104aba5d Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 31 Aug 2023 16:13:24 +1000
Subject: [PATCH 17/33] Log when each subtree is added

---
 .../finalized_state/disk_format/upgrade/add_subtrees.rs       | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
index 928bd8ddabb..ee56f257a60 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -183,6 +183,8 @@ fn write_sapling_subtree(
     upgrade_db
         .write_batch(batch)
         .expect("writing sapling note commitment subtrees should always succeed.");
+
+    info!(?height, index = ?index.0, "calculated and added sapling subtree");
 }
 
 /// Writes a Orchard note commitment subtree to `upgrade_db`.
@@ -201,4 +203,6 @@ fn write_orchard_subtree(
     upgrade_db
         .write_batch(batch)
         .expect("writing orchard note commitment subtrees should always succeed.");
+
+    info!(?height, index = ?index.0, "calculated and added orchard subtree");
 }

From 307f8313155c065681aa2a73285a42ca13a888ae Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 31 Aug 2023 16:27:54 +1000
Subject: [PATCH 18/33] If a format change is cancelled, don't mark the
 database as upgraded or do format checks

---
 .../finalized_state/disk_format/upgrade.rs    | 28 +++++++++++--------
 .../disk_format/upgrade/add_subtrees.rs       | 15 ++++++----
 2 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
index 0c5da2d1ca6..300b154892b 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
@@ -68,7 +68,8 @@ pub struct DbFormatChangeThreadHandle {
     /// A handle that can wait for the running format change thread to finish.
     ///
     /// Panics from this thread are propagated into Zebra's state service.
-    update_task: Option<Arc<JoinHandle<()>>>,
+    /// The task returns an error if the upgrade was cancelled by a shutdown.
+    update_task: Option<Arc<JoinHandle<Result<(), CancelFormatChange>>>>,
 
     /// A channel that tells the running format thread to finish early.
     /// A channel that tells the running format thread to finish early.
@@ -161,7 +162,7 @@ impl DbFormatChange {
                     initial_tip_height,
                     upgrade_db,
                     cancel_receiver,
-                );
+                )
             })
         });
 
@@ -178,7 +179,7 @@ impl DbFormatChange {
     /// Apply this format change to the database.
     ///
     /// Format changes are launched in an independent `std::thread` by `apply_format_upgrade()`.
-    /// This thread runs until the upgrade is finished.
+    /// This thread runs until the upgrade is finished or cancelled.
     ///
     /// See `apply_format_upgrade()` for details.
     fn apply_format_change(
@@ -188,7 +189,7 @@ impl DbFormatChange {
         initial_tip_height: Option<Height>,
         upgrade_db: ZebraDb,
         cancel_receiver: mpsc::Receiver<CancelFormatChange>,
-    ) {
+    ) -> Result<(), CancelFormatChange> {
         match self {
             // Perform any required upgrades, then mark the state as upgraded.
             Upgrade { .. } => self.apply_format_upgrade(
@@ -197,7 +198,7 @@ impl DbFormatChange {
                 initial_tip_height,
                 upgrade_db.clone(),
                 cancel_receiver,
-            ),
+            )?,
 
             NewlyCreated { .. } => {
                 Self::mark_as_newly_created(&config, network);
@@ -225,18 +226,21 @@ impl DbFormatChange {
         // - since this Zebra code knows how to de-duplicate trees, downgrades using this code
         //   still know how to make sure trees are unique
         Self::check_for_duplicate_trees(upgrade_db);
+
+        Ok(())
     }
 
     /// Apply any required format updates to the database.
     /// Format changes should be launched in an independent `std::thread`.
     ///
     /// If `cancel_receiver` gets a message, or its sender is dropped,
-    /// the format change stops running early.
+    /// the format change stops running early, and returns an error.
     ///
     /// See the format upgrade design docs for more details:
     /// <https://github.com/ZcashFoundation/zebra/blob/main/book/src/dev/state-db-upgrades.md#design>
     //
     // New format upgrades must be added to the *end* of this method.
+    #[allow(clippy::unwrap_in_result)]
     fn apply_format_upgrade(
         self,
         config: Config,
@@ -244,7 +248,7 @@ impl DbFormatChange {
         initial_tip_height: Option<Height>,
         db: ZebraDb,
         cancel_receiver: mpsc::Receiver<CancelFormatChange>,
-    ) {
+    ) -> Result<(), CancelFormatChange> {
         let Upgrade {
             newer_running_version,
             older_disk_version,
@@ -272,7 +276,7 @@ impl DbFormatChange {
                 "empty database is fully upgraded"
             );
 
-            return;
+            return Ok(());
         };
 
         // Note commitment tree de-duplication database upgrade task.
@@ -294,7 +298,7 @@ impl DbFormatChange {
             for (height, tree) in db.sapling_tree_by_height_range(Height(1)..=initial_tip_height) {
                 // Return early if there is a cancel signal.
                 if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
-                    return;
+                    return Err(CancelFormatChange);
                 }
 
                 // Delete any duplicate trees.
@@ -321,7 +325,7 @@ impl DbFormatChange {
             for (height, tree) in db.orchard_tree_by_height_range(Height(1)..=initial_tip_height) {
                 // Return early if there is a cancel signal.
                 if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
-                    return;
+                    return Err(CancelFormatChange);
                 }
 
                 // Delete any duplicate trees.
@@ -351,7 +355,7 @@ impl DbFormatChange {
 
         // Check if we need to add note commitment subtrees to the database.
         if older_disk_version < version_for_adding_subtrees {
-            add_subtrees::run(initial_tip_height, &db, cancel_receiver);
+            add_subtrees::run(initial_tip_height, &db, cancel_receiver)?;
 
             // Before marking the state as upgraded, check that the upgrade completed successfully.
             //
@@ -375,6 +379,8 @@ impl DbFormatChange {
         // Run the latest format upgrade code after the other upgrades are complete,
         // then mark the format as upgraded. The code should check `cancel_receiver`
         // every time it runs its inner update loop.
+
+        Ok(())
     }
 
     /// Check that note commitment trees were correctly de-duplicated.
diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
index ee56f257a60..68da2b93318 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -13,17 +13,20 @@ use crate::service::finalized_state::{
 };
 
 /// Runs disk format upgrade for adding Sapling and Orchard note commitment subtrees to database.
+///
+/// Returns `Ok` if the upgrade completed, and `Err` if it was cancelled.
+#[allow(clippy::unwrap_in_result)]
 pub fn run(
     initial_tip_height: Height,
     upgrade_db: &ZebraDb,
     cancel_receiver: mpsc::Receiver<CancelFormatChange>,
-) {
+) -> Result<(), CancelFormatChange> {
     let mut subtree_count = 0;
     let mut prev_tree: Option<_> = None;
     for (height, tree) in upgrade_db.sapling_tree_by_height_range(..=initial_tip_height) {
         // Return early if there is a cancel signal.
         if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
-            return;
+            return Err(CancelFormatChange);
         }
 
         // Empty note commitment trees can't contain subtrees.
@@ -61,7 +64,7 @@ pub fn run(
             for sapling_note_commitment in block.sapling_note_commitments() {
                 // Return early if there is a cancel signal.
                 if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
-                    return;
+                    return Err(CancelFormatChange);
                 }
 
                 sapling_nct
@@ -98,7 +101,7 @@ pub fn run(
     for (height, tree) in upgrade_db.orchard_tree_by_height_range(..=initial_tip_height) {
         // Return early if there is a cancel signal.
         if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
-            return;
+            return Err(CancelFormatChange);
         }
 
         // Empty note commitment trees can't contain subtrees.
@@ -134,7 +137,7 @@ pub fn run(
             for orchard_note_commitment in block.orchard_note_commitments() {
                 // Return early if there is a cancel signal.
                 if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
-                    return;
+                    return Err(CancelFormatChange);
                 }
 
                 orchard_nct
@@ -165,6 +168,8 @@ pub fn run(
         subtree_count += 1;
         prev_tree = Some(tree);
     }
+
+    Ok(())
 }
 
 /// Writes a Sapling note commitment subtree to `upgrade_db`.

From dee2862ea68e8a6789f900da502b61b86a46964b Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 31 Aug 2023 16:38:25 +1000
Subject: [PATCH 19/33] Log subtree progress about once every two minutes

---
 .../disk_format/upgrade/add_subtrees.rs              | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
index 68da2b93318..208999345cb 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -189,7 +189,11 @@ fn write_sapling_subtree(
         .write_batch(batch)
         .expect("writing sapling note commitment subtrees should always succeed.");
 
-    info!(?height, index = ?index.0, "calculated and added sapling subtree");
+    if index.0 % 100 == 0 {
+        info!(?height, index = ?index.0, "calculated and added sapling subtree");
+    }
+    // This log happens about once per second on recent machines with SSD disks.
+    debug!(?height, index = ?index.0, ?node, "calculated and added sapling subtree");
 }
 
 /// Writes a Orchard note commitment subtree to `upgrade_db`.
@@ -209,5 +213,9 @@ fn write_orchard_subtree(
         .write_batch(batch)
         .expect("writing orchard note commitment subtrees should always succeed.");
 
-    info!(?height, index = ?index.0, "calculated and added orchard subtree");
+    // This log happens about once per second on recent machines with SSD disks.
+    if index.0 % 100 == 0 {
+        info!(?height, index = ?index.0, "calculated and added orchard subtree");
+    }
+    debug!(?height, index = ?index.0, ?node, "calculated and added orchard subtree");
 }

From b366fd3195b67da2a0ecc3408f60b6096b437484 Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Thu, 31 Aug 2023 13:29:52 -0400
Subject: [PATCH 20/33] Adds a state validity check for subtrees upgrade

---
 .../finalized_state/disk_format/upgrade.rs    |  11 +-
 .../disk_format/upgrade/add_subtrees.rs       | 209 +++++++++++++++++-
 2 files changed, 213 insertions(+), 7 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
index 300b154892b..915ac234560 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
@@ -197,7 +197,7 @@ impl DbFormatChange {
                 network,
                 initial_tip_height,
                 upgrade_db.clone(),
-                cancel_receiver,
+                &cancel_receiver,
             )?,
 
             NewlyCreated { .. } => {
@@ -225,7 +225,8 @@ impl DbFormatChange {
         // - an empty state doesn't have any trees, so it can't have duplicate trees
         // - since this Zebra code knows how to de-duplicate trees, downgrades using this code
         //   still know how to make sure trees are unique
-        Self::check_for_duplicate_trees(upgrade_db);
+        Self::check_for_duplicate_trees(upgrade_db.clone());
+        add_subtrees::check(&upgrade_db, &cancel_receiver)?;
 
         Ok(())
     }
@@ -247,7 +248,7 @@ impl DbFormatChange {
         network: Network,
         initial_tip_height: Option<Height>,
         db: ZebraDb,
-        cancel_receiver: mpsc::Receiver<CancelFormatChange>,
+        cancel_receiver: &mpsc::Receiver<CancelFormatChange>,
     ) -> Result<(), CancelFormatChange> {
         let Upgrade {
             newer_running_version,
@@ -358,9 +359,7 @@ impl DbFormatChange {
             add_subtrees::run(initial_tip_height, &db, cancel_receiver)?;
 
             // Before marking the state as upgraded, check that the upgrade completed successfully.
-            //
-            // TODO: do this check in all the same places as check_for_duplicate_trees()
-            //Self::check_for_continuous_subtrees(db);
+            add_subtrees::check(&db, cancel_receiver)?;
 
             // Mark the database as upgraded. Zebra won't repeat the upgrade anymore once the
             // database is marked, so the upgrade MUST be complete at this point.
diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
index 208999345cb..ab9a93ca900 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -19,7 +19,7 @@ use crate::service::finalized_state::{
 pub fn run(
     initial_tip_height: Height,
     upgrade_db: &ZebraDb,
-    cancel_receiver: mpsc::Receiver<CancelFormatChange>,
+    cancel_receiver: &mpsc::Receiver<CancelFormatChange>,
 ) -> Result<(), CancelFormatChange> {
     let mut subtree_count = 0;
     let mut prev_tree: Option<_> = None;
@@ -172,6 +172,191 @@ pub fn run(
     Ok(())
 }
 
+/// Check that note commitment subtrees were correctly added.
+///
+/// # Panics
+///
+/// If a note commitment subtree is missing or incorrect.
+pub fn check(
+    upgrade_db: &ZebraDb,
+    cancel_receiver: &mpsc::Receiver<CancelFormatChange>,
+) -> Result<(), CancelFormatChange> {
+    let mut is_valid = true;
+
+    let mut subtree_count = 0;
+    let mut prev_tree: Option<_> = None;
+    for (height, tree) in upgrade_db.sapling_tree_by_height_range(..) {
+        // Return early if there is a cancel signal.
+        if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
+            return Err(CancelFormatChange);
+        }
+
+        // Empty note commitment trees can't contain subtrees.
+        let Some(end_of_block_subtree_index) = tree.subtree_index() else {
+            prev_tree = Some(tree);
+            continue;
+        };
+
+        // Blocks cannot complete multiple level 16 subtrees,
+        // the subtree index can increase by a maximum of 1 every ~20 blocks.
+        // If this block does complete a subtree, the subtree is either completed by a note before
+        // the final note (so the final note is in the next subtree), or by the final note
+        // (so the final note is the end of this subtree).
+        if end_of_block_subtree_index.0 <= subtree_count && !tree.is_complete_subtree() {
+            prev_tree = Some(tree);
+            continue;
+        }
+
+        if let Some((index, node)) = tree.completed_subtree_index_and_root() {
+            assert_eq!(
+                index.0, subtree_count,
+                "trees are inserted in order with no gaps"
+            );
+
+            if !sapling_subtree_exists(upgrade_db, index, height, node) {
+                error!(?height, ?index, "missing sapling subtree");
+                is_valid = false;
+            }
+        } else {
+            let mut prev_tree = prev_tree
+                .take()
+                .expect("should have some previous sapling frontier");
+            let sapling_nct = Arc::make_mut(&mut prev_tree);
+
+            let block = upgrade_db
+                .block(height.into())
+                .expect("height with note commitment tree should have block");
+
+            for sapling_note_commitment in block.sapling_note_commitments() {
+                // Return early if there is a cancel signal.
+                if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
+                    return Err(CancelFormatChange);
+                }
+
+                sapling_nct
+                    .append(*sapling_note_commitment)
+                    .expect("finalized notes should append successfully");
+
+                // The loop always breaks on this condition,
+                // because we checked the block has enough commitments,
+                // and that the final commitment in the block doesn't complete a subtree.
+                if sapling_nct.is_complete_subtree() {
+                    break;
+                }
+            }
+
+            let (index, node) = sapling_nct.completed_subtree_index_and_root().expect(
+                "block should have completed a subtree before its final note commitment: \
+                 already checked is_complete_subtree(),\
+                 and that the block must complete a subtree",
+            );
+
+            assert_eq!(
+                index.0, subtree_count,
+                "trees are inserted in order with no gaps"
+            );
+
+            if !sapling_subtree_exists(upgrade_db, index, height, node) {
+                error!(?height, ?index, "missing sapling subtree");
+                is_valid = false;
+            }
+        };
+
+        subtree_count += 1;
+        prev_tree = Some(tree);
+    }
+
+    let mut subtree_count = 0;
+    let mut prev_tree: Option<_> = None;
+    for (height, tree) in upgrade_db.orchard_tree_by_height_range(..) {
+        // Return early if there is a cancel signal.
+        if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
+            return Err(CancelFormatChange);
+        }
+
+        // Empty note commitment trees can't contain subtrees.
+        let Some(end_of_block_subtree_index) = tree.subtree_index() else {
+            prev_tree = Some(tree);
+            continue;
+        };
+
+        // Blocks cannot complete multiple level 16 subtrees.
+        // If a block does complete a subtree, it is either inside the block, or at the end.
+        // (See the detailed comment for Sapling.)
+        if end_of_block_subtree_index.0 <= subtree_count && !tree.is_complete_subtree() {
+            prev_tree = Some(tree);
+            continue;
+        }
+
+        if let Some((index, node)) = tree.completed_subtree_index_and_root() {
+            assert_eq!(
+                index.0, subtree_count,
+                "trees are inserted in order with no gaps"
+            );
+
+            if !orchard_subtree_exists(upgrade_db, index, height, node) {
+                error!(?height, ?index, "missing orchard subtree");
+                is_valid = false;
+            }
+        } else {
+            let mut prev_tree = prev_tree
+                .take()
+                .expect("should have some previous orchard frontier");
+            let orchard_nct = Arc::make_mut(&mut prev_tree);
+
+            let block = upgrade_db
+                .block(height.into())
+                .expect("height with note commitment tree should have block");
+
+            for orchard_note_commitment in block.orchard_note_commitments() {
+                // Return early if there is a cancel signal.
+                if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
+                    return Err(CancelFormatChange);
+                }
+
+                orchard_nct
+                    .append(*orchard_note_commitment)
+                    .expect("finalized notes should append successfully");
+
+                // The loop always breaks on this condition,
+                // because we checked the block has enough commitments,
+                // and that the final commitment in the block doesn't complete a subtree.
+                if orchard_nct.is_complete_subtree() {
+                    break;
+                }
+            }
+
+            let (index, node) = orchard_nct.completed_subtree_index_and_root().expect(
+                "block should have completed a subtree before its final note commitment: \
+                 already checked is_complete_subtree(),\
+                 and that the block must complete a subtree",
+            );
+
+            assert_eq!(
+                index.0, subtree_count,
+                "trees are inserted in order with no gaps"
+            );
+
+            if !orchard_subtree_exists(upgrade_db, index, height, node) {
+                error!(?height, ?index, "missing orchard subtree");
+                is_valid = false;
+            }
+        };
+
+        subtree_count += 1;
+        prev_tree = Some(tree);
+    }
+
+    if !is_valid {
+        panic!(
+            "found duplicate sapling or orchard trees \
+                     after running de-duplicate tree upgrade"
+        );
+    }
+
+    Ok(())
+}
+
 /// Writes a Sapling note commitment subtree to `upgrade_db`.
 fn write_sapling_subtree(
     upgrade_db: &ZebraDb,
@@ -219,3 +404,25 @@ fn write_orchard_subtree(
     }
     debug!(?height, index = ?index.0, ?node, "calculated and added orchard subtree");
 }
+
+/// Checks that a Sapling note commitment subtree exists in `upgrade_db`.
+fn sapling_subtree_exists(
+    upgrade_db: &ZebraDb,
+    index: NoteCommitmentSubtreeIndex,
+    height: Height,
+    node: sapling::tree::Node,
+) -> bool {
+    upgrade_db.sapling_subtree_by_index(index)
+        == Some(NoteCommitmentSubtree::new(index, height, node))
+}
+
+/// Checks that a Orchard note commitment subtree exists in `upgrade_db`.
+fn orchard_subtree_exists(
+    upgrade_db: &ZebraDb,
+    index: NoteCommitmentSubtreeIndex,
+    height: Height,
+    node: orchard::tree::Node,
+) -> bool {
+    upgrade_db.orchard_subtree_by_index(index)
+        == Some(NoteCommitmentSubtree::new(index, height, node))
+}

From 8fc706bc7cea1610b53cbe9d14ddaa7ff51055fa Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 1 Sep 2023 08:15:21 +1000
Subject: [PATCH 21/33] Orchard is faster, decrease log interval

---
 .../finalized_state/disk_format/upgrade/add_subtrees.rs       | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
index ab9a93ca900..d1de366d31c 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -398,10 +398,10 @@ fn write_orchard_subtree(
         .write_batch(batch)
         .expect("writing orchard note commitment subtrees should always succeed.");
 
-    // This log happens about once per second on recent machines with SSD disks.
-    if index.0 % 100 == 0 {
+    if index.0 % 300 == 0 {
         info!(?height, index = ?index.0, "calculated and added orchard subtree");
     }
+    // This log happens about 3 times per second on recent machines with SSD disks.
     debug!(?height, index = ?index.0, ?node, "calculated and added orchard subtree");
 }
 

From e2735882f7f42eb3df03296a9546e9e2e4a0f7d4 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 1 Sep 2023 11:38:09 +1000
Subject: [PATCH 22/33] Clarify subtree index docs

---
 zebra-chain/src/orchard/tree.rs | 1 +
 zebra-chain/src/sapling/tree.rs | 1 +
 2 files changed, 2 insertions(+)

diff --git a/zebra-chain/src/orchard/tree.rs b/zebra-chain/src/orchard/tree.rs
index 58b6980e4c3..2e6cc6458b1 100644
--- a/zebra-chain/src/orchard/tree.rs
+++ b/zebra-chain/src/orchard/tree.rs
@@ -364,6 +364,7 @@ impl NoteCommitmentTree {
     }
 
     /// Returns the subtree index at [`TRACKED_SUBTREE_HEIGHT`].
+    /// This is the number of complete or incomplete subtrees that are currently in the tree.
     /// Returns `None` if the tree is empty.
     #[allow(clippy::unwrap_in_result)]
     pub fn subtree_index(&self) -> Option<NoteCommitmentSubtreeIndex> {
diff --git a/zebra-chain/src/sapling/tree.rs b/zebra-chain/src/sapling/tree.rs
index 27fd5e3c680..58a2a0ca743 100644
--- a/zebra-chain/src/sapling/tree.rs
+++ b/zebra-chain/src/sapling/tree.rs
@@ -349,6 +349,7 @@ impl NoteCommitmentTree {
     }
 
     /// Returns the subtree index at [`TRACKED_SUBTREE_HEIGHT`].
+    /// This is the number of complete or incomplete subtrees that are currently in the tree.
     /// Returns `None` if the tree is empty.
     #[allow(clippy::unwrap_in_result)]
     pub fn subtree_index(&self) -> Option<NoteCommitmentSubtreeIndex> {

From 9dbb24a4fe46505598f679ecc0c8adc14c149969 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 1 Sep 2023 11:38:22 +1000
Subject: [PATCH 23/33] Move a log to the correct location

---
 .../src/service/finalized_state/disk_format/upgrade.rs | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
index 915ac234560..5c62331efa7 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
@@ -366,11 +366,6 @@ impl DbFormatChange {
             Self::mark_as_upgraded_to(&version_for_adding_subtrees, &config, network);
         }
 
-        info!(
-            ?newer_running_version,
-            "Zebra automatically upgraded the database format to:"
-        );
-
         // # New Upgrades Usually Go Here
         //
         // New code goes above this comment!
@@ -379,6 +374,11 @@ impl DbFormatChange {
         // then mark the format as upgraded. The code should check `cancel_receiver`
         // every time it runs its inner update loop.
 
+        info!(
+            ?newer_running_version,
+            "Zebra automatically upgraded the database format to:"
+        );
+
         Ok(())
     }
 

From d29576564843fcc890c724bb896d5dfd2b7b5ae4 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 1 Sep 2023 11:46:51 +1000
Subject: [PATCH 24/33] Refactor subtree upgrade to remove duplicate inverted
 loop conditions

---
 .../disk_format/upgrade/add_subtrees.rs       | 46 ++++++++++---------
 1 file changed, 24 insertions(+), 22 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
index d1de366d31c..cf77d14e767 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -36,22 +36,23 @@ pub fn run(
         };
 
         // Blocks cannot complete multiple level 16 subtrees,
-        // the subtree index can increase by a maximum of 1 every ~20 blocks.
+        // so the subtree index can increase by a maximum of 1 every ~20 blocks.
         // If this block does complete a subtree, the subtree is either completed by a note before
         // the final note (so the final note is in the next subtree), or by the final note
         // (so the final note is the end of this subtree).
-        if end_of_block_subtree_index.0 <= subtree_count && !tree.is_complete_subtree() {
-            prev_tree = Some(tree);
-            continue;
-        }
 
         if let Some((index, node)) = tree.completed_subtree_index_and_root() {
+            // If the leaf at the end of the block is the final leaf in a subtree,
+            // we already have that subtree root available in the tree.
             assert_eq!(
                 index.0, subtree_count,
                 "trees are inserted in order with no gaps"
             );
             write_sapling_subtree(upgrade_db, index, height, node);
-        } else {
+            subtree_count += 1;
+        } else if end_of_block_subtree_index.0 > subtree_count {
+            // If the leaf at the end of the block is in the next subtree,
+            // we need to calculate that subtree root based on the tree from the previous block.
             let mut prev_tree = prev_tree
                 .take()
                 .expect("should have some previous sapling frontier");
@@ -81,8 +82,7 @@ pub fn run(
 
             let (index, node) = sapling_nct.completed_subtree_index_and_root().expect(
                 "block should have completed a subtree before its final note commitment: \
-                 already checked is_complete_subtree(),\
-                 and that the block must complete a subtree",
+                 already checked is_complete_subtree(), and that the block must complete a subtree",
             );
 
             assert_eq!(
@@ -90,9 +90,9 @@ pub fn run(
                 "trees are inserted in order with no gaps"
             );
             write_sapling_subtree(upgrade_db, index, height, node);
-        };
+            subtree_count += 1;
+        }
 
-        subtree_count += 1;
         prev_tree = Some(tree);
     }
 
@@ -110,21 +110,24 @@ pub fn run(
             continue;
         };
 
-        // Blocks cannot complete multiple level 16 subtrees.
-        // If a block does complete a subtree, it is either inside the block, or at the end.
-        // (See the detailed comment for Sapling.)
-        if end_of_block_subtree_index.0 <= subtree_count && !tree.is_complete_subtree() {
-            prev_tree = Some(tree);
-            continue;
-        }
+        // Blocks cannot complete multiple level 16 subtrees,
+        // so the subtree index can increase by a maximum of 1 every ~20 blocks.
+        // If this block does complete a subtree, the subtree is either completed by a note before
+        // the final note (so the final note is in the next subtree), or by the final note
+        // (so the final note is the end of this subtree).
 
         if let Some((index, node)) = tree.completed_subtree_index_and_root() {
+            // If the leaf at the end of the block is the final leaf in a subtree,
+            // we already have that subtree root available in the tree.
             assert_eq!(
                 index.0, subtree_count,
                 "trees are inserted in order with no gaps"
             );
             write_orchard_subtree(upgrade_db, index, height, node);
-        } else {
+            subtree_count += 1;
+        } else if end_of_block_subtree_index.0 > subtree_count {
+            // If the leaf at the end of the block is in the next subtree,
+            // we need to calculate that subtree root based on the tree from the previous block.
             let mut prev_tree = prev_tree
                 .take()
                 .expect("should have some previous orchard frontier");
@@ -154,8 +157,7 @@ pub fn run(
 
             let (index, node) = orchard_nct.completed_subtree_index_and_root().expect(
                 "block should have completed a subtree before its final note commitment: \
-                 already checked is_complete_subtree(),\
-                 and that the block must complete a subtree",
+                 already checked is_complete_subtree(), and that the block must complete a subtree",
             );
 
             assert_eq!(
@@ -163,9 +165,9 @@ pub fn run(
                 "trees are inserted in order with no gaps"
             );
             write_orchard_subtree(upgrade_db, index, height, node);
-        };
+            subtree_count += 1;
+        }
 
-        subtree_count += 1;
         prev_tree = Some(tree);
     }
 

From f622de01fbfa8674072b29ad7e7222f7e1df0099 Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Thu, 31 Aug 2023 21:22:24 -0400
Subject: [PATCH 25/33] updates subtree state validity check

---
 zebra-chain/src/subtree.rs                    |   2 +-
 .../finalized_state/disk_format/upgrade.rs    |   4 +-
 .../disk_format/upgrade/add_subtrees.rs       | 326 +++++++++---------
 3 files changed, 172 insertions(+), 160 deletions(-)

diff --git a/zebra-chain/src/subtree.rs b/zebra-chain/src/subtree.rs
index 3d34a3b46c4..10cf428ea46 100644
--- a/zebra-chain/src/subtree.rs
+++ b/zebra-chain/src/subtree.rs
@@ -11,7 +11,7 @@ use crate::block::Height;
 pub const TRACKED_SUBTREE_HEIGHT: u8 = 16;
 
 /// A subtree index
-#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
 pub struct NoteCommitmentSubtreeIndex(pub u16);
 
 impl From<u16> for NoteCommitmentSubtreeIndex {
diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
index 5c62331efa7..2b941dd18bb 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
@@ -226,7 +226,7 @@ impl DbFormatChange {
         // - since this Zebra code knows how to de-duplicate trees, downgrades using this code
         //   still know how to make sure trees are unique
         Self::check_for_duplicate_trees(upgrade_db.clone());
-        add_subtrees::check(&upgrade_db, &cancel_receiver)?;
+        add_subtrees::check(&upgrade_db);
 
         Ok(())
     }
@@ -359,7 +359,7 @@ impl DbFormatChange {
             add_subtrees::run(initial_tip_height, &db, cancel_receiver)?;
 
             // Before marking the state as upgraded, check that the upgrade completed successfully.
-            add_subtrees::check(&db, cancel_receiver)?;
+            add_subtrees::check(&db);
 
             // Mark the database as upgraded. Zebra won't repeat the upgrade anymore once the
             // database is marked, so the upgrade MUST be complete at this point.
diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
index cf77d14e767..16d65520348 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -179,184 +179,218 @@ pub fn run(
 /// # Panics
 ///
 /// If a note commitment subtree is missing or incorrect.
-pub fn check(
-    upgrade_db: &ZebraDb,
-    cancel_receiver: &mpsc::Receiver<CancelFormatChange>,
-) -> Result<(), CancelFormatChange> {
-    let mut is_valid = true;
+pub fn check(db: &ZebraDb) {
+    if !check_sapling_subtrees(db) || !check_orchard_subtrees(db) {
+        panic!("missing or bad subtree(s)");
+    }
+}
 
-    let mut subtree_count = 0;
-    let mut prev_tree: Option<_> = None;
-    for (height, tree) in upgrade_db.sapling_tree_by_height_range(..) {
-        // Return early if there is a cancel signal.
-        if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
-            return Err(CancelFormatChange);
-        }
+/// Check that Sapling note commitment subtrees were correctly added.
+///
+/// # Panics
+///
+/// If a note commitment subtree is missing or incorrect.
+fn check_sapling_subtrees(db: &ZebraDb) -> bool {
+    let Some(NoteCommitmentSubtreeIndex(last_subtree_index)) = db.sapling_tree().subtree_index()
+    else {
+        return true;
+    };
 
-        // Empty note commitment trees can't contain subtrees.
-        let Some(end_of_block_subtree_index) = tree.subtree_index() else {
-            prev_tree = Some(tree);
+    let mut is_valid = true;
+    for index in 0..last_subtree_index {
+        // Check that there's a continuous range of subtrees from index [0, last_subtree_index)
+        let Some(subtree) = db.sapling_subtree_by_index(index) else {
+            warn!(index, "missing subtree");
+            is_valid = false;
             continue;
         };
 
-        // Blocks cannot complete multiple level 16 subtrees,
-        // the subtree index can increase by a maximum of 1 every ~20 blocks.
-        // If this block does complete a subtree, the subtree is either completed by a note before
-        // the final note (so the final note is in the next subtree), or by the final note
-        // (so the final note is the end of this subtree).
-        if end_of_block_subtree_index.0 <= subtree_count && !tree.is_complete_subtree() {
-            prev_tree = Some(tree);
+        // Check that there was a sapling note at the subtree's end height.
+        let Some(tree) = db.sapling_tree_by_height(&subtree.end) else {
+            warn!(?subtree.end, "missing note commitment tree at subtree completion height");
+            is_valid = false;
             continue;
-        }
+        };
 
+        // Check the index and root if the sapling note commitment tree at this height is a complete subtree.
         if let Some((index, node)) = tree.completed_subtree_index_and_root() {
-            assert_eq!(
-                index.0, subtree_count,
-                "trees are inserted in order with no gaps"
-            );
-
-            if !sapling_subtree_exists(upgrade_db, index, height, node) {
-                error!(?height, ?index, "missing sapling subtree");
+            if subtree.index != index {
+                warn!("completed subtree indexes should match");
                 is_valid = false;
             }
-        } else {
-            let mut prev_tree = prev_tree
-                .take()
-                .expect("should have some previous sapling frontier");
-            let sapling_nct = Arc::make_mut(&mut prev_tree);
-
-            let block = upgrade_db
-                .block(height.into())
-                .expect("height with note commitment tree should have block");
-
-            for sapling_note_commitment in block.sapling_note_commitments() {
-                // Return early if there is a cancel signal.
-                if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
-                    return Err(CancelFormatChange);
-                }
 
-                sapling_nct
-                    .append(*sapling_note_commitment)
-                    .expect("finalized notes should append successfully");
+            if subtree.node != node {
+                warn!("completed subtree roots should match");
+                is_valid = false;
+            }
+        }
+        // Check that the final note has a greater subtree index if it didn't complete a subtree.
+        else {
+            let Some(prev_tree) = db.sapling_tree_by_height(&subtree.end.previous()) else {
+                warn!(?subtree.end, "missing note commitment tree at subtree completion height");
+                is_valid = false;
+                continue;
+            };
+
+            let prev_subtree_index = prev_tree.subtree_index();
+            let subtree_index = tree.subtree_index();
+            if subtree_index <= prev_subtree_index {
+                warn!(
+                    ?subtree_index,
+                    ?prev_subtree_index,
+                    "note commitment tree at end height should have incremented subtree index"
+                );
+                is_valid = false;
+            }
+        }
+    }
 
-                // The loop always breaks on this condition,
-                // because we checked the block has enough commitments,
-                // and that the final commitment in the block doesn't complete a subtree.
-                if sapling_nct.is_complete_subtree() {
-                    break;
-                }
+    let mut subtree_count = 0;
+    for (index, height, tree) in db
+        .sapling_tree_by_height_range(..)
+        .filter_map(|(height, tree)| Some((tree.subtree_index()?, height, tree)))
+        .filter_map(|(subtree_index, height, tree)| {
+            if tree.is_complete_subtree() || subtree_index.0 == subtree_count + 1 {
+                subtree_count += 1;
+                Some((subtree_index, height, tree))
+            } else {
+                None
             }
+        })
+    {
+        let Some(subtree) = db.sapling_subtree_by_index(index) else {
+            warn!(?index, "missing subtree");
+            is_valid = false;
+            continue;
+        };
 
-            let (index, node) = sapling_nct.completed_subtree_index_and_root().expect(
-                "block should have completed a subtree before its final note commitment: \
-                 already checked is_complete_subtree(),\
-                 and that the block must complete a subtree",
-            );
+        if subtree.index != index {
+            warn!("completed subtree indexes should match");
+            is_valid = false;
+        }
 
-            assert_eq!(
-                index.0, subtree_count,
-                "trees are inserted in order with no gaps"
-            );
+        if subtree.end != height {
+            warn!(?subtree.end, "bad subtree end height");
+            is_valid = false;
+        }
 
-            if !sapling_subtree_exists(upgrade_db, index, height, node) {
-                error!(?height, ?index, "missing sapling subtree");
+        if let Some((_index, node)) = tree.completed_subtree_index_and_root() {
+            if subtree.node != node {
+                warn!("completed subtree roots should match");
                 is_valid = false;
             }
-        };
+        }
+    }
 
-        subtree_count += 1;
-        prev_tree = Some(tree);
+    if !is_valid {
+        warn!("missing or bad sapling subtrees");
     }
 
-    let mut subtree_count = 0;
-    let mut prev_tree: Option<_> = None;
-    for (height, tree) in upgrade_db.orchard_tree_by_height_range(..) {
-        // Return early if there is a cancel signal.
-        if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
-            return Err(CancelFormatChange);
-        }
+    is_valid
+}
 
-        // Empty note commitment trees can't contain subtrees.
-        let Some(end_of_block_subtree_index) = tree.subtree_index() else {
-            prev_tree = Some(tree);
+/// Check that Orchard note commitment subtrees were correctly added.
+///
+/// # Panics
+///
+/// If a note commitment subtree is missing or incorrect.
+fn check_orchard_subtrees(db: &ZebraDb) -> bool {
+    let Some(NoteCommitmentSubtreeIndex(last_subtree_index)) = db.orchard_tree().subtree_index()
+    else {
+        return true;
+    };
+
+    let mut is_valid = true;
+    for index in 0..last_subtree_index {
+        // Check that there's a continuous range of subtrees from index [0, last_subtree_index)
+        let Some(subtree) = db.orchard_subtree_by_index(index) else {
+            warn!(index, "missing subtree");
+            is_valid = false;
             continue;
         };
 
-        // Blocks cannot complete multiple level 16 subtrees.
-        // If a block does complete a subtree, it is either inside the block, or at the end.
-        // (See the detailed comment for Sapling.)
-        if end_of_block_subtree_index.0 <= subtree_count && !tree.is_complete_subtree() {
-            prev_tree = Some(tree);
+        // Check that there was a orchard note at the subtree's end height.
+        let Some(tree) = db.orchard_tree_by_height(&subtree.end) else {
+            warn!(?subtree.end, "missing note commitment tree at subtree completion height");
+            is_valid = false;
             continue;
-        }
+        };
 
+        // Check the index and root if the orchard note commitment tree at this height is a complete subtree.
         if let Some((index, node)) = tree.completed_subtree_index_and_root() {
-            assert_eq!(
-                index.0, subtree_count,
-                "trees are inserted in order with no gaps"
-            );
-
-            if !orchard_subtree_exists(upgrade_db, index, height, node) {
-                error!(?height, ?index, "missing orchard subtree");
+            if subtree.index != index {
+                warn!("completed subtree indexes should match");
                 is_valid = false;
             }
-        } else {
-            let mut prev_tree = prev_tree
-                .take()
-                .expect("should have some previous orchard frontier");
-            let orchard_nct = Arc::make_mut(&mut prev_tree);
-
-            let block = upgrade_db
-                .block(height.into())
-                .expect("height with note commitment tree should have block");
-
-            for orchard_note_commitment in block.orchard_note_commitments() {
-                // Return early if there is a cancel signal.
-                if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
-                    return Err(CancelFormatChange);
-                }
 
-                orchard_nct
-                    .append(*orchard_note_commitment)
-                    .expect("finalized notes should append successfully");
+            if subtree.node != node {
+                warn!("completed subtree roots should match");
+                is_valid = false;
+            }
+        }
+        // Check that the final note has a greater subtree index if it didn't complete a subtree.
+        else {
+            let Some(prev_tree) = db.orchard_tree_by_height(&subtree.end.previous()) else {
+                warn!(?subtree.end, "missing note commitment tree at subtree completion height");
+                is_valid = false;
+                continue;
+            };
+
+            let prev_subtree_index = prev_tree.subtree_index();
+            let subtree_index = tree.subtree_index();
+            if subtree_index <= prev_subtree_index {
+                warn!(
+                    ?subtree_index,
+                    ?prev_subtree_index,
+                    "note commitment tree at end height should have incremented subtree index"
+                );
+                is_valid = false;
+            }
+        }
+    }
 
-                // The loop always breaks on this condition,
-                // because we checked the block has enough commitments,
-                // and that the final commitment in the block doesn't complete a subtree.
-                if orchard_nct.is_complete_subtree() {
-                    break;
-                }
+    let mut subtree_count = 0;
+    for (index, height, tree) in db
+        .orchard_tree_by_height_range(..)
+        .filter_map(|(height, tree)| Some((tree.subtree_index()?, height, tree)))
+        .filter_map(|(subtree_index, height, tree)| {
+            if tree.is_complete_subtree() || subtree_index.0 == subtree_count + 1 {
+                subtree_count += 1;
+                Some((subtree_index, height, tree))
+            } else {
+                None
             }
+        })
+    {
+        let Some(subtree) = db.orchard_subtree_by_index(index) else {
+            warn!(?index, "missing subtree");
+            is_valid = false;
+            continue;
+        };
 
-            let (index, node) = orchard_nct.completed_subtree_index_and_root().expect(
-                "block should have completed a subtree before its final note commitment: \
-                 already checked is_complete_subtree(),\
-                 and that the block must complete a subtree",
-            );
+        if subtree.index != index {
+            warn!("completed subtree indexes should match");
+            is_valid = false;
+        }
 
-            assert_eq!(
-                index.0, subtree_count,
-                "trees are inserted in order with no gaps"
-            );
+        if subtree.end != height {
+            warn!(?subtree.end, "bad subtree end height");
+            is_valid = false;
+        }
 
-            if !orchard_subtree_exists(upgrade_db, index, height, node) {
-                error!(?height, ?index, "missing orchard subtree");
+        if let Some((_index, node)) = tree.completed_subtree_index_and_root() {
+            if subtree.node != node {
+                warn!("completed subtree roots should match");
                 is_valid = false;
             }
-        };
-
-        subtree_count += 1;
-        prev_tree = Some(tree);
+        }
     }
 
     if !is_valid {
-        panic!(
-            "found duplicate sapling or orchard trees \
-                     after running de-duplicate tree upgrade"
-        );
+        warn!("missing or bad orchard subtrees");
     }
 
-    Ok(())
+    is_valid
 }
 
 /// Writes a Sapling note commitment subtree to `upgrade_db`.
@@ -406,25 +440,3 @@ fn write_orchard_subtree(
     // This log happens about 3 times per second on recent machines with SSD disks.
     debug!(?height, index = ?index.0, ?node, "calculated and added orchard subtree");
 }
-
-/// Checks that a Sapling note commitment subtree exists in `upgrade_db`.
-fn sapling_subtree_exists(
-    upgrade_db: &ZebraDb,
-    index: NoteCommitmentSubtreeIndex,
-    height: Height,
-    node: sapling::tree::Node,
-) -> bool {
-    upgrade_db.sapling_subtree_by_index(index)
-        == Some(NoteCommitmentSubtree::new(index, height, node))
-}
-
-/// Checks that a Orchard note commitment subtree exists in `upgrade_db`.
-fn orchard_subtree_exists(
-    upgrade_db: &ZebraDb,
-    index: NoteCommitmentSubtreeIndex,
-    height: Height,
-    node: orchard::tree::Node,
-) -> bool {
-    upgrade_db.orchard_subtree_by_index(index)
-        == Some(NoteCommitmentSubtree::new(index, height, node))
-}

From 821f925da6ef8505a4230f8cf384c852c0915490 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 1 Sep 2023 12:01:21 +1000
Subject: [PATCH 26/33] Add a subtree format check when there is no upgrade

---
 .../service/finalized_state/disk_format/upgrade.rs  | 13 +++++++------
 zebra-state/src/service/finalized_state/zebra_db.rs |  5 +++--
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
index 2b941dd18bb..b4755751e55 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
@@ -25,7 +25,7 @@ use crate::{
     Config,
 };
 
-mod add_subtrees;
+pub(crate) mod add_subtrees;
 
 /// The kind of database format change we're performing.
 #[derive(Clone, Debug, Eq, PartialEq)]
@@ -220,11 +220,12 @@ impl DbFormatChange {
             }
         }
 
-        // This check should pass for all format changes:
-        // - upgrades should de-duplicate trees if needed (and they already do this check)
-        // - an empty state doesn't have any trees, so it can't have duplicate trees
-        // - since this Zebra code knows how to de-duplicate trees, downgrades using this code
-        //   still know how to make sure trees are unique
+        // These checks should pass for all format changes:
+        // - upgrades should produce a valid format (and they already do that check)
+        // - an empty state should pass all the format checks
+        // - since the running Zebra code knows how to upgrade the database to this format,
+        //   downgrades using this running code still know how to create a valid database
+        //   (unless a future upgrade breaks these format checks)
         Self::check_for_duplicate_trees(upgrade_db.clone());
         add_subtrees::check(&upgrade_db);
 
diff --git a/zebra-state/src/service/finalized_state/zebra_db.rs b/zebra-state/src/service/finalized_state/zebra_db.rs
index 479f337b0e5..68ebcc31cec 100644
--- a/zebra-state/src/service/finalized_state/zebra_db.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db.rs
@@ -19,7 +19,7 @@ use crate::{
         disk_db::DiskDb,
         disk_format::{
             block::MAX_ON_DISK_HEIGHT,
-            upgrade::{DbFormatChange, DbFormatChangeThreadHandle},
+            upgrade::{self, DbFormatChange, DbFormatChangeThreadHandle},
         },
     },
     Config,
@@ -108,9 +108,10 @@ impl ZebraDb {
             db.format_change_handle = Some(format_change_handle);
         } else {
             // If we're re-opening a previously upgraded or newly created database,
-            // the trees should already be de-duplicated.
+            // the database format should be valid.
             // (There's no format change here, so the format change checks won't run.)
             DbFormatChange::check_for_duplicate_trees(db.clone());
+            upgrade::add_subtrees::check(&db.clone());
         }
 
         db

From a442d60f9930c296e32e96fb18ddd904312847c6 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 1 Sep 2023 12:17:27 +1000
Subject: [PATCH 27/33] Fix an off-by-one error with the final subtree check

---
 .../disk_format/upgrade/add_subtrees.rs       | 24 ++++++++++++++-----
 1 file changed, 18 insertions(+), 6 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
index 16d65520348..f397fea175c 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -191,14 +191,20 @@ pub fn check(db: &ZebraDb) {
 ///
 /// If a note commitment subtree is missing or incorrect.
 fn check_sapling_subtrees(db: &ZebraDb) -> bool {
-    let Some(NoteCommitmentSubtreeIndex(last_subtree_index)) = db.sapling_tree().subtree_index()
+    let Some(NoteCommitmentSubtreeIndex(mut first_incomplete_subtree_index)) =
+        db.sapling_tree().subtree_index()
     else {
         return true;
     };
 
+    // If there are no incomplete subtrees in the tree, also expect a subtree for the final index.
+    if db.sapling_tree().is_complete_subtree() {
+        first_incomplete_subtree_index += 1;
+    }
+
     let mut is_valid = true;
-    for index in 0..last_subtree_index {
-        // Check that there's a continuous range of subtrees from index [0, last_subtree_index)
+    for index in 0..first_incomplete_subtree_index {
+        // Check that there's a continuous range of subtrees from index [0, first_incomplete_subtree_index)
         let Some(subtree) = db.sapling_subtree_by_index(index) else {
             warn!(index, "missing subtree");
             is_valid = false;
@@ -295,14 +301,20 @@ fn check_sapling_subtrees(db: &ZebraDb) -> bool {
 ///
 /// If a note commitment subtree is missing or incorrect.
 fn check_orchard_subtrees(db: &ZebraDb) -> bool {
-    let Some(NoteCommitmentSubtreeIndex(last_subtree_index)) = db.orchard_tree().subtree_index()
+    let Some(NoteCommitmentSubtreeIndex(mut first_incomplete_subtree_index)) =
+        db.orchard_tree().subtree_index()
     else {
         return true;
     };
 
+    // If there are no incomplete subtrees in the tree, also expect a subtree for the final index.
+    if db.orchard_tree().is_complete_subtree() {
+        first_incomplete_subtree_index += 1;
+    }
+
     let mut is_valid = true;
-    for index in 0..last_subtree_index {
-        // Check that there's a continuous range of subtrees from index [0, last_subtree_index)
+    for index in 0..first_incomplete_subtree_index {
+        // Check that there's a continuous range of subtrees from index [0, first_incomplete_subtree_index)
         let Some(subtree) = db.orchard_subtree_by_index(index) else {
             warn!(index, "missing subtree");
             is_valid = false;

From 7641f1b2c6560165210b1af358b67e714e06b964 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 1 Sep 2023 12:17:50 +1000
Subject: [PATCH 28/33] Use error-level logs for database format checks

---
 .../disk_format/upgrade/add_subtrees.rs       | 44 +++++++++----------
 1 file changed, 22 insertions(+), 22 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
index f397fea175c..530560c74ca 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -206,14 +206,14 @@ fn check_sapling_subtrees(db: &ZebraDb) -> bool {
     for index in 0..first_incomplete_subtree_index {
         // Check that there's a continuous range of subtrees from index [0, first_incomplete_subtree_index)
         let Some(subtree) = db.sapling_subtree_by_index(index) else {
-            warn!(index, "missing subtree");
+            error!(index, "missing subtree");
             is_valid = false;
             continue;
         };
 
         // Check that there was a sapling note at the subtree's end height.
         let Some(tree) = db.sapling_tree_by_height(&subtree.end) else {
-            warn!(?subtree.end, "missing note commitment tree at subtree completion height");
+            error!(?subtree.end, "missing note commitment tree at subtree completion height");
             is_valid = false;
             continue;
         };
@@ -221,19 +221,19 @@ fn check_sapling_subtrees(db: &ZebraDb) -> bool {
         // Check the index and root if the sapling note commitment tree at this height is a complete subtree.
         if let Some((index, node)) = tree.completed_subtree_index_and_root() {
             if subtree.index != index {
-                warn!("completed subtree indexes should match");
+                error!("completed subtree indexes should match");
                 is_valid = false;
             }
 
             if subtree.node != node {
-                warn!("completed subtree roots should match");
+                error!("completed subtree roots should match");
                 is_valid = false;
             }
         }
         // Check that the final note has a greater subtree index if it didn't complete a subtree.
         else {
             let Some(prev_tree) = db.sapling_tree_by_height(&subtree.end.previous()) else {
-                warn!(?subtree.end, "missing note commitment tree at subtree completion height");
+                error!(?subtree.end, "missing note commitment tree at subtree completion height");
                 is_valid = false;
                 continue;
             };
@@ -241,7 +241,7 @@ fn check_sapling_subtrees(db: &ZebraDb) -> bool {
             let prev_subtree_index = prev_tree.subtree_index();
             let subtree_index = tree.subtree_index();
             if subtree_index <= prev_subtree_index {
-                warn!(
+                error!(
                     ?subtree_index,
                     ?prev_subtree_index,
                     "note commitment tree at end height should have incremented subtree index"
@@ -265,31 +265,31 @@ fn check_sapling_subtrees(db: &ZebraDb) -> bool {
         })
     {
         let Some(subtree) = db.sapling_subtree_by_index(index) else {
-            warn!(?index, "missing subtree");
+            error!(?index, "missing subtree");
             is_valid = false;
             continue;
         };
 
         if subtree.index != index {
-            warn!("completed subtree indexes should match");
+            error!("completed subtree indexes should match");
             is_valid = false;
         }
 
         if subtree.end != height {
-            warn!(?subtree.end, "bad subtree end height");
+            error!(?subtree.end, "bad subtree end height");
             is_valid = false;
         }
 
         if let Some((_index, node)) = tree.completed_subtree_index_and_root() {
             if subtree.node != node {
-                warn!("completed subtree roots should match");
+                error!("completed subtree roots should match");
                 is_valid = false;
             }
         }
     }
 
     if !is_valid {
-        warn!("missing or bad sapling subtrees");
+        error!("missing or bad sapling subtrees");
     }
 
     is_valid
@@ -316,14 +316,14 @@ fn check_orchard_subtrees(db: &ZebraDb) -> bool {
     for index in 0..first_incomplete_subtree_index {
         // Check that there's a continuous range of subtrees from index [0, first_incomplete_subtree_index)
         let Some(subtree) = db.orchard_subtree_by_index(index) else {
-            warn!(index, "missing subtree");
+            error!(index, "missing subtree");
             is_valid = false;
             continue;
         };
 
         // Check that there was a orchard note at the subtree's end height.
         let Some(tree) = db.orchard_tree_by_height(&subtree.end) else {
-            warn!(?subtree.end, "missing note commitment tree at subtree completion height");
+            error!(?subtree.end, "missing note commitment tree at subtree completion height");
             is_valid = false;
             continue;
         };
@@ -331,19 +331,19 @@ fn check_orchard_subtrees(db: &ZebraDb) -> bool {
         // Check the index and root if the orchard note commitment tree at this height is a complete subtree.
         if let Some((index, node)) = tree.completed_subtree_index_and_root() {
             if subtree.index != index {
-                warn!("completed subtree indexes should match");
+                error!("completed subtree indexes should match");
                 is_valid = false;
             }
 
             if subtree.node != node {
-                warn!("completed subtree roots should match");
+                error!("completed subtree roots should match");
                 is_valid = false;
             }
         }
         // Check that the final note has a greater subtree index if it didn't complete a subtree.
         else {
             let Some(prev_tree) = db.orchard_tree_by_height(&subtree.end.previous()) else {
-                warn!(?subtree.end, "missing note commitment tree at subtree completion height");
+                error!(?subtree.end, "missing note commitment tree at subtree completion height");
                 is_valid = false;
                 continue;
             };
@@ -351,7 +351,7 @@ fn check_orchard_subtrees(db: &ZebraDb) -> bool {
             let prev_subtree_index = prev_tree.subtree_index();
             let subtree_index = tree.subtree_index();
             if subtree_index <= prev_subtree_index {
-                warn!(
+                error!(
                     ?subtree_index,
                     ?prev_subtree_index,
                     "note commitment tree at end height should have incremented subtree index"
@@ -375,31 +375,31 @@ fn check_orchard_subtrees(db: &ZebraDb) -> bool {
         })
     {
         let Some(subtree) = db.orchard_subtree_by_index(index) else {
-            warn!(?index, "missing subtree");
+            error!(?index, "missing subtree");
             is_valid = false;
             continue;
         };
 
         if subtree.index != index {
-            warn!("completed subtree indexes should match");
+            error!("completed subtree indexes should match");
             is_valid = false;
         }
 
         if subtree.end != height {
-            warn!(?subtree.end, "bad subtree end height");
+            error!(?subtree.end, "bad subtree end height");
             is_valid = false;
         }
 
         if let Some((_index, node)) = tree.completed_subtree_index_and_root() {
             if subtree.node != node {
-                warn!("completed subtree roots should match");
+                error!("completed subtree roots should match");
                 is_valid = false;
             }
         }
     }
 
     if !is_valid {
-        warn!("missing or bad orchard subtrees");
+        error!("missing or bad orchard subtrees");
     }
 
     is_valid

From 44ac0035f83ae9d53b23b47e0a6eca26643d7c99 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 1 Sep 2023 15:32:21 +1000
Subject: [PATCH 29/33] Skip format checks in tests that create invalid formats

---
 zebra-state/src/service/finalized_state.rs           |  2 +-
 zebra-state/src/service/finalized_state/zebra_db.rs  | 12 ++++++++++--
 .../finalized_state/zebra_db/block/tests/vectors.rs  |  8 ++++----
 3 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/zebra-state/src/service/finalized_state.rs b/zebra-state/src/service/finalized_state.rs
index 95d118362d6..61c8f3e0810 100644
--- a/zebra-state/src/service/finalized_state.rs
+++ b/zebra-state/src/service/finalized_state.rs
@@ -98,7 +98,7 @@ impl FinalizedState {
         network: Network,
         #[cfg(feature = "elasticsearch")] elastic_db: Option<elasticsearch::Elasticsearch>,
     ) -> Self {
-        let db = ZebraDb::new(config, network);
+        let db = ZebraDb::new(config, network, false);
 
         #[cfg(feature = "elasticsearch")]
         let new_state = Self {
diff --git a/zebra-state/src/service/finalized_state/zebra_db.rs b/zebra-state/src/service/finalized_state/zebra_db.rs
index 68ebcc31cec..109e7e88310 100644
--- a/zebra-state/src/service/finalized_state/zebra_db.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db.rs
@@ -60,7 +60,10 @@ pub struct ZebraDb {
 impl ZebraDb {
     /// Opens or creates the database at `config.path` for `network`,
     /// and returns a shared high-level typed database wrapper.
-    pub fn new(config: &Config, network: Network) -> ZebraDb {
+    ///
+    /// If `debug_skip_format_upgrades` is true, don't do any format upgrades or format checks.
+    /// This argument is only used when running tests, it is ignored in production code.
+    pub fn new(config: &Config, network: Network, debug_skip_format_upgrades: bool) -> ZebraDb {
         let running_version = database_format_version_in_code();
         let disk_version = database_format_version_on_disk(config, network)
             .expect("unable to read database format version file");
@@ -84,7 +87,12 @@ impl ZebraDb {
         // a while to start, and new blocks can be committed as soon as we return from this method.
         let initial_tip_height = db.finalized_tip_height();
 
-        // Start any required format changes.
+        // Always do format upgrades & checks in production code.
+        if cfg!(test) && debug_skip_format_upgrades {
+            return db;
+        }
+
+        // Start any required format changes, and do format checks.
         //
         // TODO: should debug_stop_at_height wait for these upgrades, or not?
         if let Some(format_change) = format_change {
diff --git a/zebra-state/src/service/finalized_state/zebra_db/block/tests/vectors.rs b/zebra-state/src/service/finalized_state/zebra_db/block/tests/vectors.rs
index 1f99b24ea6f..93db70f19a1 100644
--- a/zebra-state/src/service/finalized_state/zebra_db/block/tests/vectors.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db/block/tests/vectors.rs
@@ -26,7 +26,7 @@ use zebra_chain::{
 use zebra_test::vectors::{MAINNET_BLOCKS, TESTNET_BLOCKS};
 
 use crate::{
-    service::finalized_state::{disk_db::DiskWriteBatch, FinalizedState},
+    service::finalized_state::{disk_db::DiskWriteBatch, ZebraDb},
     CheckpointVerifiedBlock, Config,
 };
 
@@ -77,11 +77,11 @@ fn test_block_db_round_trip_with(
 ) {
     let _init_guard = zebra_test::init();
 
-    let state = FinalizedState::new(
+    let state = ZebraDb::new(
         &Config::ephemeral(),
         network,
-        #[cfg(feature = "elasticsearch")]
-        None,
+        // The raw database accesses in this test create invalid database formats.
+        true,
     );
 
     // Check that each block round-trips to the database

From d8b1a25748cb470ea2e79de3e6d0a031f335c791 Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Fri, 1 Sep 2023 13:44:31 -0400
Subject: [PATCH 30/33] fix state validity test

---
 .../disk_format/upgrade/add_subtrees.rs       | 30 +++++++++++++------
 1 file changed, 21 insertions(+), 9 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
index 530560c74ca..42565167e0a 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
@@ -180,7 +180,9 @@ pub fn run(
 ///
 /// If a note commitment subtree is missing or incorrect.
 pub fn check(db: &ZebraDb) {
-    if !check_sapling_subtrees(db) || !check_orchard_subtrees(db) {
+    let check_sapling_subtrees = check_sapling_subtrees(db);
+    let check_orchard_subtrees = check_orchard_subtrees(db);
+    if !check_sapling_subtrees || !check_orchard_subtrees {
         panic!("missing or bad subtree(s)");
     }
 }
@@ -256,7 +258,8 @@ fn check_sapling_subtrees(db: &ZebraDb) -> bool {
         .sapling_tree_by_height_range(..)
         .filter_map(|(height, tree)| Some((tree.subtree_index()?, height, tree)))
         .filter_map(|(subtree_index, height, tree)| {
-            if tree.is_complete_subtree() || subtree_index.0 == subtree_count + 1 {
+            if tree.is_complete_subtree() || subtree_index.0 > subtree_count {
+                let subtree_index = subtree_count;
                 subtree_count += 1;
                 Some((subtree_index, height, tree))
             } else {
@@ -270,13 +273,14 @@ fn check_sapling_subtrees(db: &ZebraDb) -> bool {
             continue;
         };
 
-        if subtree.index != index {
+        if subtree.index.0 != index {
             error!("completed subtree indexes should match");
             is_valid = false;
         }
 
         if subtree.end != height {
-            error!(?subtree.end, "bad subtree end height");
+            let is_complete = tree.is_complete_subtree();
+            error!(?subtree.end, ?height, ?index, ?is_complete, "bad sapling subtree end height");
             is_valid = false;
         }
 
@@ -289,7 +293,10 @@ fn check_sapling_subtrees(db: &ZebraDb) -> bool {
     }
 
     if !is_valid {
-        error!("missing or bad sapling subtrees");
+        error!(
+            ?subtree_count,
+            first_incomplete_subtree_index, "missing or bad sapling subtrees"
+        );
     }
 
     is_valid
@@ -366,7 +373,8 @@ fn check_orchard_subtrees(db: &ZebraDb) -> bool {
         .orchard_tree_by_height_range(..)
         .filter_map(|(height, tree)| Some((tree.subtree_index()?, height, tree)))
         .filter_map(|(subtree_index, height, tree)| {
-            if tree.is_complete_subtree() || subtree_index.0 == subtree_count + 1 {
+            if tree.is_complete_subtree() || subtree_index.0 > subtree_count {
+                let subtree_index = subtree_count;
                 subtree_count += 1;
                 Some((subtree_index, height, tree))
             } else {
@@ -380,13 +388,14 @@ fn check_orchard_subtrees(db: &ZebraDb) -> bool {
             continue;
         };
 
-        if subtree.index != index {
+        if subtree.index.0 != index {
             error!("completed subtree indexes should match");
             is_valid = false;
         }
 
         if subtree.end != height {
-            error!(?subtree.end, "bad subtree end height");
+            let is_complete = tree.is_complete_subtree();
+            error!(?subtree.end, ?height, ?index, ?is_complete, "bad orchard subtree end height");
             is_valid = false;
         }
 
@@ -399,7 +408,10 @@ fn check_orchard_subtrees(db: &ZebraDb) -> bool {
     }
 
     if !is_valid {
-        error!("missing or bad orchard subtrees");
+        error!(
+            ?subtree_count,
+            first_incomplete_subtree_index, "missing or bad orchard subtrees"
+        );
     }
 
     is_valid

From 1e6695a983003bfa2e9d15429dd7e20d54b473ce Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Mon, 4 Sep 2023 10:42:23 +1000
Subject: [PATCH 31/33] Add a concurrency comment to subtree by height methods

---
 zebra-state/src/service/non_finalized_state/chain.rs | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/zebra-state/src/service/non_finalized_state/chain.rs b/zebra-state/src/service/non_finalized_state/chain.rs
index 9fc68c288f4..fb3865ccec3 100644
--- a/zebra-state/src/service/non_finalized_state/chain.rs
+++ b/zebra-state/src/service/non_finalized_state/chain.rs
@@ -682,6 +682,11 @@ impl Chain {
 
     /// Returns the Sapling [`NoteCommitmentSubtree`] that was completed at a block with
     /// [`HashOrHeight`], if it exists in the non-finalized [`Chain`].
+    ///
+    /// # Concurrency
+    ///
+    /// This method should not be used to get subtrees in concurrent code by height,
+    /// because the same heights in different chain forks can have different subtrees.
     pub fn sapling_subtree(
         &self,
         hash_or_height: HashOrHeight,
@@ -872,6 +877,11 @@ impl Chain {
 
     /// Returns the Orchard [`NoteCommitmentSubtree`] that was completed at a block with
     /// [`HashOrHeight`], if it exists in the non-finalized [`Chain`].
+    ///
+    /// # Concurrency
+    ///
+    /// This method should not be used to get subtrees in concurrent code by height,
+    /// because the same heights in different chain forks can have different subtrees.
     pub fn orchard_subtree(
         &self,
         hash_or_height: HashOrHeight,

From e1abeb967c2d58e8ddd12a43a3b5b68b041b850c Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Mon, 4 Sep 2023 10:51:33 +1000
Subject: [PATCH 32/33] Add individual subtree state methods: reverts removing
 these methods in an earlier PR

---
 .../finalized_state/zebra_db/shielded.rs      | 60 +++++++++++++++++--
 1 file changed, 54 insertions(+), 6 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
index 5a13f54f161..12771ce467e 100644
--- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
@@ -180,15 +180,39 @@ impl ZebraDb {
         self.db.zs_range_iter(&sapling_trees, range)
     }
 
+    /// Returns the Sapling note commitment subtree at this `index`.
+    ///
+    /// # Correctness
+    ///
+    /// This method should not be used to get subtrees for RPC responses,
+    /// because those subtree lists require that the start subtree is present in the list.
+    /// Instead, use `sapling_subtrees_by_index()`.
+    #[allow(clippy::unwrap_in_result)]
+    pub(in super::super) fn sapling_subtree_by_index(
+        &self,
+        index: impl Into<NoteCommitmentSubtreeIndex> + Copy,
+    ) -> Option<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 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.
+    /// # Correctness
+    ///
+    /// This method is specifically designed for the `z_getsubtreesbyindex` state request.
+    /// It might not work for other RPCs or state checks.
     #[allow(clippy::unwrap_in_result)]
     pub fn sapling_subtrees_by_index(
         &self,
@@ -289,15 +313,39 @@ impl ZebraDb {
         self.db.zs_range_iter(&orchard_trees, range)
     }
 
+    /// Returns the Orchard note commitment subtree at this `index`.
+    ///
+    /// # Correctness
+    ///
+    /// This method should not be used to get subtrees for RPC responses,
+    /// because those subtree lists require that the start subtree is present in the list.
+    /// Instead, use `orchard_subtrees_by_index()`.
+    #[allow(clippy::unwrap_in_result)]
+    pub(in super::super) fn orchard_subtree_by_index(
+        &self,
+        index: impl Into<NoteCommitmentSubtreeIndex> + Copy,
+    ) -> Option<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 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.
+    /// # Correctness
+    ///
+    /// This method is specifically designed for the `z_getsubtreesbyindex` state request.
+    /// It might not work for other RPCs or state checks.
     #[allow(clippy::unwrap_in_result)]
     pub fn orchard_subtrees_by_index(
         &self,

From 40896ca571bf02fbd333f07976ea68488b8e860c Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Mon, 4 Sep 2023 10:53:32 +1000
Subject: [PATCH 33/33] fastmod "subtrees_by_index"
 "subtree_list_by_index_for_rpc"

---
 .../src/service/finalized_state/zebra_db/shielded.rs      | 8 ++++----
 zebra-state/src/service/read/tree.rs                      | 4 ++--
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
index 12771ce467e..0db9c5be92d 100644
--- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
@@ -186,7 +186,7 @@ impl ZebraDb {
     ///
     /// This method should not be used to get subtrees for RPC responses,
     /// because those subtree lists require that the start subtree is present in the list.
-    /// Instead, use `sapling_subtrees_by_index()`.
+    /// Instead, use `sapling_subtree_list_by_index_for_rpc()`.
     #[allow(clippy::unwrap_in_result)]
     pub(in super::super) fn sapling_subtree_by_index(
         &self,
@@ -214,7 +214,7 @@ impl ZebraDb {
     /// This method is specifically designed for the `z_getsubtreesbyindex` state request.
     /// It might not work for other RPCs or state checks.
     #[allow(clippy::unwrap_in_result)]
-    pub fn sapling_subtrees_by_index(
+    pub fn sapling_subtree_list_by_index_for_rpc(
         &self,
         start_index: NoteCommitmentSubtreeIndex,
         limit: Option<NoteCommitmentSubtreeIndex>,
@@ -319,7 +319,7 @@ impl ZebraDb {
     ///
     /// This method should not be used to get subtrees for RPC responses,
     /// because those subtree lists require that the start subtree is present in the list.
-    /// Instead, use `orchard_subtrees_by_index()`.
+    /// Instead, use `orchard_subtree_list_by_index_for_rpc()`.
     #[allow(clippy::unwrap_in_result)]
     pub(in super::super) fn orchard_subtree_by_index(
         &self,
@@ -347,7 +347,7 @@ impl ZebraDb {
     /// This method is specifically designed for the `z_getsubtreesbyindex` state request.
     /// It might not work for other RPCs or state checks.
     #[allow(clippy::unwrap_in_result)]
-    pub fn orchard_subtrees_by_index(
+    pub fn orchard_subtree_list_by_index_for_rpc(
         &self,
         start_index: NoteCommitmentSubtreeIndex,
         limit: Option<NoteCommitmentSubtreeIndex>,
diff --git a/zebra-state/src/service/read/tree.rs b/zebra-state/src/service/read/tree.rs
index 05bd2c4a186..7599524967f 100644
--- a/zebra-state/src/service/read/tree.rs
+++ b/zebra-state/src/service/read/tree.rs
@@ -74,7 +74,7 @@ where
     // In that case, we ignore all the trees in `chain` after the first inconsistent tree,
     // because we know they will be inconsistent as well. (It is cryptographically impossible
     // for tree roots to be equal once the leaves have diverged.)
-    let mut db_list = db.sapling_subtrees_by_index(start_index, limit);
+    let mut db_list = db.sapling_subtree_list_by_index_for_rpc(start_index, limit);
 
     // If there's no chain, then we have the complete list.
     let Some(chain) = chain else {
@@ -162,7 +162,7 @@ where
     // In that case, we ignore all the trees in `chain` after the first inconsistent tree,
     // because we know they will be inconsistent as well. (It is cryptographically impossible
     // for tree roots to be equal once the leaves have diverged.)
-    let mut db_list = db.orchard_subtrees_by_index(start_index, limit);
+    let mut db_list = db.orchard_subtree_list_by_index_for_rpc(start_index, limit);
 
     // If there's no chain, then we have the complete list.
     let Some(chain) = chain else {