Refactor tree hashing (#861)

* Pre-allocated tree hash caches

* Add SmallVec to tree hash cache

* Avoid allocation for validator.pubkey

* Avoid iterator which seems to be doing heap alloc

* Add more smallvecs

* MOAR SMALLVEC

* Move non-test code to Hash256 tree hash

* Fix byte ordering error

* Add incomplete but working merkle stream impl

* Fix zero hash error

* Add zero hash fn

* Add MerkleStream comments

* Add smallvec, tidy

* Integrate into tree hash derive

* Update ssz_types tree hash

* Don't heap alloc for mix in length

* Add byte-level streaming to MerkleStream

* Avoid recursion in write method

* Update BLS to MerkleStream

* Fix some not-compiling tests

* Remove debug profiling

* Remove code duplication

* Move beacon state tree hash to new hasher

* Fix failing tests

* Update comments

* Add some fast-paths to tree_hash::merkle_root

* Remove unncessary test

* Rename MerkleStream -> MerkleHasher

* Rename new_with_leaf_count -> with_leaves

* Tidy

* Remove NonZeroUsize

* Remove todo

* Update smallvec
This commit is contained in:
Paul Hauner
2020-03-05 08:07:27 +11:00
committed by GitHub
parent 12999fb06c
commit 7f6ae4c2f5
43 changed files with 1076 additions and 292 deletions

View File

@@ -66,6 +66,7 @@ pub enum Error {
CommitteeCacheUninitialized(Option<RelativeEpoch>),
SszTypesError(ssz_types::Error),
TreeHashCacheNotInitialized,
TreeHashError(tree_hash::Error),
CachedTreeHashError(cached_tree_hash::Error),
InvalidValidatorPubkey(ssz::DecodeError),
ValidatorRegistryShrunk,
@@ -1044,3 +1045,9 @@ impl From<cached_tree_hash::Error> for Error {
Error::CachedTreeHashError(e)
}
}
impl From<tree_hash::Error> for Error {
fn from(e: tree_hash::Error) -> Error {
Error::TreeHashError(e)
}
}

View File

@@ -3,7 +3,13 @@ use crate::{BeaconState, EthSpec, Hash256, Unsigned, Validator};
use cached_tree_hash::{int_log, CacheArena, CachedTreeHash, TreeHashCache};
use rayon::prelude::*;
use ssz_derive::{Decode, Encode};
use tree_hash::{mix_in_length, TreeHash};
use tree_hash::{mix_in_length, MerkleHasher, TreeHash};
/// The number of fields on a beacon state.
const NUM_BEACON_STATE_HASHING_FIELDS: usize = 20;
/// The number of nodes in the Merkle tree of a validator record.
const NODES_PER_VALIDATOR: usize = 15;
/// The number of validator record tree hash caches stored in each arena.
///
@@ -73,64 +79,79 @@ impl BeaconTreeHashCache {
&mut self,
state: &BeaconState<T>,
) -> Result<Hash256, Error> {
let mut leaves = vec![];
let mut hasher = MerkleHasher::with_leaves(NUM_BEACON_STATE_HASHING_FIELDS);
leaves.append(&mut state.genesis_time.tree_hash_root());
leaves.append(&mut state.slot.tree_hash_root());
leaves.append(&mut state.fork.tree_hash_root());
leaves.append(&mut state.latest_block_header.tree_hash_root());
leaves.extend_from_slice(
hasher.write(state.genesis_time.tree_hash_root().as_bytes())?;
hasher.write(state.slot.tree_hash_root().as_bytes())?;
hasher.write(state.fork.tree_hash_root().as_bytes())?;
hasher.write(state.latest_block_header.tree_hash_root().as_bytes())?;
hasher.write(
state
.block_roots
.recalculate_tree_hash_root(&mut self.fixed_arena, &mut self.block_roots)?
.as_bytes(),
);
leaves.extend_from_slice(
)?;
hasher.write(
state
.state_roots
.recalculate_tree_hash_root(&mut self.fixed_arena, &mut self.state_roots)?
.as_bytes(),
);
leaves.extend_from_slice(
)?;
hasher.write(
state
.historical_roots
.recalculate_tree_hash_root(&mut self.fixed_arena, &mut self.historical_roots)?
.as_bytes(),
);
leaves.append(&mut state.eth1_data.tree_hash_root());
leaves.append(&mut state.eth1_data_votes.tree_hash_root());
leaves.append(&mut state.eth1_deposit_index.tree_hash_root());
leaves.extend_from_slice(
)?;
hasher.write(state.eth1_data.tree_hash_root().as_bytes())?;
hasher.write(state.eth1_data_votes.tree_hash_root().as_bytes())?;
hasher.write(state.eth1_deposit_index.tree_hash_root().as_bytes())?;
hasher.write(
self.validators
.recalculate_tree_hash_root(&state.validators[..])?
.as_bytes(),
);
leaves.extend_from_slice(
)?;
hasher.write(
state
.balances
.recalculate_tree_hash_root(&mut self.balances_arena, &mut self.balances)?
.as_bytes(),
);
leaves.extend_from_slice(
)?;
hasher.write(
state
.randao_mixes
.recalculate_tree_hash_root(&mut self.fixed_arena, &mut self.randao_mixes)?
.as_bytes(),
);
leaves.extend_from_slice(
)?;
hasher.write(
state
.slashings
.recalculate_tree_hash_root(&mut self.slashings_arena, &mut self.slashings)?
.as_bytes(),
);
leaves.append(&mut state.previous_epoch_attestations.tree_hash_root());
leaves.append(&mut state.current_epoch_attestations.tree_hash_root());
leaves.append(&mut state.justification_bits.tree_hash_root());
leaves.append(&mut state.previous_justified_checkpoint.tree_hash_root());
leaves.append(&mut state.current_justified_checkpoint.tree_hash_root());
leaves.append(&mut state.finalized_checkpoint.tree_hash_root());
)?;
hasher.write(
state
.previous_epoch_attestations
.tree_hash_root()
.as_bytes(),
)?;
hasher.write(state.current_epoch_attestations.tree_hash_root().as_bytes())?;
hasher.write(state.justification_bits.tree_hash_root().as_bytes())?;
hasher.write(
state
.previous_justified_checkpoint
.tree_hash_root()
.as_bytes(),
)?;
hasher.write(
state
.current_justified_checkpoint
.tree_hash_root()
.as_bytes(),
)?;
hasher.write(state.finalized_checkpoint.tree_hash_root().as_bytes())?;
Ok(Hash256::from_slice(&tree_hash::merkle_root(&leaves, 0)))
hasher.finish().map_err(Into::into)
}
}
@@ -181,10 +202,7 @@ impl ValidatorsListTreeHashCache {
std::mem::replace(&mut self.list_arena, list_arena);
Ok(Hash256::from_slice(&mix_in_length(
list_root.as_bytes(),
validators.len(),
)))
Ok(mix_in_length(&list_root, validators.len()))
}
}
@@ -202,8 +220,22 @@ impl ParallelValidatorTreeHash {
/// Allocates the necessary memory to store all of the cached Merkle trees but does perform any
/// hashing.
fn new<E: EthSpec>(validators: &[Validator]) -> Self {
let num_arenas = (validators.len() + VALIDATORS_PER_ARENA - 1) / VALIDATORS_PER_ARENA;
let mut arenas = vec![(CacheArena::default(), vec![]); num_arenas];
let num_arenas = std::cmp::max(
1,
(validators.len() + VALIDATORS_PER_ARENA - 1) / VALIDATORS_PER_ARENA,
);
let mut arenas = (1..=num_arenas)
.map(|i| {
let num_validators = if i == num_arenas {
validators.len() % VALIDATORS_PER_ARENA
} else {
VALIDATORS_PER_ARENA
};
NODES_PER_VALIDATOR * num_validators
})
.map(|capacity| (CacheArena::with_capacity(capacity), vec![]))
.collect::<Vec<_>>();
validators.iter().enumerate().for_each(|(i, v)| {
let (arena, caches) = &mut arenas[i / VALIDATORS_PER_ARENA];
@@ -272,3 +304,16 @@ impl ParallelValidatorTreeHash {
.collect()
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn validator_node_count() {
let mut arena = CacheArena::default();
let v = Validator::default();
let _cache = v.new_tree_hash_cache(&mut arena);
assert_eq!(arena.backing_len(), NODES_PER_VALIDATOR);
}
}

View File

@@ -14,12 +14,10 @@ pub struct SigningRoot {
pub trait SignedRoot: TreeHash {
fn signing_root(&self, domain: u64) -> Hash256 {
Hash256::from_slice(
&SigningRoot {
object_root: Hash256::from_slice(&self.tree_hash_root()),
domain,
}
.tree_hash_root(),
)
SigningRoot {
object_root: self.tree_hash_root(),
domain,
}
.tree_hash_root()
}
}

View File

@@ -237,8 +237,8 @@ macro_rules! impl_ssz {
32 / 8
}
fn tree_hash_root(&self) -> Vec<u8> {
int_to_bytes::int_to_bytes32(self.0)
fn tree_hash_root(&self) -> tree_hash::Hash256 {
tree_hash::Hash256::from_slice(&int_to_bytes::int_to_fixed_bytes32(self.0))
}
}

View File

@@ -284,14 +284,14 @@ impl<T: EthSpec> TestingBeaconBlockBuilder<T> {
// Vector containing all leaves
let leaves = datas
.iter()
.map(|data| Hash256::from_slice(&data.tree_hash_root()))
.map(|data| data.tree_hash_root())
.collect::<Vec<_>>();
// Building a VarList from leaves
let deposit_data_list = VariableList::<_, U4294967296>::from(leaves.clone());
// Setting the deposit_root to be the tree_hash_root of the VarList
state.eth1_data.deposit_root = Hash256::from_slice(&deposit_data_list.tree_hash_root());
state.eth1_data.deposit_root = deposit_data_list.tree_hash_root();
// Building the merkle tree used for generating proofs
let tree = MerkleTree::create(&leaves[..], spec.deposit_contract_tree_depth as usize);

View File

@@ -39,9 +39,8 @@ macro_rules! tree_hash_tests {
let mut rng = XorShiftRng::from_seed([42; 16]);
let original = <$type>::random_for_test(&mut rng);
let result = original.tree_hash_root();
assert_eq!(result.len(), 32);
// Tree hashing should not panic.
original.tree_hash_root();
}
};
}

View File

@@ -2,10 +2,10 @@
//!
//! It makes some assumptions about the layouts and update patterns of other structs in this
//! crate, and should be updated carefully whenever those structs are changed.
use crate::{Epoch, Hash256, Validator};
use crate::{Epoch, Hash256, PublicKeyBytes, Validator};
use cached_tree_hash::{int_log, CacheArena, CachedTreeHash, Error, TreeHashCache};
use int_to_bytes::int_to_fixed_bytes32;
use tree_hash::TreeHash;
use tree_hash::merkle_root;
/// Number of struct fields on `Validator`.
const NUM_VALIDATOR_FIELDS: usize = 8;
@@ -53,7 +53,7 @@ fn process_field_by_index(
force_update: bool,
) -> bool {
match field_idx {
0 => process_vec_field(v.pubkey.tree_hash_root(), leaf, force_update),
0 => process_pubkey_bytes_field(&v.pubkey, leaf, force_update),
1 => process_slice_field(v.withdrawal_credentials.as_bytes(), leaf, force_update),
2 => process_u64_field(v.effective_balance, leaf, force_update),
3 => process_bool_field(v.slashed, leaf, force_update),
@@ -68,13 +68,13 @@ fn process_field_by_index(
}
}
fn process_vec_field(new_tree_hash: Vec<u8>, leaf: &mut Hash256, force_update: bool) -> bool {
if force_update || leaf.as_bytes() != &new_tree_hash[..] {
leaf.assign_from_slice(&new_tree_hash);
true
} else {
false
}
fn process_pubkey_bytes_field(
val: &PublicKeyBytes,
leaf: &mut Hash256,
force_update: bool,
) -> bool {
let new_tree_hash = merkle_root(val.as_slice(), 0);
process_slice_field(new_tree_hash.as_bytes(), leaf, force_update)
}
fn process_slice_field(new_tree_hash: &[u8], leaf: &mut Hash256, force_update: bool) -> bool {
@@ -106,6 +106,7 @@ mod test {
use crate::Epoch;
use rand::SeedableRng;
use rand_xorshift::XorShiftRng;
use tree_hash::TreeHash;
fn test_validator_tree_hash(v: &Validator) {
let arena = &mut CacheArena::default();
@@ -152,4 +153,13 @@ mod test {
.map(|_| Validator::random_for_test(&mut rng))
.for_each(|v| test_validator_tree_hash(&v));
}
#[test]
pub fn smallvec_size_check() {
// If this test fails we need to go and reassess the length of the `SmallVec` in
// `cached_tree_hash::TreeHashCache`. If the size of the `SmallVec` is too slow we're going
// to start doing heap allocations for each validator, this will fragment memory and slow
// us down.
assert!(NUM_VALIDATOR_FIELDS <= 8,);
}
}