change(state): Add note subtree index handling to zebra-state, but don't write them to the finalized state yet

book/src/dev/state-db-upgrades.md

@@ -92,10 +92,12 @@ We use the following rocksdb column families:
 | `sapling_nullifiers`               | `sapling::Nullifier`   | `()`                          | Create  |
 | `sapling_anchors`                  | `sapling::tree::Root`  | `()`                          | Create  |
 | `sapling_note_commitment_tree`     | `block::Height`        | `sapling::NoteCommitmentTree` | Create  |
+| `sapling_note_commitment_subtree`  | `block::Height`        | `NoteCommitmentSubtreeData`   | Create  |
 | *Orchard*                          |                        |                               |         |
 | `orchard_nullifiers`               | `orchard::Nullifier`   | `()`                          | Create  |
 | `orchard_anchors`                  | `orchard::tree::Root`  | `()`                          | Create  |
 | `orchard_note_commitment_tree`     | `block::Height`        | `orchard::NoteCommitmentTree` | Create  |
+| `orchard_note_commitment_subtree`  | `block::Height`        | `NoteCommitmentSubtreeData`   | Create  |
 | *Chain*                            |                        |                               |         |
 | `history_tree`                     | `block::Height`        | `NonEmptyHistoryTree`         | Delete  |
 | `tip_chain_value_pool`             | `()`                   | `ValueBalance`                | Update  |
@@ -118,6 +120,8 @@ Block and Transaction Data:
   used instead of a `BTreeSet<OutputLocation>` value, to improve database performance
 - `AddressTransaction`: `AddressLocation \|\| TransactionLocation`
   used instead of a `BTreeSet<TransactionLocation>` value, to improve database performance
+- `NoteCommitmentSubtreeIndex`: 16 bits, big-endian, unsigned
+- `NoteCommitmentSubtreeData<{sapling, orchard}::tree::Node>`: `Height \|\| {sapling, orchard}::tree::Node`
 
 We use big-endian encoding for keys, to allow database index prefix searches.
 
@@ -334,6 +338,9 @@ So they should not be used for consensus-critical checks.
   as a "Merkle tree frontier" which is basically a (logarithmic) subset of
   the Merkle tree nodes as required to insert new items.
 
+- The `{sapling, orchard}_note_commitment_subtree` stores the completion height and
+  root for every completed level 16 note commitment subtree, for the specific pool.
+
 - `history_tree` stores the ZIP-221 history tree state at the tip of the finalized
   state. There is always a single entry for it. The tree is stored as the set of "peaks"
   of the "Merkle mountain range" tree structure, which is what is required to

zebra-chain/src/lib.rs

@@ -34,6 +34,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/arbitrary.rs

@@ -125,14 +125,37 @@ impl Arbitrary for Flags {
     type Strategy = BoxedStrategy<Self>;
 }
 
+fn pallas_base_strat() -> BoxedStrategy<pallas::Base> {
+    (vec(any::<u8>(), 64))
+        .prop_map(|bytes| {
+            let bytes = bytes.try_into().expect("vec is the correct length");
+            pallas::Base::from_uniform_bytes(&bytes)
+        })
+        .boxed()
+}
+
 impl Arbitrary for tree::Root {
     type Parameters = ();
 
     fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
-        (vec(any::<u8>(), 64))
-            .prop_map(|bytes| {
-                let bytes = bytes.try_into().expect("vec is the correct length");
-                Self::try_from(pallas::Base::from_uniform_bytes(&bytes).to_repr())
+        pallas_base_strat()
+            .prop_map(|base| {
+                Self::try_from(base.to_repr())
+                    .expect("a valid generated Orchard note commitment tree root")
+            })
+            .boxed()
+    }
+
+    type Strategy = BoxedStrategy<Self>;
+}
+
+impl Arbitrary for tree::Node {
+    type Parameters = ();
+
+    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
+        pallas_base_strat()
+            .prop_map(|base| {
+                Self::try_from(base.to_repr())
                     .expect("a valid generated Orchard note commitment tree root")
             })
             .boxed()

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::{
-    serde_helpers, ReadZcashExt, SerializationError, ZcashDeserialize, ZcashSerialize,
+use crate::{
+    serialization::{
+        serde_helpers, ReadZcashExt, SerializationError, ZcashDeserialize, ZcashSerialize,
+    },
+    subtree::TRACKED_SUBTREE_HEIGHT,
 };
 
 pub mod legacy;
@@ -170,6 +173,33 @@ 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> {
+        <[u8; 32]>::try_from(bytes)
+            .map_err(|_| "wrong byte slice len")?
+            .try_into()
+    }
+}
+
+impl TryFrom<[u8; 32]> for Node {
+    type Error = &'static str;
+
+    fn try_from(bytes: [u8; 32]) -> Result<Self, Self::Error> {
+        Option::<pallas::Base>::from(pallas::Base::from_repr(bytes))
+            .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 +347,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::{
-    block::{Block, Height},
-    orchard, sapling, sprout,
-};
+use crate::{block::Block, orchard, sapling, sprout, subtree::NoteCommitmentSubtree};
 
 /// 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(
-            [(
-                block
-                    .coinbase_height()
-                    .expect("height was already validated"),
-                block.clone(),
-            )]
-            .into_iter()
-            .collect(),
-        )
-    }
+        let block = block.clone();
+        let height = block
+            .coinbase_height()
+            .expect("height was already validated");
 
-    /// 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
-            .values()
-            .flat_map(|block| block.transactions.iter())
+        let sprout_note_commitments: Vec<_> = block
+            .transactions
+            .iter()
             .flat_map(|tx| tx.sprout_note_commitments())
             .cloned()
             .collect();
-        let sapling_note_commitments: Vec<_> = block_list
-            .values()
-            .flat_map(|block| block.transactions.iter())
+        let sapling_note_commitments: Vec<_> = block
+            .transactions
+            .iter()
             .flat_map(|tx| tx.sapling_note_commitments())
             .cloned()
             .collect();
-        let orchard_note_commitments: Vec<_> = block_list
-            .values()
-            .flat_map(|block| block.transactions.iter())
+        let orchard_note_commitments: Vec<_> = block
+            .transactions
+            .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))
     }
 }