From e4c886e7a0b092ed28a48bb4ac064a4a5a3821b3 Mon Sep 17 00:00:00 2001 From: Oleksandr Tkachenko Date: Wed, 9 Aug 2023 09:49:13 +0000 Subject: [PATCH] perf(crypto): CRP-2019 iterative MixedHashTree::digest() --- Cargo.lock | 1 + rs/crypto/tree_hash/benches/tree_hash.rs | 18 +++++- rs/crypto/tree_hash/src/lib.rs | 63 ++++++++++++++++--- rs/crypto/tree_hash/test_utils/BUILD.bazel | 1 + rs/crypto/tree_hash/test_utils/Cargo.toml | 1 + rs/crypto/tree_hash/test_utils/src/lib.rs | 30 +++++++++ .../test_utils/tests/mixed_hash_tree.rs | 21 +++++++ rs/crypto/tree_hash/tests/tree_hash.rs | 15 +++++ 8 files changed, 141 insertions(+), 9 deletions(-) create mode 100644 rs/crypto/tree_hash/test_utils/tests/mixed_hash_tree.rs diff --git a/Cargo.lock b/Cargo.lock index 0311247065e..fb93ebb04af 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7011,6 +7011,7 @@ dependencies = [ "ic-crypto-tree-hash", "proptest", "rand 0.8.5", + "thiserror", ] [[package]] diff --git a/rs/crypto/tree_hash/benches/tree_hash.rs b/rs/crypto/tree_hash/benches/tree_hash.rs index eb98369380d..f93ce28f37f 100644 --- a/rs/crypto/tree_hash/benches/tree_hash.rs +++ b/rs/crypto/tree_hash/benches/tree_hash.rs @@ -3,7 +3,9 @@ use ic_crypto_tree_hash::{ flatmap, lookup_path, FlatMap, HashTree, HashTreeBuilder, Label, LabeledTree, LabeledTree::*, MixedHashTree, WitnessGenerator, }; -use ic_crypto_tree_hash_test_utils::hash_tree_builder_from_labeled_tree; +use ic_crypto_tree_hash_test_utils::{ + hash_tree_builder_from_labeled_tree, mixed_hash_tree_digest_recursive, +}; use ic_types_test_utils::ids::message_test_id; fn new_request_status_tree(num_subtrees: usize) -> LabeledTree> { @@ -198,9 +200,21 @@ pub fn criterion_benchmark(c: &mut Criterion) { c.benchmark_group("compute_digest"); g.bench_function(BenchmarkId::new("mixed_hash_tree", num_subtrees), |b| { - b.iter(|| black_box(mixed_hash_tree.digest())); + b.iter(|| { + black_box( + mixed_hash_tree_digest_recursive(&mixed_hash_tree) + .expect("too deep recursion"), + ) + }); }); + g.bench_function( + BenchmarkId::new("mixed_hash_tree_iterative", num_subtrees), + |b| { + b.iter(|| black_box(mixed_hash_tree.digest())); + }, + ); + g.bench_function(BenchmarkId::new("witness", num_subtrees), |b| { b.iter(|| { black_box(ic_crypto_tree_hash::recompute_digest( diff --git a/rs/crypto/tree_hash/src/lib.rs b/rs/crypto/tree_hash/src/lib.rs index efb93fd3a0f..7e34680f016 100644 --- a/rs/crypto/tree_hash/src/lib.rs +++ b/rs/crypto/tree_hash/src/lib.rs @@ -447,15 +447,64 @@ impl MixedHashTree { /// Recomputes root hash of the full tree that this mixed tree was /// constructed from. pub fn digest(&self) -> Digest { - match self { - Self::Empty => tree_hash::empty_subtree_hash(), - Self::Fork(lr) => tree_hash::compute_fork_digest(&lr.0.digest(), &lr.1.digest()), - Self::Labeled(label, subtree) => { - tree_hash::compute_node_digest(label, &subtree.digest()) + #[derive(Debug)] + enum StackItem<'a> { + Expand(&'a MixedHashTree), + Collect(&'a MixedHashTree), + } + + impl<'a> StackItem<'a> { + fn to_collect(&self) -> Self { + match self { + Self::Expand(t) => Self::Collect(t), + Self::Collect(_) => panic!("expected Expand, got Collect"), + } } - Self::Leaf(buf) => tree_hash::compute_leaf_digest(&buf[..]), - Self::Pruned(digest) => digest.clone(), } + + let mut stack: Vec> = Vec::new(); + let mut digests: Vec = Vec::new(); + + stack.push(StackItem::Expand(self)); + + while let Some(t) = stack.pop() { + match t { + StackItem::Expand(Self::Fork(lr)) => { + stack.push(t.to_collect()); + stack.push(StackItem::Expand(&lr.1)); + stack.push(StackItem::Expand(&lr.0)); + } + StackItem::Expand(Self::Labeled(_, subtree)) => { + stack.push(t.to_collect()); + stack.push(StackItem::Expand(subtree)); + } + StackItem::Collect(Self::Fork(_)) => { + let right = digests.pop().expect("bug: missing right subtree digest"); + let left = digests.pop().expect("bug: missing left subtree digest"); + digests.push(tree_hash::compute_fork_digest(&left, &right)); + } + StackItem::Collect(Self::Labeled(label, _)) => { + let subtree_digest = digests.pop().expect("bug: missing subtree digest"); + let labeled_digest = tree_hash::compute_node_digest(label, &subtree_digest); + digests.push(labeled_digest); + } + StackItem::Collect(Self::Leaf(buf)) => { + digests.push(tree_hash::compute_leaf_digest(&buf[..])) + } + StackItem::Collect(Self::Pruned(digest)) => digests.push(digest.clone()), + StackItem::Collect(Self::Empty) => digests.push(tree_hash::empty_subtree_hash()), + t /* Expand of Leaf, Pruned or Empty */ => stack.push(t.to_collect()), + } + } + + assert_eq!( + digests.len(), + 1, + "bug: reduced tree to not exactly one digest: {digests:?}" + ); + assert!(stack.is_empty(), "bug: stack is not empty: {stack:?}"); + + digests[0].clone() } /// Finds a label in a hash tree. diff --git a/rs/crypto/tree_hash/test_utils/BUILD.bazel b/rs/crypto/tree_hash/test_utils/BUILD.bazel index d5e1c56bff5..7f3350b48d2 100644 --- a/rs/crypto/tree_hash/test_utils/BUILD.bazel +++ b/rs/crypto/tree_hash/test_utils/BUILD.bazel @@ -7,6 +7,7 @@ DEPENDENCIES = [ "//rs/crypto/tree_hash", "@crate_index//:proptest", "@crate_index//:rand_0_8_4", + "@crate_index//:thiserror", ] DEV_DEPENDENCIES = [ diff --git a/rs/crypto/tree_hash/test_utils/Cargo.toml b/rs/crypto/tree_hash/test_utils/Cargo.toml index 03b3116148d..7801f113632 100644 --- a/rs/crypto/tree_hash/test_utils/Cargo.toml +++ b/rs/crypto/tree_hash/test_utils/Cargo.toml @@ -8,6 +8,7 @@ assert_matches = "1.5.0" ic-crypto-tree-hash = { path = ".." } proptest = "1.0" rand = "0.8.4" +thiserror = "1.0" [dev-dependencies] ic-crypto-test-utils-reproducible-rng = { path = "../../test_utils/reproducible_rng" } \ No newline at end of file diff --git a/rs/crypto/tree_hash/test_utils/src/lib.rs b/rs/crypto/tree_hash/test_utils/src/lib.rs index 17378eebbe1..dbe4a469120 100644 --- a/rs/crypto/tree_hash/test_utils/src/lib.rs +++ b/rs/crypto/tree_hash/test_utils/src/lib.rs @@ -455,3 +455,33 @@ pub fn compute_fork_digest(left_digest: &Digest, right_digest: &Digest) -> Diges pub fn empty_subtree_hash() -> Digest { Hasher::for_domain(DOMAIN_HASHTREE_EMPTY_SUBTREE).finalize() } + +/// This error indicates that the algorithm exceeded the recursion depth limit. +#[derive(thiserror::Error, Debug, PartialEq)] +#[error("The algorithm failed due to too deep recursion (depth={0})")] +pub struct TooDeepRecursion(pub u32); + +/// Recomputes root hash of the full tree that this mixed tree was +/// constructed from. +pub fn mixed_hash_tree_digest_recursive(tree: &MixedHashTree) -> Result { + fn digest_impl(t: &MixedHashTree, depth: u32) -> Result { + if depth as usize > MAX_HASH_TREE_DEPTH { + return Err(TooDeepRecursion(depth)); + } + let result = match t { + MixedHashTree::Empty => empty_subtree_hash(), + MixedHashTree::Fork(lr) => compute_fork_digest( + &digest_impl(&lr.0, depth + 1)?, + &digest_impl(&lr.1, depth + 1)?, + ), + MixedHashTree::Labeled(label, subtree) => { + compute_node_digest(label, &digest_impl(subtree, depth + 1)?) + } + MixedHashTree::Leaf(buf) => compute_leaf_digest(&buf[..]), + MixedHashTree::Pruned(digest) => digest.clone(), + }; + Ok(result) + } + + digest_impl(tree, 1) +} diff --git a/rs/crypto/tree_hash/test_utils/tests/mixed_hash_tree.rs b/rs/crypto/tree_hash/test_utils/tests/mixed_hash_tree.rs new file mode 100644 index 00000000000..5ccea19530f --- /dev/null +++ b/rs/crypto/tree_hash/test_utils/tests/mixed_hash_tree.rs @@ -0,0 +1,21 @@ +use assert_matches::assert_matches; +use ic_crypto_tree_hash::MixedHashTree; +use ic_crypto_tree_hash_test_utils::{mixed_hash_tree_digest_recursive, TooDeepRecursion}; + +const MAX_HASH_TREE_DEPTH: usize = 128; + +#[test] +fn mixed_hash_tree_recursive_digest_errors_on_too_deep_trees() { + let mut tree = MixedHashTree::Empty; + for _ in 1..MAX_HASH_TREE_DEPTH { + tree = MixedHashTree::Fork(Box::new((tree.clone(), MixedHashTree::Empty))); + } + + assert_matches!(mixed_hash_tree_digest_recursive(&tree), Ok(_)); + + tree = MixedHashTree::Fork(Box::new((tree.clone(), MixedHashTree::Empty))); + assert_eq!( + mixed_hash_tree_digest_recursive(&tree), + Err(TooDeepRecursion(MAX_HASH_TREE_DEPTH as u32 + 1)) + ); +} diff --git a/rs/crypto/tree_hash/tests/tree_hash.rs b/rs/crypto/tree_hash/tests/tree_hash.rs index 3a5784bb0ac..7295dacc8e8 100644 --- a/rs/crypto/tree_hash/tests/tree_hash.rs +++ b/rs/crypto/tree_hash/tests/tree_hash.rs @@ -4,6 +4,7 @@ use ic_crypto_sha2::Sha256; use ic_crypto_test_utils_reproducible_rng::reproducible_rng; use ic_crypto_tree_hash::*; use ic_crypto_tree_hash_test_utils::*; +use proptest::prelude::*; use rand::Rng; use rand::{CryptoRng, RngCore}; use std::collections::BTreeMap; @@ -1482,6 +1483,20 @@ fn tree_with_three_levels() -> HashTreeBuilderImpl { builder } +proptest! { + #[test] + fn recompute_digest_for_mixed_hash_tree_iteratively_and_recursively_produces_same_digest( + tree in arbitrary::arbitrary_well_formed_mixed_hash_tree() + ){ + let rec_or_error = mixed_hash_tree_digest_recursive(&tree); + // ignore the error case, since the iterative algorithm is infallible + if let Ok(rec) = rec_or_error { + let iter = tree.digest(); + assert_eq!(rec, iter); + } + } +} + #[test] fn witness_for_simple_path_in_a_big_tree() { // Simple path : label_b -> label_b_5 -> label_b_5_1 -> leaf