feat: optimization on merkleization of chunks by precomputed hashes (#644)

Author: Godspower-Eze

bench/ssz.exs

@@ -34,3 +34,26 @@ Benchee.run(
   warmup: 2,
   time: 5
 )
+
+## Benchmark Merkleization 
+
+list = Stream.cycle([65_535]) |> Enum.take(316)
+schema = {:list, {:int, 16}, 1024}
+packed_chunks = SszEx.pack(list, schema)
+limit = SszEx.chunk_count(schema)
+
+Benchee.run(
+  %{
+    "SszEx.merkleize_chunks" => fn {chunks, leaf_count} ->
+      SszEx.merkleize_chunks(chunks, leaf_count)
+    end,
+    "SszEx.merkleize_chunks_with_virtual_padding" => fn {chunks, leaf_count} ->
+      SszEx.merkleize_chunks_with_virtual_padding(chunks, leaf_count)
+    end
+  },
+  inputs: %{
+    "args" => {packed_chunks, limit}
+  },
+  warmup: 2,
+  time: 5
+)

lib/ssz_ex.ex

@@ -5,6 +5,7 @@ defmodule LambdaEthereumConsensus.SszEx do
   alias LambdaEthereumConsensus.Utils.BitList
   alias LambdaEthereumConsensus.Utils.BitVector
   import alias LambdaEthereumConsensus.Utils.BitVector
+  alias LambdaEthereumConsensus.Utils.ZeroHashes
 
   #################
   ### Public API
@@ -15,6 +16,7 @@ defmodule LambdaEthereumConsensus.SszEx do
   @bits_per_byte 8
   @bits_per_chunk @bytes_per_chunk * @bits_per_byte
   @zero_chunk <<0::size(@bits_per_chunk)>>
+  @zero_hashes ZeroHashes.compute_zero_hashes()
 
   @spec hash(iodata()) :: binary()
   def hash(data), do: :crypto.hash(:sha256, data)
@@ -109,7 +111,7 @@ defmodule LambdaEthereumConsensus.SszEx do
     if chunks_len > limit do
       {:error, "chunk size exceeds limit"}
     else
-      root = merkleize_chunks(chunks, limit) |> mix_in_length(len)
+      root = merkleize_chunks_with_virtual_padding(chunks, limit) |> mix_in_length(len)
       {:ok, root}
     end
   end
@@ -118,7 +120,7 @@ defmodule LambdaEthereumConsensus.SszEx do
           {:ok, Types.root()} | {:error, String.t()}
   def hash_tree_root_vector_basic_type(chunks) do
     leaf_count = chunks |> get_chunks_len() |> next_pow_of_two()
-    root = merkleize_chunks(chunks, leaf_count)
+    root = merkleize_chunks_with_virtual_padding(chunks, leaf_count)
     {:ok, root}
   end
 
@@ -164,6 +166,38 @@ defmodule LambdaEthereumConsensus.SszEx do
     end
   end
 
+  def merkleize_chunks_with_virtual_padding(chunks, leaf_count) do
+    chunks_len = chunks |> get_chunks_len()
+    power = leaf_count |> compute_pow()
+    height = power + 1
+
+    cond do
+      chunks_len == 0 ->
+        depth = height - 1
+        get_zero_hash(depth)
+
+      chunks_len == 1 and leaf_count == 1 ->
+        chunks
+
+      true ->
+        power = leaf_count |> compute_pow()
+        height = power + 1
+        layers = chunks
+        last_index = chunks_len - 1
+
+        {_, final_layer} =
+          1..(height - 1)
+          |> Enum.reverse()
+          |> Enum.reduce({last_index, layers}, fn i, {acc_last_index, acc_layers} ->
+            updated_layers = update_layers(i, height, acc_layers, acc_last_index)
+            {acc_last_index |> div(2), updated_layers}
+          end)
+
+        <<root::binary-size(@bytes_per_chunk), _::binary>> = final_layer
+        root
+    end
+  end
+
   @spec pack(boolean, :bool) :: binary()
   def pack(true, :bool), do: <<1::@bits_per_chunk-little>>
   def pack(false, :bool), do: @zero_chunk
@@ -184,6 +218,11 @@ defmodule LambdaEthereumConsensus.SszEx do
     end
   end
 
+  def chunk_count({:list, type, max_size}) do
+    size = size_of(type)
+    (max_size * size + 31) |> div(32)
+  end
+
   #################
   ### Private functions
   #################
@@ -592,16 +631,6 @@ defmodule LambdaEthereumConsensus.SszEx do
 
   defp size_of({:int, size}), do: size |> div(@bits_per_byte)
 
-  defp chunk_count({:list, {:int, size}, max_size}) do
-    size = size_of({:int, size})
-    (max_size * size + 31) |> div(32)
-  end
-
-  defp chunk_count({:list, :bool, max_size}) do
-    size = size_of(:bool)
-    (max_size * size + 31) |> div(32)
-  end
-
   defp pack_basic_type_list(list, schema) do
     list
     |> Enum.reduce(<<>>, fn x, acc ->
@@ -635,12 +664,69 @@ defmodule LambdaEthereumConsensus.SszEx do
     end
   end
 
-  defp next_pow_of_two(len) when is_integer(len) and len >= 0 do
-    if len == 0 do
-      0
-    else
-      n = ((len <<< 1) - 1) |> :math.log2() |> trunc()
-      2 ** n
+  defp next_pow_of_two(0), do: 0
+
+  defp next_pow_of_two(len) when is_integer(len) and len > 0 do
+    n = ((len <<< 1) - 1) |> compute_pow()
+    2 ** n
+  end
+
+  defp get_chunks_len(chunks) do
+    chunks |> byte_size() |> div(@bytes_per_chunk)
+  end
+
+  defp compute_pow(value) do
+    :math.log2(value) |> trunc()
+  end
+
+  defp update_layers(i, height, acc_layers, acc_last_index) do
+    0..(2 ** i - 1)
+    |> Enum.filter(fn x -> rem(x, 2) == 0 end)
+    |> Enum.reduce_while(acc_layers, fn j, acc_layers ->
+      parent_index = j |> div(2)
+      nodes = get_nodes(parent_index, i, j, height, acc_layers, acc_last_index)
+      hash_nodes_and_replace(nodes, acc_layers)
+    end)
+  end
+
+  defp get_nodes(parent_index, _i, j, _height, acc_layers, acc_last_index)
+       when j < acc_last_index do
+    start = parent_index * @bytes_per_chunk
+    stop = (j + 2) * @bytes_per_chunk
+    focus = acc_layers |> :binary.part(start, stop - start)
+    focus_len = focus |> byte_size()
+    children_index = focus_len - 2 * @bytes_per_chunk
+    <<_::binary-size(children_index), children::binary>> = focus
+
+    <<left::binary-size(@bytes_per_chunk), right::binary-size(@bytes_per_chunk)>> =
+      children
+
+    {children_index, left, right}
+  end
+
+  defp get_nodes(parent_index, i, j, height, acc_layers, acc_last_index)
+       when j == acc_last_index do
+    start = parent_index * @bytes_per_chunk
+    stop = (j + 1) * @bytes_per_chunk
+    focus = acc_layers |> :binary.part(start, stop - start)
+    focus_len = focus |> byte_size()
+    children_index = focus_len - @bytes_per_chunk
+    <<_::binary-size(children_index), left::binary>> = focus
+    depth = height - i - 1
+    right = get_zero_hash(depth)
+    {children_index, left, right}
+  end
+
+  defp get_nodes(_, _, _, _, _, _), do: :error
+
+  defp hash_nodes_and_replace(nodes, layers) do
+    case nodes do
+      :error ->
+        {:halt, layers}
+
+      {index, left, right} ->
+        hash = hash_nodes(left, right)
+        {:cont, replace_chunk(layers, index, hash)}
     end
   end
 
@@ -651,7 +737,9 @@ defmodule LambdaEthereumConsensus.SszEx do
     <<left::binary, new_chunk::binary, right::binary>>
   end
 
-  defp get_chunks_len(chunks) do
-    chunks |> byte_size() |> div(@bytes_per_chunk)
+  defp get_zero_hash(depth) do
+    offset = (depth + 1) * @bytes_per_chunk - @bytes_per_chunk
+    <<_::binary-size(offset), hash::binary-size(@bytes_per_chunk), _::binary>> = @zero_hashes
+    hash
   end
 end

lib/utils/zero_hashes.ex

@@ -0,0 +1,34 @@
+defmodule LambdaEthereumConsensus.Utils.ZeroHashes do
+  @moduledoc """
+  Precomputed zero hashes
+  """
+
+  @bits_per_byte 8
+  @bytes_per_chunk 32
+  @bits_per_chunk @bytes_per_chunk * @bits_per_byte
+  @max_merkle_tree_depth 64
+
+  def compute_zero_hashes do
+    buffer = <<0::size(@bytes_per_chunk * @max_merkle_tree_depth * @bits_per_byte)>>
+
+    0..(@max_merkle_tree_depth - 2)
+    |> Enum.reduce(buffer, fn index, acc_buffer ->
+      start = index * @bytes_per_chunk
+      stop = (index + 2) * @bytes_per_chunk
+      focus = acc_buffer |> :binary.part(start, stop - start)
+      <<left::binary-size(@bytes_per_chunk), _::binary>> = focus
+      hash = hash_nodes(left, left)
+      change_index = (index + 1) * @bytes_per_chunk
+      replace_chunk(acc_buffer, change_index, hash)
+    end)
+  end
+
+  defp hash_nodes(left, right), do: :crypto.hash(:sha256, left <> right)
+
+  defp replace_chunk(chunks, start, new_chunk) do
+    <<left::binary-size(start), _::size(@bits_per_chunk), right::binary>> =
+      chunks
+
+    <<left::binary, new_chunk::binary, right::binary>>
+  end
+end

test/unit/ssz_ex_test.exs

@@ -72,8 +72,8 @@ defmodule Unit.SSZExTest do
   end
 
   test "merklelization of chunks" do
-    ## Reference:  https://github.com/ralexstokes/ssz-rs/blob/1f94d5dfc70c86dab672e91ac46af04a5f96c342/ssz-rs/src/merkleization/mod.rs#L371
-    ##            https://github.com/ralexstokes/ssz-rs/blob/1f94d5dfc70c86dab672e91ac46af04a5f96c342/ssz-rs/src/merkleization/mod.rs#L416
+    # Reference:  https://github.com/ralexstokes/ssz-rs/blob/1f94d5dfc70c86dab672e91ac46af04a5f96c342/ssz-rs/src/merkleization/mod.rs#L371
+    #            https://github.com/ralexstokes/ssz-rs/blob/1f94d5dfc70c86dab672e91ac46af04a5f96c342/ssz-rs/src/merkleization/mod.rs#L416
     zero = <<0::256>>
 
     chunks = zero
@@ -136,13 +136,71 @@ defmodule Unit.SSZExTest do
     root = chunks |> SszEx.merkleize_chunks(2 ** 10)
     expected_value = "2647cb9e26bd83eeb0982814b2ac4d6cc4a65d0d98637f1a73a4c06d3db0e6ce"
     assert root |> Base.encode16(case: :lower) == expected_value
+  end
+
+  test "merklelization of chunks with virtual padding" do
+    zero = <<0::256>>
+
+    chunks = zero
+    root = SszEx.merkleize_chunks_with_virtual_padding(chunks, 1)
+    expected_value = "0000000000000000000000000000000000000000000000000000000000000000"
+    assert root |> Base.encode16(case: :lower) == expected_value
+
+    chunks = zero <> zero
+    root = chunks |> SszEx.merkleize_chunks_with_virtual_padding(2)
+    expected_value = "f5a5fd42d16a20302798ef6ed309979b43003d2320d9f0e8ea9831a92759fb4b"
+    assert root |> Base.encode16(case: :lower) == expected_value
+
+    ones = 0..31 |> Enum.reduce(<<>>, fn _, acc -> <<1>> <> acc end)
+
+    chunks = ones <> ones
+    root = chunks |> SszEx.merkleize_chunks_with_virtual_padding(2)
+    expected_value = "7c8975e1e60a5c8337f28edf8c33c3b180360b7279644a9bc1af3c51e6220bf5"
+    assert root |> Base.encode16(case: :lower) == expected_value
+
+    chunks = zero <> zero <> zero <> zero
+    root = chunks |> SszEx.merkleize_chunks_with_virtual_padding(4)
+    expected_value = "db56114e00fdd4c1f85c892bf35ac9a89289aaecb1ebd0a96cde606a748b5d71"
+    assert root |> Base.encode16(case: :lower) == expected_value
+
+    chunks = zero <> zero <> zero <> zero <> zero <> zero <> zero <> zero
+    root = chunks |> SszEx.merkleize_chunks_with_virtual_padding(8)
+    expected_value = "c78009fdf07fc56a11f122370658a353aaa542ed63e44c4bc15ff4cd105ab33c"
+    assert root |> Base.encode16(case: :lower) == expected_value
+
+    chunks = ones
+    root = chunks |> SszEx.merkleize_chunks_with_virtual_padding(4)
+    expected_value = "29797eded0e83376b70f2bf034cc0811ae7f1414653b1d720dfd18f74cf13309"
+    assert root |> Base.encode16(case: :lower) == expected_value
+
+    twos = 0..31 |> Enum.reduce(<<>>, fn _, acc -> <<2>> <> acc end)
+
+    chunks = twos
+    root = chunks |> SszEx.merkleize_chunks_with_virtual_padding(8)
+    expected_value = "fa4cf775712aa8a2fe5dcb5a517d19b2e9effcf58ff311b9fd8e4a7d308e6d00"
+    assert root |> Base.encode16(case: :lower) == expected_value
+
+    chunks = ones <> ones <> ones
+    root = chunks |> SszEx.merkleize_chunks_with_virtual_padding(4)
+    expected_value = "65aa94f2b59e517abd400cab655f42821374e433e41b8fe599f6bb15484adcec"
+    assert root |> Base.encode16(case: :lower) == expected_value
 
-    ## TOO HEAVY COMPUTATION!
-    # chunks = 1..70 |> Enum.reduce(<<>>, fn _, acc -> acc <> ones end)
-    # leaf_count = 9_223_372_036_854_775_808 # 2 ** 63
-    # root = chunks |> SszEx.merkleize_chunks(leaf_count)
-    # expected_value = "9317695d95b5a3b46e976b5a9cbfcfccb600accaddeda9ac867cc9669b862979"
-    # assert root |> Base.encode16(case: :lower) == expected_value
+    chunks = ones <> ones <> ones <> ones <> ones
+    root = chunks |> SszEx.merkleize_chunks_with_virtual_padding(8)
+    expected_value = "0ae67e34cba4ad2bbfea5dc39e6679b444021522d861fab00f05063c54341289"
+    assert root |> Base.encode16(case: :lower) == expected_value
+
+    chunks = ones <> ones <> ones <> ones <> ones <> ones
+    root = chunks |> SszEx.merkleize_chunks_with_virtual_padding(8)
+    expected_value = "0ef7df63c204ef203d76145627b8083c49aa7c55ebdee2967556f55a4f65a238"
+    assert root |> Base.encode16(case: :lower) == expected_value
+
+    ## Large Leaf Count
+
+    chunks = ones <> ones <> ones <> ones <> ones
+    root = chunks |> SszEx.merkleize_chunks_with_virtual_padding(2 ** 10)
+    expected_value = "2647cb9e26bd83eeb0982814b2ac4d6cc4a65d0d98637f1a73a4c06d3db0e6ce"
+    assert root |> Base.encode16(case: :lower) == expected_value
   end
 
   test "hash tree root of list" do