diff --git a/eth2/types/src/test_utils/macros.rs b/eth2/types/src/test_utils/macros.rs index 5f355bfe94..984f2962f2 100644 --- a/eth2/types/src/test_utils/macros.rs +++ b/eth2/types/src/test_utils/macros.rs @@ -44,16 +44,14 @@ macro_rules! cached_tree_hash_tests { let mut rng = XorShiftRng::from_seed([42; 16]); + // Test the original hash let original = $type::random_for_test(&mut rng); - let mut hasher = cached_tree_hash::CachedTreeHasher::new(&original).unwrap(); - assert_eq!(hasher.tree_hash_root().unwrap(), original.tree_hash_root()); + // Test the updated hash let modified = $type::random_for_test(&mut rng); - hasher.update(&modified).unwrap(); - assert_eq!(hasher.tree_hash_root().unwrap(), modified.tree_hash_root()); } }; diff --git a/eth2/utils/cached_tree_hash/src/lib.rs b/eth2/utils/cached_tree_hash/src/lib.rs index 0faf56ea22..b1397b0f4c 100644 --- a/eth2/utils/cached_tree_hash/src/lib.rs +++ b/eth2/utils/cached_tree_hash/src/lib.rs @@ -28,7 +28,7 @@ pub trait CachedTreeHash: TreeHash { #[derive(Debug, PartialEq)] pub struct CachedTreeHasher { - cache: TreeHashCache, + pub cache: TreeHashCache, } impl CachedTreeHasher { diff --git a/eth2/utils/cached_tree_hash/tests/tests.rs b/eth2/utils/cached_tree_hash/tests/tests.rs index c3392ba273..4ac6a4607c 100644 --- a/eth2/utils/cached_tree_hash/tests/tests.rs +++ b/eth2/utils/cached_tree_hash/tests/tests.rs @@ -423,6 +423,45 @@ fn test_struct_with_two_vecs() { } } +#[derive(Clone, Debug, TreeHash, CachedTreeHash)] +pub struct U64AndTwoStructs { + pub a: u64, + pub b: Inner, + pub c: Inner, +} + +#[test] +fn test_u64_and_two_structs() { + let inners = get_inners(); + + let variants = vec![ + U64AndTwoStructs { + a: 99, + b: inners[0].clone(), + c: inners[1].clone(), + }, + U64AndTwoStructs { + a: 10, + b: inners[2].clone(), + c: inners[3].clone(), + }, + U64AndTwoStructs { + a: 0, + b: inners[1].clone(), + c: inners[1].clone(), + }, + U64AndTwoStructs { + a: 0, + b: inners[1].clone(), + c: inners[1].clone(), + }, + ]; + + for v in &variants { + test_routine(v.clone(), variants.clone()); + } +} + #[derive(Clone, Debug, TreeHash, CachedTreeHash)] pub struct Inner { pub a: u64, diff --git a/eth2/utils/tree_hash_derive/src/lib.rs b/eth2/utils/tree_hash_derive/src/lib.rs index fdfe1e5c0d..3ba846307e 100644 --- a/eth2/utils/tree_hash_derive/src/lib.rs +++ b/eth2/utils/tree_hash_derive/src/lib.rs @@ -102,6 +102,8 @@ pub fn subtree_derive(input: TokenStream) -> TokenStream { // Iterate through the internal nodes, updating them if their children have changed. cache.update_internal_nodes(&overlay)?; + cache.chunk_index = overlay.next_node(); + Ok(()) } }