From 84a83a79a47b39fc901fa6d5e39a01410502bfbf Mon Sep 17 00:00:00 2001
From: Godspower Eze <61994334+Godspower-Eze@users.noreply.github.com>
Date: Tue, 30 Jan 2024 13:03:54 +0100
Subject: [PATCH] feat: merklelization of containers (#697)

---
 lib/spec/runners/ssz_generic.ex |  6 +++--
 lib/ssz_ex.ex                   | 40 +++++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 2 deletions(-)

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
     <<value::size(size)-little>> |> 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