Track anchors and note commitment trees in zebra-state (#2458)

* Tidy chain Cargo.toml

* Organize imports

* Add method to get note commitments from all Actions in Orchard shielded data

* Add method to get note commitments from all JoinSplits in Sprout JoinSplitData

* Add Request and Response variants for awaiting anchors

* Add anchors and note commitment trees to finalized state db

* Add (From|Into)Disk impls for tree::Roots and stubs for NoteCommitmentTrees

* Track anchors and note commitment trees in Chain

Append note commitments to their trees when doing update_chain_state_with,
then use the resulting Sapling and Orchard roots to pass to history_tree, and add
new roots to the anchor sets.

* Handle errors when appending to note commitment trees

* Add comments explaining why note commitment are not removed from the tree in revert_chain_state_with

* Implementing note commitments in finalized state

* Finish serialization of Orchard tree; remove old tree when updating finalize state

* Add serialization and finalized state updates for Sprout and Sapling trees

* Partially handle trees in non-finalized state. Use Option for trees in Chain

* Rebuild trees when forking; change finalized state tree getters to not require height

* Pass empty trees to tests; use empty trees by default in Chain

* Also rebuild anchor sets when forking

* Use empty tree as default in finalized state tree getters (for now)

* Use HashMultiSet for anchors in order to make pop_root() work correctly

* Reduce DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES and MAX_PARTIAL_CHAIN_BLOCKS

* Reduce DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES and MAX_PARTIAL_CHAIN_BLOCKS even more

* Apply suggestions from code review

* Add comments about order of note commitments and related methods/fields

* Don't use Option for trees

* Set DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES=1 and restore MAX_PARTIAL_CHAIN_BLOCKS

* Remove unneeded anchor set rebuilding in fork()

* Improve proptest formatting

* Add missing comparisons to eq_internal_state

* Renamed sprout::tree::NoteCommitmentTree::hash() to root()

* Improve comments

* Add asserts, add issues to TODOs

* Remove impl Default for Chain since it was only used by tests

* Improve documentation and assertions; add tree serialization tests

* Remove Sprout code, which will be moved to another branch

* Add todo! in Sprout tree append()

* Remove stub request, response *Anchor* handling for now

* Add test for validating Sapling note commitment tree using test blocks

* Increase database version (new columns added for note commitment trees and anchors)

* Update test to make sure the order of sapling_note_commitments() is being tested

* Improve comments and structure of the test

* Improve variable names again

* Rustfmt

Co-authored-by: Deirdre Connolly <deirdre@zfnd.org>
Co-authored-by: Conrado P. L. Gouvea <conradoplg@gmail.com>
Co-authored-by: Conrado Gouvea <conrado@zfnd.org>
Co-authored-by: teor <teor@riseup.net>

zebra-chain/Cargo.toml

@@ -15,10 +15,12 @@ bench = ["zebra-test"]
 [dependencies]
 aes = "0.6"
 bech32 = "0.8.1"
+bigint = "4"
 bitflags = "1.2.1"
 bitvec = "0.22"
 blake2b_simd = "0.5.11"
 blake2s_simd = "0.5.11"
+bls12_381 = "0.5.0"
 bs58 = { version = "0.4", features = ["check"] }
 byteorder = "1.4"
 chrono = { version = "0.4", features = ["serde"] }
@@ -30,6 +32,7 @@ group = "0.10"
 # Note: if updating this, also update the workspace Cargo.toml to match.
 halo2 = { git = "https://github.com/zcash/halo2.git", rev = "236115917df9db45282fec24d1e1e36f275f71ab" }
 hex = "0.4"
+incrementalmerkletree = "0.1.0"
 jubjub = "0.7.0"
 lazy_static = "1.4.0"
 rand_core = "0.6"
@@ -40,13 +43,10 @@ serde-big-array = "0.3.2"
 sha2 = { version = "0.9.5", features=["compress"] }
 subtle = "2.4"
 thiserror = "1"
+uint = "0.9.1"
 x25519-dalek = { version = "1.1", features = ["serde"] }
 zcash_history = { git = "https://github.com/zcash/librustzcash.git", rev = "0c3ed159985affa774e44d10172d4471d798a85a" }
 zcash_primitives = { git = "https://github.com/zcash/librustzcash.git", rev = "0c3ed159985affa774e44d10172d4471d798a85a" }
-bigint = "4"
-uint = "0.9.1"
-bls12_381 = "0.5.0"
-incrementalmerkletree = "0.1.0"
 
 proptest = { version = "0.10", optional = true }
 proptest-derive = { version = "0.3.0", optional = true }

zebra-chain/src/orchard/shielded_data.rs

@@ -1,5 +1,14 @@
 //! Orchard shielded data for `V5` `Transaction`s.
 
+use std::{
+    cmp::{Eq, PartialEq},
+    fmt::Debug,
+    io,
+};
+
+use byteorder::{ReadBytesExt, WriteBytesExt};
+use halo2::pasta::pallas;
+
 use crate::{
     amount::{Amount, NegativeAllowed},
     block::MAX_BLOCK_BYTES,
@@ -13,14 +22,6 @@ use crate::{
     },
 };
 
-use byteorder::{ReadBytesExt, WriteBytesExt};
-
-use std::{
-    cmp::{Eq, PartialEq},
-    fmt::Debug,
-    io,
-};
-
 /// A bundle of [`Action`] descriptions and signature data.
 #[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
 pub struct ShieldedData {
@@ -32,14 +33,15 @@ pub struct ShieldedData {
     pub shared_anchor: tree::Root,
     /// The aggregated zk-SNARK proof for all the actions in this transaction.
     pub proof: Halo2Proof,
-    /// The Orchard Actions.
+    /// The Orchard Actions, in the order they appear in the transaction.
     pub actions: AtLeastOne<AuthorizedAction>,
     /// A signature on the transaction `sighash`.
     pub binding_sig: Signature<Binding>,
 }
 
 impl ShieldedData {
-    /// Iterate over the [`Action`]s for the [`AuthorizedAction`]s in this transaction.
+    /// Iterate over the [`Action`]s for the [`AuthorizedAction`]s in this
+    /// transaction, in the order they appear in it.
     pub fn actions(&self) -> impl Iterator<Item = &Action> {
         self.actions.actions()
     }
@@ -55,6 +57,12 @@ impl ShieldedData {
     pub fn value_balance(&self) -> Amount<NegativeAllowed> {
         self.value_balance
     }
+
+    /// Collect the cm_x's for this transaction, if it contains [`Action`]s with
+    /// outputs, in the order they appear in the transaction.
+    pub fn note_commitments(&self) -> impl Iterator<Item = &pallas::Base> {
+        self.actions().map(|action| &action.cm_x)
+    }
 }
 
 impl AtLeastOne<AuthorizedAction> {

zebra-chain/src/orchard/tree.rs

@@ -111,6 +111,12 @@ impl From<Root> for [u8; 32] {
     }
 }
 
+impl From<&Root> for [u8; 32] {
+    fn from(root: &Root) -> Self {
+        (*root).into()
+    }
+}
+
 impl Hash for Root {
     fn hash<H: Hasher>(&self, state: &mut H) {
         self.0.to_bytes().hash(state)
@@ -177,15 +183,36 @@ 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,
+    {
+        self.0.to_bytes().serialize(serializer)
+    }
+}
+
+impl<'de> serde::Deserialize<'de> for Node {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        let bytes = <[u8; 32]>::deserialize(deserializer)?;
+        Option::<pallas::Base>::from(pallas::Base::from_bytes(&bytes))
+            .map(Node)
+            .ok_or_else(|| serde::de::Error::custom("invalid Pallas field element"))
+    }
+}
+
 #[allow(dead_code, missing_docs)]
-#[derive(Error, Debug, PartialEq)]
+#[derive(Error, Debug, PartialEq, Eq)]
 pub enum NoteCommitmentTreeError {
     #[error("The note commitment tree is full")]
     FullTree,
 }
 
 /// 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/sapling/shielded_data.rs

@@ -104,7 +104,7 @@ where
     pub value_balance: Amount,
 
     /// A bundle of spends and outputs, containing at least one spend or
-    /// output.
+    /// output, in the order they appear in the transaction.
     ///
     /// In V5 transactions, also contains a shared anchor, if there are any
     /// spends.
@@ -154,7 +154,8 @@ where
         /// [`Spend`]s in this `TransferData`.
         spends: AtLeastOne<Spend<AnchorV>>,
 
-        /// Maybe some outputs (can be empty).
+        /// Maybe some outputs (can be empty), in the order they appear in the
+        /// transaction.
         ///
         /// Use the [`ShieldedData::outputs`] method to get an iterator over the
         /// [`Outputs`]s in this `TransferData`.
@@ -167,7 +168,7 @@ where
     /// In Transaction::V5, if there are no spends, there must not be a shared
     /// anchor.
     JustOutputs {
-        /// At least one output.
+        /// At least one output, in the order they appear in the transaction.
         ///
         /// Use the [`ShieldedData::outputs`] method to get an iterator over the
         /// [`Outputs`]s in this `TransferData`.
@@ -205,7 +206,8 @@ where
         self.transfers.spends()
     }
 
-    /// Iterate over the [`Output`]s for this transaction.
+    /// Iterate over the [`Output`]s for this transaction, in the order they
+    /// appear in it.
     pub fn outputs(&self) -> impl Iterator<Item = &Output> {
         self.transfers.outputs()
     }
@@ -225,7 +227,8 @@ where
         self.spends().map(|spend| &spend.nullifier)
     }
 
-    /// Collect the cm_u's for this transaction, if it contains [`Output`]s.
+    /// Collect the cm_u's for this transaction, if it contains [`Output`]s,
+    /// in the order they appear in the transaction.
     pub fn note_commitments(&self) -> impl Iterator<Item = &jubjub::Fq> {
         self.outputs().map(|output| &output.cm_u)
     }

zebra-chain/src/sapling/tests/tree.rs

@@ -1,7 +1,17 @@
+use std::sync::Arc;
+
+use color_eyre::eyre;
+use eyre::Result;
 use hex::FromHex;
 
-use crate::sapling::tests::test_vectors;
-use crate::sapling::tree::*;
+use crate::block::Block;
+use crate::parameters::NetworkUpgrade;
+use crate::sapling::{self, tree::*};
+use crate::serialization::ZcashDeserializeInto;
+use crate::{parameters::Network, sapling::tests::test_vectors};
+use zebra_test::vectors::{
+    MAINNET_BLOCKS, MAINNET_FINAL_SAPLING_ROOTS, TESTNET_BLOCKS, TESTNET_FINAL_SAPLING_ROOTS,
+};
 
 #[test]
 fn empty_roots() {
@@ -41,3 +51,83 @@ fn incremental_roots() {
         );
     }
 }
+
+#[test]
+fn incremental_roots_with_blocks() -> Result<()> {
+    incremental_roots_with_blocks_for_network(Network::Mainnet)?;
+    incremental_roots_with_blocks_for_network(Network::Testnet)?;
+    Ok(())
+}
+
+fn incremental_roots_with_blocks_for_network(network: Network) -> Result<()> {
+    let (blocks, sapling_roots) = match network {
+        Network::Mainnet => (&*MAINNET_BLOCKS, &*MAINNET_FINAL_SAPLING_ROOTS),
+        Network::Testnet => (&*TESTNET_BLOCKS, &*TESTNET_FINAL_SAPLING_ROOTS),
+    };
+    let height = NetworkUpgrade::Sapling
+        .activation_height(network)
+        .unwrap()
+        .0;
+
+    // Build empty note commitment tree
+    let mut tree = sapling::tree::NoteCommitmentTree::default();
+
+    // Load Sapling activation block
+    let sapling_activation_block = Arc::new(
+        blocks
+            .get(&height)
+            .expect("test vector exists")
+            .zcash_deserialize_into::<Block>()
+            .expect("block is structurally valid"),
+    );
+
+    // Add note commitments from the Sapling activation block to the tree
+    for transaction in sapling_activation_block.transactions.iter() {
+        for sapling_note_commitment in transaction.sapling_note_commitments() {
+            tree.append(*sapling_note_commitment)
+                .expect("test vector is correct");
+        }
+    }
+
+    // Check if root of the tree of the activation block is correct
+    let sapling_activation_block_root =
+        sapling::tree::Root(**sapling_roots.get(&height).expect("test vector exists"));
+    assert_eq!(sapling_activation_block_root, tree.root());
+
+    // Load the block immediately after Sapling activation (activation + 1)
+    let block_after_sapling_activation = Arc::new(
+        blocks
+            .get(&(height + 1))
+            .expect("test vector exists")
+            .zcash_deserialize_into::<Block>()
+            .expect("block is structurally valid"),
+    );
+    let block_after_sapling_activation_root = sapling::tree::Root(
+        **sapling_roots
+            .get(&(height + 1))
+            .expect("test vector exists"),
+    );
+
+    // Add note commitments from the block after Sapling activatoin to the tree
+    let mut appended_count = 0;
+    for transaction in block_after_sapling_activation.transactions.iter() {
+        for sapling_note_commitment in transaction.sapling_note_commitments() {
+            tree.append(*sapling_note_commitment)
+                .expect("test vector is correct");
+            appended_count += 1;
+        }
+    }
+    // We also want to make sure that sapling_note_commitments() is returning
+    // the commitments in the right order. But this will only be actually tested
+    // if there are more than one note commitment in a block.
+    // In the test vectors this applies only for the block 1 in mainnet,
+    // so we make this explicit in this assert.
+    if network == Network::Mainnet {
+        assert!(appended_count > 1);
+    }
+
+    // Check if root of the second block is correct
+    assert_eq!(block_after_sapling_activation_root, tree.root());
+
+    Ok(())
+}

zebra-chain/src/sapling/tree.rs

@@ -147,15 +147,36 @@ impl From<jubjub::Fq> for Node {
     }
 }
 
+impl serde::Serialize for Node {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        self.0.serialize(serializer)
+    }
+}
+
+impl<'de> serde::Deserialize<'de> for Node {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        let bytes = <[u8; 32]>::deserialize(deserializer)?;
+        Option::<jubjub::Fq>::from(jubjub::Fq::from_bytes(&bytes))
+            .map(Node::from)
+            .ok_or_else(|| serde::de::Error::custom("invalid JubJub field element"))
+    }
+}
+
 #[allow(dead_code, missing_docs)]
-#[derive(Error, Debug, PartialEq)]
+#[derive(Error, Debug, PartialEq, Eq)]
 pub enum NoteCommitmentTreeError {
     #[error("The note commitment tree is full")]
     FullTree,
 }
 
 /// Sapling Incremental Note Commitment Tree.
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct NoteCommitmentTree {
     /// The tree represented as a Frontier.
     ///