From 5be55fb1e290eb808c6e637b97eef26d2a8a5637 Mon Sep 17 00:00:00 2001
From: avilagaston9 <gavila@fi.uba.ar>
Date: Mon, 22 Apr 2024 17:26:22 -0300
Subject: [PATCH] Add stacktrace to decode

---
 lib/ssz_ex/decode.ex      | 17 +++++++++++------
 lib/ssz_ex/encode.ex      | 11 ++++-------
 lib/ssz_ex/utils.ex       |  6 ++++++
 test/unit/ssz_ex_test.exs | 20 ++++++++------------
 4 files changed, 29 insertions(+), 25 deletions(-)

diff --git a/lib/ssz_ex/decode.ex b/lib/ssz_ex/decode.ex
index 2dd49671f..8efc49c50 100644
--- a/lib/ssz_ex/decode.ex
+++ b/lib/ssz_ex/decode.ex
@@ -39,8 +39,8 @@ defmodule SszEx.Decode do
   def decode(binary, module) when is_atom(module) do
     with {:ok, result} <-
            if(Utils.variable_size?(module),
-             do: decode_variable_container(binary, module),
-             else: decode_fixed_container(binary, module)
+             do: decode_variable_container(binary, module) |> Utils.add_trace("#{module}"),
+             else: decode_fixed_container(binary, module) |> Utils.add_trace("#{module}")
            ) do
       if exported?(module, :decode_ex, 1) do
         {:ok, module.decode_ex(result)}
@@ -337,14 +337,17 @@ defmodule SszEx.Decode do
           :ok ->
             size = next_offset - offset
             <<chunk::binary-size(size), rest::bitstring>> = rest_bytes
-            {:cont, {rest, [{key, decode(chunk, schema)} | acc_variable_parts]}}
+
+            {:cont,
+             {rest, [{key, decode(chunk, schema) |> Utils.add_trace(key)} | acc_variable_parts]}}
 
           error ->
             {:halt, {<<>>, [{key, error} | acc_variable_parts]}}
         end
 
       [{_offset, {key, schema}}], {rest_bytes, acc_variable_parts} ->
-        {:cont, {<<>>, [{key, decode(rest_bytes, schema)} | acc_variable_parts]}}
+        {:cont,
+         {<<>>, [{key, decode(rest_bytes, schema) |> Utils.add_trace(key)} | acc_variable_parts]}}
     end)
     |> then(fn {<<>>, variable_parts} ->
       flatten_container_results(variable_parts)
@@ -363,7 +366,9 @@ defmodule SszEx.Decode do
       else
         ssz_fixed_len = Utils.get_fixed_size(schema)
         <<chunk::binary-size(ssz_fixed_len), rest::bitstring>> = binary
-        {rest, [{key, decode(chunk, schema)} | fixed_parts], offsets, items_index + ssz_fixed_len}
+
+        {rest, [{key, decode(chunk, schema) |> Utils.add_trace(key)} | fixed_parts], offsets,
+         items_index + ssz_fixed_len}
       end
     end)
     |> then(fn {_rest_bytes, fixed_parts, offsets, items_index} ->
@@ -455,7 +460,7 @@ defmodule SszEx.Decode do
     case Enum.group_by(results, fn {_, {type, _}} -> type end, fn {key, {_, result}} ->
            {key, result}
          end) do
-      %{error: errors} -> {:error, errors}
+      %{error: [first_error | _rest]} -> {:error, first_error}
       summary -> {:ok, Map.get(summary, :ok, [])}
     end
   end
diff --git a/lib/ssz_ex/encode.ex b/lib/ssz_ex/encode.ex
index 71eeb24fb..e32caec3c 100644
--- a/lib/ssz_ex/encode.ex
+++ b/lib/ssz_ex/encode.ex
@@ -42,7 +42,7 @@ defmodule SszEx.Encode do
     do: encode_bitvector(value, size)
 
   def encode(container, module) when is_map(container) do
-    encode_container(container, module.schema()) |> add_trace("#{module}")
+    encode_container(container, module.schema()) |> Utils.add_trace("#{module}")
   end
 
   defp encode_int(value, size) when is_integer(value), do: {:ok, <<value::size(size)-little>>}
@@ -172,7 +172,9 @@ defmodule SszEx.Encode do
   end
 
   defp encode_schemas(tuple_values) do
-    Enum.map(tuple_values, fn {value, key, schema} -> encode(value, schema) |> add_trace(key) end)
+    Enum.map(tuple_values, fn {value, key, schema} ->
+      encode(value, schema) |> Utils.add_trace(key)
+    end)
     |> Utils.flatten_results()
   end
 
@@ -210,9 +212,4 @@ defmodule SszEx.Encode do
        }}
     end
   end
-
-  defp add_trace({:error, %Error{} = error}, module),
-    do: {:error, Error.add_trace(error, module)}
-
-  defp add_trace(value, _module), do: value
 end
diff --git a/lib/ssz_ex/utils.ex b/lib/ssz_ex/utils.ex
index 099bb7ad0..edb52f845 100644
--- a/lib/ssz_ex/utils.ex
+++ b/lib/ssz_ex/utils.ex
@@ -5,6 +5,7 @@ defmodule SszEx.Utils do
 
   alias LambdaEthereumConsensus.Utils.BitList
   alias LambdaEthereumConsensus.Utils.BitVector
+  alias SszEx.Error
 
   @allowed_uints [8, 16, 32, 64, 128, 256]
   @bits_per_byte 8
@@ -114,4 +115,9 @@ defmodule SszEx.Utils do
   def size_of(:bool), do: @bytes_per_boolean
 
   def size_of({:int, size}), do: size |> div(@bits_per_byte)
+
+  def add_trace({:error, %Error{} = error}, module),
+    do: {:error, Error.add_trace(error, module)}
+
+  def add_trace(value, _module), do: value
 end
diff --git a/test/unit/ssz_ex_test.exs b/test/unit/ssz_ex_test.exs
index 56b3e44bd..c009ea604 100644
--- a/test/unit/ssz_ex_test.exs
+++ b/test/unit/ssz_ex_test.exs
@@ -738,24 +738,20 @@ defmodule Unit.SSZExTest do
     encoded_checkpoint =
       <<0, 0, 0>>
 
-    assert SszEx.decode(encoded_checkpoint, Checkpoint) ==
-             {:error,
-              %Error{
-                message:
-                  "Invalid binary length while decoding Elixir.Types.Checkpoint. \nExpected 40. \nFound 3.\n"
-              }}
+    {:error, error} = SszEx.decode(encoded_checkpoint, Checkpoint)
+
+    assert "#{error}" ==
+             "Invalid binary length while decoding Elixir.Types.Checkpoint. \nExpected 40. \nFound 3.\nStacktrace: Elixir.Types.Checkpoint"
   end
 
   test "decode longer checkpoint" do
     encoded_checkpoint =
       <<0::size(41 * 8)>>
 
-    assert SszEx.decode(encoded_checkpoint, Checkpoint) ==
-             {:error,
-              %Error{
-                message:
-                  "Invalid binary length while decoding Elixir.Types.Checkpoint. \nExpected 40. \nFound 41.\n"
-              }}
+    {:error, error} = SszEx.decode(encoded_checkpoint, Checkpoint)
+
+    assert "#{error}" ==
+             "Invalid binary length while decoding Elixir.Types.Checkpoint. \nExpected 40. \nFound 41.\nStacktrace: Elixir.Types.Checkpoint"
   end
 
   test "hash tree root of list exceeding max size" do