0. Add note commitment subtree types to zebra-chain (#7371)

* zebra-chain changes from the subtree-boundaries branch

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

* Temporarily populate new subtree fields with None - for revert

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

* Applies suggestions from code review

* removes from_repr_unchecked methods

* simplifies loop

---------

Co-authored-by: arya2 <aryasolhi@gmail.com>

zebra-chain/src/lib.rs

@@ -44,6 +44,7 @@ pub mod sapling;
 pub mod serialization;
 pub mod shutdown;
 pub mod sprout;
+pub mod subtree;
 pub mod transaction;
 pub mod transparent;
 pub mod value_balance;

zebra-chain/src/orchard/tree.rs

@@ -18,7 +18,7 @@ use std::{
 };
 
 use bitvec::prelude::*;
-use bridgetree;
+use bridgetree::{self, NonEmptyFrontier};
 use halo2::pasta::{group::ff::PrimeField, pallas};
 use incrementalmerkletree::Hashable;
 use lazy_static::lazy_static;
@@ -27,8 +27,11 @@ use zcash_primitives::merkle_tree::{write_commitment_tree, HashSer};
 
 use super::sinsemilla::*;
 
-use crate::serialization::{
+use crate::{
-    serde_helpers, ReadZcashExt, SerializationError, ZcashDeserialize, ZcashSerialize,
+    serialization::{
+        serde_helpers, ReadZcashExt, SerializationError, ZcashDeserialize, ZcashSerialize,
+    },
+    subtree::TRACKED_SUBTREE_HEIGHT,
 };
 
 pub mod legacy;
@@ -170,6 +173,25 @@ impl ZcashDeserialize for Root {
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
 pub struct Node(pallas::Base);
 
+impl Node {
+    /// Calls `to_repr()` on inner value.
+    pub fn to_repr(&self) -> [u8; 32] {
+        self.0.to_repr()
+    }
+}
+
+impl TryFrom<&[u8]> for Node {
+    type Error = &'static str;
+
+    fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
+        Option::<pallas::Base>::from(pallas::Base::from_repr(
+            bytes.try_into().map_err(|_| "wrong byte slice len")?,
+        ))
+        .map(Node)
+        .ok_or("invalid Pallas field element")
+    }
+}
+
 /// Required to convert [`NoteCommitmentTree`] into [`SerializedTree`].
 ///
 /// Zebra stores Orchard note commitment trees as [`Frontier`][1]s while the
@@ -317,6 +339,32 @@ impl NoteCommitmentTree {
         }
     }
 
+    /// Returns true if the most recently appended leaf completes the subtree
+    pub fn is_complete_subtree(tree: &NonEmptyFrontier<Node>) -> bool {
+        tree.position()
+            .is_complete_subtree(TRACKED_SUBTREE_HEIGHT.into())
+    }
+
+    /// Returns subtree address at [`TRACKED_SUBTREE_HEIGHT`]
+    pub fn subtree_address(tree: &NonEmptyFrontier<Node>) -> incrementalmerkletree::Address {
+        incrementalmerkletree::Address::above_position(
+            TRACKED_SUBTREE_HEIGHT.into(),
+            tree.position(),
+        )
+    }
+
+    /// Returns subtree index and root if the most recently appended leaf completes the subtree
+    #[allow(clippy::unwrap_in_result)]
+    pub fn completed_subtree_index_and_root(&self) -> Option<(u16, Node)> {
+        let value = self.inner.value()?;
+        Self::is_complete_subtree(value).then_some(())?;
+        let address = Self::subtree_address(value);
+        let index = address.index().try_into().expect("should fit in u16");
+        let root = value.root(Some(TRACKED_SUBTREE_HEIGHT.into()));
+
+        Some((index, root))
+    }
+
     /// Returns the current root of the tree, used as an anchor in Orchard
     /// shielded transactions.
     pub fn root(&self) -> Root {

zebra-chain/src/parallel/tree.rs

@@ -1,13 +1,10 @@
 //! Parallel note commitment tree update methods.
 
-use std::{collections::BTreeMap, sync::Arc};
+use std::sync::Arc;
 
 use thiserror::Error;
 
-use crate::{
+use crate::{block::Block, orchard, sapling, sprout, subtree::NoteCommitmentSubtree};
-    block::{Block, Height},
-    orchard, sapling, sprout,
-};
 
 /// An argument wrapper struct for note commitment trees.
 #[derive(Clone, Debug)]
@@ -18,8 +15,14 @@ pub struct NoteCommitmentTrees {
     /// The sapling note commitment tree.
     pub sapling: Arc<sapling::tree::NoteCommitmentTree>,
 
+    /// The sapling note commitment subtree.
+    pub sapling_subtree: Option<Arc<NoteCommitmentSubtree<sapling::tree::Node>>>,
+
     /// The orchard note commitment tree.
     pub orchard: Arc<orchard::tree::NoteCommitmentTree>,
+
+    /// The orchard note commitment subtree.
+    pub orchard_subtree: Option<Arc<NoteCommitmentSubtree<orchard::tree::Node>>>,
 }
 
 /// Note commitment tree errors.
@@ -49,49 +52,34 @@ impl NoteCommitmentTrees {
         &mut self,
         block: &Arc<Block>,
     ) -> Result<(), NoteCommitmentTreeError> {
-        self.update_trees_parallel_list(
+        let block = block.clone();
-            [(
+        let height = block
-                block
+            .coinbase_height()
-                    .coinbase_height()
+            .expect("height was already validated");
-                    .expect("height was already validated"),
-                block.clone(),
-            )]
-            .into_iter()
-            .collect(),
-        )
-    }
 
-    /// Updates the note commitment trees using the transactions in `block`,
-    /// then re-calculates the cached tree roots, using parallel `rayon` threads.
-    ///
-    /// If any of the tree updates cause an error,
-    /// it will be returned at the end of the parallel batches.
-    pub fn update_trees_parallel_list(
-        &mut self,
-        block_list: BTreeMap<Height, Arc<Block>>,
-    ) -> Result<(), NoteCommitmentTreeError> {
         // Prepare arguments for parallel threads
         let NoteCommitmentTrees {
             sprout,
             sapling,
             orchard,
+            ..
         } = self.clone();
 
-        let sprout_note_commitments: Vec<_> = block_list
+        let sprout_note_commitments: Vec<_> = block
-            .values()
+            .transactions
-            .flat_map(|block| block.transactions.iter())
+            .iter()
             .flat_map(|tx| tx.sprout_note_commitments())
             .cloned()
             .collect();
-        let sapling_note_commitments: Vec<_> = block_list
+        let sapling_note_commitments: Vec<_> = block
-            .values()
+            .transactions
-            .flat_map(|block| block.transactions.iter())
+            .iter()
             .flat_map(|tx| tx.sapling_note_commitments())
             .cloned()
             .collect();
-        let orchard_note_commitments: Vec<_> = block_list
+        let orchard_note_commitments: Vec<_> = block
-            .values()
+            .transactions
-            .flat_map(|block| block.transactions.iter())
+            .iter()
             .flat_map(|tx| tx.orchard_note_commitments())
             .cloned()
             .collect();
@@ -132,12 +120,20 @@ impl NoteCommitmentTrees {
         if let Some(sprout_result) = sprout_result {
             self.sprout = sprout_result?;
         }
+
         if let Some(sapling_result) = sapling_result {
-            self.sapling = sapling_result?;
+            let (sapling, subtree_root) = sapling_result?;
-        }
+            self.sapling = sapling;
+            self.sapling_subtree =
+                subtree_root.map(|(idx, node)| NoteCommitmentSubtree::new(idx, height, node));
+        };
+
         if let Some(orchard_result) = orchard_result {
-            self.orchard = orchard_result?;
+            let (orchard, subtree_root) = orchard_result?;
-        }
+            self.orchard = orchard;
+            self.orchard_subtree =
+                subtree_root.map(|(idx, node)| NoteCommitmentSubtree::new(idx, height, node));
+        };
 
         Ok(())
     }
@@ -160,36 +156,74 @@ impl NoteCommitmentTrees {
     }
 
     /// Update the sapling note commitment tree.
+    #[allow(clippy::unwrap_in_result)]
     fn update_sapling_note_commitment_tree(
         mut sapling: Arc<sapling::tree::NoteCommitmentTree>,
         sapling_note_commitments: Vec<sapling::tree::NoteCommitmentUpdate>,
-    ) -> Result<Arc<sapling::tree::NoteCommitmentTree>, NoteCommitmentTreeError> {
+    ) -> Result<
+        (
+            Arc<sapling::tree::NoteCommitmentTree>,
+            Option<(u16, sapling::tree::Node)>,
+        ),
+        NoteCommitmentTreeError,
+    > {
         let sapling_nct = Arc::make_mut(&mut sapling);
 
+        // It is impossible for blocks to contain more than one level 16 sapling root:
+        // > [NU5 onward] nSpendsSapling, nOutputsSapling, and nActionsOrchard MUST all be less than 2^16.
+        // <https://zips.z.cash/protocol/protocol.pdf#txnconsensus>
+        //
+        // Before NU5, this limit holds due to the minimum size of Sapling outputs (948 bytes)
+        // and the maximum size of a block:
+        // > The size of a block MUST be less than or equal to 2000000 bytes.
+        // <https://zips.z.cash/protocol/protocol.pdf#blockheader>
+        // <https://zips.z.cash/protocol/protocol.pdf#txnencoding>
+        let mut subtree_root = None;
+
         for sapling_note_commitment in sapling_note_commitments {
+            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.
         let _ = sapling_nct.root();
 
-        Ok(sapling)
+        Ok((sapling, subtree_root))
     }
 
     /// Update the orchard note commitment tree.
+    #[allow(clippy::unwrap_in_result)]
     fn update_orchard_note_commitment_tree(
         mut orchard: Arc<orchard::tree::NoteCommitmentTree>,
         orchard_note_commitments: Vec<orchard::tree::NoteCommitmentUpdate>,
-    ) -> Result<Arc<orchard::tree::NoteCommitmentTree>, NoteCommitmentTreeError> {
+    ) -> Result<
+        (
+            Arc<orchard::tree::NoteCommitmentTree>,
+            Option<(u16, orchard::tree::Node)>,
+        ),
+        NoteCommitmentTreeError,
+    > {
         let orchard_nct = Arc::make_mut(&mut orchard);
 
+        // It is impossible for blocks to contain more than one level 16 orchard root:
+        // > [NU5 onward] nSpendsSapling, nOutputsSapling, and nActionsOrchard MUST all be less than 2^16.
+        // <https://zips.z.cash/protocol/protocol.pdf#txnconsensus>
+        let mut subtree_root = None;
+
         for orchard_note_commitment in orchard_note_commitments {
+            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.
         let _ = orchard_nct.root();
 
-        Ok(orchard)
+        Ok((orchard, subtree_root))
     }
 }

zebra-chain/src/sapling/tree.rs

@@ -18,7 +18,7 @@ use std::{
 };
 
 use bitvec::prelude::*;
-use bridgetree;
+use bridgetree::{self, NonEmptyFrontier};
 use incrementalmerkletree::{frontier::Frontier, Hashable};
 
 use lazy_static::lazy_static;
@@ -28,8 +28,11 @@ use zcash_primitives::merkle_tree::HashSer;
 
 use super::commitment::pedersen_hashes::pedersen_hash;
 
-use crate::serialization::{
+use crate::{
-    serde_helpers, ReadZcashExt, SerializationError, ZcashDeserialize, ZcashSerialize,
+    serialization::{
+        serde_helpers, ReadZcashExt, SerializationError, ZcashDeserialize, ZcashSerialize,
+    },
+    subtree::TRACKED_SUBTREE_HEIGHT,
 };
 
 pub mod legacy;
@@ -165,6 +168,12 @@ impl ZcashDeserialize for Root {
 #[derive(Copy, Clone, Eq, PartialEq)]
 pub struct Node([u8; 32]);
 
+impl AsRef<[u8; 32]> for Node {
+    fn as_ref(&self) -> &[u8; 32] {
+        &self.0
+    }
+}
+
 impl fmt::Debug for Node {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         f.debug_tuple("Node").field(&hex::encode(self.0)).finish()
@@ -217,6 +226,18 @@ impl From<jubjub::Fq> for Node {
     }
 }
 
+impl TryFrom<&[u8]> for Node {
+    type Error = &'static str;
+
+    fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
+        Option::<jubjub::Fq>::from(jubjub::Fq::from_bytes(
+            bytes.try_into().map_err(|_| "wrong byte slice len")?,
+        ))
+        .map(Node::from)
+        .ok_or("invalid jubjub field element")
+    }
+}
+
 impl serde::Serialize for Node {
     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
     where
@@ -311,6 +332,32 @@ impl NoteCommitmentTree {
         }
     }
 
+    /// Returns true if the most recently appended leaf completes the subtree
+    pub fn is_complete_subtree(tree: &NonEmptyFrontier<Node>) -> bool {
+        tree.position()
+            .is_complete_subtree(TRACKED_SUBTREE_HEIGHT.into())
+    }
+
+    /// Returns subtree address at [`TRACKED_SUBTREE_HEIGHT`]
+    pub fn subtree_address(tree: &NonEmptyFrontier<Node>) -> incrementalmerkletree::Address {
+        incrementalmerkletree::Address::above_position(
+            TRACKED_SUBTREE_HEIGHT.into(),
+            tree.position(),
+        )
+    }
+
+    /// Returns subtree index and root if the most recently appended leaf completes the subtree
+    #[allow(clippy::unwrap_in_result)]
+    pub fn completed_subtree_index_and_root(&self) -> Option<(u16, Node)> {
+        let value = self.inner.value()?;
+        Self::is_complete_subtree(value).then_some(())?;
+        let address = Self::subtree_address(value);
+        let index = address.index().try_into().expect("should fit in u16");
+        let root = value.root(Some(TRACKED_SUBTREE_HEIGHT.into()));
+
+        Some((index, root))
+    }
+
     /// Returns the current root of the tree, used as an anchor in Sapling
     /// shielded transactions.
     pub fn root(&self) -> Root {

zebra-chain/src/subtree.rs

@@ -0,0 +1,63 @@
+//! Struct representing Sapling/Orchard note commitment subtrees
+
+use std::sync::Arc;
+
+use crate::block::Height;
+
+/// Height at which Zebra tracks subtree roots
+pub const TRACKED_SUBTREE_HEIGHT: u8 = 16;
+
+/// A subtree index
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+pub struct NoteCommitmentSubtreeIndex(pub u16);
+
+impl From<u16> for NoteCommitmentSubtreeIndex {
+    fn from(value: u16) -> Self {
+        Self(value)
+    }
+}
+
+/// Subtree root of Sapling or Orchard note commitment tree,
+/// with its associated block height and subtree index.
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+pub struct NoteCommitmentSubtree<Node> {
+    /// Index of this subtree
+    pub index: NoteCommitmentSubtreeIndex,
+    /// End boundary of this subtree, the block height of its last leaf.
+    pub end: Height,
+    /// Root of this subtree.
+    pub node: Node,
+}
+
+impl<Node> NoteCommitmentSubtree<Node> {
+    /// Creates new [`NoteCommitmentSubtree`]
+    pub fn new(index: impl Into<NoteCommitmentSubtreeIndex>, end: Height, node: Node) -> Arc<Self> {
+        let index = index.into();
+        Arc::new(Self { index, end, node })
+    }
+}
+
+/// Subtree root of Sapling or Orchard note commitment tree, with block height, but without the subtree index.
+/// Used for database key-value serialization, where the subtree index is the key, and this struct is the value.
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+pub struct NoteCommitmentSubtreeData<Node> {
+    /// End boundary of this subtree, the block height of its last leaf.
+    pub end: Height,
+    /// Root of this subtree.
+    pub node: Node,
+}
+
+impl<Node> NoteCommitmentSubtreeData<Node> {
+    /// Creates new [`NoteCommitmentSubtreeData`]
+    pub fn new(end: Height, node: Node) -> Self {
+        Self { end, node }
+    }
+
+    /// Creates new [`NoteCommitmentSubtree`] from a [`NoteCommitmentSubtreeData`] and index
+    pub fn with_index(
+        self,
+        index: impl Into<NoteCommitmentSubtreeIndex>,
+    ) -> Arc<NoteCommitmentSubtree<Node>> {
+        NoteCommitmentSubtree::new(index, self.end, self.node)
+    }
+}

zebra-state/src/request.rs

@@ -241,7 +241,9 @@ impl Treestate {
             note_commitment_trees: NoteCommitmentTrees {
                 sprout,
                 sapling,
+                sapling_subtree: None,
                 orchard,
+                orchard_subtree: None,
             },
             history_tree,
         }

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

@@ -224,7 +224,9 @@ impl ZebraDb {
         NoteCommitmentTrees {
             sprout: self.sprout_tree(),
             sapling: self.sapling_tree(),
+            sapling_subtree: None,
             orchard: self.orchard_tree(),
+            orchard_subtree: None,
         }
     }
 }

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

@@ -1251,7 +1251,9 @@ impl Chain {
         let mut nct = NoteCommitmentTrees {
             sprout: self.sprout_note_commitment_tree(),
             sapling: self.sapling_note_commitment_tree(),
+            sapling_subtree: None,
             orchard: self.orchard_note_commitment_tree(),
+            orchard_subtree: None,
         };
 
         let mut tree_result = None;