From 40ac61dae7bcbc871f0beb1552922d14394db502 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 7 Sep 2023 10:17:33 +1000
Subject: [PATCH 01/56] Avoid manual handling of previous sapling trees by
 using iterator windows instead

---
 .../disk_format/upgrade/add_subtrees.rs       | 33 +++++++++++--------
 1 file changed, 20 insertions(+), 13 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 9d0547a5893..dbaa40d1e79 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,6 +2,8 @@
 
 use std::sync::{mpsc, Arc};
 
+use itertools::Itertools;
+
 use zebra_chain::{
     block::Height,
     orchard, sapling,
@@ -22,21 +24,31 @@ pub fn run(
     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) {
+
+    for ((_prev_height, mut prev_tree), (height, tree)) in upgrade_db
+        .sapling_tree_by_height_range(..=initial_tip_height)
+        .tuple_windows()
+    {
         // 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;
-        };
+        // Since:
+        // - the empty genesis tree is the first tree,
+        // - the next tree after genesis contains at least one note commitment,
+        // - the first window contains the genesis tree as the previous tree, and
+        // - trees are deduplicated, so each tree only appears once;
+        // we will never see an empty tree as the current `tree`.
+        let end_of_block_subtree_index = tree.subtree_index().expect(
+            "the genesis tree is the only empty tree, and it is the prev_tree in the first window",
+        );
 
-        // Blocks cannot complete multiple level 16 subtrees,
-        // so the subtree index can increase by a maximum of 1 every ~20 blocks.
+        // Due to the 2^16 limit on sapling outputs, blocks cannot complete multiple level 16
+        // subtrees. Currently, with 2MB blocks and v4/v5 sapling output sizes, 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).
@@ -53,9 +65,6 @@ pub fn run(
         } 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");
             let sapling_nct = Arc::make_mut(&mut prev_tree);
 
             let block = upgrade_db
@@ -92,8 +101,6 @@ pub fn run(
             write_sapling_subtree(upgrade_db, index, height, node);
             subtree_count += 1;
         }
-
-        prev_tree = Some(tree);
     }
 
     let mut subtree_count = 0;

From 5a11f00fee9d1d3ebec055851a7f2fa97f4ad258 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 7 Sep 2023 10:37:35 +1000
Subject: [PATCH 02/56] Avoid manual sapling subtree index handling by
 comparing prev and current subtree indexes instead

---
 .../disk_format/upgrade/add_subtrees.rs       | 40 +++++++++++--------
 1 file changed, 23 insertions(+), 17 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 dbaa40d1e79..a820be9c82f 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
@@ -23,8 +23,6 @@ pub fn run(
     upgrade_db: &ZebraDb,
     cancel_receiver: &mpsc::Receiver<CancelFormatChange>,
 ) -> Result<(), CancelFormatChange> {
-    let mut subtree_count = 0;
-
     for ((_prev_height, mut prev_tree), (height, tree)) in upgrade_db
         .sapling_tree_by_height_range(..=initial_tip_height)
         .tuple_windows()
@@ -34,6 +32,8 @@ pub fn run(
             return Err(CancelFormatChange);
         }
 
+        let before_block_subtree_index = prev_tree.subtree_index();
+
         // Empty note commitment trees can't contain subtrees.
         // Since:
         // - the empty genesis tree is the first tree,
@@ -53,18 +53,29 @@ pub fn run(
         // 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 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.
         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);
-            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.
+            continue;
+        }
+
+        // If the first block completes a subtree, then we have just written that subtree.
+        // If it doesn't, then the previous tree for every new subtree has a valid subtree index.
+        //
+        // The only subtree that doesn't have a valid index is genesis, and we've either:
+        // - just used genesis as the previous tree of a new tree in the code above, or
+        // - genesis isn't the previous tree of a new tree at all, so it can be ignored.
+        //
+        // (The first case can't happen for sapling in v4/v5 transactions, because its outputs are
+        // too large. But it might matter for future transaction formats or shielded pools.)
+        let Some(before_block_subtree_index) = before_block_subtree_index else {
+            continue;
+        };
+
+        // 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.
+        if end_of_block_subtree_index > before_block_subtree_index {
             let sapling_nct = Arc::make_mut(&mut prev_tree);
 
             let block = upgrade_db
@@ -94,12 +105,7 @@ pub fn run(
                  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"
-            );
             write_sapling_subtree(upgrade_db, index, height, node);
-            subtree_count += 1;
         }
     }
 

From effff632a29d5f51408917d99175cb8b3d51c8bc Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 7 Sep 2023 11:27:48 +1000
Subject: [PATCH 03/56] Simplify adding notes by using the exact number of
 remaining notes

---
 zebra-chain/src/sapling/tree.rs               | 36 +++++++++++++++++++
 .../disk_format/upgrade/add_subtrees.rs       | 22 ++++++------
 2 files changed, 46 insertions(+), 12 deletions(-)

diff --git a/zebra-chain/src/sapling/tree.rs b/zebra-chain/src/sapling/tree.rs
index 6112aa0cc43..87742c081eb 100644
--- a/zebra-chain/src/sapling/tree.rs
+++ b/zebra-chain/src/sapling/tree.rs
@@ -404,6 +404,42 @@ impl NoteCommitmentTree {
         Some(index)
     }
 
+    /// Returns the number of leaf nodes required to complete the subtree at
+    /// [`TRACKED_SUBTREE_HEIGHT`].
+    ///
+    /// Returns `2^TRACKED_SUBTREE_HEIGHT` if the tree is empty.
+    #[allow(clippy::unwrap_in_result)]
+    pub fn remaining_subtree_leaf_nodes(&self) -> usize {
+        let remaining = match self.frontier() {
+            // If the subtree has at least one leaf node, the remaining number of nodes can be
+            // calculated using the maximum subtree position and the current position.
+            Some(tree) => {
+                let max_position = incrementalmerkletree::Address::above_position(
+                    TRACKED_SUBTREE_HEIGHT.into(),
+                    tree.position(),
+                )
+                .max_position();
+
+                max_position - tree.position().into()
+            }
+            // If the subtree has no nodes, the remaining number of nodes is the number of nodes in
+            // a subtree.
+            None => {
+                // This position is guaranteed to be in the first subtree.
+                let first_position = 0.into();
+
+                let subtree_address = incrementalmerkletree::Address::above_position(
+                    TRACKED_SUBTREE_HEIGHT.into(),
+                    first_position,
+                );
+
+                subtree_address.position_range_end() - subtree_address.position_range_start().into()
+            }
+        };
+
+        u64::from(remaining).try_into().expect("fits in usize")
+    }
+
     /// Returns subtree index and root if the most recently appended leaf completes the subtree
     pub fn completed_subtree_index_and_root(&self) -> Option<(NoteCommitmentSubtreeIndex, Node)> {
         if !self.is_complete_subtree() {
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 a820be9c82f..8d17b0f5a63 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
@@ -77,12 +77,18 @@ pub fn run(
         // we need to calculate that subtree root based on the tree from the previous block.
         if end_of_block_subtree_index > before_block_subtree_index {
             let sapling_nct = Arc::make_mut(&mut prev_tree);
+            let remaining_notes = sapling_nct.remaining_subtree_leaf_nodes();
+
+            assert!(
+                remaining_notes > 0,
+                "just checked for complete subtrees above"
+            );
 
             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() {
+            for sapling_note_commitment in block.sapling_note_commitments().take(remaining_notes) {
                 // Return early if there is a cancel signal.
                 if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
                     return Err(CancelFormatChange);
@@ -91,19 +97,11 @@ pub fn run(
                 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",
-            );
+            let (index, node) = sapling_nct
+                .completed_subtree_index_and_root()
+                .expect("already checked that the block completed a subtree");
 
             write_sapling_subtree(upgrade_db, index, height, node);
         }

From c0e61ed1a540100e2ceb4148bbe1d41e461c7ad7 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 7 Sep 2023 11:32:37 +1000
Subject: [PATCH 04/56] Simplify by skipping the first block, because it can't
 complete a subtree

---
 .../disk_format/upgrade/add_subtrees.rs       | 60 +++++++------------
 1 file changed, 22 insertions(+), 38 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 8d17b0f5a63..6a575d0cd72 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
@@ -26,56 +26,40 @@ pub fn run(
     for ((_prev_height, mut prev_tree), (height, tree)) in upgrade_db
         .sapling_tree_by_height_range(..=initial_tip_height)
         .tuple_windows()
+        // The first block with sapling notes can't complete a subtree, see below for details.
+        .skip(1)
     {
         // Return early if there is a cancel signal.
         if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
             return Err(CancelFormatChange);
         }
 
-        let before_block_subtree_index = prev_tree.subtree_index();
-
-        // Empty note commitment trees can't contain subtrees.
-        // Since:
-        // - the empty genesis tree is the first tree,
-        // - the next tree after genesis contains at least one note commitment,
-        // - the first window contains the genesis tree as the previous tree, and
-        // - trees are deduplicated, so each tree only appears once;
-        // we will never see an empty tree as the current `tree`.
-        let end_of_block_subtree_index = tree.subtree_index().expect(
-            "the genesis tree is the only empty tree, and it is the prev_tree in the first window",
-        );
-
-        // Due to the 2^16 limit on sapling outputs, blocks cannot complete multiple level 16
-        // subtrees. Currently, with 2MB blocks and v4/v5 sapling output sizes, the subtree index
-        // can increase by a maximum of 1 every ~20 blocks.
+        // Zebra stores exactly one note commitment tree for every block with sapling notes.
+        // (It also stores the empty note commitment tree for the genesis block, but we skip that.)
         //
-        // If this block does complete a subtree, the subtree is either completed by a note before
+        // The consensus rules limit blocks to less than 2^16 sapling outputs. So a block can't
+        // complete multiple level 16 subtrees (or complete an entire subtree by itself).
+        // Currently, with 2MB blocks and v4/v5 sapling output sizes, the subtree index can
+        // increase by 1 every ~20 full blocks.
+
+        // Empty note commitment trees can't contain subtrees, so they have invalid subtree indexes.
+        // But we skip the empty genesis tree in the iterator, and all other trees must have notes.
+        let before_block_subtree_index = prev_tree
+            .subtree_index()
+            .expect("already skipped empty tree");
+        let end_of_block_subtree_index = tree.subtree_index().expect("already skipped empty tree");
+
+        // If this block completed 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 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.
         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.
             write_sapling_subtree(upgrade_db, index, height, node);
             continue;
-        }
-
-        // If the first block completes a subtree, then we have just written that subtree.
-        // If it doesn't, then the previous tree for every new subtree has a valid subtree index.
-        //
-        // The only subtree that doesn't have a valid index is genesis, and we've either:
-        // - just used genesis as the previous tree of a new tree in the code above, or
-        // - genesis isn't the previous tree of a new tree at all, so it can be ignored.
-        //
-        // (The first case can't happen for sapling in v4/v5 transactions, because its outputs are
-        // too large. But it might matter for future transaction formats or shielded pools.)
-        let Some(before_block_subtree_index) = before_block_subtree_index else {
-            continue;
-        };
-
-        // 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.
-        if end_of_block_subtree_index > before_block_subtree_index {
+        } else if end_of_block_subtree_index > before_block_subtree_index {
+            // 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 sapling_nct = Arc::make_mut(&mut prev_tree);
             let remaining_notes = sapling_nct.remaining_subtree_leaf_nodes();
 

From 3c9eb58d8df8c28c2d963aa37440ee04fff41cb8 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 7 Sep 2023 11:49:26 +1000
Subject: [PATCH 05/56] Re-use existing tree update code

---
 zebra-chain/src/parallel/tree.rs              |  2 +-
 .../disk_format/upgrade/add_subtrees.rs       | 39 ++++++++++---------
 2 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/zebra-chain/src/parallel/tree.rs b/zebra-chain/src/parallel/tree.rs
index 361b35d1058..af34a1364f1 100644
--- a/zebra-chain/src/parallel/tree.rs
+++ b/zebra-chain/src/parallel/tree.rs
@@ -146,7 +146,7 @@ impl NoteCommitmentTrees {
 
     /// Update the sapling note commitment tree.
     #[allow(clippy::unwrap_in_result)]
-    fn update_sapling_note_commitment_tree(
+    pub fn update_sapling_note_commitment_tree(
         mut sapling: Arc<sapling::tree::NoteCommitmentTree>,
         sapling_note_commitments: Vec<sapling::tree::NoteCommitmentUpdate>,
     ) -> Result<
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 6a575d0cd72..842bf0306ef 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
@@ -6,7 +6,9 @@ use itertools::Itertools;
 
 use zebra_chain::{
     block::Height,
-    orchard, sapling,
+    orchard,
+    parallel::tree::NoteCommitmentTrees,
+    sapling,
     subtree::{NoteCommitmentSubtree, NoteCommitmentSubtreeIndex},
 };
 
@@ -23,7 +25,7 @@ pub fn run(
     upgrade_db: &ZebraDb,
     cancel_receiver: &mpsc::Receiver<CancelFormatChange>,
 ) -> Result<(), CancelFormatChange> {
-    for ((_prev_height, mut prev_tree), (height, tree)) in upgrade_db
+    for ((_prev_height, prev_tree), (height, tree)) in upgrade_db
         .sapling_tree_by_height_range(..=initial_tip_height)
         .tuple_windows()
         // The first block with sapling notes can't complete a subtree, see below for details.
@@ -60,8 +62,7 @@ pub fn run(
         } else if end_of_block_subtree_index > before_block_subtree_index {
             // 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 sapling_nct = Arc::make_mut(&mut prev_tree);
-            let remaining_notes = sapling_nct.remaining_subtree_leaf_nodes();
+            let remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
 
             assert!(
                 remaining_notes > 0,
@@ -71,21 +72,21 @@ pub fn run(
             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().take(remaining_notes) {
-                // 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");
-            }
-
-            let (index, node) = sapling_nct
-                .completed_subtree_index_and_root()
-                .expect("already checked that the block completed a subtree");
+            let sapling_note_commitments = block
+                .sapling_note_commitments()
+                .take(remaining_notes)
+                .cloned()
+                .collect();
+
+            // This takes less than 1 second per tree, so we don't need to make it cancellable.
+            let (_sapling_nct, subtree) = NoteCommitmentTrees::update_sapling_note_commitment_tree(
+                prev_tree,
+                sapling_note_commitments,
+            )
+            .expect("finalized notes should append successfully");
+
+            let (index, node) =
+                subtree.expect("already checked that the block completed a subtree");
 
             write_sapling_subtree(upgrade_db, index, height, node);
         }

From bc0c674a843abe4a41c9cae5aa994c468b523a74 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 7 Sep 2023 11:50:46 +1000
Subject: [PATCH 06/56] Apply the sapling changes to orchard subtree updates

---
 zebra-chain/src/orchard/tree.rs               | 36 +++++++
 zebra-chain/src/parallel/tree.rs              |  2 +-
 .../disk_format/upgrade/add_subtrees.rs       | 94 +++++++++----------
 3 files changed, 80 insertions(+), 52 deletions(-)

diff --git a/zebra-chain/src/orchard/tree.rs b/zebra-chain/src/orchard/tree.rs
index c2adee9ba70..a3df5d4cf87 100644
--- a/zebra-chain/src/orchard/tree.rs
+++ b/zebra-chain/src/orchard/tree.rs
@@ -423,6 +423,42 @@ impl NoteCommitmentTree {
         Some(index)
     }
 
+    /// Returns the number of leaf nodes required to complete the subtree at
+    /// [`TRACKED_SUBTREE_HEIGHT`].
+    ///
+    /// Returns `2^TRACKED_SUBTREE_HEIGHT` if the tree is empty.
+    #[allow(clippy::unwrap_in_result)]
+    pub fn remaining_subtree_leaf_nodes(&self) -> usize {
+        let remaining = match self.frontier() {
+            // If the subtree has at least one leaf node, the remaining number of nodes can be
+            // calculated using the maximum subtree position and the current position.
+            Some(tree) => {
+                let max_position = incrementalmerkletree::Address::above_position(
+                    TRACKED_SUBTREE_HEIGHT.into(),
+                    tree.position(),
+                )
+                .max_position();
+
+                max_position - tree.position().into()
+            }
+            // If the subtree has no nodes, the remaining number of nodes is the number of nodes in
+            // a subtree.
+            None => {
+                // This position is guaranteed to be in the first subtree.
+                let first_position = 0.into();
+
+                let subtree_address = incrementalmerkletree::Address::above_position(
+                    TRACKED_SUBTREE_HEIGHT.into(),
+                    first_position,
+                );
+
+                subtree_address.position_range_end() - subtree_address.position_range_start().into()
+            }
+        };
+
+        u64::from(remaining).try_into().expect("fits in usize")
+    }
+
     /// Returns subtree index and root if the most recently appended leaf completes the subtree
     pub fn completed_subtree_index_and_root(&self) -> Option<(NoteCommitmentSubtreeIndex, Node)> {
         if !self.is_complete_subtree() {
diff --git a/zebra-chain/src/parallel/tree.rs b/zebra-chain/src/parallel/tree.rs
index af34a1364f1..4dfb5c298a2 100644
--- a/zebra-chain/src/parallel/tree.rs
+++ b/zebra-chain/src/parallel/tree.rs
@@ -185,7 +185,7 @@ impl NoteCommitmentTrees {
 
     /// Update the orchard note commitment tree.
     #[allow(clippy::unwrap_in_result)]
-    fn update_orchard_note_commitment_tree(
+    pub fn update_orchard_note_commitment_tree(
         mut orchard: Arc<orchard::tree::NoteCommitmentTree>,
         orchard_note_commitments: Vec<orchard::tree::NoteCommitmentUpdate>,
     ) -> Result<
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 842bf0306ef..5c68805afaa 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,6 +1,6 @@
 //! Fully populate the Sapling and Orchard note commitment subtrees for existing blocks in the database.
 
-use std::sync::{mpsc, Arc};
+use std::sync::mpsc;
 
 use itertools::Itertools;
 
@@ -92,79 +92,71 @@ 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) {
+    for ((_prev_height, prev_tree), (height, tree)) in upgrade_db
+        .orchard_tree_by_height_range(..=initial_tip_height)
+        .tuple_windows()
+        // The first block with orchard notes can't complete a subtree, see below for details.
+        .skip(1)
+    {
         // 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;
-        };
+        // Zebra stores exactly one note commitment tree for every block with orchard notes.
+        // (It also stores the empty note commitment tree for the genesis block, but we skip that.)
+        //
+        // The consensus rules limit blocks to less than 2^16 orchard outputs. So a block can't
+        // complete multiple level 16 subtrees (or complete an entire subtree by itself).
+        // Currently, with 2MB blocks and v5 orchard output sizes, the subtree index can
+        // increase by 1 every ~20 full blocks.
 
-        // 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
+        // Empty note commitment trees can't contain subtrees, so they have invalid subtree indexes.
+        // But we skip the empty genesis tree in the iterator, and all other trees must have notes.
+        let before_block_subtree_index = prev_tree
+            .subtree_index()
+            .expect("already skipped empty tree");
+        let end_of_block_subtree_index = tree.subtree_index().expect("already skipped empty tree");
+
+        // If this block completed 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);
-            subtree_count += 1;
-        } else if end_of_block_subtree_index.0 > subtree_count {
+            continue;
+        } else if end_of_block_subtree_index > before_block_subtree_index {
             // 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");
-            let orchard_nct = Arc::make_mut(&mut prev_tree);
+            let remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
+
+            assert!(
+                remaining_notes > 0,
+                "just checked for complete subtrees above"
+            );
 
             let block = upgrade_db
                 .block(height.into())
                 .expect("height with note commitment tree should have block");
+            let orchard_note_commitments = block
+                .orchard_note_commitments()
+                .take(remaining_notes)
+                .cloned()
+                .collect();
 
-            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;
-                }
-            }
+            // This takes less than 1 second per tree, so we don't need to make it cancellable.
+            let (_orchard_nct, subtree) = NoteCommitmentTrees::update_orchard_note_commitment_tree(
+                prev_tree,
+                orchard_note_commitments,
+            )
+            .expect("finalized notes should append successfully");
 
-            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",
-            );
+            let (index, node) =
+                subtree.expect("already checked that the block completed a subtree");
 
-            assert_eq!(
-                index.0, subtree_count,
-                "trees are inserted in order with no gaps"
-            );
             write_orchard_subtree(upgrade_db, index, height, node);
-            subtree_count += 1;
         }
-
-        prev_tree = Some(tree);
     }
 
     Ok(())

From 387f8bd4b240887c21a4f1406ad06ff700a83172 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 8 Sep 2023 06:28:02 +1000
Subject: [PATCH 07/56] add a reverse database column family iterator function

---
 .../src/service/finalized_state/disk_db.rs    | 105 ++++++++++++++----
 .../finalized_state/zebra_db/shielded.rs      |   4 +-
 2 files changed, 88 insertions(+), 21 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db.rs b/zebra-state/src/service/finalized_state/disk_db.rs
index eb9b95d0ac2..50a10ba4540 100644
--- a/zebra-state/src/service/finalized_state/disk_db.rs
+++ b/zebra-state/src/service/finalized_state/disk_db.rs
@@ -424,6 +424,46 @@ impl DiskDb {
     ///
     /// Holding this iterator open might delay block commit transactions.
     pub fn zs_range_iter<C, K, V, R>(&self, cf: &C, range: R) -> impl Iterator<Item = (K, V)> + '_
+    where
+        C: rocksdb::AsColumnFamilyRef,
+        K: IntoDisk + FromDisk,
+        V: FromDisk,
+        R: RangeBounds<K>,
+    {
+        self.zs_range_iter_with_direction(cf, range, false)
+    }
+
+    /// Returns a reverse iterator over the items in `cf` in `range`.
+    ///
+    /// Holding this iterator open might delay block commit transactions.
+    ///
+    /// This code is coped from `zs_range_iter()`, but with the mode reversed.
+    pub fn zs_reverse_range_iter<C, K, V, R>(
+        &self,
+        cf: &C,
+        range: R,
+    ) -> impl Iterator<Item = (K, V)> + '_
+    where
+        C: rocksdb::AsColumnFamilyRef,
+        K: IntoDisk + FromDisk,
+        V: FromDisk,
+        R: RangeBounds<K>,
+    {
+        self.zs_range_iter_with_direction(cf, range, true)
+    }
+
+    /// Returns an iterator over the items in `cf` in `range`.
+    ///
+    /// RocksDB iterators are ordered by increasing key bytes by default.
+    /// Otherwise, if `reverse` is `true`, the iterator is ordered by decreasing key bytes.
+    ///
+    /// Holding this iterator open might delay block commit transactions.
+    fn zs_range_iter_with_direction<C, K, V, R>(
+        &self,
+        cf: &C,
+        range: R,
+        reverse: bool,
+    ) -> impl Iterator<Item = (K, V)> + '_
     where
         C: rocksdb::AsColumnFamilyRef,
         K: IntoDisk + FromDisk,
@@ -444,40 +484,67 @@ impl DiskDb {
 
         let start_bound = map_to_vec(range.start_bound());
         let end_bound = map_to_vec(range.end_bound());
-        let range = (start_bound.clone(), end_bound);
-
-        let start_bound_vec =
-            if let Included(ref start_bound) | Excluded(ref start_bound) = start_bound {
-                start_bound.clone()
-            } else {
-                // Actually unused
-                Vec::new()
-            };
+        let range = (start_bound, end_bound);
 
-        let start_mode = if matches!(start_bound, Unbounded) {
-            // Unbounded iterators start at the first item
-            rocksdb::IteratorMode::Start
-        } else {
-            rocksdb::IteratorMode::From(start_bound_vec.as_slice(), rocksdb::Direction::Forward)
-        };
+        let mode = Self::zs_iter_mode(&range, reverse);
 
         // Reading multiple items from iterators has caused database hangs,
         // in previous RocksDB versions
         self.db
-            .iterator_cf(cf, start_mode)
+            .iterator_cf(cf, mode)
             .map(|result| result.expect("unexpected database failure"))
             .map(|(key, value)| (key.to_vec(), value))
-            // Skip excluded start bound and empty ranges. The `start_mode` already skips keys
-            // before the start bound.
+            // Skip excluded "from" bound and empty ranges. The `mode` already skips keys
+            // strictly before the "from" bound.
             .skip_while({
                 let range = range.clone();
                 move |(key, _value)| !range.contains(key)
             })
-            // Take until the excluded end bound is reached, or we're after the included end bound.
+            // Take until the excluded "to" bound is reached,
+            // or we're after the included "to" bound.
             .take_while(move |(key, _value)| range.contains(key))
             .map(|(key, value)| (K::from_bytes(key), V::from_bytes(value)))
     }
 
+    /// Returns the RocksDB iterator "from" mode for `range`.
+    ///
+    /// RocksDB iterators are ordered by increasing key bytes by default.
+    /// Otherwise, if `reverse` is `true`, the iterator is ordered by decreasing key bytes.
+    fn zs_iter_mode<R>(range: &R, reverse: bool) -> rocksdb::IteratorMode
+    where
+        R: RangeBounds<Vec<u8>>,
+    {
+        use std::ops::Bound::*;
+
+        let from_bound = if reverse {
+            range.end_bound()
+        } else {
+            range.start_bound()
+        };
+
+        match from_bound {
+            Unbounded => {
+                if reverse {
+                    // Reversed unbounded iterators start from the last item
+                    rocksdb::IteratorMode::End
+                } else {
+                    // Unbounded iterators start from the first item
+                    rocksdb::IteratorMode::Start
+                }
+            }
+
+            Included(bound) | Excluded(bound) => {
+                let direction = if reverse {
+                    rocksdb::Direction::Reverse
+                } else {
+                    rocksdb::Direction::Forward
+                };
+
+                rocksdb::IteratorMode::From(bound.as_slice(), direction)
+            }
+        }
+    }
+
     /// The ideal open file limit for Zebra
     const IDEAL_OPEN_FILE_LIMIT: u64 = 1024;
 
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 0db9c5be92d..f44dac0f080 100644
--- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
@@ -167,7 +167,7 @@ impl ZebraDb {
         Some(Arc::new(tree))
     }
 
-    /// Returns the Sapling note commitment trees in the supplied range.
+    /// Returns the Sapling note commitment trees in the supplied range, in increasing height order.
     #[allow(clippy::unwrap_in_result)]
     pub fn sapling_tree_by_height_range<R>(
         &self,
@@ -300,7 +300,7 @@ impl ZebraDb {
         Some(Arc::new(tree))
     }
 
-    /// Returns the Orchard note commitment trees in the supplied range.
+    /// Returns the Orchard note commitment trees in the supplied range, in increasing height order.
     #[allow(clippy::unwrap_in_result)]
     pub fn orchard_tree_by_height_range<R>(
         &self,

From c15d73404505a6e11db511104a0ba35e55039483 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 8 Sep 2023 08:35:53 +1000
Subject: [PATCH 08/56] Make skipping the lowest tree independent of iteration
 order

---
 .../finalized_state/disk_format/upgrade/add_subtrees.rs     | 6 ++++--
 1 file changed, 4 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 5c68805afaa..3ea29b3f398 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
@@ -25,11 +25,13 @@ pub fn run(
     upgrade_db: &ZebraDb,
     cancel_receiver: &mpsc::Receiver<CancelFormatChange>,
 ) -> Result<(), CancelFormatChange> {
-    for ((_prev_height, prev_tree), (height, tree)) in upgrade_db
+    for (prev_tree, height, tree) in upgrade_db
         .sapling_tree_by_height_range(..=initial_tip_height)
         .tuple_windows()
         // The first block with sapling notes can't complete a subtree, see below for details.
-        .skip(1)
+        .filter_map(|((prev_height, prev_tree), (height, tree))| {
+            prev_height.is_min().then_some((prev_tree, height, tree))
+        })
     {
         // Return early if there is a cancel signal.
         if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {

From 0695da82d4f3d79de54704ac03c95a95d5057c14 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 8 Sep 2023 08:58:48 +1000
Subject: [PATCH 09/56] Move new subtree checks into the iterator, rename to
 end_height

---
 .../disk_format/upgrade/add_subtrees.rs       | 68 ++++++++++---------
 1 file changed, 35 insertions(+), 33 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 3ea29b3f398..41214f4bc46 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
@@ -25,43 +25,45 @@ pub fn run(
     upgrade_db: &ZebraDb,
     cancel_receiver: &mpsc::Receiver<CancelFormatChange>,
 ) -> Result<(), CancelFormatChange> {
-    for (prev_tree, height, tree) in upgrade_db
+    // Zebra stores exactly one note commitment tree for every block with sapling notes.
+    // (It also stores the empty note commitment tree for the genesis block, but we skip that.)
+    //
+    // The consensus rules limit blocks to less than 2^16 sapling and 2^16 orchard outputs. So a
+    // block can't complete multiple level 16 subtrees (or complete an entire subtree by itself).
+    // Currently, with 2MB blocks and v4/v5 sapling and orchard output sizes, the subtree index can
+    // increase by 1 every ~20 full blocks.
+    //
+    // Therefore, the first block with shielded note can't complete a subtree, which means we can
+    // skip the (genesis block, first shielded block) tree pair.
+
+    // Generate a list of subtrees, with note commitment trees, end heights, and the previous tree.
+    let subtrees = upgrade_db
         .sapling_tree_by_height_range(..=initial_tip_height)
+        // The first block with sapling notes can't complete a subtree, see above for details.
+        .filter(|(height, _tree)| !height.is_min())
+        // We need both the tree and its previous tree for each shielded block.
         .tuple_windows()
-        // The first block with sapling notes can't complete a subtree, see below for details.
-        .filter_map(|((prev_height, prev_tree), (height, tree))| {
-            prev_height.is_min().then_some((prev_tree, height, tree))
-        })
-    {
+        .map(|((_prev_height, prev_tree), (end_height, tree))| (prev_tree, end_height, tree))
+        // Empty note commitment trees can't contain subtrees, so they have invalid subtree indexes.
+        // But since we skip the empty genesis tree, all trees must have valid indexes.
+        // So we don't need to unwrap the optional values for this comparison to be correct.
+        .filter(|(prev_tree, _end_height, tree)| tree.subtree_index() > prev_tree.subtree_index());
+
+    for (prev_tree, end_height, tree) in subtrees {
         // Return early if there is a cancel signal.
         if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
             return Err(CancelFormatChange);
         }
 
-        // Zebra stores exactly one note commitment tree for every block with sapling notes.
-        // (It also stores the empty note commitment tree for the genesis block, but we skip that.)
-        //
-        // The consensus rules limit blocks to less than 2^16 sapling outputs. So a block can't
-        // complete multiple level 16 subtrees (or complete an entire subtree by itself).
-        // Currently, with 2MB blocks and v4/v5 sapling output sizes, the subtree index can
-        // increase by 1 every ~20 full blocks.
-
-        // Empty note commitment trees can't contain subtrees, so they have invalid subtree indexes.
-        // But we skip the empty genesis tree in the iterator, and all other trees must have notes.
-        let before_block_subtree_index = prev_tree
-            .subtree_index()
-            .expect("already skipped empty tree");
-        let end_of_block_subtree_index = tree.subtree_index().expect("already skipped empty tree");
-
         // If this block completed 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.
-            write_sapling_subtree(upgrade_db, index, height, node);
+            write_sapling_subtree(upgrade_db, index, end_height, node);
             continue;
-        } else if end_of_block_subtree_index > before_block_subtree_index {
+        } else {
             // 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 remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
@@ -72,7 +74,7 @@ pub fn run(
             );
 
             let block = upgrade_db
-                .block(height.into())
+                .block(end_height.into())
                 .expect("height with note commitment tree should have block");
             let sapling_note_commitments = block
                 .sapling_note_commitments()
@@ -90,7 +92,7 @@ pub fn run(
             let (index, node) =
                 subtree.expect("already checked that the block completed a subtree");
 
-            write_sapling_subtree(upgrade_db, index, height, node);
+            write_sapling_subtree(upgrade_db, index, end_height, node);
         }
     }
 
@@ -411,10 +413,10 @@ fn check_orchard_subtrees(db: &ZebraDb) -> bool {
 fn write_sapling_subtree(
     upgrade_db: &ZebraDb,
     index: NoteCommitmentSubtreeIndex,
-    height: Height,
+    end_height: Height,
     node: sapling::tree::Node,
 ) {
-    let subtree = NoteCommitmentSubtree::new(index, height, node);
+    let subtree = NoteCommitmentSubtree::new(index, end_height, node);
 
     let mut batch = DiskWriteBatch::new();
 
@@ -425,20 +427,20 @@ fn write_sapling_subtree(
         .expect("writing sapling note commitment subtrees should always succeed.");
 
     if index.0 % 100 == 0 {
-        info!(?height, index = ?index.0, "calculated and added sapling subtree");
+        info!(?end_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");
+    debug!(?end_height, index = ?index.0, ?node, "calculated and added sapling subtree");
 }
 
 /// Writes a Orchard note commitment subtree to `upgrade_db`.
 fn write_orchard_subtree(
     upgrade_db: &ZebraDb,
     index: NoteCommitmentSubtreeIndex,
-    height: Height,
+    end_height: Height,
     node: orchard::tree::Node,
 ) {
-    let subtree = NoteCommitmentSubtree::new(index, height, node);
+    let subtree = NoteCommitmentSubtree::new(index, end_height, node);
 
     let mut batch = DiskWriteBatch::new();
 
@@ -449,8 +451,8 @@ fn write_orchard_subtree(
         .expect("writing orchard note commitment subtrees should always succeed.");
 
     if index.0 % 300 == 0 {
-        info!(?height, index = ?index.0, "calculated and added orchard subtree");
+        info!(?end_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");
+    debug!(?end_height, index = ?index.0, ?node, "calculated and added orchard subtree");
 }

From f0be95c36312fbb3c5b910df5446df5c460d4dcd Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 8 Sep 2023 09:05:46 +1000
Subject: [PATCH 10/56] Split subtree calculation into a new method

---
 .../disk_format/upgrade/add_subtrees.rs       | 95 +++++++++++--------
 1 file changed, 55 insertions(+), 40 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 41214f4bc46..8fd5f4647aa 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,6 +1,6 @@
 //! Fully populate the Sapling and Orchard note commitment subtrees for existing blocks in the database.
 
-use std::sync::mpsc;
+use std::sync::{mpsc, Arc};
 
 use itertools::Itertools;
 
@@ -55,45 +55,7 @@ pub fn run(
             return Err(CancelFormatChange);
         }
 
-        // If this block completed 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.
-            write_sapling_subtree(upgrade_db, index, end_height, node);
-            continue;
-        } else {
-            // 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 remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
-
-            assert!(
-                remaining_notes > 0,
-                "just checked for complete subtrees above"
-            );
-
-            let block = upgrade_db
-                .block(end_height.into())
-                .expect("height with note commitment tree should have block");
-            let sapling_note_commitments = block
-                .sapling_note_commitments()
-                .take(remaining_notes)
-                .cloned()
-                .collect();
-
-            // This takes less than 1 second per tree, so we don't need to make it cancellable.
-            let (_sapling_nct, subtree) = NoteCommitmentTrees::update_sapling_note_commitment_tree(
-                prev_tree,
-                sapling_note_commitments,
-            )
-            .expect("finalized notes should append successfully");
-
-            let (index, node) =
-                subtree.expect("already checked that the block completed a subtree");
-
-            write_sapling_subtree(upgrade_db, index, end_height, node);
-        }
+        calculate_and_write_sapling_subtree(upgrade_db, prev_tree, end_height, tree);
     }
 
     for ((_prev_height, prev_tree), (height, tree)) in upgrade_db
@@ -409,6 +371,59 @@ fn check_orchard_subtrees(db: &ZebraDb) -> bool {
     is_valid
 }
 
+/// Calculates and writes a Sapling note commitment subtree to `upgrade_db`.
+///
+/// `tree` must be a note commitment tree containing a recently completed subtree,
+/// which was not already completed in `prev_tree`.
+///
+/// # Panics
+///
+/// If `tree` does not contain a recently completed subtree.
+fn calculate_and_write_sapling_subtree(
+    upgrade_db: &ZebraDb,
+    prev_tree: Arc<sapling::tree::NoteCommitmentTree>,
+    end_height: Height,
+    tree: Arc<sapling::tree::NoteCommitmentTree>,
+) {
+    // If this block completed 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.
+        write_sapling_subtree(upgrade_db, index, end_height, node);
+    } else {
+        // 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 remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
+
+        assert!(
+            remaining_notes > 0,
+            "tree must contain a recently completed subtree"
+        );
+
+        let block = upgrade_db
+            .block(end_height.into())
+            .expect("height with note commitment tree should have block");
+        let sapling_note_commitments = block
+            .sapling_note_commitments()
+            .take(remaining_notes)
+            .cloned()
+            .collect();
+
+        // This takes less than 1 second per tree, so we don't need to make it cancellable.
+        let (_sapling_nct, subtree) = NoteCommitmentTrees::update_sapling_note_commitment_tree(
+            prev_tree,
+            sapling_note_commitments,
+        )
+        .expect("finalized notes should append successfully");
+
+        let (index, node) = subtree.expect("already checked that the block completed a subtree");
+
+        write_sapling_subtree(upgrade_db, index, end_height, node);
+    }
+}
+
 /// Writes a Sapling note commitment subtree to `upgrade_db`.
 fn write_sapling_subtree(
     upgrade_db: &ZebraDb,

From 886437054d888edbd6a9f9f1982c8978c1a74658 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 12 Sep 2023 10:09:13 +1000
Subject: [PATCH 11/56] Split the calculate and write methods

---
 .../disk_format/upgrade/add_subtrees.rs       | 35 ++++++++++---------
 1 file changed, 18 insertions(+), 17 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 8fd5f4647aa..f09f7f15dd5 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
@@ -50,12 +50,13 @@ pub fn run(
         .filter(|(prev_tree, _end_height, tree)| tree.subtree_index() > prev_tree.subtree_index());
 
     for (prev_tree, end_height, tree) in subtrees {
-        // Return early if there is a cancel signal.
+        // Return early if the upgrade is cancelled.
         if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
             return Err(CancelFormatChange);
         }
 
-        calculate_and_write_sapling_subtree(upgrade_db, prev_tree, end_height, tree);
+        let subtree = calculate_sapling_subtree(upgrade_db, prev_tree, end_height, tree);
+        write_sapling_subtree(upgrade_db, subtree);
     }
 
     for ((_prev_height, prev_tree), (height, tree)) in upgrade_db
@@ -371,27 +372,31 @@ fn check_orchard_subtrees(db: &ZebraDb) -> bool {
     is_valid
 }
 
-/// Calculates and writes a Sapling note commitment subtree to `upgrade_db`.
+/// Calculates a Sapling note commitment subtree, reading blocks from `read_db` if needed.
 ///
 /// `tree` must be a note commitment tree containing a recently completed subtree,
 /// which was not already completed in `prev_tree`.
 ///
+/// `prev_tree` is only used to rebuild the subtree, if it was completed inside the block.
+/// If the subtree was completed by the final note commitment in the block, `prev_tree` is unused.
+///
 /// # Panics
 ///
 /// If `tree` does not contain a recently completed subtree.
-fn calculate_and_write_sapling_subtree(
-    upgrade_db: &ZebraDb,
+#[must_use = "subtree should be written to the database after it is calculated"]
+fn calculate_sapling_subtree(
+    read_db: &ZebraDb,
     prev_tree: Arc<sapling::tree::NoteCommitmentTree>,
     end_height: Height,
     tree: Arc<sapling::tree::NoteCommitmentTree>,
-) {
+) -> NoteCommitmentSubtree<sapling::tree::Node> {
     // If this block completed 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.
-        write_sapling_subtree(upgrade_db, index, end_height, node);
+        NoteCommitmentSubtree::new(index, end_height, node)
     } else {
         // 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.
@@ -402,7 +407,7 @@ fn calculate_and_write_sapling_subtree(
             "tree must contain a recently completed subtree"
         );
 
-        let block = upgrade_db
+        let block = read_db
             .block(end_height.into())
             .expect("height with note commitment tree should have block");
         let sapling_note_commitments = block
@@ -420,19 +425,15 @@ fn calculate_and_write_sapling_subtree(
 
         let (index, node) = subtree.expect("already checked that the block completed a subtree");
 
-        write_sapling_subtree(upgrade_db, index, end_height, node);
+        NoteCommitmentSubtree::new(index, end_height, node)
     }
 }
 
 /// Writes a Sapling note commitment subtree to `upgrade_db`.
 fn write_sapling_subtree(
     upgrade_db: &ZebraDb,
-    index: NoteCommitmentSubtreeIndex,
-    end_height: Height,
-    node: sapling::tree::Node,
+    subtree: NoteCommitmentSubtree<sapling::tree::Node>,
 ) {
-    let subtree = NoteCommitmentSubtree::new(index, end_height, node);
-
     let mut batch = DiskWriteBatch::new();
 
     batch.insert_sapling_subtree(upgrade_db, &subtree);
@@ -441,11 +442,11 @@ fn write_sapling_subtree(
         .write_batch(batch)
         .expect("writing sapling note commitment subtrees should always succeed.");
 
-    if index.0 % 100 == 0 {
-        info!(?end_height, index = ?index.0, "calculated and added sapling subtree");
+    if subtree.index.0 % 100 == 0 {
+        info!(end_height = ?subtree.end, index = ?subtree.index.0, "calculated and added sapling subtree");
     }
     // This log happens about once per second on recent machines with SSD disks.
-    debug!(?end_height, index = ?index.0, ?node, "calculated and added sapling subtree");
+    debug!(end_height = ?subtree.end, index = ?subtree.index.0, "calculated and added sapling subtree");
 }
 
 /// Writes a Orchard note commitment subtree to `upgrade_db`.

From 538653df828feb1d4e518a06a6bd5c3484527cb6 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 12 Sep 2023 10:38:43 +1000
Subject: [PATCH 12/56] Quickly check the first subtree before running the full
 upgrade

---
 Cargo.lock                                    |  7 +++++
 zebra-state/Cargo.toml                        |  1 +
 .../src/service/finalized_state/disk_db.rs    | 14 ++++++++++
 .../disk_format/upgrade/add_subtrees.rs       | 27 +++++++++++++++++++
 4 files changed, 49 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index 908a802095a..c75c162cd95 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1795,6 +1795,12 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "hex-literal"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6fe2267d4ed49bc07b63801559be28c718ea06c4738b7a03c94df7386d2cde46"
+
 [[package]]
 name = "hmac"
 version = "0.12.1"
@@ -5723,6 +5729,7 @@ dependencies = [
  "futures",
  "halo2_proofs",
  "hex",
+ "hex-literal",
  "howudoin",
  "indexmap 2.0.0",
  "insta",
diff --git a/zebra-state/Cargo.toml b/zebra-state/Cargo.toml
index f1dadc9a10f..d091e414efa 100644
--- a/zebra-state/Cargo.toml
+++ b/zebra-state/Cargo.toml
@@ -48,6 +48,7 @@ chrono = { version = "0.4.30", default-features = false, features = ["clock", "s
 dirs = "5.0.1"
 futures = "0.3.28"
 hex = "0.4.3"
+hex-literal = "0.4.1"
 indexmap = "2.0.0"
 itertools = "0.11.0"
 lazy_static = "1.4.0"
diff --git a/zebra-state/src/service/finalized_state/disk_db.rs b/zebra-state/src/service/finalized_state/disk_db.rs
index 50a10ba4540..51af47f72ea 100644
--- a/zebra-state/src/service/finalized_state/disk_db.rs
+++ b/zebra-state/src/service/finalized_state/disk_db.rs
@@ -64,6 +64,9 @@ pub struct DiskDb {
     // This configuration cannot be modified after the database is initialized,
     // because some clones would have different values.
     //
+    /// The configured network for this database.
+    network: Network,
+
     /// The configured temporary database setting.
     ///
     /// If true, the database files are deleted on drop.
@@ -247,10 +250,15 @@ pub trait ReadDisk {
 impl PartialEq for DiskDb {
     fn eq(&self, other: &Self) -> bool {
         if self.db.path() == other.db.path() {
+            assert_eq!(
+                self.network, other.network,
+                "database with same path but different network configs",
+            );
             assert_eq!(
                 self.ephemeral, other.ephemeral,
                 "database with same path but different ephemeral configs",
             );
+
             return true;
         }
 
@@ -636,6 +644,7 @@ impl DiskDb {
                 info!("Opened Zebra state cache at {}", path.display());
 
                 let db = DiskDb {
+                    network,
                     ephemeral: config.ephemeral,
                     db: Arc::new(db),
                 };
@@ -656,6 +665,11 @@ impl DiskDb {
 
     // Accessor methods
 
+    /// Returns the configured network for this database.
+    pub fn network(&self) -> Network {
+        self.network
+    }
+
     /// Returns the `Path` where the files used by this database are located.
     pub fn path(&self) -> &Path {
         self.db.path()
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 f09f7f15dd5..32d52a1c239 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,12 +2,14 @@
 
 use std::sync::{mpsc, Arc};
 
+use hex_literal::hex;
 use itertools::Itertools;
 
 use zebra_chain::{
     block::Height,
     orchard,
     parallel::tree::NoteCommitmentTrees,
+    parameters::Network::*,
     sapling,
     subtree::{NoteCommitmentSubtree, NoteCommitmentSubtreeIndex},
 };
@@ -16,6 +18,22 @@ use crate::service::finalized_state::{
     disk_format::upgrade::CancelFormatChange, DiskWriteBatch, ZebraDb,
 };
 
+/// A quick test vector that allows us to fail an incorrect upgrade within a few seconds.
+fn first_sapling_mainnet_subtree() -> NoteCommitmentSubtree<sapling::tree::Node> {
+    // This test vector was generated using the command:
+    // ```sh
+    // zcash-cli z_getsubtreesbyindex sapling 0 1
+    // ```
+    NoteCommitmentSubtree {
+        index: 0.into(),
+        node: hex!("754bb593ea42d231a7ddf367640f09bbf59dc00f2c1d2003cc340e0c016b5b13")
+            .as_slice()
+            .try_into()
+            .expect("test vector is valid"),
+        end: Height(558822),
+    }
+}
+
 /// 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.
@@ -25,6 +43,9 @@ pub fn run(
     upgrade_db: &ZebraDb,
     cancel_receiver: &mpsc::Receiver<CancelFormatChange>,
 ) -> Result<(), CancelFormatChange> {
+    // Creating this test vector involves a cryptographic check, so only do it once.
+    let first_sapling_mainnet_subtree = first_sapling_mainnet_subtree();
+
     // Zebra stores exactly one note commitment tree for every block with sapling notes.
     // (It also stores the empty note commitment tree for the genesis block, but we skip that.)
     //
@@ -56,6 +77,12 @@ pub fn run(
         }
 
         let subtree = calculate_sapling_subtree(upgrade_db, prev_tree, end_height, tree);
+
+        // TODO: split this out into a quick check method
+        if upgrade_db.network() == Mainnet && subtree.index == first_sapling_mainnet_subtree.index {
+            assert_eq!(subtree, first_sapling_mainnet_subtree)
+        }
+
         write_sapling_subtree(upgrade_db, subtree);
     }
 

From a678989728eb5419d7fab6ef11932c7747e1089c Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 12 Sep 2023 11:22:51 +1000
Subject: [PATCH 13/56] Do the quick checks every time Zebra runs, and refactor
 slow check error handling

---
 .../finalized_state/disk_format/upgrade.rs    |   4 +
 .../disk_format/upgrade/add_subtrees.rs       | 197 +++++++++++-------
 .../src/service/finalized_state/zebra_db.rs   |   4 +
 3 files changed, 129 insertions(+), 76 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 b4755751e55..d3dab563e44 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
@@ -190,6 +190,10 @@ impl DbFormatChange {
         upgrade_db: ZebraDb,
         cancel_receiver: mpsc::Receiver<CancelFormatChange>,
     ) -> Result<(), CancelFormatChange> {
+        // These quick checks should pass for all format changes.
+        // (See the detailed comment at the end of this method.)
+        add_subtrees::quick_check(&upgrade_db);
+
         match self {
             // Perform any required upgrades, then mark the state as upgraded.
             Upgrade { .. } => self.apply_format_upgrade(
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 32d52a1c239..be20fb0454b 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
@@ -43,9 +43,6 @@ pub fn run(
     upgrade_db: &ZebraDb,
     cancel_receiver: &mpsc::Receiver<CancelFormatChange>,
 ) -> Result<(), CancelFormatChange> {
-    // Creating this test vector involves a cryptographic check, so only do it once.
-    let first_sapling_mainnet_subtree = first_sapling_mainnet_subtree();
-
     // Zebra stores exactly one note commitment tree for every block with sapling notes.
     // (It also stores the empty note commitment tree for the genesis block, but we skip that.)
     //
@@ -77,12 +74,6 @@ pub fn run(
         }
 
         let subtree = calculate_sapling_subtree(upgrade_db, prev_tree, end_height, tree);
-
-        // TODO: split this out into a quick check method
-        if upgrade_db.network() == Mainnet && subtree.index == first_sapling_mainnet_subtree.index {
-            assert_eq!(subtree, first_sapling_mainnet_subtree)
-        }
-
         write_sapling_subtree(upgrade_db, subtree);
     }
 
@@ -156,29 +147,87 @@ pub fn run(
     Ok(())
 }
 
-/// Check that note commitment subtrees were correctly added.
+/// Quickly check that the first calculated subtree is correct.
+///
+/// This allows us to fail the upgrade quickly in tests and during development,
+/// rather than waiting ~20 minutes to see if it failed.
 ///
 /// # Panics
 ///
 /// If a note commitment subtree is missing or incorrect.
-pub fn check(db: &ZebraDb) {
-    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)");
+fn quick_check(db: &ZebraDb) {
+    // We check the first sapling tree on mainnet, so skip this check if it isn't available.
+    let Some(NoteCommitmentSubtreeIndex(mut first_incomplete_subtree_index)) =
+        db.sapling_tree().subtree_index()
+    else {
+        return;
+    };
+
+    if first_incomplete_subtree_index == 0 || db.network() != Mainnet {
+        return;
+    }
+
+    let mut result = Ok(());
+
+    // Find the first complete subtree, with its note commitment tree, end height, and the previous tree.
+    let first_complete_subtree = db
+        .sapling_tree_by_height_range(..)
+        // The first block with sapling notes can't complete a subtree, see above for details.
+        .filter(|(height, _tree)| !height.is_min())
+        // We need both the tree and its previous tree for each shielded block.
+        .tuple_windows()
+        .map(|((_prev_height, prev_tree), (end_height, tree))| (prev_tree, end_height, tree))
+        // Empty note commitment trees can't contain subtrees, so they have invalid subtree indexes.
+        // But since we skip the empty genesis tree, all trees must have valid indexes.
+        // So we don't need to unwrap the optional values for this comparison to be correct.
+        .find(|(prev_tree, _end_height, tree)| tree.subtree_index() > prev_tree.subtree_index());
+
+    let Some((prev_tree, end_height, tree)) = first_complete_subtree else {
+        result = Err("iterator did not find complete subtree, but the tree has it");
+        error!(?result);
+    };
+
+    // Creating this test vector involves a cryptographic check, so only do it once.
+    let expected_subtree = first_sapling_mainnet_subtree();
+
+    let db_subtree = calculate_sapling_subtree(db, prev_tree, end_height, tree);
+
+    if db_subtree != expected_subtree {
+        result = Err("first subtree did not match expected test vector");
+        error!(?result, ?db_subtree, ?expected_subtree);
+    }
+
+    if result.is_err() {
+        panic!("incorrect first sapling subtree: {result}");
     }
 }
 
-/// Check that Sapling note commitment subtrees were correctly added.
+/// Check that note commitment subtrees were correctly added.
 ///
 /// # Panics
 ///
 /// If a note commitment subtree is missing or incorrect.
-fn check_sapling_subtrees(db: &ZebraDb) -> bool {
+pub fn check(db: &ZebraDb) {
+    let sapling_result = check_sapling_subtrees(db);
+    let orchard_result = check_orchard_subtrees(db);
+
+    if sapling_result.is_err() || orchard_result.is_err() {
+        // TODO: when the check functions are refactored so they are called from a single function,
+        //       move this panic into that function, but still log a detailed message here
+        panic!(
+            "missing or bad subtree(s): sapling: {sapling_result:?}, orchard: {orchard_result:?}"
+        );
+    }
+}
+
+/// Check that Sapling note commitment subtrees were correctly added.
+///
+/// Returns an error if a note commitment subtree is missing or incorrect.
+fn check_sapling_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
     let Some(NoteCommitmentSubtreeIndex(mut first_incomplete_subtree_index)) =
         db.sapling_tree().subtree_index()
     else {
-        return true;
+        return Ok(());
     };
 
     // If there are no incomplete subtrees in the tree, also expect a subtree for the final index.
@@ -186,51 +235,48 @@ fn check_sapling_subtrees(db: &ZebraDb) -> bool {
         first_incomplete_subtree_index += 1;
     }
 
-    let mut is_valid = true;
+    let mut result = Ok(());
     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 {
-            error!(index, "missing subtree");
-            is_valid = false;
+            result = Err("missing subtree");
+            error!(?result, index);
             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 {
-            error!(?subtree.end, "missing note commitment tree at subtree completion height");
-            is_valid = false;
+            result = Err("missing note commitment tree at subtree completion height");
+            error!(?result, ?subtree.end);
             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() {
             if subtree.index != index {
-                error!("completed subtree indexes should match");
-                is_valid = false;
+                result = Err("completed subtree indexes should match");
+                error!(?result);
             }
 
             if subtree.node != node {
-                error!("completed subtree roots should match");
-                is_valid = false;
+                result = Err("completed subtree roots should match");
+                error!(?result);
             }
         }
         // 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 {
-                error!(?subtree.end, "missing note commitment tree at subtree completion height");
-                is_valid = false;
+                result = Err("missing note commitment tree at subtree completion height");
+                error!(?result, ?subtree.end);
                 continue;
             };
 
             let prev_subtree_index = prev_tree.subtree_index();
             let subtree_index = tree.subtree_index();
             if subtree_index <= prev_subtree_index {
-                error!(
-                    ?subtree_index,
-                    ?prev_subtree_index,
-                    "note commitment tree at end height should have incremented subtree index"
-                );
-                is_valid = false;
+                result =
+                    Err("note commitment tree at end height should have incremented subtree index");
+                error!(?result, ?subtree_index, ?prev_subtree_index,);
             }
         }
     }
@@ -253,47 +299,47 @@ fn check_sapling_subtrees(db: &ZebraDb) -> bool {
     {
         // Check that there's an entry for every completed sapling subtree root in all sapling trees
         let Some(subtree) = db.sapling_subtree_by_index(index) else {
-            error!(?index, "missing subtree");
-            is_valid = false;
+            result = Err("missing subtree");
+            error!(?result, index);
             continue;
         };
 
         // Check that the subtree end height matches that in the sapling trees.
         if 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;
+            result = Err("bad sapling subtree end height");
+            error!(?result, ?subtree.end, ?height, ?index, ?is_complete, );
         }
 
         // Check the 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.node != node {
-                error!("completed subtree roots should match");
-                is_valid = false;
+                result = Err("completed subtree roots should match");
+                error!(?result);
             }
         }
     }
 
-    if !is_valid {
+    if result.is_err() {
         error!(
+            ?result,
             ?subtree_count,
-            first_incomplete_subtree_index, "missing or bad sapling subtrees"
+            first_incomplete_subtree_index,
+            "missing or bad sapling subtrees"
         );
     }
 
-    is_valid
+    result
 }
 
 /// 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 {
+/// Returns an error if a note commitment subtree is missing or incorrect.
+fn check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
     let Some(NoteCommitmentSubtreeIndex(mut first_incomplete_subtree_index)) =
         db.orchard_tree().subtree_index()
     else {
-        return true;
+        return Ok(());
     };
 
     // If there are no incomplete subtrees in the tree, also expect a subtree for the final index.
@@ -301,51 +347,48 @@ fn check_orchard_subtrees(db: &ZebraDb) -> bool {
         first_incomplete_subtree_index += 1;
     }
 
-    let mut is_valid = true;
+    let mut result = Ok(());
     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 {
-            error!(index, "missing subtree");
-            is_valid = false;
+            result = Err("missing subtree");
+            error!(?result, index);
             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 {
-            error!(?subtree.end, "missing note commitment tree at subtree completion height");
-            is_valid = false;
+            result = Err("missing note commitment tree at subtree completion height");
+            error!(?result, ?subtree.end);
             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() {
             if subtree.index != index {
-                error!("completed subtree indexes should match");
-                is_valid = false;
+                result = Err("completed subtree indexes should match");
+                error!(?result);
             }
 
             if subtree.node != node {
-                error!("completed subtree roots should match");
-                is_valid = false;
+                result = Err("completed subtree roots should match");
+                error!(?result);
             }
         }
         // 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 {
-                error!(?subtree.end, "missing note commitment tree at subtree completion height");
-                is_valid = false;
+                result = Err("missing note commitment tree at subtree completion height");
+                error!(?result, ?subtree.end);
                 continue;
             };
 
             let prev_subtree_index = prev_tree.subtree_index();
             let subtree_index = tree.subtree_index();
             if subtree_index <= prev_subtree_index {
-                error!(
-                    ?subtree_index,
-                    ?prev_subtree_index,
-                    "note commitment tree at end height should have incremented subtree index"
-                );
-                is_valid = false;
+                result =
+                    Err("note commitment tree at end height should have incremented subtree index");
+                error!(?result, ?subtree_index, ?prev_subtree_index,);
             }
         }
     }
@@ -355,7 +398,7 @@ fn check_orchard_subtrees(db: &ZebraDb) -> bool {
         .orchard_tree_by_height_range(..)
         // Exclude empty orchard tree and add subtree indexes
         .filter_map(|(height, tree)| Some((tree.subtree_index()?, height, tree)))
-        // Exclude heights that don't complete a subtree and count completed subtree
+        // Exclude heights that don't complete a subtree and count completed subtrees
         .filter_map(|(subtree_index, height, tree)| {
             if tree.is_complete_subtree() || subtree_index.0 > subtree_count {
                 let subtree_index = subtree_count;
@@ -368,35 +411,37 @@ fn check_orchard_subtrees(db: &ZebraDb) -> bool {
     {
         // Check that there's an entry for every completed orchard subtree root in all orchard trees
         let Some(subtree) = db.orchard_subtree_by_index(index) else {
-            error!(?index, "missing subtree");
-            is_valid = false;
+            result = Err("missing subtree");
+            error!(?result, index);
             continue;
         };
 
         // Check that the subtree end height matches that in the orchard trees.
         if 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;
+            result = Err("bad orchard subtree end height");
+            error!(?result, ?subtree.end, ?height, ?index, ?is_complete, );
         }
 
         // Check the 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.node != node {
-                error!("completed subtree roots should match");
-                is_valid = false;
+                result = Err("completed subtree roots should match");
+                error!(?result);
             }
         }
     }
 
-    if !is_valid {
+    if result.is_err() {
         error!(
+            ?result,
             ?subtree_count,
-            first_incomplete_subtree_index, "missing or bad orchard subtrees"
+            first_incomplete_subtree_index,
+            "missing or bad orchard subtrees"
         );
     }
 
-    is_valid
+    result
 }
 
 /// Calculates a Sapling note commitment subtree, reading blocks from `read_db` if needed.
diff --git a/zebra-state/src/service/finalized_state/zebra_db.rs b/zebra-state/src/service/finalized_state/zebra_db.rs
index 109e7e88310..341c77c2e15 100644
--- a/zebra-state/src/service/finalized_state/zebra_db.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db.rs
@@ -118,6 +118,10 @@ impl ZebraDb {
             // If we're re-opening a previously upgraded or newly created database,
             // the database format should be valid.
             // (There's no format change here, so the format change checks won't run.)
+            //
+            // Do the quick checks first, then the slower checks.
+            add_subtrees::quick_check(&upgrade_db);
+
             DbFormatChange::check_for_duplicate_trees(db.clone());
             upgrade::add_subtrees::check(&db.clone());
         }

From bebb2260d419c31782aeb44001ecab8b179ba7d7 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 12 Sep 2023 11:30:32 +1000
Subject: [PATCH 14/56] Do quick checks for orchard as well

---
 .../disk_format/upgrade/add_subtrees.rs       | 193 +++++++++++++++---
 .../src/service/finalized_state/zebra_db.rs   |   4 +-
 2 files changed, 169 insertions(+), 28 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 be20fb0454b..a7f8ac1ad87 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
@@ -18,22 +18,6 @@ use crate::service::finalized_state::{
     disk_format::upgrade::CancelFormatChange, DiskWriteBatch, ZebraDb,
 };
 
-/// A quick test vector that allows us to fail an incorrect upgrade within a few seconds.
-fn first_sapling_mainnet_subtree() -> NoteCommitmentSubtree<sapling::tree::Node> {
-    // This test vector was generated using the command:
-    // ```sh
-    // zcash-cli z_getsubtreesbyindex sapling 0 1
-    // ```
-    NoteCommitmentSubtree {
-        index: 0.into(),
-        node: hex!("754bb593ea42d231a7ddf367640f09bbf59dc00f2c1d2003cc340e0c016b5b13")
-            .as_slice()
-            .try_into()
-            .expect("test vector is valid"),
-        end: Height(558822),
-    }
-}
-
 /// 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.
@@ -155,20 +139,69 @@ pub fn run(
 /// # Panics
 ///
 /// If a note commitment subtree is missing or incorrect.
-fn quick_check(db: &ZebraDb) {
+pub fn quick_check(db: &ZebraDb) {
+    let sapling_result = quick_check_sapling_subtrees(db);
+    let orchard_result = quick_check_orchard_subtrees(db);
+
+    if sapling_result.is_err() || orchard_result.is_err() {
+        // TODO: when the check functions are refactored so they are called from a single function,
+        //       move this panic into that function, but still log a detailed message here
+        panic!(
+            "missing or bad first subtree: sapling: {sapling_result:?}, orchard: {orchard_result:?}"
+        );
+    }
+}
+
+/// A quick test vector that allows us to fail an incorrect upgrade within a few seconds.
+fn first_sapling_mainnet_subtree() -> NoteCommitmentSubtree<sapling::tree::Node> {
+    // This test vector was generated using the command:
+    // ```sh
+    // zcash-cli z_getsubtreesbyindex sapling 0 1
+    // ```
+    NoteCommitmentSubtree {
+        index: 0.into(),
+        node: hex!("754bb593ea42d231a7ddf367640f09bbf59dc00f2c1d2003cc340e0c016b5b13")
+            .as_slice()
+            .try_into()
+            .expect("test vector is valid"),
+        end: Height(558822),
+    }
+}
+
+/// A quick test vector that allows us to fail an incorrect upgrade within a few seconds.
+fn first_orchard_mainnet_subtree() -> NoteCommitmentSubtree<orchard::tree::Node> {
+    // This test vector was generated using the command:
+    // ```sh
+    // zcash-cli z_getsubtreesbyindex orchard 0 1
+    // ```
+    NoteCommitmentSubtree {
+        index: 0.into(),
+        node: hex!("d4e323b3ae0cabfb6be4087fec8c66d9a9bbfc354bf1d9588b6620448182063b")
+            .as_slice()
+            .try_into()
+            .expect("test vector is valid"),
+        end: Height(1707429),
+    }
+}
+
+/// Quickly check that the first calculated sapling subtree is correct.
+///
+/// This allows us to fail the upgrade quickly in tests and during development,
+/// rather than waiting ~20 minutes to see if it failed.
+///
+/// Returns an error if a note commitment subtree is missing or incorrect.
+fn quick_check_sapling_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
     // We check the first sapling tree on mainnet, so skip this check if it isn't available.
-    let Some(NoteCommitmentSubtreeIndex(mut first_incomplete_subtree_index)) =
+    let Some(NoteCommitmentSubtreeIndex(first_incomplete_subtree_index)) =
         db.sapling_tree().subtree_index()
     else {
-        return;
+        return Ok(());
     };
 
     if first_incomplete_subtree_index == 0 || db.network() != Mainnet {
-        return;
+        return Ok(());
     }
 
-    let mut result = Ok(());
-
     // Find the first complete subtree, with its note commitment tree, end height, and the previous tree.
     let first_complete_subtree = db
         .sapling_tree_by_height_range(..)
@@ -183,8 +216,9 @@ fn quick_check(db: &ZebraDb) {
         .find(|(prev_tree, _end_height, tree)| tree.subtree_index() > prev_tree.subtree_index());
 
     let Some((prev_tree, end_height, tree)) = first_complete_subtree else {
-        result = Err("iterator did not find complete subtree, but the tree has it");
+        let result = Err("iterator did not find complete subtree, but the tree has it");
         error!(?result);
+        return result;
     };
 
     // Creating this test vector involves a cryptographic check, so only do it once.
@@ -193,13 +227,63 @@ fn quick_check(db: &ZebraDb) {
     let db_subtree = calculate_sapling_subtree(db, prev_tree, end_height, tree);
 
     if db_subtree != expected_subtree {
-        result = Err("first subtree did not match expected test vector");
+        let result = Err("first subtree did not match expected test vector");
         error!(?result, ?db_subtree, ?expected_subtree);
+        return result;
     }
 
-    if result.is_err() {
-        panic!("incorrect first sapling subtree: {result}");
+    Ok(())
+}
+
+/// Quickly check that the first calculated orchard subtree is correct.
+///
+/// This allows us to fail the upgrade quickly in tests and during development,
+/// rather than waiting ~20 minutes to see if it failed.
+///
+/// Returns an error if a note commitment subtree is missing or incorrect.
+fn quick_check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
+    // We check the first orchard tree on mainnet, so skip this check if it isn't available.
+    let Some(NoteCommitmentSubtreeIndex(first_incomplete_subtree_index)) =
+        db.orchard_tree().subtree_index()
+    else {
+        return Ok(());
+    };
+
+    if first_incomplete_subtree_index == 0 || db.network() != Mainnet {
+        return Ok(());
+    }
+
+    // Find the first complete subtree, with its note commitment tree, end height, and the previous tree.
+    let first_complete_subtree = db
+        .orchard_tree_by_height_range(..)
+        // The first block with orchard notes can't complete a subtree, see above for details.
+        .filter(|(height, _tree)| !height.is_min())
+        // We need both the tree and its previous tree for each shielded block.
+        .tuple_windows()
+        .map(|((_prev_height, prev_tree), (end_height, tree))| (prev_tree, end_height, tree))
+        // Empty note commitment trees can't contain subtrees, so they have invalid subtree indexes.
+        // But since we skip the empty genesis tree, all trees must have valid indexes.
+        // So we don't need to unwrap the optional values for this comparison to be correct.
+        .find(|(prev_tree, _end_height, tree)| tree.subtree_index() > prev_tree.subtree_index());
+
+    let Some((prev_tree, end_height, tree)) = first_complete_subtree else {
+        let result = Err("iterator did not find complete subtree, but the tree has it");
+        error!(?result);
+        return result;
+    };
+
+    // Creating this test vector involves a cryptographic check, so only do it once.
+    let expected_subtree = first_orchard_mainnet_subtree();
+
+    let db_subtree = calculate_orchard_subtree(db, prev_tree, end_height, tree);
+
+    if db_subtree != expected_subtree {
+        let result = Err("first subtree did not match expected test vector");
+        error!(?result, ?db_subtree, ?expected_subtree);
+        return result;
     }
+
+    Ok(())
 }
 
 /// Check that note commitment subtrees were correctly added.
@@ -521,6 +605,63 @@ fn write_sapling_subtree(
     debug!(end_height = ?subtree.end, index = ?subtree.index.0, "calculated and added sapling subtree");
 }
 
+/// Calculates a Orchard note commitment subtree, reading blocks from `read_db` if needed.
+///
+/// `tree` must be a note commitment tree containing a recently completed subtree,
+/// which was not already completed in `prev_tree`.
+///
+/// `prev_tree` is only used to rebuild the subtree, if it was completed inside the block.
+/// If the subtree was completed by the final note commitment in the block, `prev_tree` is unused.
+///
+/// # Panics
+///
+/// If `tree` does not contain a recently completed subtree.
+#[must_use = "subtree should be written to the database after it is calculated"]
+fn calculate_orchard_subtree(
+    read_db: &ZebraDb,
+    prev_tree: Arc<orchard::tree::NoteCommitmentTree>,
+    end_height: Height,
+    tree: Arc<orchard::tree::NoteCommitmentTree>,
+) -> NoteCommitmentSubtree<orchard::tree::Node> {
+    // If this block completed 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.
+        NoteCommitmentSubtree::new(index, end_height, node)
+    } else {
+        // 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 remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
+
+        assert!(
+            remaining_notes > 0,
+            "tree must contain a recently completed subtree"
+        );
+
+        let block = read_db
+            .block(end_height.into())
+            .expect("height with note commitment tree should have block");
+        let orchard_note_commitments = block
+            .orchard_note_commitments()
+            .take(remaining_notes)
+            .cloned()
+            .collect();
+
+        // This takes less than 1 second per tree, so we don't need to make it cancellable.
+        let (_orchard_nct, subtree) = NoteCommitmentTrees::update_orchard_note_commitment_tree(
+            prev_tree,
+            orchard_note_commitments,
+        )
+        .expect("finalized notes should append successfully");
+
+        let (index, node) = subtree.expect("already checked that the block completed a subtree");
+
+        NoteCommitmentSubtree::new(index, end_height, node)
+    }
+}
+
 /// Writes a Orchard note commitment subtree to `upgrade_db`.
 fn write_orchard_subtree(
     upgrade_db: &ZebraDb,
diff --git a/zebra-state/src/service/finalized_state/zebra_db.rs b/zebra-state/src/service/finalized_state/zebra_db.rs
index 341c77c2e15..9a6d0f068d2 100644
--- a/zebra-state/src/service/finalized_state/zebra_db.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db.rs
@@ -120,10 +120,10 @@ impl ZebraDb {
             // (There's no format change here, so the format change checks won't run.)
             //
             // Do the quick checks first, then the slower checks.
-            add_subtrees::quick_check(&upgrade_db);
+            upgrade::add_subtrees::quick_check(&db);
 
             DbFormatChange::check_for_duplicate_trees(db.clone());
-            upgrade::add_subtrees::check(&db.clone());
+            upgrade::add_subtrees::check(&db);
         }
 
         db

From 8b54c09a3493334b0a543a0f05fc927d45eb7f68 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 12 Sep 2023 11:34:46 +1000
Subject: [PATCH 15/56] Make orchard tree upgrade match sapling upgrade code

---
 .../disk_format/upgrade/add_subtrees.rs       | 130 ++++++------------
 .../src/service/finalized_state/zebra_db.rs   |   5 +
 2 files changed, 46 insertions(+), 89 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 a7f8ac1ad87..b0301d48386 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
@@ -38,7 +38,7 @@ pub fn run(
     // Therefore, the first block with shielded note can't complete a subtree, which means we can
     // skip the (genesis block, first shielded block) tree pair.
 
-    // Generate a list of subtrees, with note commitment trees, end heights, and the previous tree.
+    // Generate a list of sapling subtree inputs: previous tree, current tree, and end height.
     let subtrees = upgrade_db
         .sapling_tree_by_height_range(..=initial_tip_height)
         // The first block with sapling notes can't complete a subtree, see above for details.
@@ -61,71 +61,27 @@ pub fn run(
         write_sapling_subtree(upgrade_db, subtree);
     }
 
-    for ((_prev_height, prev_tree), (height, tree)) in upgrade_db
+    // Generate a list of orchard subtree inputs: previous tree, current tree, and end height.
+    let subtrees = upgrade_db
         .orchard_tree_by_height_range(..=initial_tip_height)
+        // The first block with orchard notes can't complete a subtree, see above for details.
+        .filter(|(height, _tree)| !height.is_min())
+        // We need both the tree and its previous tree for each shielded block.
         .tuple_windows()
-        // The first block with orchard notes can't complete a subtree, see below for details.
-        .skip(1)
-    {
-        // Return early if there is a cancel signal.
+        .map(|((_prev_height, prev_tree), (end_height, tree))| (prev_tree, end_height, tree))
+        // Empty note commitment trees can't contain subtrees, so they have invalid subtree indexes.
+        // But since we skip the empty genesis tree, all trees must have valid indexes.
+        // So we don't need to unwrap the optional values for this comparison to be correct.
+        .filter(|(prev_tree, _end_height, tree)| tree.subtree_index() > prev_tree.subtree_index());
+
+    for (prev_tree, end_height, tree) in subtrees {
+        // Return early if the upgrade is cancelled.
         if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
             return Err(CancelFormatChange);
         }
 
-        // Zebra stores exactly one note commitment tree for every block with orchard notes.
-        // (It also stores the empty note commitment tree for the genesis block, but we skip that.)
-        //
-        // The consensus rules limit blocks to less than 2^16 orchard outputs. So a block can't
-        // complete multiple level 16 subtrees (or complete an entire subtree by itself).
-        // Currently, with 2MB blocks and v5 orchard output sizes, the subtree index can
-        // increase by 1 every ~20 full blocks.
-
-        // Empty note commitment trees can't contain subtrees, so they have invalid subtree indexes.
-        // But we skip the empty genesis tree in the iterator, and all other trees must have notes.
-        let before_block_subtree_index = prev_tree
-            .subtree_index()
-            .expect("already skipped empty tree");
-        let end_of_block_subtree_index = tree.subtree_index().expect("already skipped empty tree");
-
-        // If this block completed 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.
-            write_orchard_subtree(upgrade_db, index, height, node);
-            continue;
-        } else if end_of_block_subtree_index > before_block_subtree_index {
-            // 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 remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
-
-            assert!(
-                remaining_notes > 0,
-                "just checked for complete subtrees above"
-            );
-
-            let block = upgrade_db
-                .block(height.into())
-                .expect("height with note commitment tree should have block");
-            let orchard_note_commitments = block
-                .orchard_note_commitments()
-                .take(remaining_notes)
-                .cloned()
-                .collect();
-
-            // This takes less than 1 second per tree, so we don't need to make it cancellable.
-            let (_orchard_nct, subtree) = NoteCommitmentTrees::update_orchard_note_commitment_tree(
-                prev_tree,
-                orchard_note_commitments,
-            )
-            .expect("finalized notes should append successfully");
-
-            let (index, node) =
-                subtree.expect("already checked that the block completed a subtree");
-
-            write_orchard_subtree(upgrade_db, index, height, node);
-        }
+        let subtree = calculate_orchard_subtree(upgrade_db, prev_tree, end_height, tree);
+        write_orchard_subtree(upgrade_db, subtree);
     }
 
     Ok(())
@@ -585,26 +541,6 @@ fn calculate_sapling_subtree(
     }
 }
 
-/// Writes a Sapling note commitment subtree to `upgrade_db`.
-fn write_sapling_subtree(
-    upgrade_db: &ZebraDb,
-    subtree: NoteCommitmentSubtree<sapling::tree::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.");
-
-    if subtree.index.0 % 100 == 0 {
-        info!(end_height = ?subtree.end, index = ?subtree.index.0, "calculated and added sapling subtree");
-    }
-    // This log happens about once per second on recent machines with SSD disks.
-    debug!(end_height = ?subtree.end, index = ?subtree.index.0, "calculated and added sapling subtree");
-}
-
 /// Calculates a Orchard note commitment subtree, reading blocks from `read_db` if needed.
 ///
 /// `tree` must be a note commitment tree containing a recently completed subtree,
@@ -662,15 +598,31 @@ fn calculate_orchard_subtree(
     }
 }
 
+/// Writes a Sapling note commitment subtree to `upgrade_db`.
+fn write_sapling_subtree(
+    upgrade_db: &ZebraDb,
+    subtree: NoteCommitmentSubtree<sapling::tree::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.");
+
+    if subtree.index.0 % 100 == 0 {
+        info!(end_height = ?subtree.end, index = ?subtree.index.0, "calculated and added sapling subtree");
+    }
+    // This log happens about once per second on recent machines with SSD disks.
+    debug!(end_height = ?subtree.end, index = ?subtree.index.0, "calculated and added sapling subtree");
+}
+
 /// Writes a Orchard note commitment subtree to `upgrade_db`.
 fn write_orchard_subtree(
     upgrade_db: &ZebraDb,
-    index: NoteCommitmentSubtreeIndex,
-    end_height: Height,
-    node: orchard::tree::Node,
+    subtree: NoteCommitmentSubtree<orchard::tree::Node>,
 ) {
-    let subtree = NoteCommitmentSubtree::new(index, end_height, node);
-
     let mut batch = DiskWriteBatch::new();
 
     batch.insert_orchard_subtree(upgrade_db, &subtree);
@@ -679,9 +631,9 @@ fn write_orchard_subtree(
         .write_batch(batch)
         .expect("writing orchard note commitment subtrees should always succeed.");
 
-    if index.0 % 300 == 0 {
-        info!(?end_height, index = ?index.0, "calculated and added orchard subtree");
+    if subtree.index.0 % 100 == 0 {
+        info!(end_height = ?subtree.end, index = ?subtree.index.0, "calculated and added orchard subtree");
     }
-    // This log happens about 3 times per second on recent machines with SSD disks.
-    debug!(?end_height, index = ?index.0, ?node, "calculated and added orchard subtree");
+    // This log happens about once per second on recent machines with SSD disks.
+    debug!(end_height = ?subtree.end, index = ?subtree.index.0, "calculated and added orchard subtree");
 }
diff --git a/zebra-state/src/service/finalized_state/zebra_db.rs b/zebra-state/src/service/finalized_state/zebra_db.rs
index 9a6d0f068d2..58271bfeedb 100644
--- a/zebra-state/src/service/finalized_state/zebra_db.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db.rs
@@ -129,6 +129,11 @@ impl ZebraDb {
         db
     }
 
+    /// Returns the configured network for this database.
+    pub fn network(&self) -> Network {
+        self.db.network()
+    }
+
     /// Returns the `Path` where the files used by this database are located.
     pub fn path(&self) -> &Path {
         self.db.path()

From a63d1099aec333ade3940034063b070554ac9f79 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 12 Sep 2023 11:44:57 +1000
Subject: [PATCH 16/56] Upgrade subtrees in reverse height order

---
 .../disk_format/upgrade/add_subtrees.rs       | 22 +++++++++++++---
 .../finalized_state/zebra_db/shielded.rs      | 26 +++++++++++++++++++
 2 files changed, 44 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 b0301d48386..7387d02109c 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
@@ -20,6 +20,9 @@ use crate::service::finalized_state::{
 
 /// Runs disk format upgrade for adding Sapling and Orchard note commitment subtrees to database.
 ///
+/// Trees are added to the database in reverse height order, so that wallets can sync correctly
+/// while the upgrade is running.
+///
 /// Returns `Ok` if the upgrade completed, and `Err` if it was cancelled.
 #[allow(clippy::unwrap_in_result)]
 pub fn run(
@@ -27,6 +30,8 @@ pub fn run(
     upgrade_db: &ZebraDb,
     cancel_receiver: &mpsc::Receiver<CancelFormatChange>,
 ) -> Result<(), CancelFormatChange> {
+    // # Consensus
+    //
     // Zebra stores exactly one note commitment tree for every block with sapling notes.
     // (It also stores the empty note commitment tree for the genesis block, but we skip that.)
     //
@@ -37,15 +42,23 @@ pub fn run(
     //
     // Therefore, the first block with shielded note can't complete a subtree, which means we can
     // skip the (genesis block, first shielded block) tree pair.
+    //
+    // # Compatibility
+    //
+    // Because wallets search backwards from the chain tip, subtrees need to be added to the
+    // database in reverse height order. (Tip first, genesis last.)
+    //
+    // Otherwise, wallets that sync during the upgrade will be missing some notes.
 
     // Generate a list of sapling subtree inputs: previous tree, current tree, and end height.
     let subtrees = upgrade_db
-        .sapling_tree_by_height_range(..=initial_tip_height)
+        .sapling_tree_by_reversed_height_range(..=initial_tip_height)
         // The first block with sapling notes can't complete a subtree, see above for details.
         .filter(|(height, _tree)| !height.is_min())
         // We need both the tree and its previous tree for each shielded block.
         .tuple_windows()
-        .map(|((_prev_height, prev_tree), (end_height, tree))| (prev_tree, end_height, tree))
+        // Because the iterator is reversed, the larger tree is first.
+        .map(|((end_height, tree), (_prev_height, prev_tree))| (prev_tree, end_height, tree))
         // Empty note commitment trees can't contain subtrees, so they have invalid subtree indexes.
         // But since we skip the empty genesis tree, all trees must have valid indexes.
         // So we don't need to unwrap the optional values for this comparison to be correct.
@@ -63,12 +76,13 @@ pub fn run(
 
     // Generate a list of orchard subtree inputs: previous tree, current tree, and end height.
     let subtrees = upgrade_db
-        .orchard_tree_by_height_range(..=initial_tip_height)
+        .orchard_tree_by_reversed_height_range(..=initial_tip_height)
         // The first block with orchard notes can't complete a subtree, see above for details.
         .filter(|(height, _tree)| !height.is_min())
         // We need both the tree and its previous tree for each shielded block.
         .tuple_windows()
-        .map(|((_prev_height, prev_tree), (end_height, tree))| (prev_tree, end_height, tree))
+        // Because the iterator is reversed, the larger tree is first.
+        .map(|((end_height, tree), (_prev_height, prev_tree))| (prev_tree, end_height, tree))
         // Empty note commitment trees can't contain subtrees, so they have invalid subtree indexes.
         // But since we skip the empty genesis tree, all trees must have valid indexes.
         // So we don't need to unwrap the optional values for this comparison to be correct.
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 f44dac0f080..5f890fcc640 100644
--- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
@@ -180,6 +180,19 @@ impl ZebraDb {
         self.db.zs_range_iter(&sapling_trees, range)
     }
 
+    /// Returns the Sapling note commitment trees in the reversed range, in decreasing height order.
+    #[allow(clippy::unwrap_in_result)]
+    pub fn sapling_tree_by_reversed_height_range<R>(
+        &self,
+        range: R,
+    ) -> impl Iterator<Item = (Height, Arc<sapling::tree::NoteCommitmentTree>)> + '_
+    where
+        R: std::ops::RangeBounds<Height>,
+    {
+        let sapling_trees = self.db.cf_handle("sapling_note_commitment_tree").unwrap();
+        self.db.zs_reverse_range_iter(&sapling_trees, range)
+    }
+
     /// Returns the Sapling note commitment subtree at this `index`.
     ///
     /// # Correctness
@@ -313,6 +326,19 @@ impl ZebraDb {
         self.db.zs_range_iter(&orchard_trees, range)
     }
 
+    /// Returns the Orchard note commitment trees in the reversed range, in decreasing height order.
+    #[allow(clippy::unwrap_in_result)]
+    pub fn orchard_tree_by_reversed_height_range<R>(
+        &self,
+        range: R,
+    ) -> impl Iterator<Item = (Height, Arc<orchard::tree::NoteCommitmentTree>)> + '_
+    where
+        R: std::ops::RangeBounds<Height>,
+    {
+        let orchard_trees = self.db.cf_handle("orchard_note_commitment_tree").unwrap();
+        self.db.zs_reverse_range_iter(&orchard_trees, range)
+    }
+
     /// Returns the Orchard note commitment subtree at this `index`.
     ///
     /// # Correctness

From fa11314c6b3d88a5cfa16bcd8296a8a4160c8011 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 12 Sep 2023 11:47:57 +1000
Subject: [PATCH 17/56] Bump the database patch version so the upgrade runs
 again

---
 zebra-state/src/constants.rs                                   | 2 +-
 zebra-state/src/service/finalized_state/disk_format/upgrade.rs | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/zebra-state/src/constants.rs b/zebra-state/src/constants.rs
index 8abf3750d90..d87cced0522 100644
--- a/zebra-state/src/constants.rs
+++ b/zebra-state/src/constants.rs
@@ -52,7 +52,7 @@ 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 = 0;
+pub(crate) const DATABASE_FORMAT_PATCH_VERSION: u64 = 1;
 
 /// The name of the file containing the minor and patch database versions.
 ///
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 d3dab563e44..e3b03c94ef4 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
@@ -357,7 +357,7 @@ impl DbFormatChange {
         // Note commitment subtree creation database upgrade task.
 
         let version_for_adding_subtrees =
-            Version::parse("25.2.0").expect("Hardcoded version string should be valid.");
+            Version::parse("25.2.1").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 {

From 392f6994760c5024ca58682296cf136e8c8ed900 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 12 Sep 2023 12:17:44 +1000
Subject: [PATCH 18/56] Reset previous subtree upgrade data before doing this
 one

---
 .../finalized_state/disk_format/upgrade.rs    | 14 +++++--
 .../disk_format/upgrade/add_subtrees.rs       | 37 +++++++++++++++++++
 .../finalized_state/zebra_db/shielded.rs      | 34 +++++++++++++++++
 3 files changed, 82 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 e3b03c94ef4..45e296de3e3 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
@@ -356,11 +356,19 @@ impl DbFormatChange {
 
         // Note commitment subtree creation database upgrade task.
 
-        let version_for_adding_subtrees =
+        let latest_version_for_adding_subtrees =
             Version::parse("25.2.1").expect("Hardcoded version string should be valid.");
 
+        let first_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 {
+        if older_disk_version < latest_version_for_adding_subtrees {
+            if older_disk_version >= first_version_for_adding_subtrees {
+                // Clear previous upgrade data that might be different from this upgrade's data.
+                add_subtrees::reset(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.
@@ -368,7 +376,7 @@ 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.
-            Self::mark_as_upgraded_to(&version_for_adding_subtrees, &config, network);
+            Self::mark_as_upgraded_to(&latest_version_for_adding_subtrees, &config, network);
         }
 
         // # New Upgrades Usually Go Here
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 7387d02109c..fc7f6180dc2 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
@@ -101,6 +101,43 @@ pub fn run(
     Ok(())
 }
 
+/// Reset data from previous upgrades. This data can be complete or incomplete.
+///
+/// Returns `Ok` if the upgrade completed, and `Err` if it was cancelled.
+#[allow(clippy::unwrap_in_result)]
+pub fn reset(
+    _initial_tip_height: Height,
+    upgrade_db: &ZebraDb,
+    cancel_receiver: &mpsc::Receiver<CancelFormatChange>,
+) -> Result<(), CancelFormatChange> {
+    // Return early if the upgrade is cancelled.
+    if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
+        return Err(CancelFormatChange);
+    }
+
+    // This doesn't delete the maximum index, but the consensus rules make that subtree impossible.
+    // (Adding a note to a full note commitment tree is an error.)
+    //
+    // TODO: convert zs_delete_range() to take std::ops::RangeBounds, and delete the upper bound.
+    let mut batch = DiskWriteBatch::new();
+    batch.delete_range_sapling_subtree(upgrade_db, 0.into(), u16::MAX.into());
+    upgrade_db
+        .write_batch(batch)
+        .expect("deleting old sapling note commitment subtrees is a valid database operation");
+
+    if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
+        return Err(CancelFormatChange);
+    }
+
+    let mut batch = DiskWriteBatch::new();
+    batch.delete_range_orchard_subtree(upgrade_db, 0.into(), u16::MAX.into());
+    upgrade_db
+        .write_batch(batch)
+        .expect("deleting old orchard note commitment subtrees is a valid database operation");
+
+    Ok(())
+}
+
 /// Quickly check that the first calculated subtree is correct.
 ///
 /// This allows us to fail the upgrade quickly in tests and during development,
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 5f890fcc640..395ae8b6b99 100644
--- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
@@ -603,6 +603,23 @@ impl DiskWriteBatch {
         self.zs_delete_range(&sapling_tree_cf, from, to);
     }
 
+    /// Deletes the range of Sapling subtrees at the given [`NoteCommitmentSubtreeIndex`]es.
+    /// Doesn't delete the upper bound.
+    pub fn delete_range_sapling_subtree(
+        &mut self,
+        zebra_db: &ZebraDb,
+        from: NoteCommitmentSubtreeIndex,
+        to: NoteCommitmentSubtreeIndex,
+    ) {
+        let sapling_subtree_cf = zebra_db
+            .db
+            .cf_handle("sapling_note_commitment_subtree")
+            .unwrap();
+
+        // TODO: convert zs_delete_range() to take std::ops::RangeBounds
+        self.zs_delete_range(&sapling_subtree_cf, from, to);
+    }
+
     // Orchard tree methods
 
     /// Inserts the Orchard note commitment subtree.
@@ -638,4 +655,21 @@ impl DiskWriteBatch {
         // TODO: convert zs_delete_range() to take std::ops::RangeBounds
         self.zs_delete_range(&orchard_tree_cf, from, to);
     }
+
+    /// Deletes the range of Orchard subtrees at the given [`NoteCommitmentSubtreeIndex`]es.
+    /// Doesn't delete the upper bound.
+    pub fn delete_range_orchard_subtree(
+        &mut self,
+        zebra_db: &ZebraDb,
+        from: NoteCommitmentSubtreeIndex,
+        to: NoteCommitmentSubtreeIndex,
+    ) {
+        let orchard_subtree_cf = zebra_db
+            .db
+            .cf_handle("orchard_note_commitment_subtree")
+            .unwrap();
+
+        // TODO: convert zs_delete_range() to take std::ops::RangeBounds
+        self.zs_delete_range(&orchard_subtree_cf, from, to);
+    }
 }

From b3f4b802061fb79317948a958dd002e12ae3fe46 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 12 Sep 2023 12:38:24 +1000
Subject: [PATCH 19/56] Add extra checks to subtree calculation to diagnose
 errors

---
 .../disk_format/upgrade/add_subtrees.rs       | 72 +++++++++++++++----
 1 file changed, 60 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 fc7f6180dc2..83165356e56 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
@@ -563,12 +563,21 @@ fn calculate_sapling_subtree(
     } else {
         // 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 index = prev_tree
+            .subtree_index()
+            .expect("previous block must have a partial subtree");
         let remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
-
-        assert!(
-            remaining_notes > 0,
-            "tree must contain a recently completed subtree"
+        let is_complete = prev_tree.is_complete_subtree();
+
+        assert_eq!(
+            index.0 + 1,
+            tree.subtree_index()
+                .expect("current block must have a subtree")
+                .0,
+            "tree must have been completed by the current block"
         );
+        assert!(remaining_notes > 0, "just checked for a complete tree");
+        assert!(!is_complete, "just checked for a complete tree");
 
         let block = read_db
             .block(end_height.into())
@@ -580,13 +589,28 @@ fn calculate_sapling_subtree(
             .collect();
 
         // This takes less than 1 second per tree, so we don't need to make it cancellable.
-        let (_sapling_nct, subtree) = NoteCommitmentTrees::update_sapling_note_commitment_tree(
+        let (sapling_nct, subtree) = NoteCommitmentTrees::update_sapling_note_commitment_tree(
             prev_tree,
             sapling_note_commitments,
         )
         .expect("finalized notes should append successfully");
 
-        let (index, node) = subtree.expect("already checked that the block completed a subtree");
+        let (index, node) = subtree.unwrap_or_else(|| {
+            panic!(
+                "already checked that the block completed a subtree:\n\
+                 updated subtree:\n\
+                 index: {:?}\n\
+                 remaining notes: {}\n\
+                 is complete: {}\n\
+                 original subtree:\n\
+                 index: {index:?}\n\
+                 remaining notes: {remaining_notes}\n\
+                 is complete: {is_complete}\n",
+                sapling_nct.subtree_index(),
+                sapling_nct.remaining_subtree_leaf_nodes(),
+                sapling_nct.is_complete_subtree(),
+            )
+        });
 
         NoteCommitmentSubtree::new(index, end_height, node)
     }
@@ -620,12 +644,21 @@ fn calculate_orchard_subtree(
     } else {
         // 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 index = prev_tree
+            .subtree_index()
+            .expect("previous block must have a partial subtree");
         let remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
-
-        assert!(
-            remaining_notes > 0,
-            "tree must contain a recently completed subtree"
+        let is_complete = prev_tree.is_complete_subtree();
+
+        assert_eq!(
+            index.0 + 1,
+            tree.subtree_index()
+                .expect("current block must have a subtree")
+                .0,
+            "tree must have been completed by the current block"
         );
+        assert!(remaining_notes > 0, "just checked for a complete tree");
+        assert!(!is_complete, "just checked for a complete tree");
 
         let block = read_db
             .block(end_height.into())
@@ -637,13 +670,28 @@ fn calculate_orchard_subtree(
             .collect();
 
         // This takes less than 1 second per tree, so we don't need to make it cancellable.
-        let (_orchard_nct, subtree) = NoteCommitmentTrees::update_orchard_note_commitment_tree(
+        let (orchard_nct, subtree) = NoteCommitmentTrees::update_orchard_note_commitment_tree(
             prev_tree,
             orchard_note_commitments,
         )
         .expect("finalized notes should append successfully");
 
-        let (index, node) = subtree.expect("already checked that the block completed a subtree");
+        let (index, node) = subtree.unwrap_or_else(|| {
+            panic!(
+                "already checked that the block completed a subtree:\n\
+                 updated subtree:\n\
+                 index: {:?}\n\
+                 remaining notes: {}\n\
+                 is complete: {}\n\
+                 original subtree:\n\
+                 index: {index:?}\n\
+                 remaining notes: {remaining_notes}\n\
+                 is complete: {is_complete}\n",
+                orchard_nct.subtree_index(),
+                orchard_nct.remaining_subtree_leaf_nodes(),
+                orchard_nct.is_complete_subtree(),
+            )
+        });
 
         NoteCommitmentSubtree::new(index, end_height, node)
     }

From 294c9a972a582ba6364542ab4af021afc56646e7 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 12 Sep 2023 12:57:52 +1000
Subject: [PATCH 20/56] Use correct heights for subtrees completed at the end
 of a block

---
 zebra-chain/src/parallel/tree.rs | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/zebra-chain/src/parallel/tree.rs b/zebra-chain/src/parallel/tree.rs
index 4dfb5c298a2..a31855b3277 100644
--- a/zebra-chain/src/parallel/tree.rs
+++ b/zebra-chain/src/parallel/tree.rs
@@ -170,11 +170,14 @@ impl NoteCommitmentTrees {
         let mut subtree_root = None;
 
         for sapling_note_commitment in sapling_note_commitments {
+            sapling_nct.append(sapling_note_commitment)?;
+
+            // Subtrees end heights come from the blocks they are completed in,
+            // so we check for new subtrees after appending the note.
+            // (If we check before, subtrees at the end of blocks have the wrong heights.)
             if let Some(index_and_node) = sapling_nct.completed_subtree_index_and_root() {
                 subtree_root = Some(index_and_node);
             }
-
-            sapling_nct.append(sapling_note_commitment)?;
         }
 
         // Re-calculate and cache the tree root.
@@ -203,11 +206,14 @@ impl NoteCommitmentTrees {
         let mut subtree_root = None;
 
         for orchard_note_commitment in orchard_note_commitments {
+            orchard_nct.append(orchard_note_commitment)?;
+
+            // Subtrees end heights come from the blocks they are completed in,
+            // so we check for new subtrees after appending the note.
+            // (If we check before, subtrees at the end of blocks have the wrong heights.)
             if let Some(index_and_node) = orchard_nct.completed_subtree_index_and_root() {
                 subtree_root = Some(index_and_node);
             }
-
-            orchard_nct.append(orchard_note_commitment)?;
         }
 
         // Re-calculate and cache the tree root.

From f018c5de5f82f3c9611752ff605b0f9ea70d0f62 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 12 Sep 2023 13:24:13 +1000
Subject: [PATCH 21/56] Add even more checks to diagnose issues

---
 .../disk_format/upgrade/add_subtrees.rs         | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 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 83165356e56..35290298a74 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
@@ -574,10 +574,21 @@ fn calculate_sapling_subtree(
             tree.subtree_index()
                 .expect("current block must have a subtree")
                 .0,
-            "tree must have been completed by the current block"
+            "subtree must have been completed by the current block"
         );
-        assert!(remaining_notes > 0, "just checked for a complete tree");
-        assert!(!is_complete, "just checked for a complete tree");
+        assert!(
+            tree.remaining_subtree_leaf_nodes() > 0,
+            "just checked for a complete subtree in the current block"
+        );
+        assert!(
+            !tree.is_complete_subtree(),
+            "just checked for a complete tree"
+        );
+
+        // If the previous tree was complete, then the current tree can't also complete a subtree,
+        // due to the consensus limit on the number of outputs in a block.
+        assert!(remaining_notes > 0, "the caller should supply a complete subtree, so the previous tree can't also be complete");
+        assert!(!is_complete, "the caller should supply a complete subtree, so the previous tree can't also be complete");
 
         let block = read_db
             .block(end_height.into())

From 2a47cc0d1efba1cbd89fbddd3aefe77060dcaada Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 12 Sep 2023 13:33:08 +1000
Subject: [PATCH 22/56] Instrument upgrade methods to improve diagnostics

---
 .../finalized_state/disk_format/upgrade/add_subtrees.rs      | 5 +++++
 1 file changed, 5 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 35290298a74..9c5d8bdcd61 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
@@ -4,6 +4,7 @@ use std::sync::{mpsc, Arc};
 
 use hex_literal::hex;
 use itertools::Itertools;
+use tracing::instrument;
 
 use zebra_chain::{
     block::Height,
@@ -25,6 +26,7 @@ use crate::service::finalized_state::{
 ///
 /// Returns `Ok` if the upgrade completed, and `Err` if it was cancelled.
 #[allow(clippy::unwrap_in_result)]
+#[instrument(skip(upgrade_db, cancel_receiver))]
 pub fn run(
     initial_tip_height: Height,
     upgrade_db: &ZebraDb,
@@ -105,6 +107,7 @@ pub fn run(
 ///
 /// Returns `Ok` if the upgrade completed, and `Err` if it was cancelled.
 #[allow(clippy::unwrap_in_result)]
+#[instrument(skip(upgrade_db, cancel_receiver))]
 pub fn reset(
     _initial_tip_height: Height,
     upgrade_db: &ZebraDb,
@@ -547,6 +550,7 @@ fn check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
 ///
 /// If `tree` does not contain a recently completed subtree.
 #[must_use = "subtree should be written to the database after it is calculated"]
+#[instrument(skip(read_db, prev_tree, tree))]
 fn calculate_sapling_subtree(
     read_db: &ZebraDb,
     prev_tree: Arc<sapling::tree::NoteCommitmentTree>,
@@ -639,6 +643,7 @@ fn calculate_sapling_subtree(
 ///
 /// If `tree` does not contain a recently completed subtree.
 #[must_use = "subtree should be written to the database after it is calculated"]
+#[instrument(skip(read_db, prev_tree, tree))]
 fn calculate_orchard_subtree(
     read_db: &ZebraDb,
     prev_tree: Arc<orchard::tree::NoteCommitmentTree>,

From 5799da8807453b4a23fa57bbb9150e62673a2956 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 12 Sep 2023 14:17:01 +1000
Subject: [PATCH 23/56] Prevent modification of re-used trees

---
 zebra-chain/src/parallel/tree.rs              |  3 +++
 .../disk_format/upgrade/add_subtrees.rs       | 27 ++++++++++++++++---
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/zebra-chain/src/parallel/tree.rs b/zebra-chain/src/parallel/tree.rs
index a31855b3277..88eb6a3b94d 100644
--- a/zebra-chain/src/parallel/tree.rs
+++ b/zebra-chain/src/parallel/tree.rs
@@ -128,6 +128,7 @@ impl NoteCommitmentTrees {
     }
 
     /// Update the sprout note commitment tree.
+    /// This method modifies the tree inside the `Arc`, if the `Arc` only has one reference.
     fn update_sprout_note_commitment_tree(
         mut sprout: Arc<sprout::tree::NoteCommitmentTree>,
         sprout_note_commitments: Vec<sprout::NoteCommitment>,
@@ -145,6 +146,7 @@ impl NoteCommitmentTrees {
     }
 
     /// Update the sapling note commitment tree.
+    /// This method modifies the tree inside the `Arc`, if the `Arc` only has one reference.
     #[allow(clippy::unwrap_in_result)]
     pub fn update_sapling_note_commitment_tree(
         mut sapling: Arc<sapling::tree::NoteCommitmentTree>,
@@ -187,6 +189,7 @@ impl NoteCommitmentTrees {
     }
 
     /// Update the orchard note commitment tree.
+    /// This method modifies the tree inside the `Arc`, if the `Arc` only has one reference.
     #[allow(clippy::unwrap_in_result)]
     pub fn update_orchard_note_commitment_tree(
         mut orchard: Arc<orchard::tree::NoteCommitmentTree>,
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 9c5d8bdcd61..e9ba8dc4619 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
@@ -603,6 +603,11 @@ fn calculate_sapling_subtree(
             .cloned()
             .collect();
 
+        // The update method modifies the tree inside the Arc if it only has one reference.
+        // This is usually ok, but during reverse iteration, we want re-use the original previous
+        // tree as the next current tree. Cloning the Arc preserves the original tree data.
+        let _force_inner_clone_of_prev_tree = prev_tree.clone();
+
         // This takes less than 1 second per tree, so we don't need to make it cancellable.
         let (sapling_nct, subtree) = NoteCommitmentTrees::update_sapling_note_commitment_tree(
             prev_tree,
@@ -671,10 +676,21 @@ fn calculate_orchard_subtree(
             tree.subtree_index()
                 .expect("current block must have a subtree")
                 .0,
-            "tree must have been completed by the current block"
+            "subtree must have been completed by the current block"
+        );
+        assert!(
+            tree.remaining_subtree_leaf_nodes() > 0,
+            "just checked for a complete subtree in the current block"
+        );
+        assert!(
+            !tree.is_complete_subtree(),
+            "just checked for a complete tree"
         );
-        assert!(remaining_notes > 0, "just checked for a complete tree");
-        assert!(!is_complete, "just checked for a complete tree");
+
+        // If the previous tree was complete, then the current tree can't also complete a subtree,
+        // due to the consensus limit on the number of outputs in a block.
+        assert!(remaining_notes > 0, "the caller should supply a complete subtree, so the previous tree can't also be complete");
+        assert!(!is_complete, "the caller should supply a complete subtree, so the previous tree can't also be complete");
 
         let block = read_db
             .block(end_height.into())
@@ -685,6 +701,11 @@ fn calculate_orchard_subtree(
             .cloned()
             .collect();
 
+        // The update method modifies the tree inside the Arc if it only has one reference.
+        // This is usually ok, but during reverse iteration, we want re-use the original previous
+        // tree as the next current tree. Cloning the Arc preserves the original tree data.
+        let _force_inner_clone_of_prev_tree = prev_tree.clone();
+
         // This takes less than 1 second per tree, so we don't need to make it cancellable.
         let (orchard_nct, subtree) = NoteCommitmentTrees::update_orchard_note_commitment_tree(
             prev_tree,

From fd72b1fbdfbf8721538d157bd3bd4ff2b457b20a Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 12 Sep 2023 14:43:23 +1000
Subject: [PATCH 24/56] Debug with subtree positions as well

---
 zebra-chain/src/orchard/tree.rs               |  12 ++
 zebra-chain/src/sapling/tree.rs               |  12 ++
 zebra-chain/src/subtree.rs                    |   8 +-
 .../disk_format/upgrade/add_subtrees.rs       | 202 ++++++++++++++----
 4 files changed, 193 insertions(+), 41 deletions(-)

diff --git a/zebra-chain/src/orchard/tree.rs b/zebra-chain/src/orchard/tree.rs
index a3df5d4cf87..5b2470c74a3 100644
--- a/zebra-chain/src/orchard/tree.rs
+++ b/zebra-chain/src/orchard/tree.rs
@@ -394,6 +394,18 @@ impl NoteCommitmentTree {
         self.inner.value()
     }
 
+    /// Returns the position of the most recently appended leaf in the tree.
+    ///
+    /// This method is used for debugging, use `incrementalmerkletree::Address` for tree operations.
+    pub fn position(&self) -> Option<u64> {
+        let Some(tree) = self.frontier() else {
+            // An empty tree doesn't have a previous leaf.
+            return None;
+        };
+
+        Some(tree.position().into())
+    }
+
     /// Returns true if the most recently appended leaf completes the subtree
     pub fn is_complete_subtree(&self) -> bool {
         let Some(tree) = self.frontier() else {
diff --git a/zebra-chain/src/sapling/tree.rs b/zebra-chain/src/sapling/tree.rs
index 87742c081eb..060bed5bc68 100644
--- a/zebra-chain/src/sapling/tree.rs
+++ b/zebra-chain/src/sapling/tree.rs
@@ -375,6 +375,18 @@ impl NoteCommitmentTree {
         self.inner.value()
     }
 
+    /// Returns the position of the most recently appended leaf in the tree.
+    ///
+    /// This method is used for debugging, use `incrementalmerkletree::Address` for tree operations.
+    pub fn position(&self) -> Option<u64> {
+        let Some(tree) = self.frontier() else {
+            // An empty tree doesn't have a previous leaf.
+            return None;
+        };
+
+        Some(tree.position().into())
+    }
+
     /// Returns true if the most recently appended leaf completes the subtree
     pub fn is_complete_subtree(&self) -> bool {
         let Some(tree) = self.frontier() else {
diff --git a/zebra-chain/src/subtree.rs b/zebra-chain/src/subtree.rs
index 4a9ef87688b..52e3ef3d7d0 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::num::TryFromIntError;
+use std::{fmt, num::TryFromIntError};
 
 use serde::{Deserialize, Serialize};
 
@@ -19,6 +19,12 @@ pub const TRACKED_SUBTREE_HEIGHT: u8 = 16;
 #[serde(transparent)]
 pub struct NoteCommitmentSubtreeIndex(pub u16);
 
+impl fmt::Display for NoteCommitmentSubtreeIndex {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(&self.0.to_string())
+    }
+}
+
 impl From<u16> for NoteCommitmentSubtreeIndex {
     fn from(value: u16) -> Self {
         Self(value)
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 e9ba8dc4619..71e359a54f8 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
@@ -567,39 +567,96 @@ fn calculate_sapling_subtree(
     } else {
         // 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 index = prev_tree
+        let prev_position = prev_tree
+            .position()
+            .expect("previous block must have a partial subtree");
+
+        let prev_index = prev_tree
             .subtree_index()
             .expect("previous block must have a partial subtree");
-        let remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
-        let is_complete = prev_tree.is_complete_subtree();
+        let prev_remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
+        let prev_is_complete = prev_tree.is_complete_subtree();
+
+        let current_position = tree.position().expect("current block must have a subtree");
+
+        let current_index = tree
+            .subtree_index()
+            .expect("current block must have a subtree");
+        let current_remaining_notes = tree.remaining_subtree_leaf_nodes();
 
         assert_eq!(
-            index.0 + 1,
-            tree.subtree_index()
-                .expect("current block must have a subtree")
-                .0,
-            "subtree must have been completed by the current block"
+            prev_index.0 + 1,
+            current_index.0,
+            "subtree must have been completed by the current block:\n\
+             previous subtree:\n\
+             index: {prev_index}\n\
+             position: {prev_position}\n\
+             remaining: {prev_remaining_notes}\n\
+             current subtree:\n\
+             index: {current_index}\n\
+             position: {current_position}\n\
+             remaining: {current_remaining_notes}"
         );
         assert!(
-            tree.remaining_subtree_leaf_nodes() > 0,
-            "just checked for a complete subtree in the current block"
+            current_remaining_notes > 0,
+            "just checked for a complete subtree in the current block:\n\
+             previous subtree:\n\
+             index: {prev_index}\n\
+             position: {prev_position}\n\
+             remaining: {prev_remaining_notes}\n\
+             current subtree:\n\
+             index: {current_index}\n\
+             position: {current_position}\n\
+             remaining: {current_remaining_notes}"
         );
         assert!(
             !tree.is_complete_subtree(),
-            "just checked for a complete tree"
+            "just checked for a complete subtree in the current block:\n\
+             previous subtree:\n\
+             index: {prev_index}\n\
+             position: {prev_position}\n\
+             remaining: {prev_remaining_notes}\n\
+             current subtree:\n\
+             index: {current_index}\n\
+             position: {current_position}\n\
+             remaining: {current_remaining_notes}"
         );
 
         // If the previous tree was complete, then the current tree can't also complete a subtree,
         // due to the consensus limit on the number of outputs in a block.
-        assert!(remaining_notes > 0, "the caller should supply a complete subtree, so the previous tree can't also be complete");
-        assert!(!is_complete, "the caller should supply a complete subtree, so the previous tree can't also be complete");
+        assert!(
+            prev_remaining_notes > 0,
+            "the caller should supply a complete subtree, \
+             so the previous tree can't also be complete:\n\
+             previous subtree:\n\
+             index: {prev_index}\n\
+             position: {prev_position}\n\
+             remaining: {prev_remaining_notes}\n\
+             current subtree:\n\
+             index: {current_index}\n\
+             position: {current_position}\n\
+             remaining: {current_remaining_notes}"
+        );
+        assert!(
+            !prev_is_complete,
+            "the caller should supply a complete subtree, \
+             so the previous tree can't also be complete:\n\
+             previous subtree:\n\
+             index: {prev_index}\n\
+             position: {prev_position}\n\
+             remaining: {prev_remaining_notes}\n\
+             current subtree:\n\
+             index: {current_index}\n\
+             position: {current_position}\n\
+             remaining: {current_remaining_notes}"
+        );
 
         let block = read_db
             .block(end_height.into())
             .expect("height with note commitment tree should have block");
         let sapling_note_commitments = block
             .sapling_note_commitments()
-            .take(remaining_notes)
+            .take(prev_remaining_notes)
             .cloned()
             .collect();
 
@@ -620,15 +677,19 @@ fn calculate_sapling_subtree(
                 "already checked that the block completed a subtree:\n\
                  updated subtree:\n\
                  index: {:?}\n\
+                 position: {:?}\n\
                  remaining notes: {}\n\
-                 is complete: {}\n\
-                 original subtree:\n\
-                 index: {index:?}\n\
-                 remaining notes: {remaining_notes}\n\
-                 is complete: {is_complete}\n",
+                 original previous subtree:\n\
+                 index: {prev_index}\n\
+                 position: {prev_position}\n\
+                 remaining: {prev_remaining_notes}\n\
+                 original current subtree:\n\
+                 index: {current_index}\n\
+                 position: {current_position}\n\
+                 remaining: {current_remaining_notes}",
                 sapling_nct.subtree_index(),
+                sapling_nct.position(),
                 sapling_nct.remaining_subtree_leaf_nodes(),
-                sapling_nct.is_complete_subtree(),
             )
         });
 
@@ -665,39 +726,96 @@ fn calculate_orchard_subtree(
     } else {
         // 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 index = prev_tree
+        let prev_position = prev_tree
+            .position()
+            .expect("previous block must have a partial subtree");
+
+        let prev_index = prev_tree
             .subtree_index()
             .expect("previous block must have a partial subtree");
-        let remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
-        let is_complete = prev_tree.is_complete_subtree();
+        let prev_remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
+        let prev_is_complete = prev_tree.is_complete_subtree();
+
+        let current_position = tree.position().expect("current block must have a subtree");
+
+        let current_index = tree
+            .subtree_index()
+            .expect("current block must have a subtree");
+        let current_remaining_notes = tree.remaining_subtree_leaf_nodes();
 
         assert_eq!(
-            index.0 + 1,
-            tree.subtree_index()
-                .expect("current block must have a subtree")
-                .0,
-            "subtree must have been completed by the current block"
+            prev_index.0 + 1,
+            current_index.0,
+            "subtree must have been completed by the current block:\n\
+             previous subtree:\n\
+             index: {prev_index}\n\
+             position: {prev_position}\n\
+             remaining: {prev_remaining_notes}\n\
+             current subtree:\n\
+             index: {current_index}\n\
+             position: {current_position}\n\
+             remaining: {current_remaining_notes}"
         );
         assert!(
-            tree.remaining_subtree_leaf_nodes() > 0,
-            "just checked for a complete subtree in the current block"
+            current_remaining_notes > 0,
+            "just checked for a complete subtree in the current block:\n\
+             previous subtree:\n\
+             index: {prev_index}\n\
+             position: {prev_position}\n\
+             remaining: {prev_remaining_notes}\n\
+             current subtree:\n\
+             index: {current_index}\n\
+             position: {current_position}\n\
+             remaining: {current_remaining_notes}"
         );
         assert!(
             !tree.is_complete_subtree(),
-            "just checked for a complete tree"
+            "just checked for a complete subtree in the current block:\n\
+             previous subtree:\n\
+             index: {prev_index}\n\
+             position: {prev_position}\n\
+             remaining: {prev_remaining_notes}\n\
+             current subtree:\n\
+             index: {current_index}\n\
+             position: {current_position}\n\
+             remaining: {current_remaining_notes}"
         );
 
         // If the previous tree was complete, then the current tree can't also complete a subtree,
         // due to the consensus limit on the number of outputs in a block.
-        assert!(remaining_notes > 0, "the caller should supply a complete subtree, so the previous tree can't also be complete");
-        assert!(!is_complete, "the caller should supply a complete subtree, so the previous tree can't also be complete");
+        assert!(
+            prev_remaining_notes > 0,
+            "the caller should supply a complete subtree, \
+             so the previous tree can't also be complete:\n\
+             previous subtree:\n\
+             index: {prev_index}\n\
+             position: {prev_position}\n\
+             remaining: {prev_remaining_notes}\n\
+             current subtree:\n\
+             index: {current_index}\n\
+             position: {current_position}\n\
+             remaining: {current_remaining_notes}"
+        );
+        assert!(
+            !prev_is_complete,
+            "the caller should supply a complete subtree, \
+             so the previous tree can't also be complete:\n\
+             previous subtree:\n\
+             index: {prev_index}\n\
+             position: {prev_position}\n\
+             remaining: {prev_remaining_notes}\n\
+             current subtree:\n\
+             index: {current_index}\n\
+             position: {current_position}\n\
+             remaining: {current_remaining_notes}"
+        );
 
         let block = read_db
             .block(end_height.into())
             .expect("height with note commitment tree should have block");
         let orchard_note_commitments = block
             .orchard_note_commitments()
-            .take(remaining_notes)
+            .take(prev_remaining_notes)
             .cloned()
             .collect();
 
@@ -718,15 +836,19 @@ fn calculate_orchard_subtree(
                 "already checked that the block completed a subtree:\n\
                  updated subtree:\n\
                  index: {:?}\n\
+                 position: {:?}\n\
                  remaining notes: {}\n\
-                 is complete: {}\n\
-                 original subtree:\n\
-                 index: {index:?}\n\
-                 remaining notes: {remaining_notes}\n\
-                 is complete: {is_complete}\n",
+                 original previous subtree:\n\
+                 index: {prev_index}\n\
+                 position: {prev_position}\n\
+                 remaining: {prev_remaining_notes}\n\
+                 original current subtree:\n\
+                 index: {current_index}\n\
+                 position: {current_position}\n\
+                 remaining: {current_remaining_notes}",
                 orchard_nct.subtree_index(),
+                orchard_nct.position(),
                 orchard_nct.remaining_subtree_leaf_nodes(),
-                orchard_nct.is_complete_subtree(),
             )
         });
 

From 5e6bcd7a8511e828e3a2ef4d452056e1f00c466e Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 12 Sep 2023 15:05:50 +1000
Subject: [PATCH 25/56] Fix an off-by-one error with completed subtrees

---
 .../disk_format/upgrade/add_subtrees.rs       | 114 +++++++++++-------
 1 file changed, 68 insertions(+), 46 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 71e359a54f8..9129e534eaa 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
@@ -52,7 +52,7 @@ pub fn run(
     //
     // Otherwise, wallets that sync during the upgrade will be missing some notes.
 
-    // Generate a list of sapling subtree inputs: previous tree, current tree, and end height.
+    // Generate a list of sapling subtree inputs: previous and current trees, and their end heights.
     let subtrees = upgrade_db
         .sapling_tree_by_reversed_height_range(..=initial_tip_height)
         // The first block with sapling notes can't complete a subtree, see above for details.
@@ -60,23 +60,28 @@ pub fn run(
         // We need both the tree and its previous tree for each shielded block.
         .tuple_windows()
         // Because the iterator is reversed, the larger tree is first.
-        .map(|((end_height, tree), (_prev_height, prev_tree))| (prev_tree, end_height, tree))
+        .map(|((end_height, tree), (prev_end_height, prev_tree))| {
+            (prev_end_height, prev_tree, end_height, tree)
+        })
         // Empty note commitment trees can't contain subtrees, so they have invalid subtree indexes.
         // But since we skip the empty genesis tree, all trees must have valid indexes.
         // So we don't need to unwrap the optional values for this comparison to be correct.
-        .filter(|(prev_tree, _end_height, tree)| tree.subtree_index() > prev_tree.subtree_index());
+        .filter(|(_prev_end_height, prev_tree, _end_height, tree)| {
+            tree.subtree_index() > prev_tree.subtree_index()
+        });
 
-    for (prev_tree, end_height, tree) in subtrees {
+    for (prev_end_height, prev_tree, end_height, tree) in subtrees {
         // Return early if the upgrade is cancelled.
         if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
             return Err(CancelFormatChange);
         }
 
-        let subtree = calculate_sapling_subtree(upgrade_db, prev_tree, end_height, tree);
+        let subtree =
+            calculate_sapling_subtree(upgrade_db, prev_end_height, prev_tree, end_height, tree);
         write_sapling_subtree(upgrade_db, subtree);
     }
 
-    // Generate a list of orchard subtree inputs: previous tree, current tree, and end height.
+    // Generate a list of orchard subtree inputs: previous and current trees, and their end heights.
     let subtrees = upgrade_db
         .orchard_tree_by_reversed_height_range(..=initial_tip_height)
         // The first block with orchard notes can't complete a subtree, see above for details.
@@ -84,19 +89,24 @@ pub fn run(
         // We need both the tree and its previous tree for each shielded block.
         .tuple_windows()
         // Because the iterator is reversed, the larger tree is first.
-        .map(|((end_height, tree), (_prev_height, prev_tree))| (prev_tree, end_height, tree))
+        .map(|((end_height, tree), (prev_end_height, prev_tree))| {
+            (prev_end_height, prev_tree, end_height, tree)
+        })
         // Empty note commitment trees can't contain subtrees, so they have invalid subtree indexes.
         // But since we skip the empty genesis tree, all trees must have valid indexes.
         // So we don't need to unwrap the optional values for this comparison to be correct.
-        .filter(|(prev_tree, _end_height, tree)| tree.subtree_index() > prev_tree.subtree_index());
+        .filter(|(_prev_end_height, prev_tree, _end_height, tree)| {
+            tree.subtree_index() > prev_tree.subtree_index()
+        });
 
-    for (prev_tree, end_height, tree) in subtrees {
+    for (prev_end_height, prev_tree, end_height, tree) in subtrees {
         // Return early if the upgrade is cancelled.
         if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
             return Err(CancelFormatChange);
         }
 
-        let subtree = calculate_orchard_subtree(upgrade_db, prev_tree, end_height, tree);
+        let subtree =
+            calculate_orchard_subtree(upgrade_db, prev_end_height, prev_tree, end_height, tree);
         write_orchard_subtree(upgrade_db, subtree);
     }
 
@@ -219,13 +229,17 @@ fn quick_check_sapling_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
         .filter(|(height, _tree)| !height.is_min())
         // We need both the tree and its previous tree for each shielded block.
         .tuple_windows()
-        .map(|((_prev_height, prev_tree), (end_height, tree))| (prev_tree, end_height, tree))
+        .map(|((prev_end_height, prev_tree), (end_height, tree))| {
+            (prev_end_height, prev_tree, end_height, tree)
+        })
         // Empty note commitment trees can't contain subtrees, so they have invalid subtree indexes.
         // But since we skip the empty genesis tree, all trees must have valid indexes.
         // So we don't need to unwrap the optional values for this comparison to be correct.
-        .find(|(prev_tree, _end_height, tree)| tree.subtree_index() > prev_tree.subtree_index());
+        .find(|(_prev_end_height, prev_tree, _end_height, tree)| {
+            tree.subtree_index() > prev_tree.subtree_index()
+        });
 
-    let Some((prev_tree, end_height, tree)) = first_complete_subtree else {
+    let Some((prev_end_height, prev_tree, end_height, tree)) = first_complete_subtree else {
         let result = Err("iterator did not find complete subtree, but the tree has it");
         error!(?result);
         return result;
@@ -234,7 +248,7 @@ fn quick_check_sapling_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
     // Creating this test vector involves a cryptographic check, so only do it once.
     let expected_subtree = first_sapling_mainnet_subtree();
 
-    let db_subtree = calculate_sapling_subtree(db, prev_tree, end_height, tree);
+    let db_subtree = calculate_sapling_subtree(db, prev_end_height, prev_tree, end_height, tree);
 
     if db_subtree != expected_subtree {
         let result = Err("first subtree did not match expected test vector");
@@ -270,13 +284,17 @@ fn quick_check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
         .filter(|(height, _tree)| !height.is_min())
         // We need both the tree and its previous tree for each shielded block.
         .tuple_windows()
-        .map(|((_prev_height, prev_tree), (end_height, tree))| (prev_tree, end_height, tree))
+        .map(|((prev_end_height, prev_tree), (end_height, tree))| {
+            (prev_end_height, prev_tree, end_height, tree)
+        })
         // Empty note commitment trees can't contain subtrees, so they have invalid subtree indexes.
         // But since we skip the empty genesis tree, all trees must have valid indexes.
         // So we don't need to unwrap the optional values for this comparison to be correct.
-        .find(|(prev_tree, _end_height, tree)| tree.subtree_index() > prev_tree.subtree_index());
+        .find(|(_prev_end_height, prev_tree, _end_height, tree)| {
+            tree.subtree_index() > prev_tree.subtree_index()
+        });
 
-    let Some((prev_tree, end_height, tree)) = first_complete_subtree else {
+    let Some((prev_end_height, prev_tree, end_height, tree)) = first_complete_subtree else {
         let result = Err("iterator did not find complete subtree, but the tree has it");
         error!(?result);
         return result;
@@ -285,7 +303,7 @@ fn quick_check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
     // Creating this test vector involves a cryptographic check, so only do it once.
     let expected_subtree = first_orchard_mainnet_subtree();
 
-    let db_subtree = calculate_orchard_subtree(db, prev_tree, end_height, tree);
+    let db_subtree = calculate_orchard_subtree(db, prev_end_height, prev_tree, end_height, tree);
 
     if db_subtree != expected_subtree {
         let result = Err("first subtree did not match expected test vector");
@@ -540,37 +558,38 @@ fn check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
 
 /// Calculates a Sapling note commitment subtree, reading blocks from `read_db` if needed.
 ///
-/// `tree` must be a note commitment tree containing a recently completed subtree,
-/// which was not already completed in `prev_tree`.
+/// `tree` must be a note commitment tree that starts a new subtree.
+/// The subtree can be completed at the end of `prev_tree`, or before the end of `tree`.
 ///
-/// `prev_tree` is only used to rebuild the subtree, if it was completed inside the block.
-/// If the subtree was completed by the final note commitment in the block, `prev_tree` is unused.
+/// `prev_tree` is only used to rebuild the subtree if it was completed before the end of `tree`.
 ///
 /// # Panics
 ///
-/// If `tree` does not contain a recently completed subtree.
+/// If `tree` does not start a new subtree.
 #[must_use = "subtree should be written to the database after it is calculated"]
 #[instrument(skip(read_db, prev_tree, tree))]
 fn calculate_sapling_subtree(
     read_db: &ZebraDb,
+    prev_end_height: Height,
     prev_tree: Arc<sapling::tree::NoteCommitmentTree>,
     end_height: Height,
     tree: Arc<sapling::tree::NoteCommitmentTree>,
 ) -> NoteCommitmentSubtree<sapling::tree::Node> {
-    // If this block completed 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,
+    // If this block starts a new subtree, that subtree is either completed by:
+    // - a note in this block (so the final note is before the end of this block), or
+    // - the final note in the previous block.
+    if let Some((index, node)) = prev_tree.completed_subtree_index_and_root() {
+        // If the leaf at the end of the previous block is the final leaf in a subtree,
         // we already have that subtree root available in the tree.
-        NoteCommitmentSubtree::new(index, end_height, node)
+        NoteCommitmentSubtree::new(index, prev_end_height, node)
     } else {
-        // 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.
+        // If the leaf at the end of this block is in the next subtree,
+        // we need to calculate the subtree root based on the tree from the previous block.
+
+        // TODO: move these checks into a separate function
         let prev_position = prev_tree
             .position()
             .expect("previous block must have a partial subtree");
-
         let prev_index = prev_tree
             .subtree_index()
             .expect("previous block must have a partial subtree");
@@ -651,6 +670,7 @@ fn calculate_sapling_subtree(
              remaining: {current_remaining_notes}"
         );
 
+        // Get the missing notes needed to complete the subtree.
         let block = read_db
             .block(end_height.into())
             .expect("height with note commitment tree should have block");
@@ -699,37 +719,38 @@ fn calculate_sapling_subtree(
 
 /// Calculates a Orchard note commitment subtree, reading blocks from `read_db` if needed.
 ///
-/// `tree` must be a note commitment tree containing a recently completed subtree,
-/// which was not already completed in `prev_tree`.
+/// `tree` must be a note commitment tree that starts a new subtree.
+/// The subtree can be completed at the end of `prev_tree`, or before the end of `tree`.
 ///
-/// `prev_tree` is only used to rebuild the subtree, if it was completed inside the block.
-/// If the subtree was completed by the final note commitment in the block, `prev_tree` is unused.
+/// `prev_tree` is only used to rebuild the subtree if it was completed before the end of `tree`.
 ///
 /// # Panics
 ///
-/// If `tree` does not contain a recently completed subtree.
+/// If `tree` does not start a new subtree.
 #[must_use = "subtree should be written to the database after it is calculated"]
 #[instrument(skip(read_db, prev_tree, tree))]
 fn calculate_orchard_subtree(
     read_db: &ZebraDb,
+    prev_end_height: Height,
     prev_tree: Arc<orchard::tree::NoteCommitmentTree>,
     end_height: Height,
     tree: Arc<orchard::tree::NoteCommitmentTree>,
 ) -> NoteCommitmentSubtree<orchard::tree::Node> {
-    // If this block completed 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,
+    // If this block starts a new subtree, that subtree is either completed by:
+    // - a note in this block (so the final note is before the end of this block), or
+    // - the final note in the previous block.
+    if let Some((index, node)) = prev_tree.completed_subtree_index_and_root() {
+        // If the leaf at the end of the previous block is the final leaf in a subtree,
         // we already have that subtree root available in the tree.
-        NoteCommitmentSubtree::new(index, end_height, node)
+        NoteCommitmentSubtree::new(index, prev_end_height, node)
     } else {
-        // 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.
+        // If the leaf at the end of this block is in the next subtree,
+        // we need to calculate the subtree root based on the tree from the previous block.
+
+        // TODO: move these checks into a separate function
         let prev_position = prev_tree
             .position()
             .expect("previous block must have a partial subtree");
-
         let prev_index = prev_tree
             .subtree_index()
             .expect("previous block must have a partial subtree");
@@ -810,6 +831,7 @@ fn calculate_orchard_subtree(
              remaining: {current_remaining_notes}"
         );
 
+        // Get the missing notes needed to complete the subtree.
         let block = read_db
             .block(end_height.into())
             .expect("height with note commitment tree should have block");

From cd1ecc5d89074a0ec81f177263d925f6cfd4e6af Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 15 Sep 2023 09:14:44 +1000
Subject: [PATCH 26/56] Fix typos and confusing comments

Co-authored-by: Marek <mail@marek.onl>
---
 zebra-state/src/service/finalized_state/disk_db.rs            | 2 +-
 .../finalized_state/disk_format/upgrade/add_subtrees.rs       | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db.rs b/zebra-state/src/service/finalized_state/disk_db.rs
index 51af47f72ea..3772fb7a789 100644
--- a/zebra-state/src/service/finalized_state/disk_db.rs
+++ b/zebra-state/src/service/finalized_state/disk_db.rs
@@ -445,7 +445,7 @@ impl DiskDb {
     ///
     /// Holding this iterator open might delay block commit transactions.
     ///
-    /// This code is coped from `zs_range_iter()`, but with the mode reversed.
+    /// This code is copied from `zs_range_iter()`, but with the mode reversed.
     pub fn zs_reverse_range_iter<C, K, V, R>(
         &self,
         cf: &C,
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 9129e534eaa..520576d2b40 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
@@ -40,7 +40,7 @@ pub fn run(
     // The consensus rules limit blocks to less than 2^16 sapling and 2^16 orchard outputs. So a
     // block can't complete multiple level 16 subtrees (or complete an entire subtree by itself).
     // Currently, with 2MB blocks and v4/v5 sapling and orchard output sizes, the subtree index can
-    // increase by 1 every ~20 full blocks.
+    // increase by at most 1 every ~20 blocks.
     //
     // Therefore, the first block with shielded note can't complete a subtree, which means we can
     // skip the (genesis block, first shielded block) tree pair.
@@ -898,7 +898,7 @@ fn write_sapling_subtree(
     debug!(end_height = ?subtree.end, index = ?subtree.index.0, "calculated and added sapling subtree");
 }
 
-/// Writes a Orchard note commitment subtree to `upgrade_db`.
+/// Writes an Orchard note commitment subtree to `upgrade_db`.
 fn write_orchard_subtree(
     upgrade_db: &ZebraDb,
     subtree: NoteCommitmentSubtree<orchard::tree::Node>,

From 540933f7541f28e88a70f33fb7e3e4f5aa91572b Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 15 Sep 2023 09:52:59 +1000
Subject: [PATCH 27/56] Fix mistaken previous tree handling and end tree
 comments

---
 .../disk_format/upgrade/add_subtrees.rs       | 76 +++++++++----------
 1 file changed, 36 insertions(+), 40 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 520576d2b40..2f83fbb283b 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
@@ -70,14 +70,13 @@ pub fn run(
             tree.subtree_index() > prev_tree.subtree_index()
         });
 
-    for (prev_end_height, prev_tree, end_height, tree) in subtrees {
+    for (_prev_end_height, prev_tree, end_height, tree) in subtrees {
         // Return early if the upgrade is cancelled.
         if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
             return Err(CancelFormatChange);
         }
 
-        let subtree =
-            calculate_sapling_subtree(upgrade_db, prev_end_height, prev_tree, end_height, tree);
+        let subtree = calculate_sapling_subtree(upgrade_db, prev_tree, end_height, tree);
         write_sapling_subtree(upgrade_db, subtree);
     }
 
@@ -99,14 +98,13 @@ pub fn run(
             tree.subtree_index() > prev_tree.subtree_index()
         });
 
-    for (prev_end_height, prev_tree, end_height, tree) in subtrees {
+    for (_prev_end_height, prev_tree, end_height, tree) in subtrees {
         // Return early if the upgrade is cancelled.
         if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
             return Err(CancelFormatChange);
         }
 
-        let subtree =
-            calculate_orchard_subtree(upgrade_db, prev_end_height, prev_tree, end_height, tree);
+        let subtree = calculate_orchard_subtree(upgrade_db, prev_tree, end_height, tree);
         write_orchard_subtree(upgrade_db, subtree);
     }
 
@@ -239,7 +237,7 @@ fn quick_check_sapling_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
             tree.subtree_index() > prev_tree.subtree_index()
         });
 
-    let Some((prev_end_height, prev_tree, end_height, tree)) = first_complete_subtree else {
+    let Some((_prev_end_height, prev_tree, end_height, tree)) = first_complete_subtree else {
         let result = Err("iterator did not find complete subtree, but the tree has it");
         error!(?result);
         return result;
@@ -248,7 +246,7 @@ fn quick_check_sapling_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
     // Creating this test vector involves a cryptographic check, so only do it once.
     let expected_subtree = first_sapling_mainnet_subtree();
 
-    let db_subtree = calculate_sapling_subtree(db, prev_end_height, prev_tree, end_height, tree);
+    let db_subtree = calculate_sapling_subtree(db, prev_tree, end_height, tree);
 
     if db_subtree != expected_subtree {
         let result = Err("first subtree did not match expected test vector");
@@ -294,7 +292,7 @@ fn quick_check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
             tree.subtree_index() > prev_tree.subtree_index()
         });
 
-    let Some((prev_end_height, prev_tree, end_height, tree)) = first_complete_subtree else {
+    let Some((_prev_end_height, prev_tree, end_height, tree)) = first_complete_subtree else {
         let result = Err("iterator did not find complete subtree, but the tree has it");
         error!(?result);
         return result;
@@ -303,7 +301,7 @@ fn quick_check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
     // Creating this test vector involves a cryptographic check, so only do it once.
     let expected_subtree = first_orchard_mainnet_subtree();
 
-    let db_subtree = calculate_orchard_subtree(db, prev_end_height, prev_tree, end_height, tree);
+    let db_subtree = calculate_orchard_subtree(db, prev_tree, end_height, tree);
 
     if db_subtree != expected_subtree {
         let result = Err("first subtree did not match expected test vector");
@@ -558,33 +556,32 @@ fn check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
 
 /// Calculates a Sapling note commitment subtree, reading blocks from `read_db` if needed.
 ///
-/// `tree` must be a note commitment tree that starts a new subtree.
-/// The subtree can be completed at the end of `prev_tree`, or before the end of `tree`.
+/// The subtree must be completed by a note commitment in the block at `end_height`.
+/// `tree` is the tree for that block, and `prev_tree` is the tree for the previous block.
 ///
-/// `prev_tree` is only used to rebuild the subtree if it was completed before the end of `tree`.
+/// `prev_tree` is only used to rebuild the subtree if it was completed without using the last
+/// note commitment in the block at `end_height`.
 ///
 /// # Panics
 ///
-/// If `tree` does not start a new subtree.
+/// If a subtree is not completed by a note commitment in the block at `end_height`.
 #[must_use = "subtree should be written to the database after it is calculated"]
 #[instrument(skip(read_db, prev_tree, tree))]
 fn calculate_sapling_subtree(
     read_db: &ZebraDb,
-    prev_end_height: Height,
     prev_tree: Arc<sapling::tree::NoteCommitmentTree>,
     end_height: Height,
     tree: Arc<sapling::tree::NoteCommitmentTree>,
 ) -> NoteCommitmentSubtree<sapling::tree::Node> {
-    // If this block starts a new subtree, that subtree is either completed by:
-    // - a note in this block (so the final note is before the end of this block), or
-    // - the final note in the previous block.
-    if let Some((index, node)) = prev_tree.completed_subtree_index_and_root() {
-        // If the leaf at the end of the previous block is the final leaf in a subtree,
+    // If a subtree is completed by a note commentment in the block at `end_height`,
+    // then that subtree can be completed in two different ways:
+    if let Some((index, node)) = tree.completed_subtree_index_and_root() {
+        // If the subtree is completed by the last note commitment in that block,
         // we already have that subtree root available in the tree.
-        NoteCommitmentSubtree::new(index, prev_end_height, node)
+        NoteCommitmentSubtree::new(index, end_height, node)
     } else {
-        // If the leaf at the end of this block is in the next subtree,
-        // we need to calculate the subtree root based on the tree from the previous block.
+        // If the subtree is completed without using the last note commitment in the block,
+        // we need to calculate the subtree root, starting with the tree from the previous block.
 
         // TODO: move these checks into a separate function
         let prev_position = prev_tree
@@ -597,11 +594,11 @@ fn calculate_sapling_subtree(
         let prev_is_complete = prev_tree.is_complete_subtree();
 
         let current_position = tree.position().expect("current block must have a subtree");
-
         let current_index = tree
             .subtree_index()
             .expect("current block must have a subtree");
         let current_remaining_notes = tree.remaining_subtree_leaf_nodes();
+        let current_is_complete = tree.is_complete_subtree();
 
         assert_eq!(
             prev_index.0 + 1,
@@ -629,7 +626,7 @@ fn calculate_sapling_subtree(
              remaining: {current_remaining_notes}"
         );
         assert!(
-            !tree.is_complete_subtree(),
+            !current_is_complete,
             "just checked for a complete subtree in the current block:\n\
              previous subtree:\n\
              index: {prev_index}\n\
@@ -719,33 +716,32 @@ fn calculate_sapling_subtree(
 
 /// Calculates a Orchard note commitment subtree, reading blocks from `read_db` if needed.
 ///
-/// `tree` must be a note commitment tree that starts a new subtree.
-/// The subtree can be completed at the end of `prev_tree`, or before the end of `tree`.
+/// The subtree must be completed by a note commitment in the block at `end_height`.
+/// `tree` is the tree for that block, and `prev_tree` is the tree for the previous block.
 ///
-/// `prev_tree` is only used to rebuild the subtree if it was completed before the end of `tree`.
+/// `prev_tree` is only used to rebuild the subtree if it was completed without using the last
+/// note commitment in the block at `end_height`.
 ///
 /// # Panics
 ///
-/// If `tree` does not start a new subtree.
+/// If a subtree is not completed by a note commitment in the block at `end_height`.
 #[must_use = "subtree should be written to the database after it is calculated"]
 #[instrument(skip(read_db, prev_tree, tree))]
 fn calculate_orchard_subtree(
     read_db: &ZebraDb,
-    prev_end_height: Height,
     prev_tree: Arc<orchard::tree::NoteCommitmentTree>,
     end_height: Height,
     tree: Arc<orchard::tree::NoteCommitmentTree>,
 ) -> NoteCommitmentSubtree<orchard::tree::Node> {
-    // If this block starts a new subtree, that subtree is either completed by:
-    // - a note in this block (so the final note is before the end of this block), or
-    // - the final note in the previous block.
-    if let Some((index, node)) = prev_tree.completed_subtree_index_and_root() {
-        // If the leaf at the end of the previous block is the final leaf in a subtree,
+    // If a subtree is completed by a note commentment in the block at `end_height`,
+    // then that subtree can be completed in two different ways:
+    if let Some((index, node)) = tree.completed_subtree_index_and_root() {
+        // If the subtree is completed by the last note commitment in that block,
         // we already have that subtree root available in the tree.
-        NoteCommitmentSubtree::new(index, prev_end_height, node)
+        NoteCommitmentSubtree::new(index, end_height, node)
     } else {
-        // If the leaf at the end of this block is in the next subtree,
-        // we need to calculate the subtree root based on the tree from the previous block.
+        // If the subtree is completed without using the last note commitment in the block,
+        // we need to calculate the subtree root, starting with the tree from the previous block.
 
         // TODO: move these checks into a separate function
         let prev_position = prev_tree
@@ -758,11 +754,11 @@ fn calculate_orchard_subtree(
         let prev_is_complete = prev_tree.is_complete_subtree();
 
         let current_position = tree.position().expect("current block must have a subtree");
-
         let current_index = tree
             .subtree_index()
             .expect("current block must have a subtree");
         let current_remaining_notes = tree.remaining_subtree_leaf_nodes();
+        let current_is_complete = tree.is_complete_subtree();
 
         assert_eq!(
             prev_index.0 + 1,
@@ -790,7 +786,7 @@ fn calculate_orchard_subtree(
              remaining: {current_remaining_notes}"
         );
         assert!(
-            !tree.is_complete_subtree(),
+            !current_is_complete,
             "just checked for a complete subtree in the current block:\n\
              previous subtree:\n\
              index: {prev_index}\n\

From c4382bd0bc98eb74e35a20422d7d8a7f046b68df Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 15 Sep 2023 09:55:48 +1000
Subject: [PATCH 28/56] Remove unnecessary subtraction in remaining leaves calc

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

diff --git a/zebra-chain/src/orchard/tree.rs b/zebra-chain/src/orchard/tree.rs
index 5b2470c74a3..5fa6407ed92 100644
--- a/zebra-chain/src/orchard/tree.rs
+++ b/zebra-chain/src/orchard/tree.rs
@@ -464,7 +464,13 @@ impl NoteCommitmentTree {
                     first_position,
                 );
 
-                subtree_address.position_range_end() - subtree_address.position_range_start().into()
+                assert_eq!(
+                    subtree_address.position_range_start(),
+                    0.into(),
+                    "address is not in the first subtree"
+                );
+
+                subtree_address.position_range_end()
             }
         };
 
diff --git a/zebra-chain/src/sapling/tree.rs b/zebra-chain/src/sapling/tree.rs
index 060bed5bc68..f2702fea2d9 100644
--- a/zebra-chain/src/sapling/tree.rs
+++ b/zebra-chain/src/sapling/tree.rs
@@ -445,7 +445,13 @@ impl NoteCommitmentTree {
                     first_position,
                 );
 
-                subtree_address.position_range_end() - subtree_address.position_range_start().into()
+                assert_eq!(
+                    subtree_address.position_range_start(),
+                    0.into(),
+                    "address is not in the first subtree"
+                );
+
+                subtree_address.position_range_end()
             }
         };
 

From 1f5f58d2d0e1449234f299851e9e0814293d9c8d Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 15 Sep 2023 10:04:49 +1000
Subject: [PATCH 29/56] Log heights when assertions fail

---
 .../disk_format/upgrade/add_subtrees.rs       | 48 +++++++++++++++----
 1 file changed, 38 insertions(+), 10 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 2f83fbb283b..accafd9cad4 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
@@ -70,13 +70,14 @@ pub fn run(
             tree.subtree_index() > prev_tree.subtree_index()
         });
 
-    for (_prev_end_height, prev_tree, end_height, tree) in subtrees {
+    for (prev_end_height, prev_tree, end_height, tree) in subtrees {
         // Return early if the upgrade is cancelled.
         if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
             return Err(CancelFormatChange);
         }
 
-        let subtree = calculate_sapling_subtree(upgrade_db, prev_tree, end_height, tree);
+        let subtree =
+            calculate_sapling_subtree(upgrade_db, prev_end_height, prev_tree, end_height, tree);
         write_sapling_subtree(upgrade_db, subtree);
     }
 
@@ -98,13 +99,14 @@ pub fn run(
             tree.subtree_index() > prev_tree.subtree_index()
         });
 
-    for (_prev_end_height, prev_tree, end_height, tree) in subtrees {
+    for (prev_end_height, prev_tree, end_height, tree) in subtrees {
         // Return early if the upgrade is cancelled.
         if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
             return Err(CancelFormatChange);
         }
 
-        let subtree = calculate_orchard_subtree(upgrade_db, prev_tree, end_height, tree);
+        let subtree =
+            calculate_orchard_subtree(upgrade_db, prev_end_height, prev_tree, end_height, tree);
         write_orchard_subtree(upgrade_db, subtree);
     }
 
@@ -237,7 +239,7 @@ fn quick_check_sapling_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
             tree.subtree_index() > prev_tree.subtree_index()
         });
 
-    let Some((_prev_end_height, prev_tree, end_height, tree)) = first_complete_subtree else {
+    let Some((prev_end_height, prev_tree, end_height, tree)) = first_complete_subtree else {
         let result = Err("iterator did not find complete subtree, but the tree has it");
         error!(?result);
         return result;
@@ -246,7 +248,7 @@ fn quick_check_sapling_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
     // Creating this test vector involves a cryptographic check, so only do it once.
     let expected_subtree = first_sapling_mainnet_subtree();
 
-    let db_subtree = calculate_sapling_subtree(db, prev_tree, end_height, tree);
+    let db_subtree = calculate_sapling_subtree(db, prev_end_height, prev_tree, end_height, tree);
 
     if db_subtree != expected_subtree {
         let result = Err("first subtree did not match expected test vector");
@@ -292,7 +294,7 @@ fn quick_check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
             tree.subtree_index() > prev_tree.subtree_index()
         });
 
-    let Some((_prev_end_height, prev_tree, end_height, tree)) = first_complete_subtree else {
+    let Some((prev_end_height, prev_tree, end_height, tree)) = first_complete_subtree else {
         let result = Err("iterator did not find complete subtree, but the tree has it");
         error!(?result);
         return result;
@@ -301,7 +303,7 @@ fn quick_check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
     // Creating this test vector involves a cryptographic check, so only do it once.
     let expected_subtree = first_orchard_mainnet_subtree();
 
-    let db_subtree = calculate_orchard_subtree(db, prev_tree, end_height, tree);
+    let db_subtree = calculate_orchard_subtree(db, prev_end_height, prev_tree, end_height, tree);
 
     if db_subtree != expected_subtree {
         let result = Err("first subtree did not match expected test vector");
@@ -554,7 +556,7 @@ fn check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
     result
 }
 
-/// Calculates a Sapling note commitment subtree, reading blocks from `read_db` if needed.
+/// Calculates a note commitment subtree for Sapling, reading blocks from `read_db` if needed.
 ///
 /// The subtree must be completed by a note commitment in the block at `end_height`.
 /// `tree` is the tree for that block, and `prev_tree` is the tree for the previous block.
@@ -569,6 +571,7 @@ fn check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
 #[instrument(skip(read_db, prev_tree, tree))]
 fn calculate_sapling_subtree(
     read_db: &ZebraDb,
+    prev_end_height: Height,
     prev_tree: Arc<sapling::tree::NoteCommitmentTree>,
     end_height: Height,
     tree: Arc<sapling::tree::NoteCommitmentTree>,
@@ -605,10 +608,12 @@ fn calculate_sapling_subtree(
             current_index.0,
             "subtree must have been completed by the current block:\n\
              previous subtree:\n\
+             height: {prev_end_height:?}\n\
              index: {prev_index}\n\
              position: {prev_position}\n\
              remaining: {prev_remaining_notes}\n\
              current subtree:\n\
+             height: {end_height:?}\n\
              index: {current_index}\n\
              position: {current_position}\n\
              remaining: {current_remaining_notes}"
@@ -617,10 +622,12 @@ fn calculate_sapling_subtree(
             current_remaining_notes > 0,
             "just checked for a complete subtree in the current block:\n\
              previous subtree:\n\
+             height: {prev_end_height:?}\n\
              index: {prev_index}\n\
              position: {prev_position}\n\
              remaining: {prev_remaining_notes}\n\
              current subtree:\n\
+             height: {end_height:?}\n\
              index: {current_index}\n\
              position: {current_position}\n\
              remaining: {current_remaining_notes}"
@@ -629,10 +636,12 @@ fn calculate_sapling_subtree(
             !current_is_complete,
             "just checked for a complete subtree in the current block:\n\
              previous subtree:\n\
+             height: {prev_end_height:?}\n\
              index: {prev_index}\n\
              position: {prev_position}\n\
              remaining: {prev_remaining_notes}\n\
              current subtree:\n\
+             height: {end_height:?}\n\
              index: {current_index}\n\
              position: {current_position}\n\
              remaining: {current_remaining_notes}"
@@ -645,10 +654,12 @@ fn calculate_sapling_subtree(
             "the caller should supply a complete subtree, \
              so the previous tree can't also be complete:\n\
              previous subtree:\n\
+             height: {prev_end_height:?}\n\
              index: {prev_index}\n\
              position: {prev_position}\n\
              remaining: {prev_remaining_notes}\n\
              current subtree:\n\
+             height: {end_height:?}\n\
              index: {current_index}\n\
              position: {current_position}\n\
              remaining: {current_remaining_notes}"
@@ -658,10 +669,12 @@ fn calculate_sapling_subtree(
             "the caller should supply a complete subtree, \
              so the previous tree can't also be complete:\n\
              previous subtree:\n\
+             height: {prev_end_height:?}\n\
              index: {prev_index}\n\
              position: {prev_position}\n\
              remaining: {prev_remaining_notes}\n\
              current subtree:\n\
+             height: {end_height:?}\n\
              index: {current_index}\n\
              position: {current_position}\n\
              remaining: {current_remaining_notes}"
@@ -697,10 +710,12 @@ fn calculate_sapling_subtree(
                  position: {:?}\n\
                  remaining notes: {}\n\
                  original previous subtree:\n\
+                 height: {prev_end_height:?}\n\
                  index: {prev_index}\n\
                  position: {prev_position}\n\
                  remaining: {prev_remaining_notes}\n\
                  original current subtree:\n\
+                 height: {end_height:?}\n\
                  index: {current_index}\n\
                  position: {current_position}\n\
                  remaining: {current_remaining_notes}",
@@ -714,7 +729,7 @@ fn calculate_sapling_subtree(
     }
 }
 
-/// Calculates a Orchard note commitment subtree, reading blocks from `read_db` if needed.
+/// Calculates a note commitment subtree for Orchard, reading blocks from `read_db` if needed.
 ///
 /// The subtree must be completed by a note commitment in the block at `end_height`.
 /// `tree` is the tree for that block, and `prev_tree` is the tree for the previous block.
@@ -729,6 +744,7 @@ fn calculate_sapling_subtree(
 #[instrument(skip(read_db, prev_tree, tree))]
 fn calculate_orchard_subtree(
     read_db: &ZebraDb,
+    prev_end_height: Height,
     prev_tree: Arc<orchard::tree::NoteCommitmentTree>,
     end_height: Height,
     tree: Arc<orchard::tree::NoteCommitmentTree>,
@@ -765,10 +781,12 @@ fn calculate_orchard_subtree(
             current_index.0,
             "subtree must have been completed by the current block:\n\
              previous subtree:\n\
+             height: {prev_end_height:?}\n\
              index: {prev_index}\n\
              position: {prev_position}\n\
              remaining: {prev_remaining_notes}\n\
              current subtree:\n\
+             height: {end_height:?}\n\
              index: {current_index}\n\
              position: {current_position}\n\
              remaining: {current_remaining_notes}"
@@ -777,10 +795,12 @@ fn calculate_orchard_subtree(
             current_remaining_notes > 0,
             "just checked for a complete subtree in the current block:\n\
              previous subtree:\n\
+             height: {prev_end_height:?}\n\
              index: {prev_index}\n\
              position: {prev_position}\n\
              remaining: {prev_remaining_notes}\n\
              current subtree:\n\
+             height: {end_height:?}\n\
              index: {current_index}\n\
              position: {current_position}\n\
              remaining: {current_remaining_notes}"
@@ -789,10 +809,12 @@ fn calculate_orchard_subtree(
             !current_is_complete,
             "just checked for a complete subtree in the current block:\n\
              previous subtree:\n\
+             height: {prev_end_height:?}\n\
              index: {prev_index}\n\
              position: {prev_position}\n\
              remaining: {prev_remaining_notes}\n\
              current subtree:\n\
+             height: {end_height:?}\n\
              index: {current_index}\n\
              position: {current_position}\n\
              remaining: {current_remaining_notes}"
@@ -805,10 +827,12 @@ fn calculate_orchard_subtree(
             "the caller should supply a complete subtree, \
              so the previous tree can't also be complete:\n\
              previous subtree:\n\
+             height: {prev_end_height:?}\n\
              index: {prev_index}\n\
              position: {prev_position}\n\
              remaining: {prev_remaining_notes}\n\
              current subtree:\n\
+             height: {end_height:?}\n\
              index: {current_index}\n\
              position: {current_position}\n\
              remaining: {current_remaining_notes}"
@@ -818,10 +842,12 @@ fn calculate_orchard_subtree(
             "the caller should supply a complete subtree, \
              so the previous tree can't also be complete:\n\
              previous subtree:\n\
+             height: {prev_end_height:?}\n\
              index: {prev_index}\n\
              position: {prev_position}\n\
              remaining: {prev_remaining_notes}\n\
              current subtree:\n\
+             height: {end_height:?}\n\
              index: {current_index}\n\
              position: {current_position}\n\
              remaining: {current_remaining_notes}"
@@ -857,10 +883,12 @@ fn calculate_orchard_subtree(
                  position: {:?}\n\
                  remaining notes: {}\n\
                  original previous subtree:\n\
+                 height: {prev_end_height:?}\n\
                  index: {prev_index}\n\
                  position: {prev_position}\n\
                  remaining: {prev_remaining_notes}\n\
                  original current subtree:\n\
+                 height: {end_height:?}\n\
                  index: {current_index}\n\
                  position: {current_position}\n\
                  remaining: {current_remaining_notes}",

From bc259c38a637cbbfd13c7954f9e4abb0c72bb81a Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 15 Sep 2023 10:25:23 +1000
Subject: [PATCH 30/56] Fix new subtree detection filter

---
 .../disk_format/upgrade/add_subtrees.rs       | 30 ++++++++++++++-----
 1 file changed, 22 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 accafd9cad4..e4d1af06b87 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
@@ -63,11 +63,18 @@ pub fn run(
         .map(|((end_height, tree), (prev_end_height, prev_tree))| {
             (prev_end_height, prev_tree, end_height, tree)
         })
-        // Empty note commitment trees can't contain subtrees, so they have invalid subtree indexes.
-        // But since we skip the empty genesis tree, all trees must have valid indexes.
-        // So we don't need to unwrap the optional values for this comparison to be correct.
+        // Find new subtrees.
         .filter(|(_prev_end_height, prev_tree, _end_height, tree)| {
-            tree.subtree_index() > prev_tree.subtree_index()
+            // The subtree is completed by the last note commitment in this block.
+            tree.is_complete_subtree() ||
+            // The subtree is completed by a note commitment before the last one in this block.
+            // We need to exclude subtrees completed at the end of the previous block
+            // because they also increase the subtree index.
+            //
+            // Empty note commitment trees can't contain subtrees, so they have invalid subtree
+            // indexes. But we have already skipped the empty genesis tree, so all trees have valid
+            // indexes. So this comparison is correct even if we don't unwrap the indexes.
+            (tree.subtree_index() > prev_tree.subtree_index() && !prev_tree.is_complete_subtree())
         });
 
     for (prev_end_height, prev_tree, end_height, tree) in subtrees {
@@ -92,11 +99,18 @@ pub fn run(
         .map(|((end_height, tree), (prev_end_height, prev_tree))| {
             (prev_end_height, prev_tree, end_height, tree)
         })
-        // Empty note commitment trees can't contain subtrees, so they have invalid subtree indexes.
-        // But since we skip the empty genesis tree, all trees must have valid indexes.
-        // So we don't need to unwrap the optional values for this comparison to be correct.
+        // Find new subtrees.
         .filter(|(_prev_end_height, prev_tree, _end_height, tree)| {
-            tree.subtree_index() > prev_tree.subtree_index()
+            // The subtree is completed by the last note commitment in this block.
+            tree.is_complete_subtree() ||
+            // The subtree is completed by a note commitment before the last one in this block.
+            // We need to exclude subtrees completed at the end of the previous block
+            // because they also increase the subtree index.
+            //
+            // Empty note commitment trees can't contain subtrees, so they have invalid subtree
+            // indexes. But we have already skipped the empty genesis tree, so all trees have valid
+            // indexes. So this comparison is correct even if we don't unwrap the indexes.
+            (tree.subtree_index() > prev_tree.subtree_index() && !prev_tree.is_complete_subtree())
         });
 
     for (prev_end_height, prev_tree, end_height, tree) in subtrees {

From a86891d360669ff8d2af9ec916f94856a9afbccd Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 15 Sep 2023 10:49:12 +1000
Subject: [PATCH 31/56] Move new subtree check into a method, cleanup unused
 code

---
 zebra-chain/src/orchard/tree.rs               | 41 +++++++++++++++++++
 zebra-chain/src/sapling/tree.rs               | 41 +++++++++++++++++++
 .../disk_format/upgrade/add_subtrees.rs       | 38 ++++-------------
 3 files changed, 90 insertions(+), 30 deletions(-)

diff --git a/zebra-chain/src/orchard/tree.rs b/zebra-chain/src/orchard/tree.rs
index 5fa6407ed92..d0c2146915f 100644
--- a/zebra-chain/src/orchard/tree.rs
+++ b/zebra-chain/src/orchard/tree.rs
@@ -406,6 +406,47 @@ impl NoteCommitmentTree {
         Some(tree.position().into())
     }
 
+    /// Returns true if this tree has a new subtree, when compared with `prev_tree`.
+    pub fn contains_new_subtree(&self, prev_tree: &Self) -> bool {
+        // A new subtree was completed by the last note commitment in this tree.
+        if self.is_complete_subtree() {
+            return true;
+        }
+
+        // Use -1 for the index of the subtree with no notes, so the comparisons are valid.
+        let index = self.subtree_index().map_or(-1, |index| i32::from(index.0));
+        let prev_index = prev_tree
+            .subtree_index()
+            .map_or(-1, |index| i32::from(index.0));
+
+        // If the index is equal or lower, there can't be any new subtrees.
+        if index <= prev_index {
+            return false;
+        }
+
+        // This calculation can't overflow, because we're using i32 for u16 values.
+        let index_difference = index - prev_index;
+
+        // There are multiple new subtrees.
+        if index_difference > 1 {
+            return true;
+        }
+
+        // There is one new subtree somewhere in the trees. It is either:
+        // - a new subtree at the end of the previous tree, or
+        // - a new subtree in this tree (but not at the end).
+        //
+        // This happens because the subtree index only increases when the first note is added to
+        // the new subtree. So we need to exclude subtrees completed by the last note commitment in
+        // the previous tree.
+        if prev_tree.is_complete_subtree() {
+            return false;
+        }
+
+        // A new subtree was completed by a note commitment that isn't in the previous tree.
+        true
+    }
+
     /// Returns true if the most recently appended leaf completes the subtree
     pub fn is_complete_subtree(&self) -> bool {
         let Some(tree) = self.frontier() else {
diff --git a/zebra-chain/src/sapling/tree.rs b/zebra-chain/src/sapling/tree.rs
index f2702fea2d9..5ef07b2df94 100644
--- a/zebra-chain/src/sapling/tree.rs
+++ b/zebra-chain/src/sapling/tree.rs
@@ -387,6 +387,47 @@ impl NoteCommitmentTree {
         Some(tree.position().into())
     }
 
+    /// Returns true if this tree has a new subtree, when compared with `prev_tree`.
+    pub fn contains_new_subtree(&self, prev_tree: &Self) -> bool {
+        // A new subtree was completed by the last note commitment in this tree.
+        if self.is_complete_subtree() {
+            return true;
+        }
+
+        // Use -1 for the index of the subtree with no notes, so the comparisons are valid.
+        let index = self.subtree_index().map_or(-1, |index| i32::from(index.0));
+        let prev_index = prev_tree
+            .subtree_index()
+            .map_or(-1, |index| i32::from(index.0));
+
+        // If the index is equal or lower, there can't be any new subtrees.
+        if index <= prev_index {
+            return false;
+        }
+
+        // This calculation can't overflow, because we're using i32 for u16 values.
+        let index_difference = index - prev_index;
+
+        // There are multiple new subtrees.
+        if index_difference > 1 {
+            return true;
+        }
+
+        // There is one new subtree somewhere in the trees. It is either:
+        // - a new subtree at the end of the previous tree, or
+        // - a new subtree in this tree (but not at the end).
+        //
+        // This happens because the subtree index only increases when the first note is added to
+        // the new subtree. So we need to exclude subtrees completed by the last note commitment in
+        // the previous tree.
+        if prev_tree.is_complete_subtree() {
+            return false;
+        }
+
+        // A new subtree was completed by a note commitment that isn't in the previous tree.
+        true
+    }
+
     /// Returns true if the most recently appended leaf completes the subtree
     pub fn is_complete_subtree(&self) -> bool {
         let Some(tree) = self.frontier() else {
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 e4d1af06b87..f54a88a0709 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
@@ -65,16 +65,7 @@ pub fn run(
         })
         // Find new subtrees.
         .filter(|(_prev_end_height, prev_tree, _end_height, tree)| {
-            // The subtree is completed by the last note commitment in this block.
-            tree.is_complete_subtree() ||
-            // The subtree is completed by a note commitment before the last one in this block.
-            // We need to exclude subtrees completed at the end of the previous block
-            // because they also increase the subtree index.
-            //
-            // Empty note commitment trees can't contain subtrees, so they have invalid subtree
-            // indexes. But we have already skipped the empty genesis tree, so all trees have valid
-            // indexes. So this comparison is correct even if we don't unwrap the indexes.
-            (tree.subtree_index() > prev_tree.subtree_index() && !prev_tree.is_complete_subtree())
+            tree.contains_new_subtree(prev_tree)
         });
 
     for (prev_end_height, prev_tree, end_height, tree) in subtrees {
@@ -101,16 +92,7 @@ pub fn run(
         })
         // Find new subtrees.
         .filter(|(_prev_end_height, prev_tree, _end_height, tree)| {
-            // The subtree is completed by the last note commitment in this block.
-            tree.is_complete_subtree() ||
-            // The subtree is completed by a note commitment before the last one in this block.
-            // We need to exclude subtrees completed at the end of the previous block
-            // because they also increase the subtree index.
-            //
-            // Empty note commitment trees can't contain subtrees, so they have invalid subtree
-            // indexes. But we have already skipped the empty genesis tree, so all trees have valid
-            // indexes. So this comparison is correct even if we don't unwrap the indexes.
-            (tree.subtree_index() > prev_tree.subtree_index() && !prev_tree.is_complete_subtree())
+            tree.contains_new_subtree(prev_tree)
         });
 
     for (prev_end_height, prev_tree, end_height, tree) in subtrees {
@@ -695,6 +677,9 @@ fn calculate_sapling_subtree(
         );
 
         // Get the missing notes needed to complete the subtree.
+        //
+        // TODO: consider just reading the block's transactions from the database file,
+        //       because we don't use the block header data at all.
         let block = read_db
             .block(end_height.into())
             .expect("height with note commitment tree should have block");
@@ -704,11 +689,6 @@ fn calculate_sapling_subtree(
             .cloned()
             .collect();
 
-        // The update method modifies the tree inside the Arc if it only has one reference.
-        // This is usually ok, but during reverse iteration, we want re-use the original previous
-        // tree as the next current tree. Cloning the Arc preserves the original tree data.
-        let _force_inner_clone_of_prev_tree = prev_tree.clone();
-
         // This takes less than 1 second per tree, so we don't need to make it cancellable.
         let (sapling_nct, subtree) = NoteCommitmentTrees::update_sapling_note_commitment_tree(
             prev_tree,
@@ -868,6 +848,9 @@ fn calculate_orchard_subtree(
         );
 
         // Get the missing notes needed to complete the subtree.
+        //
+        // TODO: consider just reading the block's transactions from the database file,
+        //       because we don't use the block header data at all.
         let block = read_db
             .block(end_height.into())
             .expect("height with note commitment tree should have block");
@@ -877,11 +860,6 @@ fn calculate_orchard_subtree(
             .cloned()
             .collect();
 
-        // The update method modifies the tree inside the Arc if it only has one reference.
-        // This is usually ok, but during reverse iteration, we want re-use the original previous
-        // tree as the next current tree. Cloning the Arc preserves the original tree data.
-        let _force_inner_clone_of_prev_tree = prev_tree.clone();
-
         // This takes less than 1 second per tree, so we don't need to make it cancellable.
         let (orchard_nct, subtree) = NoteCommitmentTrees::update_orchard_note_commitment_tree(
             prev_tree,

From 58ee6f80bddb0f8186c2f339c9cff683701d8ad3 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 15 Sep 2023 10:52:19 +1000
Subject: [PATCH 32/56] Remove redundant assertions

---
 .../disk_format/upgrade/add_subtrees.rs       | 130 +-----------------
 1 file changed, 2 insertions(+), 128 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 f54a88a0709..4c4fb2eb4b8 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
@@ -582,7 +582,7 @@ fn calculate_sapling_subtree(
         // If the subtree is completed without using the last note commitment in the block,
         // we need to calculate the subtree root, starting with the tree from the previous block.
 
-        // TODO: move these checks into a separate function
+        // TODO: move the assertion/panic log string formatting into a separate function?
         let prev_position = prev_tree
             .position()
             .expect("previous block must have a partial subtree");
@@ -590,14 +590,12 @@ fn calculate_sapling_subtree(
             .subtree_index()
             .expect("previous block must have a partial subtree");
         let prev_remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
-        let prev_is_complete = prev_tree.is_complete_subtree();
 
         let current_position = tree.position().expect("current block must have a subtree");
         let current_index = tree
             .subtree_index()
             .expect("current block must have a subtree");
         let current_remaining_notes = tree.remaining_subtree_leaf_nodes();
-        let current_is_complete = tree.is_complete_subtree();
 
         assert_eq!(
             prev_index.0 + 1,
@@ -614,67 +612,6 @@ fn calculate_sapling_subtree(
              position: {current_position}\n\
              remaining: {current_remaining_notes}"
         );
-        assert!(
-            current_remaining_notes > 0,
-            "just checked for a complete subtree in the current block:\n\
-             previous subtree:\n\
-             height: {prev_end_height:?}\n\
-             index: {prev_index}\n\
-             position: {prev_position}\n\
-             remaining: {prev_remaining_notes}\n\
-             current subtree:\n\
-             height: {end_height:?}\n\
-             index: {current_index}\n\
-             position: {current_position}\n\
-             remaining: {current_remaining_notes}"
-        );
-        assert!(
-            !current_is_complete,
-            "just checked for a complete subtree in the current block:\n\
-             previous subtree:\n\
-             height: {prev_end_height:?}\n\
-             index: {prev_index}\n\
-             position: {prev_position}\n\
-             remaining: {prev_remaining_notes}\n\
-             current subtree:\n\
-             height: {end_height:?}\n\
-             index: {current_index}\n\
-             position: {current_position}\n\
-             remaining: {current_remaining_notes}"
-        );
-
-        // If the previous tree was complete, then the current tree can't also complete a subtree,
-        // due to the consensus limit on the number of outputs in a block.
-        assert!(
-            prev_remaining_notes > 0,
-            "the caller should supply a complete subtree, \
-             so the previous tree can't also be complete:\n\
-             previous subtree:\n\
-             height: {prev_end_height:?}\n\
-             index: {prev_index}\n\
-             position: {prev_position}\n\
-             remaining: {prev_remaining_notes}\n\
-             current subtree:\n\
-             height: {end_height:?}\n\
-             index: {current_index}\n\
-             position: {current_position}\n\
-             remaining: {current_remaining_notes}"
-        );
-        assert!(
-            !prev_is_complete,
-            "the caller should supply a complete subtree, \
-             so the previous tree can't also be complete:\n\
-             previous subtree:\n\
-             height: {prev_end_height:?}\n\
-             index: {prev_index}\n\
-             position: {prev_position}\n\
-             remaining: {prev_remaining_notes}\n\
-             current subtree:\n\
-             height: {end_height:?}\n\
-             index: {current_index}\n\
-             position: {current_position}\n\
-             remaining: {current_remaining_notes}"
-        );
 
         // Get the missing notes needed to complete the subtree.
         //
@@ -753,7 +690,7 @@ fn calculate_orchard_subtree(
         // If the subtree is completed without using the last note commitment in the block,
         // we need to calculate the subtree root, starting with the tree from the previous block.
 
-        // TODO: move these checks into a separate function
+        // TODO: move the assertion/panic log string formatting into a separate function?
         let prev_position = prev_tree
             .position()
             .expect("previous block must have a partial subtree");
@@ -761,14 +698,12 @@ fn calculate_orchard_subtree(
             .subtree_index()
             .expect("previous block must have a partial subtree");
         let prev_remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
-        let prev_is_complete = prev_tree.is_complete_subtree();
 
         let current_position = tree.position().expect("current block must have a subtree");
         let current_index = tree
             .subtree_index()
             .expect("current block must have a subtree");
         let current_remaining_notes = tree.remaining_subtree_leaf_nodes();
-        let current_is_complete = tree.is_complete_subtree();
 
         assert_eq!(
             prev_index.0 + 1,
@@ -785,67 +720,6 @@ fn calculate_orchard_subtree(
              position: {current_position}\n\
              remaining: {current_remaining_notes}"
         );
-        assert!(
-            current_remaining_notes > 0,
-            "just checked for a complete subtree in the current block:\n\
-             previous subtree:\n\
-             height: {prev_end_height:?}\n\
-             index: {prev_index}\n\
-             position: {prev_position}\n\
-             remaining: {prev_remaining_notes}\n\
-             current subtree:\n\
-             height: {end_height:?}\n\
-             index: {current_index}\n\
-             position: {current_position}\n\
-             remaining: {current_remaining_notes}"
-        );
-        assert!(
-            !current_is_complete,
-            "just checked for a complete subtree in the current block:\n\
-             previous subtree:\n\
-             height: {prev_end_height:?}\n\
-             index: {prev_index}\n\
-             position: {prev_position}\n\
-             remaining: {prev_remaining_notes}\n\
-             current subtree:\n\
-             height: {end_height:?}\n\
-             index: {current_index}\n\
-             position: {current_position}\n\
-             remaining: {current_remaining_notes}"
-        );
-
-        // If the previous tree was complete, then the current tree can't also complete a subtree,
-        // due to the consensus limit on the number of outputs in a block.
-        assert!(
-            prev_remaining_notes > 0,
-            "the caller should supply a complete subtree, \
-             so the previous tree can't also be complete:\n\
-             previous subtree:\n\
-             height: {prev_end_height:?}\n\
-             index: {prev_index}\n\
-             position: {prev_position}\n\
-             remaining: {prev_remaining_notes}\n\
-             current subtree:\n\
-             height: {end_height:?}\n\
-             index: {current_index}\n\
-             position: {current_position}\n\
-             remaining: {current_remaining_notes}"
-        );
-        assert!(
-            !prev_is_complete,
-            "the caller should supply a complete subtree, \
-             so the previous tree can't also be complete:\n\
-             previous subtree:\n\
-             height: {prev_end_height:?}\n\
-             index: {prev_index}\n\
-             position: {prev_position}\n\
-             remaining: {prev_remaining_notes}\n\
-             current subtree:\n\
-             height: {end_height:?}\n\
-             index: {current_index}\n\
-             position: {current_position}\n\
-             remaining: {current_remaining_notes}"
-        );
 
         // Get the missing notes needed to complete the subtree.
         //

From 0146f9f725598fd73fa58306f3a1312641393cd6 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 15 Sep 2023 16:30:19 +1000
Subject: [PATCH 33/56] Wait for subtree upgrade before testing RPCs

---
 zebra-state/src/constants.rs                  |   8 ++
 zebra-state/src/lib.rs                        |   3 +
 .../finalized_state/disk_format/upgrade.rs    |  43 ++++----
 .../common/lightwalletd/wallet_grpc_test.rs   | 104 +++++++++++-------
 4 files changed, 96 insertions(+), 62 deletions(-)

diff --git a/zebra-state/src/constants.rs b/zebra-state/src/constants.rs
index d87cced0522..07177f11cfa 100644
--- a/zebra-state/src/constants.rs
+++ b/zebra-state/src/constants.rs
@@ -2,6 +2,7 @@
 
 use lazy_static::lazy_static;
 use regex::Regex;
+use semver::Version;
 
 // For doc comment links
 #[allow(unused_imports)]
@@ -54,6 +55,13 @@ pub(crate) const DATABASE_FORMAT_MINOR_VERSION: u64 = 2;
 /// significant format compatibility fix.
 pub(crate) const DATABASE_FORMAT_PATCH_VERSION: u64 = 1;
 
+/// Returns the highest database version that modifies the subtree index format.
+///
+/// This version is used by tests to wait for the subtree upgrade to finish.
+pub fn latest_version_for_adding_subtrees() -> Version {
+    Version::parse("25.2.1").expect("Hardcoded version string should be valid.")
+}
+
 /// The name of the file containing the minor and patch database versions.
 ///
 /// Use [`Config::version_file_path()`] to get the path to this file.
diff --git a/zebra-state/src/lib.rs b/zebra-state/src/lib.rs
index eaab5fe083b..d1076c68b95 100644
--- a/zebra-state/src/lib.rs
+++ b/zebra-state/src/lib.rs
@@ -73,4 +73,7 @@ pub use service::{
 #[cfg(any(test, feature = "proptest-impl"))]
 pub use config::write_database_format_version_to_disk;
 
+#[cfg(any(test, feature = "proptest-impl"))]
+pub use constants::latest_version_for_adding_subtrees;
+
 pub(crate) use request::ContextuallyVerifiedBlock;
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 45e296de3e3..11f0c13f653 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
@@ -19,7 +19,7 @@ use DbFormatChange::*;
 
 use crate::{
     config::write_database_format_version_to_disk,
-    constants::DATABASE_FORMAT_VERSION,
+    constants::{latest_version_for_adding_subtrees, DATABASE_FORMAT_VERSION},
     database_format_version_in_code, database_format_version_on_disk,
     service::finalized_state::{DiskWriteBatch, ZebraDb},
     Config,
@@ -90,7 +90,7 @@ impl DbFormatChange {
     pub fn new(running_version: Version, disk_version: Option<Version>) -> Option<Self> {
         let Some(disk_version) = disk_version else {
             info!(
-                ?running_version,
+                %running_version,
                 "creating new database with the current format"
             );
 
@@ -100,8 +100,8 @@ impl DbFormatChange {
         match disk_version.cmp(&running_version) {
             Ordering::Less => {
                 info!(
-                    ?running_version,
-                    ?disk_version,
+                    %running_version,
+                    %disk_version,
                     "trying to open older database format: launching upgrade task"
                 );
 
@@ -112,8 +112,8 @@ impl DbFormatChange {
             }
             Ordering::Greater => {
                 info!(
-                    ?running_version,
-                    ?disk_version,
+                    %running_version,
+                    %disk_version,
                     "trying to open newer database format: data should be compatible"
                 );
 
@@ -123,7 +123,7 @@ impl DbFormatChange {
                 })
             }
             Ordering::Equal => {
-                info!(?running_version, "trying to open current database format");
+                info!(%running_version, "trying to open current database format");
 
                 None
             }
@@ -269,16 +269,16 @@ impl DbFormatChange {
         let Some(initial_tip_height) = initial_tip_height else {
             // If the database is empty, then the RocksDb format doesn't need any changes.
             info!(
-                ?newer_running_version,
-                ?older_disk_version,
+                %newer_running_version,
+                %older_disk_version,
                 "marking empty database as upgraded"
             );
 
             Self::mark_as_upgraded_to(&database_format_version_in_code(), &config, network);
 
             info!(
-                ?newer_running_version,
-                ?older_disk_version,
+                %newer_running_version,
+                %older_disk_version,
                 "empty database is fully upgraded"
             );
 
@@ -356,9 +356,7 @@ impl DbFormatChange {
 
         // Note commitment subtree creation database upgrade task.
 
-        let latest_version_for_adding_subtrees =
-            Version::parse("25.2.1").expect("Hardcoded version string should be valid.");
-
+        let latest_version_for_adding_subtrees = latest_version_for_adding_subtrees();
         let first_version_for_adding_subtrees =
             Version::parse("25.2.0").expect("Hardcoded version string should be valid.");
 
@@ -388,7 +386,7 @@ impl DbFormatChange {
         // every time it runs its inner update loop.
 
         info!(
-            ?newer_running_version,
+            %newer_running_version,
             "Zebra automatically upgraded the database format to:"
         );
 
@@ -486,8 +484,8 @@ impl DbFormatChange {
             .expect("unable to write database format version file to disk");
 
         info!(
-            ?running_version,
-            ?disk_version,
+            %running_version,
+            disk_version = %disk_version.map_or("None".to_string(), |version| version.to_string()),
             "marked database format as newly created"
         );
     }
@@ -547,9 +545,10 @@ impl DbFormatChange {
             .expect("unable to write database format version file to disk");
 
         info!(
-            ?running_version,
-            ?format_upgrade_version,
-            ?disk_version,
+            %running_version,
+            %disk_version,
+            // The wallet_grpc_test needs this to be the last field, so the regex matches correctly
+            %format_upgrade_version,
             "marked database format as upgraded"
         );
     }
@@ -586,8 +585,8 @@ impl DbFormatChange {
             .expect("unable to write database format version file to disk");
 
         info!(
-            ?running_version,
-            ?disk_version,
+            %running_version,
+            %disk_version,
             "marked database format as downgraded"
         );
     }
diff --git a/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs b/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
index 351c14843a8..9834199df48 100644
--- a/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
+++ b/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
@@ -97,6 +97,18 @@ pub async fn run() -> Result<()> {
 
     let zebra_rpc_address = zebra_rpc_address.expect("lightwalletd test must have RPC port");
 
+    tracing::info!(
+        ?test_type,
+        "launched zebrad, waiting for zebrad to open the state database..."
+    );
+    // Zebra logs one of these lines on startup, depending on the disk and running formats.
+    let state_format_change = zebrad.expect_stdout_line_matches(
+        "(creating new database with the current format)|\
+         (trying to open older database format)|\
+         (trying to open newer database format)|\
+         (trying to open current database format)",
+    )?;
+
     tracing::info!(
         ?test_type,
         ?zebra_rpc_address,
@@ -339,6 +351,58 @@ pub async fn run() -> Result<()> {
         *zebra_test::vectors::SAPLING_TREESTATE_MAINNET_419201_STRING
     );
 
+    // Call `GetAddressUtxos` with the ZF funding stream address that will always have utxos
+    let utxos = rpc_client
+        .get_address_utxos(GetAddressUtxosArg {
+            addresses: vec!["t3dvVE3SQEi7kqNzwrfNePxZ1d4hUyztBA1".to_string()],
+            start_height: 1,
+            max_entries: 1,
+        })
+        .await?
+        .into_inner();
+
+    // As we requested one entry we should get a response of length 1
+    assert_eq!(utxos.address_utxos.len(), 1);
+
+    // Call `GetAddressUtxosStream` with the ZF funding stream address that will always have utxos
+    let mut utxos_zf = rpc_client
+        .get_address_utxos_stream(GetAddressUtxosArg {
+            addresses: vec!["t3dvVE3SQEi7kqNzwrfNePxZ1d4hUyztBA1".to_string()],
+            start_height: 1,
+            max_entries: 2,
+        })
+        .await?
+        .into_inner();
+
+    let mut counter = 0;
+    while let Some(_utxos) = utxos_zf.message().await? {
+        counter += 1;
+    }
+    // As we are in a "in sync" chain we know there are more than 2 utxos for this address
+    // but we will receive the max of 2 from the stream response because we used a limit of 2 `max_entries`.
+    assert_eq!(2, counter);
+
+    // Call `GetLightdInfo`
+    let lightd_info = rpc_client.get_lightd_info(Empty {}).await?.into_inner();
+
+    // Make sure the subversion field is zebra the user agent
+    assert_eq!(
+        lightd_info.zcashd_subversion,
+        zebrad::application::user_agent()
+    );
+
+    // Before we call `z_getsubtreesbyindex`, we might need to wait for a database upgrade.
+    if state_format_change.contains("launching upgrade task") {
+        let latest_version_for_adding_subtrees = latest_version_for_adding_subtrees();
+        tracing::info!(
+            ?test_type,
+            ?zebra_rpc_address,
+            latest_version_for_adding_subtrees,
+            "waiting for zebrad subtree state upgrade..."
+        );
+        zebrad.expect_stdout_line_matches(&format!("marked database format as upgraded.*format_upgrade_version.*=.*{latest_version_for_adding_subtrees}"))?;
+    }
+
     // Call `z_getsubtreesbyindex` separately for
 
     // ... Sapling.
@@ -411,45 +475,5 @@ pub async fn run() -> Result<()> {
     }
     assert_eq!(counter, 2);
 
-    // Call `GetAddressUtxos` with the ZF funding stream address that will always have utxos
-    let utxos = rpc_client
-        .get_address_utxos(GetAddressUtxosArg {
-            addresses: vec!["t3dvVE3SQEi7kqNzwrfNePxZ1d4hUyztBA1".to_string()],
-            start_height: 1,
-            max_entries: 1,
-        })
-        .await?
-        .into_inner();
-
-    // As we requested one entry we should get a response of length 1
-    assert_eq!(utxos.address_utxos.len(), 1);
-
-    // Call `GetAddressUtxosStream` with the ZF funding stream address that will always have utxos
-    let mut utxos_zf = rpc_client
-        .get_address_utxos_stream(GetAddressUtxosArg {
-            addresses: vec!["t3dvVE3SQEi7kqNzwrfNePxZ1d4hUyztBA1".to_string()],
-            start_height: 1,
-            max_entries: 2,
-        })
-        .await?
-        .into_inner();
-
-    let mut counter = 0;
-    while let Some(_utxos) = utxos_zf.message().await? {
-        counter += 1;
-    }
-    // As we are in a "in sync" chain we know there are more than 2 utxos for this address
-    // but we will receive the max of 2 from the stream response because we used a limit of 2 `max_entries`.
-    assert_eq!(2, counter);
-
-    // Call `GetLightdInfo`
-    let lightd_info = rpc_client.get_lightd_info(Empty {}).await?.into_inner();
-
-    // Make sure the subversion field is zebra the user agent
-    assert_eq!(
-        lightd_info.zcashd_subversion,
-        zebrad::application::user_agent()
-    );
-
     Ok(())
 }

From 1d889aa3a831b70cdab0ec49640be2a80314ac3f Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 15 Sep 2023 16:41:13 +1000
Subject: [PATCH 34/56] Fix subtree search in quick check

---
 .../disk_format/upgrade/add_subtrees.rs               | 11 +++++++----
 1 file changed, 7 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 4c4fb2eb4b8..8b3856cd7e3 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
@@ -207,7 +207,7 @@ fn first_orchard_mainnet_subtree() -> NoteCommitmentSubtree<orchard::tree::Node>
 ///
 /// Returns an error if a note commitment subtree is missing or incorrect.
 fn quick_check_sapling_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
-    // We check the first sapling tree on mainnet, so skip this check if it isn't available.
+    // We check the first sapling subtree on mainnet, so skip this check if it isn't available.
     let Some(NoteCommitmentSubtreeIndex(first_incomplete_subtree_index)) =
         db.sapling_tree().subtree_index()
     else {
@@ -232,7 +232,7 @@ fn quick_check_sapling_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
         // But since we skip the empty genesis tree, all trees must have valid indexes.
         // So we don't need to unwrap the optional values for this comparison to be correct.
         .find(|(_prev_end_height, prev_tree, _end_height, tree)| {
-            tree.subtree_index() > prev_tree.subtree_index()
+            tree.contains_new_subtree(prev_tree)
         });
 
     let Some((prev_end_height, prev_tree, end_height, tree)) = first_complete_subtree else {
@@ -262,7 +262,7 @@ fn quick_check_sapling_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
 ///
 /// Returns an error if a note commitment subtree is missing or incorrect.
 fn quick_check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
-    // We check the first orchard tree on mainnet, so skip this check if it isn't available.
+    // We check the first orchard subtree on mainnet, so skip this check if it isn't available.
     let Some(NoteCommitmentSubtreeIndex(first_incomplete_subtree_index)) =
         db.orchard_tree().subtree_index()
     else {
@@ -287,7 +287,7 @@ fn quick_check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
         // But since we skip the empty genesis tree, all trees must have valid indexes.
         // So we don't need to unwrap the optional values for this comparison to be correct.
         .find(|(_prev_end_height, prev_tree, _end_height, tree)| {
-            tree.subtree_index() > prev_tree.subtree_index()
+            tree.contains_new_subtree(prev_tree)
         });
 
     let Some((prev_end_height, prev_tree, end_height, tree)) = first_complete_subtree else {
@@ -316,6 +316,9 @@ fn quick_check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
 ///
 /// If a note commitment subtree is missing or incorrect.
 pub fn check(db: &ZebraDb) {
+    // This check is partly redundant, but we want to make sure it's never missed.
+    quick_check(db);
+
     let sapling_result = check_sapling_subtrees(db);
     let orchard_result = check_orchard_subtrees(db);
 

From a9558be21401eb23f0079ef0f6a3e5086dba16e5 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 15 Sep 2023 16:54:59 +1000
Subject: [PATCH 35/56] Temporarily upgrade subtrees in forward height order

---
 .../src/service/finalized_state/disk_db.rs    | 106 ++++--------------
 .../disk_format/upgrade/add_subtrees.rs       |  17 +--
 .../finalized_state/zebra_db/shielded.rs      |  26 -----
 3 files changed, 23 insertions(+), 126 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db.rs b/zebra-state/src/service/finalized_state/disk_db.rs
index 3772fb7a789..d001c87d157 100644
--- a/zebra-state/src/service/finalized_state/disk_db.rs
+++ b/zebra-state/src/service/finalized_state/disk_db.rs
@@ -258,7 +258,6 @@ impl PartialEq for DiskDb {
                 self.ephemeral, other.ephemeral,
                 "database with same path but different ephemeral configs",
             );
-
             return true;
         }
 
@@ -432,46 +431,6 @@ impl DiskDb {
     ///
     /// Holding this iterator open might delay block commit transactions.
     pub fn zs_range_iter<C, K, V, R>(&self, cf: &C, range: R) -> impl Iterator<Item = (K, V)> + '_
-    where
-        C: rocksdb::AsColumnFamilyRef,
-        K: IntoDisk + FromDisk,
-        V: FromDisk,
-        R: RangeBounds<K>,
-    {
-        self.zs_range_iter_with_direction(cf, range, false)
-    }
-
-    /// Returns a reverse iterator over the items in `cf` in `range`.
-    ///
-    /// Holding this iterator open might delay block commit transactions.
-    ///
-    /// This code is copied from `zs_range_iter()`, but with the mode reversed.
-    pub fn zs_reverse_range_iter<C, K, V, R>(
-        &self,
-        cf: &C,
-        range: R,
-    ) -> impl Iterator<Item = (K, V)> + '_
-    where
-        C: rocksdb::AsColumnFamilyRef,
-        K: IntoDisk + FromDisk,
-        V: FromDisk,
-        R: RangeBounds<K>,
-    {
-        self.zs_range_iter_with_direction(cf, range, true)
-    }
-
-    /// Returns an iterator over the items in `cf` in `range`.
-    ///
-    /// RocksDB iterators are ordered by increasing key bytes by default.
-    /// Otherwise, if `reverse` is `true`, the iterator is ordered by decreasing key bytes.
-    ///
-    /// Holding this iterator open might delay block commit transactions.
-    fn zs_range_iter_with_direction<C, K, V, R>(
-        &self,
-        cf: &C,
-        range: R,
-        reverse: bool,
-    ) -> impl Iterator<Item = (K, V)> + '_
     where
         C: rocksdb::AsColumnFamilyRef,
         K: IntoDisk + FromDisk,
@@ -492,67 +451,40 @@ impl DiskDb {
 
         let start_bound = map_to_vec(range.start_bound());
         let end_bound = map_to_vec(range.end_bound());
-        let range = (start_bound, end_bound);
+        let range = (start_bound.clone(), end_bound);
+
+        let start_bound_vec =
+            if let Included(ref start_bound) | Excluded(ref start_bound) = start_bound {
+                start_bound.clone()
+            } else {
+                // Actually unused
+                Vec::new()
+            };
 
-        let mode = Self::zs_iter_mode(&range, reverse);
+        let start_mode = if matches!(start_bound, Unbounded) {
+            // Unbounded iterators start at the first item
+            rocksdb::IteratorMode::Start
+        } else {
+            rocksdb::IteratorMode::From(start_bound_vec.as_slice(), rocksdb::Direction::Forward)
+        };
 
         // Reading multiple items from iterators has caused database hangs,
         // in previous RocksDB versions
         self.db
-            .iterator_cf(cf, mode)
+            .iterator_cf(cf, start_mode)
             .map(|result| result.expect("unexpected database failure"))
             .map(|(key, value)| (key.to_vec(), value))
-            // Skip excluded "from" bound and empty ranges. The `mode` already skips keys
-            // strictly before the "from" bound.
+            // Skip excluded start bound and empty ranges. The `start_mode` already skips keys
+            // before the start bound.
             .skip_while({
                 let range = range.clone();
                 move |(key, _value)| !range.contains(key)
             })
-            // Take until the excluded "to" bound is reached,
-            // or we're after the included "to" bound.
+            // Take until the excluded end bound is reached, or we're after the included end bound.
             .take_while(move |(key, _value)| range.contains(key))
             .map(|(key, value)| (K::from_bytes(key), V::from_bytes(value)))
     }
 
-    /// Returns the RocksDB iterator "from" mode for `range`.
-    ///
-    /// RocksDB iterators are ordered by increasing key bytes by default.
-    /// Otherwise, if `reverse` is `true`, the iterator is ordered by decreasing key bytes.
-    fn zs_iter_mode<R>(range: &R, reverse: bool) -> rocksdb::IteratorMode
-    where
-        R: RangeBounds<Vec<u8>>,
-    {
-        use std::ops::Bound::*;
-
-        let from_bound = if reverse {
-            range.end_bound()
-        } else {
-            range.start_bound()
-        };
-
-        match from_bound {
-            Unbounded => {
-                if reverse {
-                    // Reversed unbounded iterators start from the last item
-                    rocksdb::IteratorMode::End
-                } else {
-                    // Unbounded iterators start from the first item
-                    rocksdb::IteratorMode::Start
-                }
-            }
-
-            Included(bound) | Excluded(bound) => {
-                let direction = if reverse {
-                    rocksdb::Direction::Reverse
-                } else {
-                    rocksdb::Direction::Forward
-                };
-
-                rocksdb::IteratorMode::From(bound.as_slice(), direction)
-            }
-        }
-    }
-
     /// The ideal open file limit for Zebra
     const IDEAL_OPEN_FILE_LIMIT: u64 = 1024;
 
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 8b3856cd7e3..e659317426b 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
@@ -44,23 +44,15 @@ pub fn run(
     //
     // Therefore, the first block with shielded note can't complete a subtree, which means we can
     // skip the (genesis block, first shielded block) tree pair.
-    //
-    // # Compatibility
-    //
-    // Because wallets search backwards from the chain tip, subtrees need to be added to the
-    // database in reverse height order. (Tip first, genesis last.)
-    //
-    // Otherwise, wallets that sync during the upgrade will be missing some notes.
 
     // Generate a list of sapling subtree inputs: previous and current trees, and their end heights.
     let subtrees = upgrade_db
-        .sapling_tree_by_reversed_height_range(..=initial_tip_height)
+        .sapling_tree_by_height_range(..=initial_tip_height)
         // The first block with sapling notes can't complete a subtree, see above for details.
         .filter(|(height, _tree)| !height.is_min())
         // We need both the tree and its previous tree for each shielded block.
         .tuple_windows()
-        // Because the iterator is reversed, the larger tree is first.
-        .map(|((end_height, tree), (prev_end_height, prev_tree))| {
+        .map(|((prev_end_height, prev_tree), (end_height, tree))| {
             (prev_end_height, prev_tree, end_height, tree)
         })
         // Find new subtrees.
@@ -81,13 +73,12 @@ pub fn run(
 
     // Generate a list of orchard subtree inputs: previous and current trees, and their end heights.
     let subtrees = upgrade_db
-        .orchard_tree_by_reversed_height_range(..=initial_tip_height)
+        .orchard_tree_by_height_range(..=initial_tip_height)
         // The first block with orchard notes can't complete a subtree, see above for details.
         .filter(|(height, _tree)| !height.is_min())
         // We need both the tree and its previous tree for each shielded block.
         .tuple_windows()
-        // Because the iterator is reversed, the larger tree is first.
-        .map(|((end_height, tree), (prev_end_height, prev_tree))| {
+        .map(|((prev_end_height, prev_tree), (end_height, tree))| {
             (prev_end_height, prev_tree, end_height, tree)
         })
         // Find new subtrees.
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 395ae8b6b99..aaca367019f 100644
--- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
@@ -180,19 +180,6 @@ impl ZebraDb {
         self.db.zs_range_iter(&sapling_trees, range)
     }
 
-    /// Returns the Sapling note commitment trees in the reversed range, in decreasing height order.
-    #[allow(clippy::unwrap_in_result)]
-    pub fn sapling_tree_by_reversed_height_range<R>(
-        &self,
-        range: R,
-    ) -> impl Iterator<Item = (Height, Arc<sapling::tree::NoteCommitmentTree>)> + '_
-    where
-        R: std::ops::RangeBounds<Height>,
-    {
-        let sapling_trees = self.db.cf_handle("sapling_note_commitment_tree").unwrap();
-        self.db.zs_reverse_range_iter(&sapling_trees, range)
-    }
-
     /// Returns the Sapling note commitment subtree at this `index`.
     ///
     /// # Correctness
@@ -326,19 +313,6 @@ impl ZebraDb {
         self.db.zs_range_iter(&orchard_trees, range)
     }
 
-    /// Returns the Orchard note commitment trees in the reversed range, in decreasing height order.
-    #[allow(clippy::unwrap_in_result)]
-    pub fn orchard_tree_by_reversed_height_range<R>(
-        &self,
-        range: R,
-    ) -> impl Iterator<Item = (Height, Arc<orchard::tree::NoteCommitmentTree>)> + '_
-    where
-        R: std::ops::RangeBounds<Height>,
-    {
-        let orchard_trees = self.db.cf_handle("orchard_note_commitment_tree").unwrap();
-        self.db.zs_reverse_range_iter(&orchard_trees, range)
-    }
-
     /// Returns the Orchard note commitment subtree at this `index`.
     ///
     /// # Correctness

From 74db92343d727cc35c901b691f8c67e121d2d236 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 15 Sep 2023 16:55:15 +1000
Subject: [PATCH 36/56] Clarify some comments

---
 .../src/service/finalized_state/disk_format/upgrade.rs        | 4 ++--
 1 file changed, 2 insertions(+), 2 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 11f0c13f653..c084b4385aa 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
@@ -360,10 +360,10 @@ impl DbFormatChange {
         let first_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.
+        // Check if we need to add or fix note commitment subtrees in the database.
         if older_disk_version < latest_version_for_adding_subtrees {
             if older_disk_version >= first_version_for_adding_subtrees {
-                // Clear previous upgrade data that might be different from this upgrade's data.
+                // Clear previous upgrade data, because it was incorrect.
                 add_subtrees::reset(initial_tip_height, &db, cancel_receiver)?;
             }
 

From 2f8885873c0190e581ccbbedc883f8776ffae742 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 15 Sep 2023 16:58:11 +1000
Subject: [PATCH 37/56] Fix missing test imports

---
 zebrad/tests/common/lightwalletd/wallet_grpc_test.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs b/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
index 9834199df48..6e883723e71 100644
--- a/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
+++ b/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
@@ -43,6 +43,7 @@ use zebra_chain::{
     parameters::NetworkUpgrade::{Nu5, Sapling},
     serialization::ZcashDeserializeInto,
 };
+use zebra_state::latest_version_for_adding_subtrees;
 
 use crate::common::{
     launch::spawn_zebrad_for_rpc,

From d457eaf9eddb1c7699c3718cf74c572ceb8edddd Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 15 Sep 2023 17:11:09 +1000
Subject: [PATCH 38/56] Fix subtree logging

---
 zebrad/tests/common/lightwalletd/wallet_grpc_test.rs | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs b/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
index 6e883723e71..ff55b6f8efa 100644
--- a/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
+++ b/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
@@ -398,7 +398,8 @@ pub async fn run() -> Result<()> {
         tracing::info!(
             ?test_type,
             ?zebra_rpc_address,
-            latest_version_for_adding_subtrees,
+            %state_format_change,
+            %latest_version_for_adding_subtrees,
             "waiting for zebrad subtree state upgrade..."
         );
         zebrad.expect_stdout_line_matches(&format!("marked database format as upgraded.*format_upgrade_version.*=.*{latest_version_for_adding_subtrees}"))?;

From 1d8f45803821305fce8b4741aad7db4776fb6761 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 15 Sep 2023 17:18:33 +1000
Subject: [PATCH 39/56] Add a comment about a potential hang with future
 upgrades

---
 zebrad/tests/common/lightwalletd/wallet_grpc_test.rs | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs b/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
index ff55b6f8efa..59f21dd34dc 100644
--- a/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
+++ b/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
@@ -402,6 +402,11 @@ pub async fn run() -> Result<()> {
             %latest_version_for_adding_subtrees,
             "waiting for zebrad subtree state upgrade..."
         );
+        //
+        // TODO: this line will hang if the state upgrade finishes before the subtree tests start.
+        // But that is unlikely with the 25.2 upgrade, because it takes 20+ minutes.
+        // If it happens for a later upgrade, this code can be removed, as long as all the
+        // lightwalletd cached states are version 25.2.2 or later.
         zebrad.expect_stdout_line_matches(&format!("marked database format as upgraded.*format_upgrade_version.*=.*{latest_version_for_adding_subtrees}"))?;
     }
 

From 6d1e494d6b4c9b9398696bf4cbfd7358e6fd049c Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 15 Sep 2023 17:20:45 +1000
Subject: [PATCH 40/56] Fix zebrad var ownership

---
 zebrad/tests/common/lightwalletd/wallet_grpc_test.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs b/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
index 59f21dd34dc..7f2ea84286a 100644
--- a/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
+++ b/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
@@ -132,7 +132,7 @@ pub async fn run() -> Result<()> {
         "spawned lightwalletd connected to zebrad, waiting for them both to sync...",
     );
 
-    let (_lightwalletd, _zebrad) = wait_for_zebrad_and_lightwalletd_sync(
+    let (_lightwalletd, mut zebrad) = wait_for_zebrad_and_lightwalletd_sync(
         lightwalletd,
         lightwalletd_rpc_port,
         zebrad,

From c6b8aa7bbed758b11d0628b42ea2864cba3bd872 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Mon, 18 Sep 2023 05:59:29 +1000
Subject: [PATCH 41/56] Log more info when add_subtrees.rs fails

---
 .../disk_format/upgrade/add_subtrees.rs       | 42 +++++++++++++++++--
 1 file changed, 38 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 e659317426b..9c82cb8f06d 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
@@ -579,13 +579,30 @@ fn calculate_sapling_subtree(
         // TODO: move the assertion/panic log string formatting into a separate function?
         let prev_position = prev_tree
             .position()
-            .expect("previous block must have a partial subtree");
+            .expect(&format!(
+                "previous block must have a partial subtree:\n\
+                previous subtree:\n\
+                height: {prev_end_height:?}\n\
+                current subtree:\n\
+                height: {end_height:?}"
+            ));
         let prev_index = prev_tree
             .subtree_index()
             .expect("previous block must have a partial subtree");
         let prev_remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
 
-        let current_position = tree.position().expect("current block must have a subtree");
+        let current_position = tree
+            .position()
+            .expect(&format!(
+                "current block must have a subtree:\n\
+                previous subtree:\n\
+                height: {prev_end_height:?}\n\
+                index: {prev_index}\n\
+                position: {prev_position}\n\
+                remaining: {prev_remaining_notes}\n\
+                current subtree:\n\
+                height: {end_height:?}"
+            ));
         let current_index = tree
             .subtree_index()
             .expect("current block must have a subtree");
@@ -687,13 +704,30 @@ fn calculate_orchard_subtree(
         // TODO: move the assertion/panic log string formatting into a separate function?
         let prev_position = prev_tree
             .position()
-            .expect("previous block must have a partial subtree");
+            .expect(&format!(
+                "previous block must have a partial subtree:\n\
+                previous subtree:\n\
+                height: {prev_end_height:?}\n\
+                current subtree:\n\
+                height: {end_height:?}"
+            ));
         let prev_index = prev_tree
             .subtree_index()
             .expect("previous block must have a partial subtree");
         let prev_remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
 
-        let current_position = tree.position().expect("current block must have a subtree");
+        let current_position = tree
+            .position()
+            .expect(&format!(
+                "current block must have a subtree:\n\
+                previous subtree:\n\
+                height: {prev_end_height:?}\n\
+                index: {prev_index}\n\
+                position: {prev_position}\n\
+                remaining: {prev_remaining_notes}\n\
+                current subtree:\n\
+                height: {end_height:?}"
+            ));
         let current_index = tree
             .subtree_index()
             .expect("current block must have a subtree");

From 5ee803adcd60170634f7d1b3f31adedea928ce5f Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Mon, 18 Sep 2023 06:16:00 +1000
Subject: [PATCH 42/56] cargo fmt --all

---
 .../disk_format/upgrade/add_subtrees.rs       | 32 +++++++------------
 1 file changed, 12 insertions(+), 20 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 9c82cb8f06d..3171aacc347 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
@@ -577,24 +577,20 @@ fn calculate_sapling_subtree(
         // we need to calculate the subtree root, starting with the tree from the previous block.
 
         // TODO: move the assertion/panic log string formatting into a separate function?
-        let prev_position = prev_tree
-            .position()
-            .expect(&format!(
-                "previous block must have a partial subtree:\n\
+        let prev_position = prev_tree.position().expect(&format!(
+            "previous block must have a partial subtree:\n\
                 previous subtree:\n\
                 height: {prev_end_height:?}\n\
                 current subtree:\n\
                 height: {end_height:?}"
-            ));
+        ));
         let prev_index = prev_tree
             .subtree_index()
             .expect("previous block must have a partial subtree");
         let prev_remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
 
-        let current_position = tree
-            .position()
-            .expect(&format!(
-                "current block must have a subtree:\n\
+        let current_position = tree.position().expect(&format!(
+            "current block must have a subtree:\n\
                 previous subtree:\n\
                 height: {prev_end_height:?}\n\
                 index: {prev_index}\n\
@@ -602,7 +598,7 @@ fn calculate_sapling_subtree(
                 remaining: {prev_remaining_notes}\n\
                 current subtree:\n\
                 height: {end_height:?}"
-            ));
+        ));
         let current_index = tree
             .subtree_index()
             .expect("current block must have a subtree");
@@ -702,24 +698,20 @@ fn calculate_orchard_subtree(
         // we need to calculate the subtree root, starting with the tree from the previous block.
 
         // TODO: move the assertion/panic log string formatting into a separate function?
-        let prev_position = prev_tree
-            .position()
-            .expect(&format!(
-                "previous block must have a partial subtree:\n\
+        let prev_position = prev_tree.position().expect(&format!(
+            "previous block must have a partial subtree:\n\
                 previous subtree:\n\
                 height: {prev_end_height:?}\n\
                 current subtree:\n\
                 height: {end_height:?}"
-            ));
+        ));
         let prev_index = prev_tree
             .subtree_index()
             .expect("previous block must have a partial subtree");
         let prev_remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
 
-        let current_position = tree
-            .position()
-            .expect(&format!(
-                "current block must have a subtree:\n\
+        let current_position = tree.position().expect(&format!(
+            "current block must have a subtree:\n\
                 previous subtree:\n\
                 height: {prev_end_height:?}\n\
                 index: {prev_index}\n\
@@ -727,7 +719,7 @@ fn calculate_orchard_subtree(
                 remaining: {prev_remaining_notes}\n\
                 current subtree:\n\
                 height: {end_height:?}"
-            ));
+        ));
         let current_index = tree
             .subtree_index()
             .expect("current block must have a subtree");

From 96ace9016e86d5d2117c10f0ca92a9250592edf5 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Mon, 18 Sep 2023 06:24:25 +1000
Subject: [PATCH 43/56] Fix unrelated clippy::unnecessary_unwrap

---
 zebra-chain/src/block/arbitrary.rs | 22 ++++++++++------------
 1 file changed, 10 insertions(+), 12 deletions(-)

diff --git a/zebra-chain/src/block/arbitrary.rs b/zebra-chain/src/block/arbitrary.rs
index 36734c86b24..dc06d4ff43e 100644
--- a/zebra-chain/src/block/arbitrary.rs
+++ b/zebra-chain/src/block/arbitrary.rs
@@ -517,7 +517,16 @@ impl Block {
                         }
                     }
                     // update history tree for the next block
-                    if history_tree.is_none() {
+                    if let Some(history_tree) = history_tree.as_mut() {
+                        history_tree
+                            .push(
+                                current.network,
+                                Arc::new(block.clone()),
+                                sapling_tree.root(),
+                                orchard_tree.root(),
+                            )
+                            .unwrap();
+                    } else {
                         history_tree = Some(
                             HistoryTree::from_block(
                                 current.network,
@@ -527,17 +536,6 @@ impl Block {
                             )
                             .unwrap(),
                         );
-                    } else {
-                        history_tree
-                            .as_mut()
-                            .unwrap()
-                            .push(
-                                current.network,
-                                Arc::new(block.clone()),
-                                sapling_tree.root(),
-                                orchard_tree.root(),
-                            )
-                            .unwrap();
                     }
                 }
 

From d33f497da98cfbf2c522d648e418f64254172d51 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Mon, 18 Sep 2023 06:27:45 +1000
Subject: [PATCH 44/56] cargo clippy --fix --all-features --all-targets; cargo
 fmt --all

---
 .../disk_format/upgrade/add_subtrees.rs       | 32 ++++++++++++-------
 1 file changed, 20 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 3171aacc347..6598c4e72ea 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
@@ -577,20 +577,23 @@ fn calculate_sapling_subtree(
         // we need to calculate the subtree root, starting with the tree from the previous block.
 
         // TODO: move the assertion/panic log string formatting into a separate function?
-        let prev_position = prev_tree.position().expect(&format!(
-            "previous block must have a partial subtree:\n\
+        let prev_position = prev_tree.position().unwrap_or_else(|| {
+            panic!(
+                "previous block must have a partial subtree:\n\
                 previous subtree:\n\
                 height: {prev_end_height:?}\n\
                 current subtree:\n\
                 height: {end_height:?}"
-        ));
+            )
+        });
         let prev_index = prev_tree
             .subtree_index()
             .expect("previous block must have a partial subtree");
         let prev_remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
 
-        let current_position = tree.position().expect(&format!(
-            "current block must have a subtree:\n\
+        let current_position = tree.position().unwrap_or_else(|| {
+            panic!(
+                "current block must have a subtree:\n\
                 previous subtree:\n\
                 height: {prev_end_height:?}\n\
                 index: {prev_index}\n\
@@ -598,7 +601,8 @@ fn calculate_sapling_subtree(
                 remaining: {prev_remaining_notes}\n\
                 current subtree:\n\
                 height: {end_height:?}"
-        ));
+            )
+        });
         let current_index = tree
             .subtree_index()
             .expect("current block must have a subtree");
@@ -698,20 +702,23 @@ fn calculate_orchard_subtree(
         // we need to calculate the subtree root, starting with the tree from the previous block.
 
         // TODO: move the assertion/panic log string formatting into a separate function?
-        let prev_position = prev_tree.position().expect(&format!(
-            "previous block must have a partial subtree:\n\
+        let prev_position = prev_tree.position().unwrap_or_else(|| {
+            panic!(
+                "previous block must have a partial subtree:\n\
                 previous subtree:\n\
                 height: {prev_end_height:?}\n\
                 current subtree:\n\
                 height: {end_height:?}"
-        ));
+            )
+        });
         let prev_index = prev_tree
             .subtree_index()
             .expect("previous block must have a partial subtree");
         let prev_remaining_notes = prev_tree.remaining_subtree_leaf_nodes();
 
-        let current_position = tree.position().expect(&format!(
-            "current block must have a subtree:\n\
+        let current_position = tree.position().unwrap_or_else(|| {
+            panic!(
+                "current block must have a subtree:\n\
                 previous subtree:\n\
                 height: {prev_end_height:?}\n\
                 index: {prev_index}\n\
@@ -719,7 +726,8 @@ fn calculate_orchard_subtree(
                 remaining: {prev_remaining_notes}\n\
                 current subtree:\n\
                 height: {end_height:?}"
-        ));
+            )
+        });
         let current_index = tree
             .subtree_index()
             .expect("current block must have a subtree");

From 7b0baf884914e79c05dd29754b04809b412c593d Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Mon, 18 Sep 2023 07:47:41 +1000
Subject: [PATCH 45/56] Stop the quick check depending on tree de-duplication

---
 zebra-chain/src/orchard/tree.rs               | 11 +++++++---
 zebra-chain/src/sapling/tree.rs               | 11 +++++++---
 .../disk_format/upgrade/add_subtrees.rs       | 21 ++-----------------
 3 files changed, 18 insertions(+), 25 deletions(-)

diff --git a/zebra-chain/src/orchard/tree.rs b/zebra-chain/src/orchard/tree.rs
index d0c2146915f..6992a7ba9c1 100644
--- a/zebra-chain/src/orchard/tree.rs
+++ b/zebra-chain/src/orchard/tree.rs
@@ -406,7 +406,7 @@ impl NoteCommitmentTree {
         Some(tree.position().into())
     }
 
-    /// Returns true if this tree has a new subtree, when compared with `prev_tree`.
+    /// Returns true if this tree has at least one new subtree, when compared with `prev_tree`.
     pub fn contains_new_subtree(&self, prev_tree: &Self) -> bool {
         // A new subtree was completed by the last note commitment in this tree.
         if self.is_complete_subtree() {
@@ -427,11 +427,13 @@ impl NoteCommitmentTree {
         // This calculation can't overflow, because we're using i32 for u16 values.
         let index_difference = index - prev_index;
 
-        // There are multiple new subtrees.
+        // There are multiple new subtrees, or one new subtree and a spurious index difference.
         if index_difference > 1 {
             return true;
         }
 
+        // Check for spurious index differences.
+        //
         // There is one new subtree somewhere in the trees. It is either:
         // - a new subtree at the end of the previous tree, or
         // - a new subtree in this tree (but not at the end).
@@ -439,7 +441,10 @@ impl NoteCommitmentTree {
         // This happens because the subtree index only increases when the first note is added to
         // the new subtree. So we need to exclude subtrees completed by the last note commitment in
         // the previous tree.
-        if prev_tree.is_complete_subtree() {
+        //
+        // We also need to exclude empty previous subtrees, because the index changes to zero when
+        // the first note is added, but a subtree wasn't completed.
+        if prev_tree.is_complete_subtree() || prev_index == -1 {
             return false;
         }
 
diff --git a/zebra-chain/src/sapling/tree.rs b/zebra-chain/src/sapling/tree.rs
index 5ef07b2df94..093cdb2e6aa 100644
--- a/zebra-chain/src/sapling/tree.rs
+++ b/zebra-chain/src/sapling/tree.rs
@@ -387,7 +387,7 @@ impl NoteCommitmentTree {
         Some(tree.position().into())
     }
 
-    /// Returns true if this tree has a new subtree, when compared with `prev_tree`.
+    /// Returns true if this tree has at least one new subtree, when compared with `prev_tree`.
     pub fn contains_new_subtree(&self, prev_tree: &Self) -> bool {
         // A new subtree was completed by the last note commitment in this tree.
         if self.is_complete_subtree() {
@@ -408,11 +408,13 @@ impl NoteCommitmentTree {
         // This calculation can't overflow, because we're using i32 for u16 values.
         let index_difference = index - prev_index;
 
-        // There are multiple new subtrees.
+        // There are multiple new subtrees, or one new subtree and a spurious index difference.
         if index_difference > 1 {
             return true;
         }
 
+        // Check for spurious index differences.
+        //
         // There is one new subtree somewhere in the trees. It is either:
         // - a new subtree at the end of the previous tree, or
         // - a new subtree in this tree (but not at the end).
@@ -420,7 +422,10 @@ impl NoteCommitmentTree {
         // This happens because the subtree index only increases when the first note is added to
         // the new subtree. So we need to exclude subtrees completed by the last note commitment in
         // the previous tree.
-        if prev_tree.is_complete_subtree() {
+        //
+        // We also need to exclude empty previous subtrees, because the index changes to zero when
+        // the first note is added, but a subtree wasn't completed.
+        if prev_tree.is_complete_subtree() || prev_index == -1 {
             return false;
         }
 
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 6598c4e72ea..792430803f1 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
@@ -41,15 +41,10 @@ pub fn run(
     // block can't complete multiple level 16 subtrees (or complete an entire subtree by itself).
     // Currently, with 2MB blocks and v4/v5 sapling and orchard output sizes, the subtree index can
     // increase by at most 1 every ~20 blocks.
-    //
-    // Therefore, the first block with shielded note can't complete a subtree, which means we can
-    // skip the (genesis block, first shielded block) tree pair.
 
     // Generate a list of sapling subtree inputs: previous and current trees, and their end heights.
     let subtrees = upgrade_db
         .sapling_tree_by_height_range(..=initial_tip_height)
-        // The first block with sapling notes can't complete a subtree, see above for details.
-        .filter(|(height, _tree)| !height.is_min())
         // We need both the tree and its previous tree for each shielded block.
         .tuple_windows()
         .map(|((prev_end_height, prev_tree), (end_height, tree))| {
@@ -74,8 +69,6 @@ pub fn run(
     // Generate a list of orchard subtree inputs: previous and current trees, and their end heights.
     let subtrees = upgrade_db
         .orchard_tree_by_height_range(..=initial_tip_height)
-        // The first block with orchard notes can't complete a subtree, see above for details.
-        .filter(|(height, _tree)| !height.is_min())
         // We need both the tree and its previous tree for each shielded block.
         .tuple_windows()
         .map(|((prev_end_height, prev_tree), (end_height, tree))| {
@@ -209,19 +202,14 @@ fn quick_check_sapling_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
         return Ok(());
     }
 
-    // Find the first complete subtree, with its note commitment tree, end height, and the previous tree.
+    // Find the first complete subtree: previous and current trees, and their end heights.
     let first_complete_subtree = db
         .sapling_tree_by_height_range(..)
-        // The first block with sapling notes can't complete a subtree, see above for details.
-        .filter(|(height, _tree)| !height.is_min())
         // We need both the tree and its previous tree for each shielded block.
         .tuple_windows()
         .map(|((prev_end_height, prev_tree), (end_height, tree))| {
             (prev_end_height, prev_tree, end_height, tree)
         })
-        // Empty note commitment trees can't contain subtrees, so they have invalid subtree indexes.
-        // But since we skip the empty genesis tree, all trees must have valid indexes.
-        // So we don't need to unwrap the optional values for this comparison to be correct.
         .find(|(_prev_end_height, prev_tree, _end_height, tree)| {
             tree.contains_new_subtree(prev_tree)
         });
@@ -264,19 +252,14 @@ fn quick_check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
         return Ok(());
     }
 
-    // Find the first complete subtree, with its note commitment tree, end height, and the previous tree.
+    // Find the first complete subtree: previous and current trees, and their end heights.
     let first_complete_subtree = db
         .orchard_tree_by_height_range(..)
-        // The first block with orchard notes can't complete a subtree, see above for details.
-        .filter(|(height, _tree)| !height.is_min())
         // We need both the tree and its previous tree for each shielded block.
         .tuple_windows()
         .map(|((prev_end_height, prev_tree), (end_height, tree))| {
             (prev_end_height, prev_tree, end_height, tree)
         })
-        // Empty note commitment trees can't contain subtrees, so they have invalid subtree indexes.
-        // But since we skip the empty genesis tree, all trees must have valid indexes.
-        // So we don't need to unwrap the optional values for this comparison to be correct.
         .find(|(_prev_end_height, prev_tree, _end_height, tree)| {
             tree.contains_new_subtree(prev_tree)
         });

From e3b014a1018da13d44843789498b88a979c10b45 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Mon, 18 Sep 2023 08:34:57 +1000
Subject: [PATCH 46/56] Refactor waiting for the upgrade into functions

---
 .../finalized_state/disk_format/upgrade.rs    |  3 +-
 zebrad/tests/common/cached_state.rs           | 55 +++++++++++++++++++
 .../common/lightwalletd/wallet_grpc_test.rs   | 42 +++++---------
 3 files changed, 71 insertions(+), 29 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 c084b4385aa..e2903a4855f 100644
--- a/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
+++ b/zebra-state/src/service/finalized_state/disk_format/upgrade.rs
@@ -547,7 +547,8 @@ impl DbFormatChange {
         info!(
             %running_version,
             %disk_version,
-            // The wallet_grpc_test needs this to be the last field, so the regex matches correctly
+            // wait_for_state_version_upgrade() needs this to be the last field,
+            // so the regex matches correctly
             %format_upgrade_version,
             "marked database format as upgraded"
         );
diff --git a/zebrad/tests/common/cached_state.rs b/zebrad/tests/common/cached_state.rs
index 284add1c0df..3e293321636 100644
--- a/zebrad/tests/common/cached_state.rs
+++ b/zebrad/tests/common/cached_state.rs
@@ -11,6 +11,7 @@ use std::{
 };
 
 use color_eyre::eyre::{eyre, Result};
+use semver::Version;
 use tower::{util::BoxService, Service};
 
 use zebra_chain::{
@@ -21,6 +22,7 @@ use zebra_chain::{
 };
 use zebra_node_services::rpc_client::RpcRequestClient;
 use zebra_state::{ChainTipChange, LatestChainTip, MAX_BLOCK_REORG_HEIGHT};
+use zebra_test::command::TestChild;
 
 use crate::common::{
     launch::spawn_zebrad_for_rpc,
@@ -35,6 +37,59 @@ pub const ZEBRA_CACHED_STATE_DIR: &str = "ZEBRA_CACHED_STATE_DIR";
 pub type BoxStateService =
     BoxService<zebra_state::Request, zebra_state::Response, zebra_state::BoxError>;
 
+/// Waits for the startup logs generated by the cached state version checks.
+/// Returns the state version log message.
+///
+/// This function should be called immediately after launching `zebrad`.
+#[tracing::instrument(skip(zebrad))]
+pub fn wait_for_state_version_message<T>(zebrad: &mut TestChild<T>) -> Result<String> {
+    tracing::info!(
+        zebrad = ?zebrad.cmd,
+        "launched zebrad, waiting for zebrad to open the state database..."
+    );
+
+    // Zebra logs one of these lines on startup, depending on the disk and running formats.
+    zebrad.expect_stdout_line_matches(
+        "(creating new database with the current format)|\
+         (trying to open older database format)|\
+         (trying to open newer database format)|\
+         (trying to open current database format)",
+    )
+}
+
+/// Waits for the `required_version` state upgrade to complete, if needed.
+///
+/// This function should be called with the output of [`wait_for_state_version_message()`].
+#[tracing::instrument(skip(zebrad))]
+pub fn wait_for_state_version_upgrade<T>(
+    zebrad: &mut TestChild<T>,
+    state_version_message: &str,
+    required_version: Version,
+) -> Result<()> {
+    if state_version_message.contains("launching upgrade task") {
+        tracing::info!(
+            zebrad = ?zebrad.cmd,
+            %state_version_message,
+            %required_version,
+            "waiting for zebrad state upgrade..."
+        );
+
+        let upgrade_message = zebrad.expect_stdout_line_matches(&format!(
+            "marked database format as upgraded.*format_upgrade_version.*=.*{required_version}"
+        ))?;
+
+        tracing::info!(
+            zebrad = ?zebrad.cmd,
+            %state_version_message,
+            %required_version,
+            %upgrade_message,
+            "zebrad state has been upgraded"
+        );
+    }
+
+    Ok(())
+}
+
 /// Starts a state service using the provided `cache_dir` as the directory with the chain state.
 #[tracing::instrument(skip(cache_dir))]
 pub async fn start_state_service_with_cache_dir(
diff --git a/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs b/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
index 7f2ea84286a..33c566d0cf5 100644
--- a/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
+++ b/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
@@ -46,6 +46,7 @@ use zebra_chain::{
 use zebra_state::latest_version_for_adding_subtrees;
 
 use crate::common::{
+    cached_state::{wait_for_state_version_message, wait_for_state_version_upgrade},
     launch::spawn_zebrad_for_rpc,
     lightwalletd::{
         can_spawn_lightwalletd_for_rpc, spawn_lightwalletd_for_rpc,
@@ -98,17 +99,8 @@ pub async fn run() -> Result<()> {
 
     let zebra_rpc_address = zebra_rpc_address.expect("lightwalletd test must have RPC port");
 
-    tracing::info!(
-        ?test_type,
-        "launched zebrad, waiting for zebrad to open the state database..."
-    );
-    // Zebra logs one of these lines on startup, depending on the disk and running formats.
-    let state_format_change = zebrad.expect_stdout_line_matches(
-        "(creating new database with the current format)|\
-         (trying to open older database format)|\
-         (trying to open newer database format)|\
-         (trying to open current database format)",
-    )?;
+    // Store the state version message so we can wait for the upgrade later if needed.
+    let state_version_message = wait_for_state_version_message(&mut zebrad)?;
 
     tracing::info!(
         ?test_type,
@@ -393,24 +385,18 @@ pub async fn run() -> Result<()> {
     );
 
     // Before we call `z_getsubtreesbyindex`, we might need to wait for a database upgrade.
-    if state_format_change.contains("launching upgrade task") {
-        let latest_version_for_adding_subtrees = latest_version_for_adding_subtrees();
-        tracing::info!(
-            ?test_type,
-            ?zebra_rpc_address,
-            %state_format_change,
-            %latest_version_for_adding_subtrees,
-            "waiting for zebrad subtree state upgrade..."
-        );
-        //
-        // TODO: this line will hang if the state upgrade finishes before the subtree tests start.
-        // But that is unlikely with the 25.2 upgrade, because it takes 20+ minutes.
-        // If it happens for a later upgrade, this code can be removed, as long as all the
-        // lightwalletd cached states are version 25.2.2 or later.
-        zebrad.expect_stdout_line_matches(&format!("marked database format as upgraded.*format_upgrade_version.*=.*{latest_version_for_adding_subtrees}"))?;
-    }
+    //
+    // TODO: this line will hang if the state upgrade finishes before the subtree tests start.
+    // But that is unlikely with the 25.2 upgrade, because it takes 20+ minutes.
+    // If it happens for a later upgrade, this code can be removed, as long as all the
+    // lightwalletd cached states are version 25.2.2 or later.
+    wait_for_state_version_upgrade(
+        &mut zebrad,
+        &state_version_message,
+        latest_version_for_adding_subtrees(),
+    )?;
 
-    // Call `z_getsubtreesbyindex` separately for
+    // Call `z_getsubtreesbyindex` separately for...
 
     // ... Sapling.
     let mut subtrees = rpc_client

From be6a33ee6598f7cd86cbd1da7d961a606c88d07b Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Mon, 18 Sep 2023 08:48:53 +1000
Subject: [PATCH 47/56] Wait for state upgrades whenever the cached state is
 updated

---
 zebrad/tests/acceptance.rs                    | 32 ++++++++++++++++++-
 zebrad/tests/common/checkpoints.rs            | 20 +++++++++++-
 .../common/lightwalletd/wallet_grpc_test.rs   |  4 +--
 3 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/zebrad/tests/acceptance.rs b/zebrad/tests/acceptance.rs
index 17abe47d1c4..f241ca67c90 100644
--- a/zebrad/tests/acceptance.rs
+++ b/zebrad/tests/acceptance.rs
@@ -184,6 +184,8 @@ use common::{
     test_type::TestType::{self, *},
 };
 
+use crate::common::cached_state::{wait_for_state_version_message, wait_for_state_version_upgrade};
+
 /// The maximum amount of time that we allow the creation of a future to block the `tokio` executor.
 ///
 /// This should be larger than the amount of time between thread time slices on a busy test VM.
@@ -1780,6 +1782,9 @@ fn lightwalletd_integration_test(test_type: TestType) -> Result<()> {
         return Ok(());
     };
 
+    // Store the state version message so we can wait for the upgrade later if needed.
+    let state_version_message = wait_for_state_version_message(&mut zebrad)?;
+
     if test_type.needs_zebra_cached_state() {
         zebrad
             .expect_stdout_line_matches(r"loaded Zebra state cache .*tip.*=.*Height\([0-9]{7}\)")?;
@@ -1875,6 +1880,7 @@ fn lightwalletd_integration_test(test_type: TestType) -> Result<()> {
         None
     };
 
+    // Wait for zebrad and lightwalletd to sync, if needed.
     let (mut zebrad, lightwalletd) = if test_type.needs_zebra_cached_state() {
         if let Some((lightwalletd, lightwalletd_rpc_port)) = lightwalletd_and_port {
             #[cfg(feature = "lightwalletd-grpc-tests")]
@@ -1886,7 +1892,7 @@ fn lightwalletd_integration_test(test_type: TestType) -> Result<()> {
                     "waiting for zebrad and lightwalletd to sync...",
                 );
 
-                let (lightwalletd, zebrad) = wait_for_zebrad_and_lightwalletd_sync(
+                let (lightwalletd, mut zebrad) = wait_for_zebrad_and_lightwalletd_sync(
                     lightwalletd,
                     lightwalletd_rpc_port,
                     zebrad,
@@ -1897,6 +1903,18 @@ fn lightwalletd_integration_test(test_type: TestType) -> Result<()> {
                     use_internet_connection,
                 )?;
 
+                // Before we write a cached state image, wait for a database upgrade.
+                //
+                // TODO: this line will hang if the state upgrade finishes before zebra is synced.
+                // But that is unlikely with the 25.2 upgrade, because it takes 20+ minutes.
+                // If it happens for a later upgrade, this code can be moved earlier in the test,
+                // as long as all the cached states are version 25.2.2 or later.
+                wait_for_state_version_upgrade(
+                    &mut zebrad,
+                    &state_version_message,
+                    database_format_version_in_code(),
+                )?;
+
                 (zebrad, Some(lightwalletd))
             }
 
@@ -1912,6 +1930,18 @@ fn lightwalletd_integration_test(test_type: TestType) -> Result<()> {
             tracing::info!(?test_type, "waiting for zebrad to sync to the tip");
             zebrad.expect_stdout_line_matches(SYNC_FINISHED_REGEX)?;
 
+            // Before we write a cached state image, wait for a database upgrade.
+            //
+            // TODO: this line will hang if the state upgrade finishes before zebra is synced.
+            // But that is unlikely with the 25.2 upgrade, because it takes 20+ minutes.
+            // If it happens for a later upgrade, this code can be moved earlier in the test,
+            // as long as all the cached states are version 25.2.2 or later.
+            wait_for_state_version_upgrade(
+                &mut zebrad,
+                &state_version_message,
+                database_format_version_in_code(),
+            )?;
+
             (zebrad, None)
         }
     } else {
diff --git a/zebrad/tests/common/checkpoints.rs b/zebrad/tests/common/checkpoints.rs
index cc5e6be40f9..77ad12e4430 100644
--- a/zebrad/tests/common/checkpoints.rs
+++ b/zebrad/tests/common/checkpoints.rs
@@ -20,6 +20,7 @@ use zebra_chain::{
 };
 use zebra_consensus::MAX_CHECKPOINT_HEIGHT_GAP;
 use zebra_node_services::rpc_client::RpcRequestClient;
+use zebra_state::database_format_version_in_code;
 use zebra_test::{
     args,
     command::{Arguments, TestDirExt, NO_MATCHES_REGEX_ITER},
@@ -27,6 +28,7 @@ use zebra_test::{
 };
 
 use crate::common::{
+    cached_state::{wait_for_state_version_message, wait_for_state_version_upgrade},
     launch::spawn_zebrad_for_rpc,
     sync::{CHECKPOINT_VERIFIER_REGEX, SYNC_FINISHED_REGEX},
     test_type::TestType::*,
@@ -77,6 +79,9 @@ pub async fn run(network: Network) -> Result<()> {
         return Ok(());
     };
 
+    // Store the state version message so we can wait for the upgrade later if needed.
+    let state_version_message = wait_for_state_version_message(&mut zebrad)?;
+
     let zebra_rpc_address = zebra_rpc_address.expect("zebra_checkpoints test must have RPC port");
 
     tracing::info!(
@@ -147,7 +152,7 @@ pub async fn run(network: Network) -> Result<()> {
     );
     println!("\n\n");
 
-    let (_zebra_checkpoints, _zebrad) = wait_for_zebra_checkpoints_generation(
+    let (_zebra_checkpoints, mut zebrad) = wait_for_zebra_checkpoints_generation(
         zebra_checkpoints,
         zebrad,
         zebra_tip_height,
@@ -163,6 +168,19 @@ pub async fn run(network: Network) -> Result<()> {
         "finished generating Zebra checkpoints",
     );
 
+    // Before we write a cached state image, wait for a database upgrade.
+    // Currently we only write an image for testnet, but waiting also improves test coverage.
+    //
+    // TODO: this line will hang if the state upgrade finishes before zebra is synced.
+    // But that is unlikely with the 25.2 upgrade, because it takes 20+ minutes.
+    // If it happens for a later upgrade, this code can be moved earlier in the test,
+    // as long as all the cached states are version 25.2.2 or later.
+    wait_for_state_version_upgrade(
+        &mut zebrad,
+        &state_version_message,
+        database_format_version_in_code(),
+    )?;
+
     Ok(())
 }
 
diff --git a/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs b/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
index 33c566d0cf5..8aad372105e 100644
--- a/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
+++ b/zebrad/tests/common/lightwalletd/wallet_grpc_test.rs
@@ -388,8 +388,8 @@ pub async fn run() -> Result<()> {
     //
     // TODO: this line will hang if the state upgrade finishes before the subtree tests start.
     // But that is unlikely with the 25.2 upgrade, because it takes 20+ minutes.
-    // If it happens for a later upgrade, this code can be removed, as long as all the
-    // lightwalletd cached states are version 25.2.2 or later.
+    // If it happens for a later upgrade, this code can be moved earlier in the test,
+    // as long as all the cached states are version 25.2.2 or later.
     wait_for_state_version_upgrade(
         &mut zebrad,
         &state_version_message,

From 171998b70a57b2542fe7561c307d9b3fc53e7cca Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Mon, 18 Sep 2023 11:04:46 +1000
Subject: [PATCH 48/56] Wait for the testnet upgrade in the right place

---
 zebrad/tests/common/checkpoints.rs | 33 +++++++++++++++---------------
 1 file changed, 17 insertions(+), 16 deletions(-)

diff --git a/zebrad/tests/common/checkpoints.rs b/zebrad/tests/common/checkpoints.rs
index 77ad12e4430..f2337a4d679 100644
--- a/zebrad/tests/common/checkpoints.rs
+++ b/zebrad/tests/common/checkpoints.rs
@@ -15,7 +15,7 @@ use tempfile::TempDir;
 
 use zebra_chain::{
     block::{Height, HeightDiff, TryIntoHeight},
-    parameters::Network,
+    parameters::Network::{self, *},
     transparent::MIN_TRANSPARENT_COINBASE_MATURITY,
 };
 use zebra_consensus::MAX_CHECKPOINT_HEIGHT_GAP;
@@ -79,8 +79,22 @@ pub async fn run(network: Network) -> Result<()> {
         return Ok(());
     };
 
-    // Store the state version message so we can wait for the upgrade later if needed.
-    let state_version_message = wait_for_state_version_message(&mut zebrad)?;
+    // Wait for the upgrade if needed.
+    // Currently we only write an image for testnet, which is quick.
+    // (Mainnet would need to wait at the end of this function, if the upgrade is long.)
+    if network == Testnet {
+        let state_version_message = wait_for_state_version_message(&mut zebrad)?;
+
+        // Before we write a cached state image, wait for a database upgrade.
+        //
+        // TODO: this line will hang if the state upgrade is slower than the RPC server spawn.
+        // But that is unlikely, because both 25.1 and 25.2 are quick on testnet.
+        wait_for_state_version_upgrade(
+            &mut zebrad,
+            &state_version_message,
+            database_format_version_in_code(),
+        )?;
+    }
 
     let zebra_rpc_address = zebra_rpc_address.expect("zebra_checkpoints test must have RPC port");
 
@@ -168,19 +182,6 @@ pub async fn run(network: Network) -> Result<()> {
         "finished generating Zebra checkpoints",
     );
 
-    // Before we write a cached state image, wait for a database upgrade.
-    // Currently we only write an image for testnet, but waiting also improves test coverage.
-    //
-    // TODO: this line will hang if the state upgrade finishes before zebra is synced.
-    // But that is unlikely with the 25.2 upgrade, because it takes 20+ minutes.
-    // If it happens for a later upgrade, this code can be moved earlier in the test,
-    // as long as all the cached states are version 25.2.2 or later.
-    wait_for_state_version_upgrade(
-        &mut zebrad,
-        &state_version_message,
-        database_format_version_in_code(),
-    )?;
-
     Ok(())
 }
 

From 82bf50a37f5da716957e54fa85c205eee555869e Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Mon, 18 Sep 2023 11:10:01 +1000
Subject: [PATCH 49/56] Fix unused variable

---
 zebrad/tests/common/checkpoints.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/zebrad/tests/common/checkpoints.rs b/zebrad/tests/common/checkpoints.rs
index f2337a4d679..c43b5ca7e92 100644
--- a/zebrad/tests/common/checkpoints.rs
+++ b/zebrad/tests/common/checkpoints.rs
@@ -166,7 +166,7 @@ pub async fn run(network: Network) -> Result<()> {
     );
     println!("\n\n");
 
-    let (_zebra_checkpoints, mut zebrad) = wait_for_zebra_checkpoints_generation(
+    let (_zebra_checkpoints, _zebrad) = wait_for_zebra_checkpoints_generation(
         zebra_checkpoints,
         zebrad,
         zebra_tip_height,

From 5a1deb1865071977df36dcb10e64c9e328f8c322 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 19 Sep 2023 08:14:03 +1000
Subject: [PATCH 50/56] Fix a subtree detection bug and comments

---
 zebra-chain/src/orchard/tree.rs                 | 17 +++++------------
 zebra-chain/src/sapling/tree.rs                 | 17 +++++------------
 .../disk_format/upgrade/add_subtrees.rs         |  8 ++++----
 3 files changed, 14 insertions(+), 28 deletions(-)

diff --git a/zebra-chain/src/orchard/tree.rs b/zebra-chain/src/orchard/tree.rs
index 6992a7ba9c1..0c5b7531f36 100644
--- a/zebra-chain/src/orchard/tree.rs
+++ b/zebra-chain/src/orchard/tree.rs
@@ -408,11 +408,6 @@ impl NoteCommitmentTree {
 
     /// Returns true if this tree has at least one new subtree, when compared with `prev_tree`.
     pub fn contains_new_subtree(&self, prev_tree: &Self) -> bool {
-        // A new subtree was completed by the last note commitment in this tree.
-        if self.is_complete_subtree() {
-            return true;
-        }
-
         // Use -1 for the index of the subtree with no notes, so the comparisons are valid.
         let index = self.subtree_index().map_or(-1, |index| i32::from(index.0));
         let prev_index = prev_tree
@@ -438,9 +433,9 @@ impl NoteCommitmentTree {
         // - a new subtree at the end of the previous tree, or
         // - a new subtree in this tree (but not at the end).
         //
-        // This happens because the subtree index only increases when the first note is added to
-        // the new subtree. So we need to exclude subtrees completed by the last note commitment in
-        // the previous tree.
+        // Spurious index differences happen because the subtree index only increases when the
+        // first note is added to the new subtree. So we need to exclude subtrees completed by the
+        // last note commitment in the previous tree.
         //
         // We also need to exclude empty previous subtrees, because the index changes to zero when
         // the first note is added, but a subtree wasn't completed.
@@ -502,12 +497,10 @@ impl NoteCommitmentTree {
             // If the subtree has no nodes, the remaining number of nodes is the number of nodes in
             // a subtree.
             None => {
-                // This position is guaranteed to be in the first subtree.
-                let first_position = 0.into();
-
                 let subtree_address = incrementalmerkletree::Address::above_position(
                     TRACKED_SUBTREE_HEIGHT.into(),
-                    first_position,
+                    // This position is guaranteed to be in the first subtree.
+                    0.into(),
                 );
 
                 assert_eq!(
diff --git a/zebra-chain/src/sapling/tree.rs b/zebra-chain/src/sapling/tree.rs
index 093cdb2e6aa..ade7a445c6d 100644
--- a/zebra-chain/src/sapling/tree.rs
+++ b/zebra-chain/src/sapling/tree.rs
@@ -389,11 +389,6 @@ impl NoteCommitmentTree {
 
     /// Returns true if this tree has at least one new subtree, when compared with `prev_tree`.
     pub fn contains_new_subtree(&self, prev_tree: &Self) -> bool {
-        // A new subtree was completed by the last note commitment in this tree.
-        if self.is_complete_subtree() {
-            return true;
-        }
-
         // Use -1 for the index of the subtree with no notes, so the comparisons are valid.
         let index = self.subtree_index().map_or(-1, |index| i32::from(index.0));
         let prev_index = prev_tree
@@ -419,9 +414,9 @@ impl NoteCommitmentTree {
         // - a new subtree at the end of the previous tree, or
         // - a new subtree in this tree (but not at the end).
         //
-        // This happens because the subtree index only increases when the first note is added to
-        // the new subtree. So we need to exclude subtrees completed by the last note commitment in
-        // the previous tree.
+        // Spurious index differences happen because the subtree index only increases when the
+        // first note is added to the new subtree. So we need to exclude subtrees completed by the
+        // last note commitment in the previous tree.
         //
         // We also need to exclude empty previous subtrees, because the index changes to zero when
         // the first note is added, but a subtree wasn't completed.
@@ -483,12 +478,10 @@ impl NoteCommitmentTree {
             // If the subtree has no nodes, the remaining number of nodes is the number of nodes in
             // a subtree.
             None => {
-                // This position is guaranteed to be in the first subtree.
-                let first_position = 0.into();
-
                 let subtree_address = incrementalmerkletree::Address::above_position(
                     TRACKED_SUBTREE_HEIGHT.into(),
-                    first_position,
+                    // This position is guaranteed to be in the first subtree.
+                    0.into(),
                 );
 
                 assert_eq!(
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 792430803f1..7bd25d28c3e 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
@@ -351,7 +351,7 @@ fn check_sapling_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
         // 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 {
-                result = Err("missing note commitment tree at subtree completion height");
+                result = Err("missing note commitment tree below subtree completion height");
                 error!(?result, ?subtree.end);
                 continue;
             };
@@ -463,7 +463,7 @@ fn check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
         // 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 {
-                result = Err("missing note commitment tree at subtree completion height");
+                result = Err("missing note commitment tree below subtree completion height");
                 error!(?result, ?subtree.end);
                 continue;
             };
@@ -549,7 +549,7 @@ fn calculate_sapling_subtree(
     end_height: Height,
     tree: Arc<sapling::tree::NoteCommitmentTree>,
 ) -> NoteCommitmentSubtree<sapling::tree::Node> {
-    // If a subtree is completed by a note commentment in the block at `end_height`,
+    // If a subtree is completed by a note commitment in the block at `end_height`,
     // then that subtree can be completed in two different ways:
     if let Some((index, node)) = tree.completed_subtree_index_and_root() {
         // If the subtree is completed by the last note commitment in that block,
@@ -674,7 +674,7 @@ fn calculate_orchard_subtree(
     end_height: Height,
     tree: Arc<orchard::tree::NoteCommitmentTree>,
 ) -> NoteCommitmentSubtree<orchard::tree::Node> {
-    // If a subtree is completed by a note commentment in the block at `end_height`,
+    // If a subtree is completed by a note commitment in the block at `end_height`,
     // then that subtree can be completed in two different ways:
     if let Some((index, node)) = tree.completed_subtree_index_and_root() {
         // If the subtree is completed by the last note commitment in that block,

From c2064043776a11ef45fbe98d17ffc55e2be31f36 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 19 Sep 2023 08:14:28 +1000
Subject: [PATCH 51/56] Remove an early reference to reverse direction

---
 .../finalized_state/disk_format/upgrade/add_subtrees.rs        | 3 ---
 1 file changed, 3 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 7bd25d28c3e..fe9b1b583fe 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
@@ -21,9 +21,6 @@ use crate::service::finalized_state::{
 
 /// Runs disk format upgrade for adding Sapling and Orchard note commitment subtrees to database.
 ///
-/// Trees are added to the database in reverse height order, so that wallets can sync correctly
-/// while the upgrade is running.
-///
 /// Returns `Ok` if the upgrade completed, and `Err` if it was cancelled.
 #[allow(clippy::unwrap_in_result)]
 #[instrument(skip(upgrade_db, cancel_receiver))]

From 41eef096793dca9ec4d94318f117ec2cae675ad5 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 19 Sep 2023 12:27:57 +1000
Subject: [PATCH 52/56] Stop skipping subtrees completed at the end of blocks

---
 zebra-chain/src/orchard/tree.rs                |  9 +++++++--
 zebra-chain/src/sapling/tree.rs                |  9 +++++++--
 .../disk_format/upgrade/add_subtrees.rs        | 18 ++++++++++++------
 3 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/zebra-chain/src/orchard/tree.rs b/zebra-chain/src/orchard/tree.rs
index 0c5b7531f36..74a4201eecf 100644
--- a/zebra-chain/src/orchard/tree.rs
+++ b/zebra-chain/src/orchard/tree.rs
@@ -414,11 +414,16 @@ impl NoteCommitmentTree {
             .subtree_index()
             .map_or(-1, |index| i32::from(index.0));
 
-        // If the index is equal or lower, there can't be any new subtrees.
-        if index <= prev_index {
+        // There can't be any new subtrees if the current index is strictly lower.
+        if index < prev_index {
             return false;
         }
 
+        // If the indexes are equal, there can only be a new subtree if `self` just completed it.
+        if index == prev_index && self.is_complete_subtree() {
+            return true;
+        }
+
         // This calculation can't overflow, because we're using i32 for u16 values.
         let index_difference = index - prev_index;
 
diff --git a/zebra-chain/src/sapling/tree.rs b/zebra-chain/src/sapling/tree.rs
index ade7a445c6d..8434fa0c721 100644
--- a/zebra-chain/src/sapling/tree.rs
+++ b/zebra-chain/src/sapling/tree.rs
@@ -395,11 +395,16 @@ impl NoteCommitmentTree {
             .subtree_index()
             .map_or(-1, |index| i32::from(index.0));
 
-        // If the index is equal or lower, there can't be any new subtrees.
-        if index <= prev_index {
+        // There can't be any new subtrees if the current index is strictly lower.
+        if index < prev_index {
             return false;
         }
 
+        // If the indexes are equal, there can only be a new subtree if `self` just completed it.
+        if index == prev_index && self.is_complete_subtree() {
+            return true;
+        }
+
         // This calculation can't overflow, because we're using i32 for u16 values.
         let index_difference = index - prev_index;
 
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 fe9b1b583fe..fc5af751df7 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,13 +189,16 @@ fn first_orchard_mainnet_subtree() -> NoteCommitmentSubtree<orchard::tree::Node>
 /// Returns an error if a note commitment subtree is missing or incorrect.
 fn quick_check_sapling_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
     // We check the first sapling subtree on mainnet, so skip this check if it isn't available.
-    let Some(NoteCommitmentSubtreeIndex(first_incomplete_subtree_index)) =
-        db.sapling_tree().subtree_index()
+    if db.network() != Mainnet {
+        return Ok(());
+    }
+
+    let Some(NoteCommitmentSubtreeIndex(tip_subtree_index)) = db.sapling_tree().subtree_index()
     else {
         return Ok(());
     };
 
-    if first_incomplete_subtree_index == 0 || db.network() != Mainnet {
+    if tip_subtree_index == 0 && !db.sapling_tree().is_complete_subtree() {
         return Ok(());
     }
 
@@ -239,13 +242,16 @@ fn quick_check_sapling_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
 /// Returns an error if a note commitment subtree is missing or incorrect.
 fn quick_check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
     // We check the first orchard subtree on mainnet, so skip this check if it isn't available.
-    let Some(NoteCommitmentSubtreeIndex(first_incomplete_subtree_index)) =
-        db.orchard_tree().subtree_index()
+    if db.network() != Mainnet {
+        return Ok(());
+    }
+
+    let Some(NoteCommitmentSubtreeIndex(tip_subtree_index)) = db.orchard_tree().subtree_index()
     else {
         return Ok(());
     };
 
-    if first_incomplete_subtree_index == 0 || db.network() != Mainnet {
+    if tip_subtree_index == 0 && !db.orchard_tree().is_complete_subtree() {
         return Ok(());
     }
 

From 54f2bba601b55575b5748f2636ab5aafd54ca433 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 19 Sep 2023 12:45:08 +1000
Subject: [PATCH 53/56] Actually fix new subtree code

---
 zebra-chain/src/orchard/tree.rs | 26 +++++++++++++++++---------
 zebra-chain/src/sapling/tree.rs | 26 +++++++++++++++++---------
 2 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/zebra-chain/src/orchard/tree.rs b/zebra-chain/src/orchard/tree.rs
index 74a4201eecf..2bbe4872c98 100644
--- a/zebra-chain/src/orchard/tree.rs
+++ b/zebra-chain/src/orchard/tree.rs
@@ -414,25 +414,33 @@ impl NoteCommitmentTree {
             .subtree_index()
             .map_or(-1, |index| i32::from(index.0));
 
+        // This calculation can't overflow, because we're using i32 for u16 values.
+        let index_difference = index - prev_index;
+
+        // There are 4 cases we need to handle:
+        // - lower index: never a new subtree
+        // - equal index: sometimes a new subtree
+        // - next index: sometimes a new subtree
+        // - greater than the next index: always a new subtree
+        //
+        // To simplify the function, we deal with the simple cases first.
+
         // There can't be any new subtrees if the current index is strictly lower.
         if index < prev_index {
             return false;
         }
 
-        // If the indexes are equal, there can only be a new subtree if `self` just completed it.
-        if index == prev_index && self.is_complete_subtree() {
+        // There is at least one new subtree, even if there is a spurious index difference.
+        if index_difference > 1 {
             return true;
         }
 
-        // This calculation can't overflow, because we're using i32 for u16 values.
-        let index_difference = index - prev_index;
-
-        // There are multiple new subtrees, or one new subtree and a spurious index difference.
-        if index_difference > 1 {
-            return true;
+        // If the indexes are equal, there can only be a new subtree if `self` just completed it.
+        if index == prev_index {
+            return self.is_complete_subtree();
         }
 
-        // Check for spurious index differences.
+        // If `self` is the next index, check for spurious index differences.
         //
         // There is one new subtree somewhere in the trees. It is either:
         // - a new subtree at the end of the previous tree, or
diff --git a/zebra-chain/src/sapling/tree.rs b/zebra-chain/src/sapling/tree.rs
index 8434fa0c721..c940b94d242 100644
--- a/zebra-chain/src/sapling/tree.rs
+++ b/zebra-chain/src/sapling/tree.rs
@@ -395,25 +395,33 @@ impl NoteCommitmentTree {
             .subtree_index()
             .map_or(-1, |index| i32::from(index.0));
 
+        // This calculation can't overflow, because we're using i32 for u16 values.
+        let index_difference = index - prev_index;
+
+        // There are 4 cases we need to handle:
+        // - lower index: never a new subtree
+        // - equal index: sometimes a new subtree
+        // - next index: sometimes a new subtree
+        // - greater than the next index: always a new subtree
+        //
+        // To simplify the function, we deal with the simple cases first.
+
         // There can't be any new subtrees if the current index is strictly lower.
         if index < prev_index {
             return false;
         }
 
-        // If the indexes are equal, there can only be a new subtree if `self` just completed it.
-        if index == prev_index && self.is_complete_subtree() {
+        // There is at least one new subtree, even if there is a spurious index difference.
+        if index_difference > 1 {
             return true;
         }
 
-        // This calculation can't overflow, because we're using i32 for u16 values.
-        let index_difference = index - prev_index;
-
-        // There are multiple new subtrees, or one new subtree and a spurious index difference.
-        if index_difference > 1 {
-            return true;
+        // If the indexes are equal, there can only be a new subtree if `self` just completed it.
+        if index == prev_index {
+            return self.is_complete_subtree();
         }
 
-        // Check for spurious index differences.
+        // If `self` is the next index, check for spurious index differences.
         //
         // There is one new subtree somewhere in the trees. It is either:
         // - a new subtree at the end of the previous tree, or

From d1df2ab601603a657f065201554ed116605ace1d Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 15 Sep 2023 16:59:22 +1000
Subject: [PATCH 54/56] Upgrade subtrees in reverse height order

Reverts "Temporarily upgrade subtrees in forward height order"
This reverts commit a9558be21401eb23f0079ef0f6a3e5086dba16e5.
---
 .../src/service/finalized_state/disk_db.rs    | 106 ++++++++++++++----
 .../disk_format/upgrade/add_subtrees.rs       |  17 ++-
 .../finalized_state/zebra_db/shielded.rs      |  26 +++++
 3 files changed, 126 insertions(+), 23 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db.rs b/zebra-state/src/service/finalized_state/disk_db.rs
index d001c87d157..3772fb7a789 100644
--- a/zebra-state/src/service/finalized_state/disk_db.rs
+++ b/zebra-state/src/service/finalized_state/disk_db.rs
@@ -258,6 +258,7 @@ impl PartialEq for DiskDb {
                 self.ephemeral, other.ephemeral,
                 "database with same path but different ephemeral configs",
             );
+
             return true;
         }
 
@@ -431,6 +432,46 @@ impl DiskDb {
     ///
     /// Holding this iterator open might delay block commit transactions.
     pub fn zs_range_iter<C, K, V, R>(&self, cf: &C, range: R) -> impl Iterator<Item = (K, V)> + '_
+    where
+        C: rocksdb::AsColumnFamilyRef,
+        K: IntoDisk + FromDisk,
+        V: FromDisk,
+        R: RangeBounds<K>,
+    {
+        self.zs_range_iter_with_direction(cf, range, false)
+    }
+
+    /// Returns a reverse iterator over the items in `cf` in `range`.
+    ///
+    /// Holding this iterator open might delay block commit transactions.
+    ///
+    /// This code is copied from `zs_range_iter()`, but with the mode reversed.
+    pub fn zs_reverse_range_iter<C, K, V, R>(
+        &self,
+        cf: &C,
+        range: R,
+    ) -> impl Iterator<Item = (K, V)> + '_
+    where
+        C: rocksdb::AsColumnFamilyRef,
+        K: IntoDisk + FromDisk,
+        V: FromDisk,
+        R: RangeBounds<K>,
+    {
+        self.zs_range_iter_with_direction(cf, range, true)
+    }
+
+    /// Returns an iterator over the items in `cf` in `range`.
+    ///
+    /// RocksDB iterators are ordered by increasing key bytes by default.
+    /// Otherwise, if `reverse` is `true`, the iterator is ordered by decreasing key bytes.
+    ///
+    /// Holding this iterator open might delay block commit transactions.
+    fn zs_range_iter_with_direction<C, K, V, R>(
+        &self,
+        cf: &C,
+        range: R,
+        reverse: bool,
+    ) -> impl Iterator<Item = (K, V)> + '_
     where
         C: rocksdb::AsColumnFamilyRef,
         K: IntoDisk + FromDisk,
@@ -451,40 +492,67 @@ impl DiskDb {
 
         let start_bound = map_to_vec(range.start_bound());
         let end_bound = map_to_vec(range.end_bound());
-        let range = (start_bound.clone(), end_bound);
-
-        let start_bound_vec =
-            if let Included(ref start_bound) | Excluded(ref start_bound) = start_bound {
-                start_bound.clone()
-            } else {
-                // Actually unused
-                Vec::new()
-            };
+        let range = (start_bound, end_bound);
 
-        let start_mode = if matches!(start_bound, Unbounded) {
-            // Unbounded iterators start at the first item
-            rocksdb::IteratorMode::Start
-        } else {
-            rocksdb::IteratorMode::From(start_bound_vec.as_slice(), rocksdb::Direction::Forward)
-        };
+        let mode = Self::zs_iter_mode(&range, reverse);
 
         // Reading multiple items from iterators has caused database hangs,
         // in previous RocksDB versions
         self.db
-            .iterator_cf(cf, start_mode)
+            .iterator_cf(cf, mode)
             .map(|result| result.expect("unexpected database failure"))
             .map(|(key, value)| (key.to_vec(), value))
-            // Skip excluded start bound and empty ranges. The `start_mode` already skips keys
-            // before the start bound.
+            // Skip excluded "from" bound and empty ranges. The `mode` already skips keys
+            // strictly before the "from" bound.
             .skip_while({
                 let range = range.clone();
                 move |(key, _value)| !range.contains(key)
             })
-            // Take until the excluded end bound is reached, or we're after the included end bound.
+            // Take until the excluded "to" bound is reached,
+            // or we're after the included "to" bound.
             .take_while(move |(key, _value)| range.contains(key))
             .map(|(key, value)| (K::from_bytes(key), V::from_bytes(value)))
     }
 
+    /// Returns the RocksDB iterator "from" mode for `range`.
+    ///
+    /// RocksDB iterators are ordered by increasing key bytes by default.
+    /// Otherwise, if `reverse` is `true`, the iterator is ordered by decreasing key bytes.
+    fn zs_iter_mode<R>(range: &R, reverse: bool) -> rocksdb::IteratorMode
+    where
+        R: RangeBounds<Vec<u8>>,
+    {
+        use std::ops::Bound::*;
+
+        let from_bound = if reverse {
+            range.end_bound()
+        } else {
+            range.start_bound()
+        };
+
+        match from_bound {
+            Unbounded => {
+                if reverse {
+                    // Reversed unbounded iterators start from the last item
+                    rocksdb::IteratorMode::End
+                } else {
+                    // Unbounded iterators start from the first item
+                    rocksdb::IteratorMode::Start
+                }
+            }
+
+            Included(bound) | Excluded(bound) => {
+                let direction = if reverse {
+                    rocksdb::Direction::Reverse
+                } else {
+                    rocksdb::Direction::Forward
+                };
+
+                rocksdb::IteratorMode::From(bound.as_slice(), direction)
+            }
+        }
+    }
+
     /// The ideal open file limit for Zebra
     const IDEAL_OPEN_FILE_LIMIT: u64 = 1024;
 
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 fc5af751df7..2cd806b64c8 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
@@ -38,13 +38,21 @@ pub fn run(
     // block can't complete multiple level 16 subtrees (or complete an entire subtree by itself).
     // Currently, with 2MB blocks and v4/v5 sapling and orchard output sizes, the subtree index can
     // increase by at most 1 every ~20 blocks.
+    //
+    // # Compatibility
+    //
+    // Because wallets search backwards from the chain tip, subtrees need to be added to the
+    // database in reverse height order. (Tip first, genesis last.)
+    //
+    // Otherwise, wallets that sync during the upgrade will be missing some notes.
 
     // Generate a list of sapling subtree inputs: previous and current trees, and their end heights.
     let subtrees = upgrade_db
-        .sapling_tree_by_height_range(..=initial_tip_height)
+        .sapling_tree_by_reversed_height_range(..=initial_tip_height)
         // We need both the tree and its previous tree for each shielded block.
         .tuple_windows()
-        .map(|((prev_end_height, prev_tree), (end_height, tree))| {
+        // Because the iterator is reversed, the larger tree is first.
+        .map(|((end_height, tree), (prev_end_height, prev_tree))| {
             (prev_end_height, prev_tree, end_height, tree)
         })
         // Find new subtrees.
@@ -65,10 +73,11 @@ pub fn run(
 
     // Generate a list of orchard subtree inputs: previous and current trees, and their end heights.
     let subtrees = upgrade_db
-        .orchard_tree_by_height_range(..=initial_tip_height)
+        .orchard_tree_by_reversed_height_range(..=initial_tip_height)
         // We need both the tree and its previous tree for each shielded block.
         .tuple_windows()
-        .map(|((prev_end_height, prev_tree), (end_height, tree))| {
+        // Because the iterator is reversed, the larger tree is first.
+        .map(|((end_height, tree), (prev_end_height, prev_tree))| {
             (prev_end_height, prev_tree, end_height, tree)
         })
         // Find new subtrees.
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 aaca367019f..395ae8b6b99 100644
--- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
@@ -180,6 +180,19 @@ impl ZebraDb {
         self.db.zs_range_iter(&sapling_trees, range)
     }
 
+    /// Returns the Sapling note commitment trees in the reversed range, in decreasing height order.
+    #[allow(clippy::unwrap_in_result)]
+    pub fn sapling_tree_by_reversed_height_range<R>(
+        &self,
+        range: R,
+    ) -> impl Iterator<Item = (Height, Arc<sapling::tree::NoteCommitmentTree>)> + '_
+    where
+        R: std::ops::RangeBounds<Height>,
+    {
+        let sapling_trees = self.db.cf_handle("sapling_note_commitment_tree").unwrap();
+        self.db.zs_reverse_range_iter(&sapling_trees, range)
+    }
+
     /// Returns the Sapling note commitment subtree at this `index`.
     ///
     /// # Correctness
@@ -313,6 +326,19 @@ impl ZebraDb {
         self.db.zs_range_iter(&orchard_trees, range)
     }
 
+    /// Returns the Orchard note commitment trees in the reversed range, in decreasing height order.
+    #[allow(clippy::unwrap_in_result)]
+    pub fn orchard_tree_by_reversed_height_range<R>(
+        &self,
+        range: R,
+    ) -> impl Iterator<Item = (Height, Arc<orchard::tree::NoteCommitmentTree>)> + '_
+    where
+        R: std::ops::RangeBounds<Height>,
+    {
+        let orchard_trees = self.db.cf_handle("orchard_note_commitment_tree").unwrap();
+        self.db.zs_reverse_range_iter(&orchard_trees, range)
+    }
+
     /// Returns the Orchard note commitment subtree at this `index`.
     ///
     /// # Correctness

From 8f4f6cb8dac3d333ca2dd9826c7290fdd26f06b6 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Fri, 15 Sep 2023 17:01:14 +1000
Subject: [PATCH 55/56] Bump the database patch version to re-run the upgrade
 (for testing)

---
 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 07177f11cfa..b5a06851638 100644
--- a/zebra-state/src/constants.rs
+++ b/zebra-state/src/constants.rs
@@ -53,13 +53,13 @@ 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 = 2;
 
 /// Returns the highest database version that modifies the subtree index format.
 ///
 /// This version is used by tests to wait for the subtree upgrade to finish.
 pub fn latest_version_for_adding_subtrees() -> Version {
-    Version::parse("25.2.1").expect("Hardcoded version string should be valid.")
+    Version::parse("25.2.2").expect("Hardcoded version string should be valid.")
 }
 
 /// The name of the file containing the minor and patch database versions.

From 8e8f7e6df502133c9b81b7c1fb100c0a1988c38d Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 19 Sep 2023 08:14:44 +1000
Subject: [PATCH 56/56] Revert "Remove an early reference to reverse direction"

This reverts commit c2064043776a11ef45fbe98d17ffc55e2be31f36.
---
 .../finalized_state/disk_format/upgrade/add_subtrees.rs        | 3 +++
 1 file changed, 3 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 2cd806b64c8..fafff57f4d9 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
@@ -21,6 +21,9 @@ use crate::service::finalized_state::{
 
 /// Runs disk format upgrade for adding Sapling and Orchard note commitment subtrees to database.
 ///
+/// Trees are added to the database in reverse height order, so that wallets can sync correctly
+/// while the upgrade is running.
+///
 /// Returns `Ok` if the upgrade completed, and `Err` if it was cancelled.
 #[allow(clippy::unwrap_in_result)]
 #[instrument(skip(upgrade_db, cancel_receiver))]