diff --git a/tree_hash/tests/tests.rs b/tree_hash/tests/tests.rs index 8531548..0012e6b 100644 --- a/tree_hash/tests/tests.rs +++ b/tree_hash/tests/tests.rs @@ -1,6 +1,6 @@ use alloy_primitives::{Address, U128, U160, U256}; use ssz_derive::Encode; -use tree_hash::{Hash256, MerkleHasher, PackedEncoding, TreeHash, BYTES_PER_CHUNK}; +use tree_hash::{merkle_root, Hash256, MerkleHasher, PackedEncoding, TreeHash, BYTES_PER_CHUNK}; use tree_hash_derive::TreeHash; #[derive(Encode)] @@ -127,6 +127,131 @@ fn variable_union() { ); } +/// Computes the expected root of a container by merkleizing its field roots, one leaf per field. +fn container_root(field_roots: &[Hash256]) -> Hash256 { + let mut leaves = Vec::with_capacity(field_roots.len() * BYTES_PER_CHUNK); + for root in field_roots { + leaves.extend_from_slice(root.as_slice()); + } + merkle_root(&leaves, field_roots.len()) +} + +#[derive(TreeHash)] +struct Basic { + a: u8, + b: u64, + c: Hash256, +} + +#[test] +fn struct_basic() { + let x = Basic { + a: 1, + b: 2, + c: Hash256::repeat_byte(3), + }; + assert_eq!( + x.tree_hash_root(), + container_root(&[ + x.a.tree_hash_root(), + x.b.tree_hash_root(), + x.c.tree_hash_root(), + ]) + ); +} + +#[test] +fn struct_is_container_type() { + assert_eq!(Basic::tree_hash_type(), tree_hash::TreeHashType::Container); +} + +#[derive(TreeHash)] +struct Single { + only: Hash256, +} + +#[test] +fn struct_single_field() { + // A one-field container hashes to that field's root unchanged. + let x = Single { + only: Hash256::repeat_byte(7), + }; + assert_eq!(x.tree_hash_root(), x.only.tree_hash_root()); + assert_eq!( + x.tree_hash_root(), + container_root(&[x.only.tree_hash_root()]) + ); +} + +#[derive(TreeHash)] +struct WithSkip { + a: u8, + // Never read; only present to prove `skip_hashing` excludes it from the root. + #[allow(dead_code)] + #[tree_hash(skip_hashing)] + b: u64, + c: u8, +} + +#[test] +fn struct_skip_hashing() { + let x = WithSkip { a: 1, b: 999, c: 3 }; + + // The skipped field is absent from the tree: the root is a two-leaf container of `a` and `c`. + assert_eq!( + x.tree_hash_root(), + container_root(&[x.a.tree_hash_root(), x.c.tree_hash_root()]) + ); + + // Mutating only the skipped field must not change the root. + let y = WithSkip { + a: 1, + b: 12345, + c: 3, + }; + assert_eq!(x.tree_hash_root(), y.tree_hash_root()); +} + +#[derive(TreeHash)] +struct Nested { + inner: Basic, + tag: u8, +} + +#[test] +fn struct_nested() { + let x = Nested { + inner: Basic { + a: 1, + b: 2, + c: Hash256::repeat_byte(3), + }, + tag: 9, + }; + assert_eq!( + x.tree_hash_root(), + container_root(&[x.inner.tree_hash_root(), x.tag.tree_hash_root()]) + ); +} + +#[derive(TreeHash)] +struct Generic { + value: T, + count: u64, +} + +#[test] +fn struct_generic() { + let x = Generic { + value: 42u16, + count: 7, + }; + assert_eq!( + x.tree_hash_root(), + container_root(&[x.value.tree_hash_root(), x.count.tree_hash_root()]) + ); +} + /// Test that the packed encodings for different types are equal. #[test] fn packed_encoding_example() {