diff --git a/eth2/utils/ssz/src/cached_tree_hash/resize.rs b/eth2/utils/ssz/src/cached_tree_hash/resize.rs index 2cdce4827..3c2d2c407 100644 --- a/eth2/utils/ssz/src/cached_tree_hash/resize.rs +++ b/eth2/utils/ssz/src/cached_tree_hash/resize.rs @@ -26,7 +26,7 @@ pub fn grow_merkle_cache( // (e.g., the case where there are subtrees as leaves). // // If we're not on a leaf level, the number of nodes is fixed and known. - let (byte_slice, flag_slice) = if i == leaf_level { + let (old_byte_slice, old_flag_slice) = if i == leaf_level { ( old_bytes.get(first_byte_at_height(i)..)?, old_flags.get(first_node_at_height(i)..)?, @@ -38,14 +38,25 @@ pub fn grow_merkle_cache( ) }; - bytes - .get_mut(byte_range_at_height(i + to_height - from_height))? - .get_mut(0..byte_slice.len())? - .copy_from_slice(byte_slice); - flags - .get_mut(node_range_at_height(i + to_height - from_height))? - .get_mut(0..flag_slice.len())? - .copy_from_slice(flag_slice); + let new_i = i + to_height - from_height; + let (new_byte_slice, new_flag_slice) = if i == leaf_level { + ( + bytes.get_mut(first_byte_at_height(new_i)..)?, + flags.get_mut(first_node_at_height(new_i)..)?, + ) + } else { + ( + bytes.get_mut(byte_range_at_height(new_i))?, + flags.get_mut(node_range_at_height(new_i))?, + ) + }; + + new_byte_slice + .get_mut(0..old_byte_slice.len())? + .copy_from_slice(old_byte_slice); + new_flag_slice + .get_mut(0..old_flag_slice.len())? + .copy_from_slice(old_flag_slice); } Some((bytes, flags)) diff --git a/eth2/utils/ssz/src/cached_tree_hash/tests.rs b/eth2/utils/ssz/src/cached_tree_hash/tests.rs index c402fd15b..fb6ed9080 100644 --- a/eth2/utils/ssz/src/cached_tree_hash/tests.rs +++ b/eth2/utils/ssz/src/cached_tree_hash/tests.rs @@ -472,9 +472,10 @@ fn test_inner_vec_modifications(original: Vec, modified: Vec, refe } let num_leaves = leaves.len() / HASHSIZE; - let mut expected = merkleize(leaves); - expected.splice(3 * HASHSIZE.., full_bytes); + + let num_internal_nodes = num_leaves.next_power_of_two() - 1; + expected.splice(num_internal_nodes * HASHSIZE.., full_bytes); for _ in num_leaves..num_leaves.next_power_of_two() { expected.append(&mut vec![0; HASHSIZE]); @@ -589,6 +590,48 @@ fn lengthened_vec_of_inner_within_power_of_two_boundary() { test_inner_vec_modifications(original, modified, reference_vec); } +#[test] +fn lengthened_vec_of_inner_outside_power_of_two_boundary() { + let original = vec![ + Inner { + a: 0, + b: 1, + c: 2, + d: 3, + }, + Inner { + a: 4, + b: 5, + c: 6, + d: 7, + }, + Inner { + a: 8, + b: 9, + c: 10, + d: 11, + }, + Inner { + a: 12, + b: 13, + c: 14, + d: 15, + }, + ]; + + let mut modified = original.clone(); + modified.push(Inner { + a: 16, + b: 17, + c: 18, + d: 19, + }); + + let reference_vec: Vec = (0..20).collect(); + + test_inner_vec_modifications(original, modified, reference_vec); +} + #[test] fn vec_of_inner_builds() { let numbers: Vec = (0..12).collect();