change(state): Add note subtree index handling to zebra-state, but do…

…n't write them to the finalized state yet (#7334)

* 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

* adds subtrees to zebra-state

* uses split_at, from_repr, & updates state-db-upgrades.md

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

Co-authored-by: teor <[email protected]>

* renames partial_subtree to subtree_data

* tests that subtree serialization format

* adds raw data format serialization round-trip test

* decrements minor version and skips inserting subtrees in db

---------

Co-authored-by: teor <[email protected]>

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/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

@@ -184,11 +184,19 @@ 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")
+        <[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")
     }
 }
 

zebra-chain/src/sapling/arbitrary.rs

@@ -119,16 +119,30 @@ fn spendauth_verification_key_bytes() -> impl Strategy<Value = ValidatingKey> {
     })
 }
 
+fn jubjub_base_strat() -> BoxedStrategy<jubjub::Base> {
+    (vec(any::<u8>(), 64))
+        .prop_map(|bytes| {
+            let bytes = bytes.try_into().expect("vec is the correct length");
+            jubjub::Base::from_bytes_wide(&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");
-                tree::Root(jubjub::Base::from_bytes_wide(&bytes))
-            })
-            .boxed()
+        jubjub_base_strat().prop_map(tree::Root).boxed()
+    }
+
+    type Strategy = BoxedStrategy<Self>;
+}
+
+impl Arbitrary for tree::Node {
+    type Parameters = ();
+
+    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
+        jubjub_base_strat().prop_map(tree::Node::from).boxed()
     }
 
     type Strategy = BoxedStrategy<Self>;

zebra-chain/src/subtree.rs

@@ -2,6 +2,9 @@
 
 use std::sync::Arc;
 
+#[cfg(any(test, feature = "proptest-impl"))]
+use proptest_derive::Arbitrary;
+
 use crate::block::Height;
 
 /// Height at which Zebra tracks subtree roots
@@ -35,11 +38,17 @@ impl<Node> NoteCommitmentSubtree<Node> {
         let index = index.into();
         Arc::new(Self { index, end, node })
     }
+
+    /// Converts struct to [`NoteCommitmentSubtreeData`].
+    pub fn into_data(self) -> NoteCommitmentSubtreeData<Node> {
+        NoteCommitmentSubtreeData::new(self.end, self.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)]
+#[cfg_attr(any(test, feature = "proptest-impl"), derive(Arbitrary))]
 pub struct NoteCommitmentSubtreeData<Node> {
     /// End boundary of this subtree, the block height of its last leaf.
     pub end: Height,

zebra-state/src/request.rs

@@ -15,6 +15,7 @@ use zebra_chain::{
     sapling,
     serialization::SerializationError,
     sprout,
+    subtree::NoteCommitmentSubtree,
     transaction::{self, UnminedTx},
     transparent::{self, utxos_from_ordered_utxos},
     value_balance::{ValueBalance, ValueBalanceError},
@@ -235,15 +236,17 @@ impl Treestate {
         sprout: Arc<sprout::tree::NoteCommitmentTree>,
         sapling: Arc<sapling::tree::NoteCommitmentTree>,
         orchard: Arc<orchard::tree::NoteCommitmentTree>,
+        sapling_subtree: Option<Arc<NoteCommitmentSubtree<sapling::tree::Node>>>,
+        orchard_subtree: Option<Arc<NoteCommitmentSubtree<orchard::tree::Node>>>,
         history_tree: Arc<HistoryTree>,
     ) -> Self {
         Self {
             note_commitment_trees: NoteCommitmentTrees {
                 sprout,
                 sapling,
-                sapling_subtree: None,
+                sapling_subtree,
                 orchard,
-                orchard_subtree: None,
+                orchard_subtree,
             },
             history_tree,
         }

zebra-state/src/service/finalized_state/disk_db.rs

@@ -517,10 +517,12 @@ impl DiskDb {
         "sapling_nullifiers",
         "sapling_anchors",
         "sapling_note_commitment_tree",
+        "sapling_note_commitment_subtree",
         // Orchard
         "orchard_nullifiers",
         "orchard_anchors",
         "orchard_note_commitment_tree",
+        "orchard_note_commitment_subtree",
         // Chain
         "history_tree",
         "tip_chain_value_pool",

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

@@ -7,10 +7,16 @@
 
 use bincode::Options;
 
-use zebra_chain::{orchard, sapling, sprout};
+use zebra_chain::{
+    block::Height,
+    orchard, sapling, sprout,
+    subtree::{NoteCommitmentSubtreeData, NoteCommitmentSubtreeIndex},
+};
 
 use crate::service::finalized_state::disk_format::{FromDisk, IntoDisk};
 
+use super::block::HEIGHT_DISK_BYTES;
+
 impl IntoDisk for sprout::Nullifier {
     type Bytes = [u8; 32];
 
@@ -74,6 +80,14 @@ impl IntoDisk for orchard::tree::Root {
     }
 }
 
+impl IntoDisk for NoteCommitmentSubtreeIndex {
+    type Bytes = [u8; 2];
+
+    fn as_bytes(&self) -> Self::Bytes {
+        self.0.to_be_bytes()
+    }
+}
+
 impl FromDisk for orchard::tree::Root {
     fn from_bytes(bytes: impl AsRef<[u8]>) -> Self {
         let array: [u8; 32] = bytes.as_ref().try_into().unwrap();
@@ -140,3 +154,49 @@ impl FromDisk for orchard::tree::NoteCommitmentTree {
             .expect("deserialization format should match the serialization format used by IntoDisk")
     }
 }
+
+impl IntoDisk for sapling::tree::Node {
+    type Bytes = Vec<u8>;
+
+    fn as_bytes(&self) -> Self::Bytes {
+        self.as_ref().to_vec()
+    }
+}
+
+impl IntoDisk for orchard::tree::Node {
+    type Bytes = Vec<u8>;
+
+    fn as_bytes(&self) -> Self::Bytes {
+        self.to_repr().to_vec()
+    }
+}
+
+impl<Node: IntoDisk<Bytes = Vec<u8>>> IntoDisk for NoteCommitmentSubtreeData<Node> {
+    type Bytes = Vec<u8>;
+
+    fn as_bytes(&self) -> Self::Bytes {
+        [self.end.as_bytes().to_vec(), self.node.as_bytes()].concat()
+    }
+}
+
+impl FromDisk for sapling::tree::Node {
+    fn from_bytes(bytes: impl AsRef<[u8]>) -> Self {
+        Self::try_from(bytes.as_ref()).expect("trusted data should deserialize successfully")
+    }
+}
+
+impl FromDisk for orchard::tree::Node {
+    fn from_bytes(bytes: impl AsRef<[u8]>) -> Self {
+        Self::try_from(bytes.as_ref()).expect("trusted data should deserialize successfully")
+    }
+}
+
+impl<Node: FromDisk> FromDisk for NoteCommitmentSubtreeData<Node> {
+    fn from_bytes(disk_bytes: impl AsRef<[u8]>) -> Self {
+        let (height_bytes, node_bytes) = disk_bytes.as_ref().split_at(HEIGHT_DISK_BYTES);
+        Self::new(
+            Height::from_bytes(height_bytes),
+            Node::from_bytes(node_bytes),
+        )
+    }
+}

zebra-state/src/service/finalized_state/disk_format/tests/prop.rs

@@ -6,6 +6,7 @@ use zebra_chain::{
     amount::{Amount, NonNegative},
     block::{self, Height},
     orchard, sapling, sprout,
+    subtree::NoteCommitmentSubtreeData,
     transaction::{self, Transaction},
     transparent,
     value_balance::ValueBalance,
@@ -361,6 +362,16 @@ fn roundtrip_sapling_tree_root() {
     proptest!(|(val in any::<sapling::tree::Root>())| assert_value_properties(val));
 }
 
+#[test]
+fn roundtrip_sapling_subtree_data() {
+    let _init_guard = zebra_test::init();
+
+    proptest!(|(mut val in any::<NoteCommitmentSubtreeData<sapling::tree::Node>>())| {
+        val.end = val.end.clamp(Height(0), MAX_ON_DISK_HEIGHT);
+        assert_value_properties(val)
+    });
+}
+
 // TODO: test note commitment tree round-trip, after implementing proptest::Arbitrary
 
 // Orchard
@@ -436,6 +447,16 @@ fn roundtrip_orchard_tree_root() {
     proptest!(|(val in any::<orchard::tree::Root>())| assert_value_properties(val));
 }
 
+#[test]
+fn roundtrip_orchard_subtree_data() {
+    let _init_guard = zebra_test::init();
+
+    proptest!(|(mut val in any::<NoteCommitmentSubtreeData<orchard::tree::Node>>())| {
+        val.end = val.end.clamp(Height(0), MAX_ON_DISK_HEIGHT);
+        assert_value_properties(val)
+    });
+}
+
 // TODO: test note commitment tree round-trip, after implementing proptest::Arbitrary
 
 // Chain

zebra-state/src/service/finalized_state/disk_format/tests/snapshots/column_family_names.snap

@@ -1,6 +1,6 @@
 ---
 source: zebra-state/src/service/finalized_state/disk_format/tests/snapshot.rs
-assertion_line: 72
+assertion_line: 81
 expression: cf_names
 ---
 [
@@ -12,9 +12,11 @@ expression: cf_names
   "height_by_hash",
   "history_tree",
   "orchard_anchors",
+  "orchard_note_commitment_subtree",
   "orchard_note_commitment_tree",
   "orchard_nullifiers",
   "sapling_anchors",
+  "sapling_note_commitment_subtree",
   "sapling_note_commitment_tree",
   "sapling_nullifiers",
   "sprout_anchors",

zebra-state/src/service/finalized_state/disk_format/tests/snapshots/empty_column_families@mainnet_0.snap

@@ -1,14 +1,15 @@
 ---
 source: zebra-state/src/service/finalized_state/disk_format/tests/snapshot.rs
-assertion_line: 154
 expression: empty_column_families
 ---
 [
   "balance_by_transparent_addr: no entries",
   "history_tree: no entries",
   "orchard_anchors: no entries",
+  "orchard_note_commitment_subtree: no entries",
   "orchard_nullifiers: no entries",
   "sapling_anchors: no entries",
+  "sapling_note_commitment_subtree: no entries",
   "sapling_nullifiers: no entries",
   "sprout_anchors: no entries",
   "sprout_nullifiers: no entries",

zebra-state/src/service/finalized_state/disk_format/tests/snapshots/empty_column_families@mainnet_1.snap

@@ -4,7 +4,9 @@ expression: empty_column_families
 ---
 [
   "history_tree: no entries",
+  "orchard_note_commitment_subtree: no entries",
   "orchard_nullifiers: no entries",
+  "sapling_note_commitment_subtree: no entries",
   "sapling_nullifiers: no entries",
   "sprout_nullifiers: no entries",
 ]

zebra-state/src/service/finalized_state/disk_format/tests/snapshots/empty_column_families@mainnet_2.snap

@@ -4,7 +4,9 @@ expression: empty_column_families
 ---
 [
   "history_tree: no entries",
+  "orchard_note_commitment_subtree: no entries",
   "orchard_nullifiers: no entries",
+  "sapling_note_commitment_subtree: no entries",
   "sapling_nullifiers: no entries",
   "sprout_nullifiers: no entries",
 ]