Cache note commitment tree roots (#2584)

* Cache note commitment tree roots

* Add comments to cached root fields

* Apply suggestions from code review

Co-authored-by: teor <teor@riseup.net>

Co-authored-by: teor <teor@riseup.net>

zebra-chain/src/orchard/tree.rs

@@ -15,6 +15,7 @@
 #![allow(dead_code)]
 
 use std::{
+    cell::Cell,
     convert::TryFrom,
     fmt,
     hash::{Hash, Hasher},
@@ -221,6 +222,24 @@ pub struct NoteCommitmentTree {
     /// has non-empty nodes. Upper (near root) empty nodes of the branch are not
     /// stored.
     inner: bridgetree::Frontier<Node, { MERKLE_DEPTH as u8 }>,
+    /// A cached root of the tree.
+    ///
+    /// Every time the root is computed by [`Self::root`] it is cached here,
+    /// and the cached value will be returned by [`Self::root`] until the tree is
+    /// changed by [`Self::append`]. This greatly increases performance
+    /// because it avoids recomputing the root when the tree does not change
+    /// between blocks. In the finalized state, the tree is read from
+    /// disk for every block processed, which would also require recomputing
+    /// the root even if it has not changed (note that the cached root is
+    /// serialized with the tree). This is particularly important since we decided
+    /// to instantiate the trees from the genesis block, for simplicity.
+    ///
+    /// [`Cell`] offers interior mutability (it works even with a non-mutable
+    /// reference to the tree) but it prevents the tree (and anything that uses it)
+    /// from being shared between threads. If this ever becomes an issue we can
+    /// leave caching to the callers (which requires much more code), or replace
+    /// `Cell` with `Arc<Mutex<_>>` (and be careful of deadlocks and async code.)
+    cached_root: Cell<Option<Root>>,
 }
 
 impl NoteCommitmentTree {
@@ -233,6 +252,8 @@ impl NoteCommitmentTree {
     /// Returns an error if the tree is full.
     pub fn append(&mut self, cm_x: pallas::Base) -> Result<(), NoteCommitmentTreeError> {
         if self.inner.append(&cm_x.into()) {
+            // Invalidate cached root
+            self.cached_root.replace(None);
             Ok(())
         } else {
             Err(NoteCommitmentTreeError::FullTree)
@@ -242,7 +263,16 @@ impl NoteCommitmentTree {
     /// Returns the current root of the tree, used as an anchor in Orchard
     /// shielded transactions.
     pub fn root(&self) -> Root {
-        Root(self.inner.root().0)
+        match self.cached_root.get() {
+            // Return cached root
+            Some(root) => root,
+            None => {
+                // Compute root and cache it
+                let root = Root(self.inner.root().0);
+                self.cached_root.replace(Some(root));
+                root
+            }
+        }
     }
 
     /// Get the Pallas-based Sinsemilla hash / root node of this merkle tree of
@@ -274,6 +304,7 @@ impl Default for NoteCommitmentTree {
     fn default() -> Self {
         Self {
             inner: bridgetree::Frontier::new(),
+            cached_root: Default::default(),
         }
     }
 }

zebra-chain/src/sapling/tree.rs

@@ -13,7 +13,7 @@
 #![allow(clippy::unit_arg)]
 #![allow(dead_code)]
 
-use std::fmt;
+use std::{cell::Cell, fmt};
 
 use bitvec::prelude::*;
 use incrementalmerkletree::{bridgetree, Frontier};
@@ -185,6 +185,24 @@ pub struct NoteCommitmentTree {
     /// has non-empty nodes. Upper (near root) empty nodes of the branch are not
     /// stored.
     inner: bridgetree::Frontier<Node, { MERKLE_DEPTH as u8 }>,
+    /// A cached root of the tree.
+    ///
+    /// Every time the root is computed by [`Self::root`] it is cached here,
+    /// and the cached value will be returned by [`Self::root`] until the tree is
+    /// changed by [`Self::append`]. This greatly increases performance
+    /// because it avoids recomputing the root when the tree does not change
+    /// between blocks. In the finalized state, the tree is read from
+    /// disk for every block processed, which would also require recomputing
+    /// the root even if it has not changed (note that the cached root is
+    /// serialized with the tree). This is particularly important since we decided
+    /// to instantiate the trees from the genesis block, for simplicity.
+    ///
+    /// [`Cell`] offers interior mutability (it works even with a non-mutable
+    /// reference to the tree) but it prevents the tree (and anything that uses it)
+    /// from being shared between threads. If this ever becomes an issue we can
+    /// leave caching to the callers (which requires much more code), or replace
+    /// `Cell` with `Arc<Mutex<_>>` (and be careful of deadlocks and async code.)
+    cached_root: Cell<Option<Root>>,
 }
 
 impl NoteCommitmentTree {
@@ -197,6 +215,8 @@ impl NoteCommitmentTree {
     /// Returns an error if the tree is full.
     pub fn append(&mut self, cm_u: jubjub::Fq) -> Result<(), NoteCommitmentTreeError> {
         if self.inner.append(&cm_u.into()) {
+            // Invalidate cached root
+            self.cached_root.replace(None);
             Ok(())
         } else {
             Err(NoteCommitmentTreeError::FullTree)
@@ -206,7 +226,16 @@ impl NoteCommitmentTree {
     /// Returns the current root of the tree, used as an anchor in Sapling
     /// shielded transactions.
     pub fn root(&self) -> Root {
-        Root(self.inner.root().0)
+        match self.cached_root.get() {
+            // Return cached root
+            Some(root) => root,
+            None => {
+                // Compute root and cache it
+                let root = Root(self.inner.root().0);
+                self.cached_root.replace(Some(root));
+                root
+            }
+        }
     }
 
     /// Get the Jubjub-based Pedersen hash of root node of this merkle tree of
@@ -238,6 +267,7 @@ impl Default for NoteCommitmentTree {
     fn default() -> Self {
         Self {
             inner: bridgetree::Frontier::new(),
+            cached_root: Default::default(),
         }
     }
 }

zebra-state/src/constants.rs

@@ -25,7 +25,7 @@ pub const MIN_TRANSPARENT_COINBASE_MATURITY: u32 = 100;
 pub const MAX_BLOCK_REORG_HEIGHT: u32 = MIN_TRANSPARENT_COINBASE_MATURITY - 1;
 
 /// The database format version, incremented each time the database format changes.
-pub const DATABASE_FORMAT_VERSION: u32 = 8;
+pub const DATABASE_FORMAT_VERSION: u32 = 9;
 
 /// The maximum number of blocks to check for NU5 transactions,
 /// before we assume we are on a pre-NU5 legacy chain.

zebra-state/src/service/finalized_state/tests/vectors.rs

@@ -18,17 +18,24 @@ fn sapling_note_commitment_tree_serialization() {
         "7c3ea01a6e3a3d90cf59cd789e467044b5cd78eb2c84cc6816f960746d0e036c",
     ];
 
-    for cm_u_hex in hex_commitments {
+    for (idx, cm_u_hex) in hex_commitments.iter().enumerate() {
         let bytes = <[u8; 32]>::from_hex(cm_u_hex).unwrap();
 
         let cm_u = jubjub::Fq::from_bytes(&bytes).unwrap();
         incremental_tree.append(cm_u).unwrap();
+        if idx % 2 == 0 {
+            // Cache the root half of the time to make sure it works in both cases
+            let _ = incremental_tree.root();
+        }
     }
 
+    // Make sure the last root is cached
+    let _ = incremental_tree.root();
+
     // This test vector was generated by the code itself.
     // The purpose of this test is to make sure the serialization format does
     // not change by accident.
-    let expected_serialized_tree_hex = "0102007c3ea01a6e3a3d90cf59cd789e467044b5cd78eb2c84cc6816f960746d0e036c0162324ff2c329e99193a74d28a585a3c167a93bf41a255135529c913bd9b1e666";
+    let expected_serialized_tree_hex = "0102007c3ea01a6e3a3d90cf59cd789e467044b5cd78eb2c84cc6816f960746d0e036c0162324ff2c329e99193a74d28a585a3c167a93bf41a255135529c913bd9b1e66601ddaa1ab86de5c153993414f34ba97e9674c459dfadde112b89eeeafa0e5a204c";
     let serialized_tree = incremental_tree.as_bytes();
     assert_eq!(hex::encode(&serialized_tree), expected_serialized_tree_hex);
 
@@ -62,15 +69,22 @@ fn orchard_note_commitment_tree_serialization() {
         ],
     ];
 
-    for cm_x_bytes in &commitments {
+    for (idx, cm_x_bytes) in commitments.iter().enumerate() {
         let cm_x = pallas::Base::from_bytes(cm_x_bytes).unwrap();
         incremental_tree.append(cm_x).unwrap();
+        if idx % 2 == 0 {
+            // Cache the root half of the time to make sure it works in both cases
+            let _ = incremental_tree.root();
+        }
     }
 
+    // Make sure the last root is cached
+    let _ = incremental_tree.root();
+
     // This test vector was generated by the code itself.
     // The purpose of this test is to make sure the serialization format does
     // not change by accident.
-    let expected_serialized_tree_hex = "010200ee9488053a30c596b43014105d3477e6f578c89240d1d1ee1743b77bb6adc40a01a34b69a4e4d9ccf954d46e5da1004d361a5497f511aeb4d481d23c0be1778133";
+    let expected_serialized_tree_hex = "010200ee9488053a30c596b43014105d3477e6f578c89240d1d1ee1743b77bb6adc40a01a34b69a4e4d9ccf954d46e5da1004d361a5497f511aeb4d481d23c0be177813301a0be6dab19bc2c65d8299258c16e14d48ec4d4959568c6412aa85763c222a702";
     let serialized_tree = incremental_tree.as_bytes();
     assert_eq!(hex::encode(&serialized_tree), expected_serialized_tree_hex);