feat: implement ssz encode/decode container (#463)

lib/container.ex

@@ -0,0 +1,11 @@
+defmodule LambdaEthereumConsensus.Container do
+  @moduledoc """
+    Container for SSZ
+  """
+
+  @doc """
+  List of ordered {key, schema} tuples. 
+    It specifies both the serialization order and the schema for each key in the map.
+  """
+  @callback schema() :: [{atom, any}]
+end

lib/ssz_ex.ex

@@ -14,20 +14,29 @@ defmodule LambdaEthereumConsensus.SszEx do
       else: encode_fixed_size_list(list, basic_type, size)
   end
 
+  def encode(value, {:bytes, _}), do: {:ok, value}
+
+  def encode(container, module) when is_map(container),
+    do: encode_container(container, module.schema())
+
   def decode(binary, :bool), do: decode_bool(binary)
   def decode(binary, {:int, size}), do: decode_uint(binary, size)
+  def decode(value, {:bytes, _}), do: {:ok, value}
 
   def decode(binary, {:list, basic_type, size}) do
     if variable_size?(basic_type),
       do: decode_variable_list(binary, basic_type, size),
       else: decode_list(binary, basic_type, size)
   end
 
+  def decode(binary, module) when is_atom(module), do: decode_container(binary, module)
+
   #################
   ### Private functions
   #################
   @bytes_per_length_offset 4
   @bits_per_byte 8
+  @offset_bits 32
 
   defp encode_int(value, size) when is_integer(value), do: {:ok, <<value::size(size)-little>>}
   defp encode_bool(true), do: {:ok, "\x01"}
@@ -184,6 +193,123 @@ defmodule LambdaEthereumConsensus.SszEx do
     end
   end
 
+  defp encode_container(container, schemas) do
+    {fixed_size_values, fixed_length, variable_values} = analyze_schemas(container, schemas)
+
+    with {:ok, variable_parts} <- encode_schemas(variable_values),
+         offsets = calculate_offsets(variable_parts, fixed_length),
+         variable_length =
+           Enum.reduce(variable_parts, 0, fn part, acc -> byte_size(part) + acc end),
+         :ok <- check_length(fixed_length, variable_length),
+         {:ok, fixed_parts} <-
+           replace_offsets(fixed_size_values, offsets)
+           |> encode_schemas do
+      (fixed_parts ++ variable_parts)
+      |> Enum.join()
+      |> then(&{:ok, &1})
+    end
+  end
+
+  defp analyze_schemas(container, schemas) do
+    schemas
+    |> Enum.reduce({[], 0, []}, fn {key, schema},
+                                   {acc_fixed_size_values, acc_fixed_length, acc_variable_values} ->
+      value = Map.fetch!(container, key)
+
+      if variable_size?(schema) do
+        {[:offset | acc_fixed_size_values], @bytes_per_length_offset + acc_fixed_length,
+         [{value, schema} | acc_variable_values]}
+      else
+        {[{value, schema} | acc_fixed_size_values], acc_fixed_length + get_fixed_size(schema),
+         acc_variable_values}
+      end
+    end)
+  end
+
+  defp encode_schemas(tuple_values) do
+    Enum.map(tuple_values, fn {value, schema} -> encode(value, schema) end)
+    |> flatten_results()
+  end
+
+  defp calculate_offsets(variable_parts, fixed_length) do
+    {offsets, _} =
+      Enum.reduce(variable_parts, {[], fixed_length}, fn element, {res, acc} ->
+        {[{acc, {:int, 32}} | res], byte_size(element) + acc}
+      end)
+
+    offsets
+  end
+
+  defp replace_offsets(fixed_size_values, offsets) do
+    {fixed_size_values, _} =
+      Enum.reduce(fixed_size_values, {[], offsets}, &replace_offset/2)
+
+    fixed_size_values
+  end
+
+  defp replace_offset(:offset, {acc_fixed_list, [offset | rest_offsets]}),
+    do: {[offset | acc_fixed_list], rest_offsets}
+
+  defp replace_offset(element, {acc_fixed_list, acc_offsets_list}),
+    do: {[element | acc_fixed_list], acc_offsets_list}
+
+  defp decode_container(binary, module) do
+    schemas = module.schema()
+    fixed_length = get_fixed_length(schemas)
+    <<fixed_binary::binary-size(fixed_length), variable_binary::bitstring>> = binary
+
+    with {:ok, fixed_parts, offsets} <- decode_fixed_section(fixed_binary, schemas, fixed_length),
+         {:ok, variable_parts} <- decode_variable_section(variable_binary, offsets) do
+      {:ok, struct!(module, fixed_parts ++ variable_parts)}
+    end
+  end
+
+  defp decode_variable_section(binary, offsets) do
+    offsets
+    |> Enum.chunk_every(2, 1)
+    |> Enum.reduce({binary, []}, fn
+      [{offset, {key, schema}}, {next_offset, _}], {rest_bytes, acc_variable_parts} ->
+        size = next_offset - offset
+        <<chunk::binary-size(size), rest::bitstring>> = rest_bytes
+        {rest, [{key, decode(chunk, schema)} | acc_variable_parts]}
+
+      [{_offset, {key, schema}}], {rest_bytes, acc_variable_parts} ->
+        {<<>>, [{key, decode(rest_bytes, schema)} | acc_variable_parts]}
+    end)
+    |> then(fn {<<>>, variable_parts} ->
+      flatten_container_results(variable_parts)
+    end)
+  end
+
+  defp decode_fixed_section(binary, schemas, fixed_length) do
+    schemas
+    |> Enum.reduce({binary, [], []}, fn {key, schema}, {binary, fixed_parts, offsets} ->
+      if variable_size?(schema) do
+        <<offset::integer-size(@offset_bits)-little, rest::bitstring>> = binary
+        {rest, fixed_parts, [{offset - fixed_length, {key, schema}} | offsets]}
+      else
+        ssz_fixed_len = get_fixed_size(schema)
+        <<chunk::binary-size(ssz_fixed_len), rest::bitstring>> = binary
+        {rest, [{key, decode(chunk, schema)} | fixed_parts], offsets}
+      end
+    end)
+    |> then(fn {_rest_bytes, fixed_parts, offsets} ->
+      Tuple.append(flatten_container_results(fixed_parts), Enum.reverse(offsets))
+    end)
+  end
+
+  defp get_fixed_length(schemas) do
+    schemas
+    |> Stream.map(fn {_key, schema} ->
+      if variable_size?(schema) do
+        @bytes_per_length_offset
+      else
+        get_fixed_size(schema)
+      end
+    end)
+    |> Enum.sum()
+  end
+
   # https://notes.ethereum.org/ruKvDXl6QOW3gnqVYb8ezA?view
   defp sanitize_offset(offset, previous_offset, num_bytes, num_fixed_bytes) do
     cond do
@@ -227,6 +353,15 @@ defmodule LambdaEthereumConsensus.SszEx do
     end
   end
 
+  defp flatten_container_results(results) do
+    case Enum.group_by(results, fn {_, {type, _}} -> type end, fn {key, {_, result}} ->
+           {key, result}
+         end) do
+      %{error: errors} -> {:error, errors}
+      summary -> {:ok, Map.get(summary, :ok, [])}
+    end
+  end
+
   defp check_length(fixed_lengths, total_byte_size) do
     if fixed_lengths + total_byte_size <
          2 ** (@bytes_per_length_offset * @bits_per_byte) do
@@ -238,8 +373,10 @@ defmodule LambdaEthereumConsensus.SszEx do
 
   defp get_fixed_size(:bool), do: 1
   defp get_fixed_size({:int, size}), do: div(size, @bits_per_byte)
+  defp get_fixed_size({:bytes, size}), do: size
 
   defp variable_size?({:list, _, _}), do: true
   defp variable_size?(:bool), do: false
   defp variable_size?({:int, _}), do: false
+  defp variable_size?({:bytes, _}), do: false
 end

lib/ssz_types/beacon_chain/sync_committee.ex

@@ -4,6 +4,8 @@ defmodule SszTypes.SyncCommittee do
   Related definitions in `native/ssz_nif/src/types/`.
   """
 
+  @behaviour LambdaEthereumConsensus.Container
+
   fields = [
     :pubkeys,
     :aggregate_pubkey
@@ -16,4 +18,12 @@ defmodule SszTypes.SyncCommittee do
           pubkeys: list(SszTypes.bls_pubkey()),
           aggregate_pubkey: SszTypes.bls_pubkey()
         }
+
+  @impl LambdaEthereumConsensus.Container
+  def schema do
+    [
+      {:pubkeys, {:list, {:bytes, 48}, 100}},
+      {:aggregate_pubkey, {:bytes, 48}}
+    ]
+  end
 end

lib/ssz_types/beacon_chain/validator.ex

@@ -3,6 +3,7 @@ defmodule SszTypes.Validator do
   Struct definition for `Validator`.
   Related definitions in `native/ssz_nif/src/types/`.
   """
+  @behaviour LambdaEthereumConsensus.Container
 
   @eth1_address_withdrawal_prefix <<0x01>>
 
@@ -66,4 +67,18 @@ defmodule SszTypes.Validator do
     has_excess_balance = balance > max_effective_balance
     has_eth1_withdrawal_credential(validator) && has_max_effective_balance && has_excess_balance
   end
+
+  @impl LambdaEthereumConsensus.Container
+  def schema do
+    [
+      {:pubkey, {:bytes, 48}},
+      {:withdrawal_credentials, {:bytes, 32}},
+      {:effective_balance, {:int, 64}},
+      {:slashed, :bool},
+      {:activation_eligibility_epoch, {:int, 64}},
+      {:activation_epoch, {:int, 64}},
+      {:exit_epoch, {:int, 64}},
+      {:withdrawable_epoch, {:int, 64}}
+    ]
+  end
 end

test/unit/ssz_ex_test.exs

@@ -63,4 +63,74 @@ defmodule Unit.SSZExTest do
     # empty list
     assert_roundtrip(<<>>, [], {:list, {:int, 32}, 6})
   end
+
+  test "serialize and deserialize container only with fixed parts" do
+    validator = %SszTypes.Validator{
+      pubkey:
+        <<166, 144, 240, 158, 185, 117, 206, 31, 49, 45, 247, 53, 183, 95, 32, 20, 57, 245, 54,
+          60, 97, 78, 24, 81, 227, 157, 191, 150, 163, 202, 1, 72, 46, 131, 80, 54, 55, 203, 11,
+          160, 206, 88, 144, 58, 231, 142, 94, 235>>,
+      withdrawal_credentials:
+        <<31, 83, 167, 245, 158, 202, 157, 114, 98, 134, 215, 52, 106, 152, 108, 188, 15, 122, 21,
+          35, 113, 166, 17, 202, 159, 46, 180, 113, 98, 99, 233, 2>>,
+      effective_balance: 2_281_329_295_298_915_107,
+      slashed: false,
+      activation_eligibility_epoch: 8_916_476_893_047_043_501,
+      activation_epoch: 11_765_006_084_061_081_232,
+      exit_epoch: 14_221_179_644_044_541_938,
+      withdrawable_epoch: 11_813_934_873_299_048_632
+    }
+
+    serialized =
+      <<166, 144, 240, 158, 185, 117, 206, 31, 49, 45, 247, 53, 183, 95, 32, 20, 57, 245, 54, 60,
+        97, 78, 24, 81, 227, 157, 191, 150, 163, 202, 1, 72, 46, 131, 80, 54, 55, 203, 11, 160,
+        206, 88, 144, 58, 231, 142, 94, 235, 31, 83, 167, 245, 158, 202, 157, 114, 98, 134, 215,
+        52, 106, 152, 108, 188, 15, 122, 21, 35, 113, 166, 17, 202, 159, 46, 180, 113, 98, 99,
+        233, 2, 35, 235, 251, 53, 232, 232, 168, 31, 0, 173, 53, 12, 34, 126, 176, 189, 123, 144,
+        46, 197, 36, 179, 178, 69, 163, 242, 127, 74, 10, 138, 199, 91, 197, 184, 216, 150, 162,
+        44, 135, 243, 163>>
+
+    assert_roundtrip(serialized, validator, SszTypes.Validator)
+  end
+
+  test "serialize and deserialize variable container" do
+    pubkey1 =
+      <<166, 144, 240, 158, 185, 117, 206, 31, 49, 45, 247, 53, 183, 95, 32, 20, 57, 245, 54, 60,
+        97, 78, 24, 81, 227, 157, 191, 150, 163, 202, 1, 72, 46, 131, 80, 54, 55, 203, 11, 160,
+        206, 88, 144, 58, 231, 142, 94, 235>>
+
+    pubkey2 =
+      <<180, 144, 240, 158, 185, 117, 206, 31, 49, 45, 247, 53, 183, 95, 32, 20, 57, 245, 54, 60,
+        97, 78, 24, 81, 227, 157, 191, 150, 163, 202, 1, 72, 46, 131, 80, 54, 55, 203, 11, 160,
+        206, 88, 144, 58, 231, 142, 94, 235>>
+
+    pubkey3 =
+      <<190, 144, 240, 158, 185, 117, 206, 31, 49, 45, 247, 53, 183, 95, 32, 20, 57, 245, 54, 60,
+        97, 78, 24, 81, 227, 157, 191, 150, 163, 202, 1, 72, 46, 131, 80, 54, 55, 203, 11, 160,
+        206, 88, 144, 58, 231, 142, 94, 235>>
+
+    pubkey4 =
+      <<200, 144, 240, 158, 185, 117, 206, 31, 49, 45, 247, 53, 183, 95, 32, 20, 57, 245, 54, 60,
+        97, 78, 24, 81, 227, 157, 191, 150, 163, 202, 1, 72, 46, 131, 80, 54, 55, 203, 11, 160,
+        206, 88, 144, 58, 231, 142, 94, 235>>
+
+    sync = %SszTypes.SyncCommittee{
+      pubkeys: [pubkey1, pubkey2, pubkey3],
+      aggregate_pubkey: pubkey4
+    }
+
+    serialized =
+      <<52, 0, 0, 0, 200, 144, 240, 158, 185, 117, 206, 31, 49, 45, 247, 53, 183, 95, 32, 20, 57,
+        245, 54, 60, 97, 78, 24, 81, 227, 157, 191, 150, 163, 202, 1, 72, 46, 131, 80, 54, 55,
+        203, 11, 160, 206, 88, 144, 58, 231, 142, 94, 235, 166, 144, 240, 158, 185, 117, 206, 31,
+        49, 45, 247, 53, 183, 95, 32, 20, 57, 245, 54, 60, 97, 78, 24, 81, 227, 157, 191, 150,
+        163, 202, 1, 72, 46, 131, 80, 54, 55, 203, 11, 160, 206, 88, 144, 58, 231, 142, 94, 235,
+        180, 144, 240, 158, 185, 117, 206, 31, 49, 45, 247, 53, 183, 95, 32, 20, 57, 245, 54, 60,
+        97, 78, 24, 81, 227, 157, 191, 150, 163, 202, 1, 72, 46, 131, 80, 54, 55, 203, 11, 160,
+        206, 88, 144, 58, 231, 142, 94, 235, 190, 144, 240, 158, 185, 117, 206, 31, 49, 45, 247,
+        53, 183, 95, 32, 20, 57, 245, 54, 60, 97, 78, 24, 81, 227, 157, 191, 150, 163, 202, 1, 72,
+        46, 131, 80, 54, 55, 203, 11, 160, 206, 88, 144, 58, 231, 142, 94, 235>>
+
+    assert_roundtrip(serialized, sync, SszTypes.SyncCommittee)
+  end
 end