[Perf] Faster block tree (de)serialization

console/collections/Cargo.toml

@@ -44,6 +44,9 @@ workspace = true
 [dev-dependencies.snarkvm-console-network]
 path = "../network"
 
+[dev-dependencies.bincode]
+workspace = true
+
 [dev-dependencies.criterion]
 workspace = true
 

console/collections/benches/merkle_tree.rs

@@ -16,15 +16,20 @@
 #[macro_use]
 extern crate criterion;
 
+use snarkvm_console_algorithms::{BHP512, BHP1024};
+use snarkvm_console_collections::merkle_tree::MerkleTreeState;
 use snarkvm_console_network::{
+    BHP_512,
+    BHP_1024,
+    BHPMerkleTree,
     MainnetV0,
     Network,
     prelude::{Rng, TestRng, ToBits, Uniform},
 };
 use snarkvm_console_types::Field;
 
 use criterion::{BatchSize, BenchmarkId, Criterion};
-use std::collections::BTreeMap;
+use std::{collections::BTreeMap, time::Duration};
 
 const DEPTH: u8 = 32;
 const MAX_INSTANTIATED_DEPTH: u8 = 8;
@@ -33,6 +38,10 @@ const NUM_LEAVES: &[usize] = &[1, 100, 10_000];
 const APPEND_SIZES: &[usize] = &[1, 100, 10_000];
 const UPDATE_SIZES: &[usize] = &[1, 100, 1_000];
 
+/// The tree sizes used by the `MerkleTreeState` benchmarks; these reach further
+/// than `NUM_LEAVES`, as caching a tree is most interesting for large trees.
+const STATE_NUM_LEAVES: &[usize] = &[1, 100, 10_000, 1_000_000];
+
 /// Generates the specified number of random Merkle tree leaves.
 macro_rules! generate_leaves {
     ($num_leaves:expr, $rng:expr) => {{ (0..$num_leaves).map(|_| Field::<MainnetV0>::rand($rng).to_bits_le()).collect::<Vec<_>>() }};
@@ -175,9 +184,110 @@ fn update_vs_update_many(c: &mut Criterion) {
     }
 }
 
+/// A Merkle tree as it was cached before [`MerkleTreeState`] was introduced, i.e. the
+/// whole tree, hashers included.
+///
+/// Note: `bincode` encodes a tuple and a struct identically - as the concatenation of
+/// their fields - so this decodes the legacy payload byte for byte.
+type LegacyCachedTree<'a> = (BHP1024<MainnetV0>, BHP512<MainnetV0>, MerkleTreeState<'a, MainnetV0>);
+
+/// Produces the payload that caching the given tree used to write, for comparison.
+fn legacy_payload(merkle_tree: &BHPMerkleTree<MainnetV0, DEPTH>) -> Vec<u8> {
+    bincode::serialize(&(&*BHP_1024, &*BHP_512, merkle_tree.to_state())).unwrap()
+}
+
+/// Compares recreating a Merkle tree from a cached [`MerkleTreeState`] against
+/// constructing the same tree from scratch, across a range of tree sizes.
+///
+/// Run with `cargo bench --bench merkle_tree -- MerkleTreeState`.
+fn state(c: &mut Criterion) {
+    let mut group = c.benchmark_group("MerkleTreeState");
+    let mut rng = TestRng::default();
+    // Accumulate leaves in a vector to avoid recomputing across iterations.
+    let leaves = generate_leaves!(*STATE_NUM_LEAVES.last().unwrap(), &mut rng);
+
+    for num_leaves in STATE_NUM_LEAVES {
+        // Construct a Merkle tree with the specified number of leaves, and cache its state.
+        let merkle_tree = MainnetV0::merkle_tree_bhp::<DEPTH>(&leaves[..*num_leaves]).unwrap();
+        let cached = bincode::serialize(&merkle_tree.to_state()).unwrap();
+        // Report the size of the cached payload against that of the legacy one; the former
+        // should be proportional to the tree, i.e. it must not contain the hashers.
+        println!(
+            "MerkleTreeState/{num_leaves} leaves: {} bytes cached ({} bytes in the legacy format)",
+            cached.len(),
+            legacy_payload(&merkle_tree).len()
+        );
+
+        // For reference: constructing the tree from scratch, i.e. not using a cache at all.
+        group.bench_with_input(BenchmarkId::new("from_scratch", num_leaves), num_leaves, |b, num_leaves| {
+            b.iter(|| MainnetV0::merkle_tree_bhp::<DEPTH>(&leaves[..*num_leaves]).unwrap())
+        });
+        // Serializing the state of the tree, i.e. the cost of writing the cache.
+        group.bench_with_input(BenchmarkId::new("serialize", num_leaves), &merkle_tree, |b, merkle_tree| {
+            b.iter(|| bincode::serialize(&merkle_tree.to_state()).unwrap())
+        });
+        // Deserializing the state, i.e. the decoding half of the cost of reading the cache.
+        group.bench_with_input(BenchmarkId::new("deserialize", num_leaves), &cached, |b, cached| {
+            b.iter(|| bincode::deserialize::<MerkleTreeState<'_, MainnetV0>>(cached).unwrap())
+        });
+        // Recreating the tree from a deserialized state, i.e. the validating half.
+        group.bench_with_input(BenchmarkId::new("from_state", num_leaves), &cached, |b, cached| {
+            b.iter_batched(
+                || bincode::deserialize::<MerkleTreeState<'_, MainnetV0>>(cached).unwrap(),
+                |state| MainnetV0::merkle_tree_bhp_from_state::<DEPTH>(state).unwrap(),
+                BatchSize::LargeInput,
+            )
+        });
+        // Both halves together, i.e. what a node pays to load a cached tree.
+        group.bench_with_input(BenchmarkId::new("deserialize_and_recreate", num_leaves), &cached, |b, cached| {
+            b.iter(|| {
+                let state = bincode::deserialize::<MerkleTreeState<'_, MainnetV0>>(cached).unwrap();
+                MainnetV0::merkle_tree_bhp_from_state::<DEPTH>(state).unwrap()
+            })
+        });
+    }
+
+    group.finish();
+}
+
+/// Measures loading a tree from a legacy cache payload, i.e. one that carries the hashers.
+///
+/// This is kept apart from [`state`] because a single iteration takes tens of seconds: the
+/// hashers hold hundreds of thousands of group elements, and deserializing each one costs a
+/// subgroup check, i.e. a full scalar multiplication. That cost is also invariant in the
+/// size of the tree, which is why this is measured for a single tree size only.
+///
+/// Run with `cargo bench --bench merkle_tree -- LegacyMerkleTreeCache` (slow).
+fn legacy_state(c: &mut Criterion) {
+    let mut group = c.benchmark_group("LegacyMerkleTreeCache");
+    let mut rng = TestRng::default();
+
+    let num_leaves = STATE_NUM_LEAVES[0];
+    let leaves = generate_leaves!(num_leaves, &mut rng);
+    let merkle_tree = MainnetV0::merkle_tree_bhp::<DEPTH>(&leaves).unwrap();
+    let legacy = legacy_payload(&merkle_tree);
+    println!("LegacyMerkleTreeCache/{num_leaves} leaves: {} bytes cached", legacy.len());
+
+    group.bench_with_input(BenchmarkId::new("deserialize", num_leaves), &legacy, |b, legacy| {
+        b.iter(|| bincode::deserialize::<LegacyCachedTree<'_>>(legacy).unwrap())
+    });
+
+    group.finish();
+}
+
 criterion_group! {
     name = merkle_tree;
     config = Criterion::default().sample_size(10);
     targets = new, append, update, update_many, update_vs_update_many
 }
-criterion_main!(merkle_tree);
+criterion_group! {
+    name = merkle_tree_state;
+    config = Criterion::default().sample_size(10).warm_up_time(Duration::from_secs(1));
+    targets = state
+}
+criterion_group! {
+    name = legacy_merkle_tree_cache;
+    config = Criterion::default().sample_size(10).warm_up_time(Duration::from_secs(1));
+    targets = legacy_state
+}
+criterion_main!(merkle_tree, merkle_tree_state, legacy_merkle_tree_cache);

console/collections/src/merkle_tree/mod.rs

@@ -31,7 +31,7 @@ use locktick::parking_lot::Mutex;
 #[cfg(not(feature = "locktick"))]
 use parking_lot::Mutex;
 use serde::{Deserialize, Serialize};
-use std::{collections::BTreeMap, mem};
+use std::{borrow::Cow, collections::BTreeMap, mem};
 
 #[cfg(not(feature = "serial"))]
 use rayon::prelude::*;
@@ -50,8 +50,6 @@ use rayon::prelude::*;
 /// Padding levels are then added as needed to reach the full `DEPTH`, each of
 /// which is constructed by hashing the root of the previous level together with
 /// `e`.
-#[derive(Deserialize, Serialize)]
-#[serde(bound = "E: Serialize + DeserializeOwned, LH: Serialize + DeserializeOwned, PH: Serialize + DeserializeOwned")]
 pub struct MerkleTree<E: Environment, LH: LeafHash<Hash = PH::Hash>, PH: PathHash<Hash = Field<E>>, const DEPTH: u8> {
     /// The leaf hasher for the Merkle tree.
     leaf_hasher: LH,
@@ -66,10 +64,30 @@ pub struct MerkleTree<E: Environment, LH: LeafHash<Hash = PH::Hash>, PH: PathHas
     /// The number of hashed leaves in the tree.
     number_of_leaves: usize,
     /// An optimization: the previous tree allocation reused in prepare_append.
-    #[serde(skip)]
     preserved_tree_allocation: Mutex<Option<Vec<PH::Hash>>>,
 }
 
+/// The contents of a [`MerkleTree`], sans its hashers.
+///
+/// This is the serializable form of a Merkle tree, intended for caching one on
+/// disk. Note that a [`MerkleTree`] itself is deliberately **not** serializable:
+/// its hashers hold precomputed bases (tens of MiBs of group elements for the
+/// BHP hashers), and deserializing those is orders of magnitude more expensive
+/// than setting them up from scratch - each group element costs a subgroup check,
+/// i.e. a full scalar multiplication.
+#[derive(Clone, Debug, Deserialize, Serialize)]
+#[serde(bound = "E: Serialize + DeserializeOwned")]
+pub struct MerkleTreeState<'a, E: Environment> {
+    /// The computed root of the full Merkle tree.
+    root: Field<E>,
+    /// The internal hashes, from root to hashed leaves, of the full Merkle tree.
+    tree: Cow<'a, [Field<E>]>,
+    /// The canonical empty hash.
+    empty_hash: Field<E>,
+    /// The number of hashed leaves in the tree.
+    number_of_leaves: usize,
+}
+
 impl<E: Environment, LH: LeafHash<Hash = PH::Hash>, PH: PathHash<Hash = Field<E>>, const DEPTH: u8> Clone
     for MerkleTree<E, LH, PH, DEPTH>
 {
@@ -178,6 +196,80 @@ impl<E: Environment, LH: LeafHash<Hash = PH::Hash>, PH: PathHash<Hash = Field<E>
         })
     }
 
+    /// Returns the contents of the Merkle tree, sans its hashers.
+    ///
+    /// This borrows from the tree, so it is cheap even for very large trees; use
+    /// [`Self::from_state`] to recreate the tree from the returned state.
+    pub fn to_state(&self) -> MerkleTreeState<'_, E> {
+        MerkleTreeState {
+            root: self.root,
+            tree: Cow::Borrowed(&self.tree),
+            empty_hash: self.empty_hash,
+            number_of_leaves: self.number_of_leaves,
+        }
+    }
+
+    /// Recreates a Merkle tree from the given state, using the given hashers.
+    ///
+    /// The state is checked for internal consistency, which includes recomputing
+    /// the root from the topmost node; since only the padding levels are hashed,
+    /// this is cheap. Note that this cannot attest that the tree corresponds to
+    /// any particular set of leaves, so the caller is still expected to check the
+    /// resulting root against a trusted value.
+    pub fn from_state(leaf_hasher: &LH, path_hasher: &PH, state: MerkleTreeState<'_, E>) -> Result<Self> {
+        // Ensure the Merkle tree depth is greater than 0.
+        ensure!(DEPTH > 0, "Merkle tree depth must be greater than 0");
+        // Ensure the Merkle tree depth is less than or equal to 64.
+        ensure!(DEPTH <= 64u8, "Merkle tree depth must be less than or equal to 64");
+
+        let MerkleTreeState { root, tree, empty_hash, number_of_leaves } = state;
+        // Note: this is a no-op if the state was deserialized, as opposed to borrowed.
+        let tree = tree.into_owned();
+
+        // Ensure the empty hash matches the one produced by the given path hasher; a
+        // mismatch means that the state was produced for a different network or hasher.
+        ensure!(empty_hash == path_hasher.hash_empty()?, "The Merkle tree state has an invalid empty hash");
+
+        // Compute the maximum number of leaves.
+        let max_leaves = match number_of_leaves.checked_next_power_of_two() {
+            Some(num_leaves) => num_leaves,
+            None => bail!("Integer overflow when computing the maximum number of leaves in the Merkle tree"),
+        };
+        // Compute the number of nodes.
+        let num_nodes = max_leaves - 1;
+        // Compute the number of padded levels.
+        let padding_depth = DEPTH - tree_depth::<DEPTH>(max_leaves + num_nodes)?;
+
+        // Ensure the tree contains exactly as many nodes as its number of leaves implies.
+        let minimum_tree_size = std::cmp::max(
+            1,
+            num_nodes + number_of_leaves + if number_of_leaves > 1 { number_of_leaves % 2 } else { 0 },
+        );
+        ensure!(
+            tree.len() == minimum_tree_size,
+            "The Merkle tree state contains {} nodes, expected {minimum_tree_size} for {number_of_leaves} leaves",
+            tree.len()
+        );
+
+        // Recompute the root hash, by iterating from the root level up to `DEPTH`.
+        let mut root_hash = tree[0];
+        for _ in 0..padding_depth {
+            // Update the root hash, by hashing the current root hash with the empty hash.
+            root_hash = path_hasher.hash_children(&root_hash, &empty_hash)?;
+        }
+        ensure!(root_hash == root, "The Merkle tree state has an invalid root");
+
+        Ok(Self {
+            leaf_hasher: leaf_hasher.clone(),
+            path_hasher: path_hasher.clone(),
+            root,
+            tree,
+            empty_hash,
+            number_of_leaves,
+            preserved_tree_allocation: Default::default(),
+        })
+    }
+
     #[inline]
     /// Returns a new Merkle tree with the given new leaves appended to it.
     pub fn prepare_append(&self, new_leaves: &[LH::Leaf]) -> Result<Self> {

console/collections/src/merkle_tree/tests/mod.rs

@@ -17,6 +17,7 @@ use super::*;
 
 mod append;
 mod remove;
+mod state;
 mod test_print;
 mod update;
 mod update_many;