feat: add bitlist and bitvector ssz support

lib/ssz_ex.ex

@@ -2,6 +2,7 @@ defmodule LambdaEthereumConsensus.SszEx do
   @moduledoc """
     SSZ library in Elixir
   """
+  alias LambdaEthereumConsensus.Utils.BitVector
 
   #################
   ### Public API
@@ -23,6 +24,12 @@ defmodule LambdaEthereumConsensus.SszEx do
 
   def encode(value, {:bytes, _}), do: {:ok, value}
 
+  def encode(value, {:bitlist, max_size}) when is_bitstring(value),
+    do: encode_bitlist(value, max_size)
+
+  def encode(value, {:bitvector, size}) when is_bitstring(value),
+    do: encode_bitvector(value, size)
+
   def encode(container, module) when is_map(container),
     do: encode_container(container, module.schema())
 
@@ -36,6 +43,12 @@ defmodule LambdaEthereumConsensus.SszEx do
       else: decode_list(binary, basic_type, size)
   end
 
+  def decode(value, {:bitlist, max_size}) when is_bitstring(value),
+    do: decode_bitlist(value, max_size)
+
+  def decode(value, {:bitvector, size}) when is_bitstring(value),
+    do: decode_bitvector(value, size)
+
   def decode(binary, module) when is_atom(module), do: decode_container(binary, module)
 
   @spec hash_tree_root!(boolean, atom) :: SszTypes.root()
@@ -72,6 +85,103 @@ defmodule LambdaEthereumConsensus.SszEx do
     |> flatten_results_by(&Enum.join/1)
   end
 
+  defp decode_bitlist(bit_list, max_size) do
+    num_bytes = byte_size(bit_list)
+    num_bits = num_bytes * 8
+    padding = num_bits - bit_size(bit_list)
+    formatted_bit_list = <<0::size(padding), bit_list::bitstring>>
+
+    with {:ok, len} <- get_bitlist_len(formatted_bit_list) do
+      cond do
+        div(len, 8) + 1 != num_bytes ->
+          {:error, "InvalidByteCount"}
+
+        len > max_size ->
+          {:error, "OutOfBounds"}
+
+        true ->
+          bytes_for_bit_len = div(len, 8) + 1
+          bit_size = bit_size(formatted_bit_list)
+          index = rem(len, 8)
+          skip = bit_size - index - 1
+          <<pre::bitstring-size(skip), _::size(1), rest::bitstring>> = formatted_bit_list
+
+          <<pre::binary-size(bytes_for_bit_len - 1), rest::bitstring>> =
+            <<pre::bitstring, 0::1, rest::bitstring>>
+
+          remain_bits_size = len - bit_size(pre)
+          <<_::size(8 - remain_bits_size), rest::bitstring>> = rest
+          formatted = <<pre::bitstring, rest::bitstring>>
+
+          if byte_size(formatted) > bytes_for_bit_len do
+            <<truncated::binary-size(bytes_for_bit_len), _rest::bitstring>> = formatted
+            {:ok, truncated}
+          else
+            {:ok, formatted}
+          end
+      end
+    end
+  end
+
+  defp get_bitlist_len(bit_list) do
+    :binary.bin_to_list(bit_list)
+    |> Enum.with_index()
+    |> Enum.reverse()
+    |> Enum.find_value({:error, "missing_length_information"}, fn {byte, byte_index} ->
+      if for(<<(bit::1 <- <<byte>>)>>, do: bit) |> Enum.any?(&(&1 == 1)) do
+        {:ok, {byte, byte_index}}
+      else
+        {:error, "missing_length_information"}
+      end
+    end)
+    |> case do
+      {:ok, {byte, byte_index}} ->
+        leading_zeros =
+          for(<<(bit::1 <- <<byte>>)>>, do: bit)
+          |> Enum.take_while(&(&1 == 0))
+          |> Enum.count()
+
+        {:ok, byte_index * 8 + 7 - leading_zeros}
+
+      error ->
+        error
+    end
+  end
+
+  defp decode_bitvector(bit_vector, size) when bit_size(bit_vector) == size,
+    do: {:ok, BitVector.new(bit_vector, size)}
+
+  defp decode_bitvector(_bit_vector, _size), do: {:error, "invalid bit_vector length"}
+
+  defp encode_bitlist(bit_list, max_size) do
+    len = bit_size(bit_list)
+    bytes_len = byte_size(bit_list)
+    bytes_for_bit_len = div(len, 8) + 1
+
+    cond do
+      len > max_size ->
+        {:error, "ExcessBits"}
+
+      bytes_for_bit_len > bytes_len ->
+        resize = (bytes_for_bit_len - bytes_len) * @bits_per_byte
+        {:ok, <<bit_list::bitstring, 0::size(resize - 1), 1::1>>}
+
+      bytes_len > 1 ->
+        <<pre::binary-size(bytes_len - 1), last::bitstring>> = bit_list
+        {:ok, <<pre::bitstring, 0::size(8 - bit_size(last) - 1), last::bitstring, 1::1>>}
+
+      true ->
+        resized_bitlist = <<0::1, bit_list::bitstring>>
+        skip = bit_size(resized_bitlist) - len - 1
+        <<pre::bitstring-size(skip), _::size(1), rest::bitstring>> = resized_bitlist
+        padding = 8 - bit_size(pre) - 1 - bit_size(rest)
+        {:ok, <<0::size(padding), pre::bitstring, 1::1, rest::bitstring>>}
+    end
+  end
+
+  defp encode_bitvector(bit_vector, size) when bit_size(bit_vector) == size, do: {:ok, bit_vector}
+  defp encode_bitvector(_bit_vector, _size), do: {:error, "invalid bit_vector length"}
+
   defp encode_variable_size_list(list, _basic_type, max_size) when length(list) > max_size,
     do: {:error, "invalid max_size of list"}
 

test/unit/ssz_ex_test.exs

@@ -187,4 +187,62 @@ defmodule Unit.SSZExTest do
 
     assert_roundtrip(serialized, sync, SszTypes.SyncCommittee)
   end
+
+  test "serialize and deserialize bitlist" do
+    encoded_bytes = <<160, 92, 1>>
+    assert {:ok, decoded_bytes} = SszEx.decode(encoded_bytes, {:bitlist, 16})
+    assert {:ok, ^encoded_bytes} = SszEx.encode(decoded_bytes, {:bitlist, 16})
+
+    encoded_bytes = <<255, 1>>
+    assert {:ok, decoded_bytes} = SszEx.decode(encoded_bytes, {:bitlist, 16})
+    assert {:ok, ^encoded_bytes} = SszEx.encode(decoded_bytes, {:bitlist, 16})
+
+    encoded_bytes = <<31>>
+    assert {:ok, decoded_bytes} = SszEx.decode(encoded_bytes, {:bitlist, 16})
+    assert {:ok, ^encoded_bytes} = SszEx.encode(decoded_bytes, {:bitlist, 16})
+
+    encoded_bytes = <<1>>
+    assert {:ok, decoded_bytes} = SszEx.decode(encoded_bytes, {:bitlist, 31})
+    assert {:ok, ^encoded_bytes} = SszEx.encode(decoded_bytes, {:bitlist, 31})
+
+    encoded_bytes = <<106, 141, 117, 7>>
+    assert {:ok, decoded_bytes} = SszEx.decode(encoded_bytes, {:bitlist, 31})
+    assert {:ok, ^encoded_bytes} = SszEx.encode(decoded_bytes, {:bitlist, 31})
+
+    truncated_encoded_bytes = <<11::4>>
+    expected_encoded_bytes = <<11>>
+    assert {:ok, decoded_bytes} = SszEx.decode(truncated_encoded_bytes, {:bitlist, 31})
+    assert {:ok, ^expected_encoded_bytes} = SszEx.encode(decoded_bytes, {:bitlist, 31})
+
+    truncated_encoded_bytes = <<10::4>>
+    expected_encoded_bytes = <<10>>
+    assert {:ok, decoded_bytes} = SszEx.decode(truncated_encoded_bytes, {:bitlist, 31})
+    assert {:ok, ^expected_encoded_bytes} = SszEx.encode(decoded_bytes, {:bitlist, 31})
+
+    encoded_bytes = <<7>>
+    assert {:error, _msg} = SszEx.decode(encoded_bytes, {:bitlist, 1})
+
+    encoded_bytes = <<124, 3>>
+    assert {:error, _msg} = SszEx.decode(encoded_bytes, {:bitlist, 1})
+
+    encoded_bytes = <<0>>
+    assert {:error, _msg} = SszEx.decode(encoded_bytes, {:bitlist, 1})
+    assert {:error, _msg} = SszEx.decode(encoded_bytes, {:bitlist, 16})
+  end
+
+  test "serialize and deserialize bitvector" do
+    encoded_bytes = <<255, 255>>
+    assert {:ok, decoded_bytes} = SszEx.decode(encoded_bytes, {:bitvector, 16})
+    assert {:ok, ^encoded_bytes} = SszEx.encode(decoded_bytes, {:bitvector, 16})
+
+    encoded_bytes = <<0, 0>>
+    assert {:ok, decoded_bytes} = SszEx.decode(encoded_bytes, {:bitvector, 16})
+    assert {:ok, ^encoded_bytes} = SszEx.encode(decoded_bytes, {:bitvector, 16})
+
+    encoded_bytes = <<255, 255, 255, 255, 1>>
+    assert {:error, _msg} = SszEx.decode(encoded_bytes, {:bitvector, 33})
+
+    encoded_bytes = <<0>>
+    assert {:error, _msg} = SszEx.decode(encoded_bytes, {:bitvector, 9})
+  end
 end