Implementing note commitments in finalized state

zebra-chain/src/orchard/tree.rs

@@ -183,6 +183,24 @@ impl From<pallas::Base> for Node {
     }
 }
 
+impl serde::Serialize for Node {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        serializer.serialize_bytes(&self.0.to_bytes())
+    }
+}
+
+impl<'de> serde::Deserialize<'de> for Node {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        todo!()
+    }
+}
+
 #[allow(dead_code, missing_docs)]
 #[derive(Error, Debug, PartialEq, Eq)]
 pub enum NoteCommitmentTreeError {
@@ -191,7 +209,7 @@ pub enum NoteCommitmentTreeError {
 }
 
 /// Orchard Incremental Note Commitment Tree
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct NoteCommitmentTree {
     /// The tree represented as a Frontier.
     ///

zebra-chain/src/transaction.rs

@@ -1,5 +1,6 @@
 //! Transactions and transaction-related structures.
 
+use halo2::pasta::pallas;
 use serde::{Deserialize, Serialize};
 
 mod hash;
@@ -383,6 +384,47 @@ impl Transaction {
         }
     }
 
+    /// Access the Sprout note commitments in this transaction, regardless of version.
+    pub fn sprout_note_commitments(
+        &self,
+    ) -> Box<dyn Iterator<Item = &sprout::commitment::NoteCommitment> + '_> {
+        // This function returns a boxed iterator because the different
+        // transaction variants end up having different iterator types
+        // (we could extract bctv and groth as separate iterators, then chain
+        // them together, but that would be much harder to read and maintain)
+        match self {
+            // JoinSplits with Bctv14 Proofs
+            Transaction::V2 {
+                joinsplit_data: Some(joinsplit_data),
+                ..
+            }
+            | Transaction::V3 {
+                joinsplit_data: Some(joinsplit_data),
+                ..
+            } => Box::new(joinsplit_data.note_commitments()),
+            // JoinSplits with Groth Proofs
+            Transaction::V4 {
+                joinsplit_data: Some(joinsplit_data),
+                ..
+            } => Box::new(joinsplit_data.note_commitments()),
+            // No JoinSplits
+            Transaction::V1 { .. }
+            | Transaction::V2 {
+                joinsplit_data: None,
+                ..
+            }
+            | Transaction::V3 {
+                joinsplit_data: None,
+                ..
+            }
+            | Transaction::V4 {
+                joinsplit_data: None,
+                ..
+            }
+            | Transaction::V5 { .. } => Box::new(std::iter::empty()),
+        }
+    }
+
     // sapling
 
     /// Iterate over the sapling [`Spend`](sapling::Spend)s for this transaction,
@@ -481,6 +523,36 @@ impl Transaction {
         }
     }
 
+    /// Access the note commitments in this transaction, regardless of version.
+    pub fn sapling_note_commitments(&self) -> Box<dyn Iterator<Item = &jubjub::Fq> + '_> {
+        // This function returns a boxed iterator because the different
+        // transaction variants end up having different iterator types
+        match self {
+            // Spends with Groth Proofs
+            Transaction::V4 {
+                sapling_shielded_data: Some(sapling_shielded_data),
+                ..
+            } => Box::new(sapling_shielded_data.note_commitments()),
+            Transaction::V5 {
+                sapling_shielded_data: Some(sapling_shielded_data),
+                ..
+            } => Box::new(sapling_shielded_data.note_commitments()),
+
+            // No Spends
+            Transaction::V1 { .. }
+            | Transaction::V2 { .. }
+            | Transaction::V3 { .. }
+            | Transaction::V4 {
+                sapling_shielded_data: None,
+                ..
+            }
+            | Transaction::V5 {
+                sapling_shielded_data: None,
+                ..
+            } => Box::new(std::iter::empty()),
+        }
+    }
+
     /// Return if the transaction has any Sapling shielded data.
     pub fn has_sapling_shielded_data(&self) -> bool {
         match self {
@@ -534,6 +606,15 @@ impl Transaction {
             .flatten()
     }
 
+    /// Access the note commitments in this transaction, if there are any,
+    /// regardless of version.
+    pub fn orchard_note_commitments(&self) -> impl Iterator<Item = &pallas::Base> {
+        self.orchard_shielded_data()
+            .into_iter()
+            .map(orchard::ShieldedData::note_commitments)
+            .flatten()
+    }
+
     /// Access the [`orchard::Flags`] in this transaction, if there is any,
     /// regardless of version.
     pub fn orchard_flags(&self) -> Option<orchard::shielded_data::Flags> {

zebra-state/src/service/finalized_state.rs

@@ -9,6 +9,7 @@ use std::{collections::HashMap, convert::TryInto, path::Path, sync::Arc};
 
 use zebra_chain::{
     block::{self, Block},
+    orchard,
     parameters::{Network, GENESIS_PREVIOUS_BLOCK_HASH},
     sprout,
     transaction::{self, Transaction},
@@ -224,7 +225,7 @@ impl FinalizedState {
         let sprout_note_commitment_tree = self.db.cf_handle("sprout_note_commitment_tree").unwrap();
         let sapling_note_commitment_tree =
             self.db.cf_handle("sapling_note_commitment_tree").unwrap();
-        let orchard_note_commitment_tree =
+        let orchard_note_commitment_tree_cf =
             self.db.cf_handle("orchard_note_commitment_tree").unwrap();
 
         // Assert that callers (including unit tests) get the chain order correct
@@ -258,6 +259,14 @@ impl FinalizedState {
             );
         }
 
+        let mut orchard_note_commitment_tree = match self.finalized_tip_height() {
+            Some(tip_height) => self
+                .orchard_note_commitment_tree(tip_height)
+                .expect("Orchard tree must exist for the finalized tip"),
+            None => Default::default(),
+        };
+        // TODO: other trees
+
         // We use a closure so we can use an early return for control flow in
         // the genesis case
         let prepare_commit = || -> rocksdb::WriteBatch {
@@ -274,6 +283,12 @@ impl FinalizedState {
             // (There is one such zero-valued output, on each of Testnet and Mainnet .)"
             // https://zips.z.cash/protocol/protocol.pdf#txnconsensus
             if block.header.previous_block_hash == GENESIS_PREVIOUS_BLOCK_HASH {
+                // TODO: do this or only add the tree when Orchard activates?
+                batch.zs_insert(
+                    orchard_note_commitment_tree_cf,
+                    height,
+                    orchard_note_commitment_tree,
+                );
                 return batch;
             }
 
@@ -319,8 +334,30 @@ impl FinalizedState {
                 for orchard_nullifier in transaction.orchard_nullifiers() {
                     batch.zs_insert(orchard_nullifiers, orchard_nullifier, ());
                 }
+
+                // Update the note commitment trees
+                for orchard_note_commitment in transaction.orchard_note_commitments() {
+                    orchard_note_commitment_tree
+                        .append(*orchard_note_commitment)
+                        .expect("must work since it was already appended before in the non-finalized state");
+                }
+                // TODO: other trees
             }
 
+            // Compute the new anchors and index them
+            let orchard_anchor = orchard_note_commitment_tree.root();
+            batch.zs_insert(orchard_anchors, height, orchard_anchor);
+            // TODO: other trees
+
+            // Update the note commitment trees
+            // TODO: delete the previous trees
+            batch.zs_insert(
+                orchard_note_commitment_tree_cf,
+                height,
+                orchard_note_commitment_tree,
+            );
+            // TODO: other trees
+
             batch
         };
 
@@ -418,6 +455,17 @@ impl FinalizedState {
             })
     }
 
+    /// Returns the Orchard note commitment tree for a given `block::Height`
+    /// if it is present.
+    pub fn orchard_note_commitment_tree(
+        &self,
+        height: block::Height,
+    ) -> Option<orchard::tree::NoteCommitmentTree> {
+        let orchard_note_commitment_tree =
+            self.db.cf_handle("orchard_note_commitment_tree").unwrap();
+        self.db.zs_get(orchard_note_commitment_tree, &height)
+    }
+
     /// If the database is `ephemeral`, delete it.
     fn delete_ephemeral(&self) {
         if self.ephemeral {

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

@@ -283,11 +283,12 @@ impl FromDisk for sapling::tree::NoteCommitmentTree {
         unimplemented!()
     }
 }
+
 impl IntoDisk for orchard::tree::NoteCommitmentTree {
     type Bytes = Vec<u8>;
 
     fn as_bytes(&self) -> Self::Bytes {
-        unimplemented!();
+        unimplemented!()
     }
 }