fix(state): Use correct end heights for end of block subtrees during the full sync

Cargo.lock

@@ -5740,6 +5740,7 @@ dependencies = [
  "futures",
  "halo2_proofs",
  "hex",
+ "hex-literal",
  "howudoin",
  "indexmap 2.0.0",
  "insta",

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();
                     }
                 }
 

zebra-chain/src/orchard/tree.rs

@@ -394,6 +394,72 @@ 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 this tree has at least one new subtree, when compared with `prev_tree`.
+    pub fn contains_new_subtree(&self, prev_tree: &Self) -> bool {
+        // 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));
+
+        // 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;
+        }
+
+        // There is at least one new subtree, even if there is 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();
+        }
+
+        // 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
+        // - a new subtree in this tree (but not at the end).
+        //
+        // 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.
+        if prev_tree.is_complete_subtree() || prev_index == -1 {
+            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 {
@@ -423,6 +489,46 @@ 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 => {
+                let subtree_address = incrementalmerkletree::Address::above_position(
+                    TRACKED_SUBTREE_HEIGHT.into(),
+                    // This position is guaranteed to be in the first subtree.
+                    0.into(),
+                );
+
+                assert_eq!(
+                    subtree_address.position_range_start(),
+                    0.into(),
+                    "address is not in the first subtree"
+                );
+
+                subtree_address.position_range_end()
+            }
+        };
+
+        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() {

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,8 +146,9 @@ 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)]
-    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<
@@ -170,11 +172,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.
@@ -184,8 +189,9 @@ 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)]
-    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<
@@ -203,11 +209,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.

zebra-chain/src/sapling/tree.rs

@@ -375,6 +375,72 @@ 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 this tree has at least one new subtree, when compared with `prev_tree`.
+    pub fn contains_new_subtree(&self, prev_tree: &Self) -> bool {
+        // 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));
+
+        // 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;
+        }
+
+        // There is at least one new subtree, even if there is 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();
+        }
+
+        // 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
+        // - a new subtree in this tree (but not at the end).
+        //
+        // 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.
+        if prev_tree.is_complete_subtree() || prev_index == -1 {
+            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 {
@@ -404,6 +470,46 @@ 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 => {
+                let subtree_address = incrementalmerkletree::Address::above_position(
+                    TRACKED_SUBTREE_HEIGHT.into(),
+                    // This position is guaranteed to be in the first subtree.
+                    0.into(),
+                );
+
+                assert_eq!(
+                    subtree_address.position_range_start(),
+                    0.into(),
+                    "address is not in the first subtree"
+                );
+
+                subtree_address.position_range_end()
+            }
+        };
+
+        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() {

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)

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"