diff --git a/lib/spec/runners/ssz_generic.ex b/lib/spec/runners/ssz_generic.ex index be0bc3ea8..96de18061 100644 --- a/lib/spec/runners/ssz_generic.ex +++ b/lib/spec/runners/ssz_generic.ex @@ -69,13 +69,15 @@ defmodule SszGenericTestRunner do {:container, module}, real_serialized, real_deserialized, - _hash_tree_root + expected_hash_tree_root ) do real_struct = struct!(module, real_deserialized) {:ok, deserialized} = SszEx.decode(real_serialized, module) assert deserialized == real_struct {:ok, serialized} = SszEx.encode(real_struct, module) assert serialized == real_serialized + actual_hash_tree_root = SszEx.hash_tree_root!(real_struct, module) + assert actual_hash_tree_root == expected_hash_tree_root end defp assert_ssz( @@ -92,7 +94,7 @@ defmodule SszGenericTestRunner do assert serialized == real_serialized - {:ok, actual_hash_tree_root} = SszEx.hash_tree_root(real_deserialized, schema) + actual_hash_tree_root = SszEx.hash_tree_root!(real_deserialized, schema) assert actual_hash_tree_root == expected_hash_tree_root end diff --git a/lib/ssz_ex.ex b/lib/ssz_ex.ex index b7df120fc..fee004df7 100644 --- a/lib/ssz_ex.ex +++ b/lib/ssz_ex.ex @@ -75,6 +75,41 @@ defmodule LambdaEthereumConsensus.SszEx do @spec hash_tree_root!(non_neg_integer, {:int, non_neg_integer}) :: Types.root() def hash_tree_root!(value, {:int, size}), do: pack(value, {:int, size}) + @spec hash_tree_root!(binary, {:bytes, non_neg_integer}) :: Types.root() + def hash_tree_root!(value, {:bytes, size}) do + packed_chunks = pack(value, {:bytes, size}) + leaf_count = packed_chunks |> get_chunks_len() |> next_pow_of_two() + root = merkleize_chunks_with_virtual_padding(packed_chunks, leaf_count) + root + end + + @spec hash_tree_root!(list(), {:list, any, non_neg_integer}) :: Types.root() + def hash_tree_root!(list, {:list, type, size}) do + {:ok, root} = hash_tree_root(list, {:list, type, size}) + root + end + + @spec hash_tree_root!(list(), {:vector, any, non_neg_integer}) :: Types.root() + def hash_tree_root!(vector, {:vector, type, size}) do + {:ok, root} = hash_tree_root(vector, {:vector, type, size}) + root + end + + @spec hash_tree_root!(struct(), atom()) :: Types.root() + def hash_tree_root!(container, module) when is_map(container) do + chunks = + module.schema() + |> Enum.reduce(<<>>, fn {key, schema}, acc_root -> + value = container |> Map.get(key) + root = hash_tree_root!(value, schema) + acc_root <> root + end) + + leaf_count = chunks |> get_chunks_len() |> next_pow_of_two() + root = merkleize_chunks_with_virtual_padding(chunks, leaf_count) + root + end + @spec hash_tree_root(list(), {:list, any, non_neg_integer}) :: {:ok, Types.root()} | {:error, String.t()} def hash_tree_root(list, {:list, type, size}) do @@ -207,6 +242,11 @@ defmodule LambdaEthereumConsensus.SszEx do <> |> pack_bytes() end + @spec pack(binary, {:bytes, non_neg_integer}) :: binary() + def pack(value, {:bytes, _size}) do + value |> pack_bytes() + end + @spec pack(list(), {:list | :vector, any, non_neg_integer}) :: binary() | :error def pack(list, {type, schema, _}) when type in [:vector, :list] do if variable_size?(schema) do