Implement Fast Sequentially-Updated Merkle Tree

lib/examples/e_merkle_tree.ex

@@ -12,7 +12,7 @@ defmodule Examples.EMerkleTree do
       empty_tree()
       |> MerkleTree.add([:crypto.hash(:sha256, "a")])
 
-    assert Enum.at(Map.get(new_tree.nodes, 1), 0) ==
+    assert Map.get(Map.get(new_tree.nodes, 1), 0) ==
              :crypto.hash(
                :sha256,
                :crypto.hash(:sha256, "a") <>
@@ -47,6 +47,15 @@ defmodule Examples.EMerkleTree do
     {frontiers, root}
   end
 
+  def generate_a_proof_wrongly() do
+    tree = expand_merkle_tree()
+
+    assert nil ==
+             MerkleTree.generate_proof(tree, :crypto.hash(:sha256, "d"))
+
+    :ok
+  end
+
   def verify_a_proof() do
     {frontiers, root} = generate_a_proof()
     MerkleTree.verify_proof(:crypto.hash(:sha256, "b"), frontiers, root)
@@ -65,8 +74,14 @@ defmodule Examples.EMerkleTree do
 
     new_tree = MerkleTree.add(empty_tree(), leaves)
 
+    sorted_leaves =
+      Map.get(new_tree.nodes, 0)
+      |> Map.to_list()
+      |> Enum.sort(fn {ind1, _}, {ind2, _} -> ind1 < ind2 end)
+      |> Enum.map(&elem(&1, 1))
+
     # Check that the leaves are as expected
-    assert List.starts_with?(Map.get(new_tree.nodes, 0), leaves)
+    assert List.starts_with?(sorted_leaves, leaves)
 
     # Check that index is as expected
     assert new_tree.next_index == leaves_length

lib/local_domain/structures/merkle_tree.ex

@@ -8,13 +8,17 @@ defmodule Anoma.LocalDomain.MerkleTree do
   import Bitwise
   use TypedStruct
 
+  alias __MODULE__
+
   typedstruct enforce: true do
-    field(:nodes, %{integer() => list(binary())},
+    # map from levels to the map from index to the node
+    field(:nodes, %{integer() => %{integer() => binary()}},
       default: %{
-        0 => [:crypto.hash(:sha256, "EMPTY")]
+        0 => %{0 => :crypto.hash(:sha256, "EMPTY")}
       }
     )
 
+    field(:empty_nodes, %{integer() => binary()})
     field(:next_index, non_neg_integer(), default: 0)
     field(:capacity, non_neg_integer(), default: 1)
   end
@@ -28,66 +32,92 @@ defmodule Anoma.LocalDomain.MerkleTree do
   end
 
   def new() do
-    %Anoma.LocalDomain.MerkleTree{}
+    # Assume we have a tree at most of depth 32
+    empty_nodes =
+      for i <- 1..31, reduce: %{0 => :crypto.hash(:sha256, "EMPTY")} do
+        acc ->
+          previous_empty_hash = Map.get(acc, i - 1)
+
+          Map.put(
+            acc,
+            i,
+            :crypto.hash(
+              :sha256,
+              previous_empty_hash <> previous_empty_hash
+            )
+          )
+      end
+
+    %Anoma.LocalDomain.MerkleTree{empty_nodes: empty_nodes}
   end
 
   def depth(tree) do
-    length(Map.keys(tree.nodes)) - 1
+    tree.capacity |> :math.log2() |> trunc()
   end
 
   def root(tree) do
-    Enum.at(Map.get(tree.nodes, depth(tree)), 0)
-  end
-
-  def generate_empty_branch(n) do
-    Enum.map(1..n, fn _ -> empty() end)
+    Map.get(Map.get(tree.nodes, depth(tree)), 0)
   end
 
   def add(tree, values) do
-    {new_leaves, next_index, capacity} =
-      Enum.reduce(
-        values,
-        {Map.get(tree.nodes, 0), tree.next_index, tree.capacity},
-        fn leaf, {new_leaves, next_index, new_capacity} ->
-          add_leaf(new_leaves, next_index, new_capacity, leaf)
-        end
-      )
-
-    %Anoma.LocalDomain.MerkleTree{
-      nodes:
-        Map.merge(%{0 => new_leaves}, calculate_nodes(new_leaves, 1)),
-      next_index: next_index,
-      capacity: capacity
-    }
+    Enum.reduce(
+      values,
+      tree,
+      fn leaf, acc_tree ->
+        add_leaf(acc_tree, leaf)
+      end
+    )
   end
 
   def generate_proof(tree, leaf) do
-    {frontiers, root} =
-      for i <- 0..(depth(tree) - 1), reduce: {[], leaf} do
-        {acc, leaf} ->
-          leaves = Map.get(tree.nodes, i)
-
-          leaf_index =
-            leaves
-            |> Enum.find_index(&(&1 == leaf))
-
-          is_left = (leaf_index &&& 1) == 0
-
-          if is_left do
-            neighbour = Enum.at(leaves, leaf_index + 1)
+    leaves = Map.get(tree.nodes, 0)
+
+    # Get the index of the leaf
+    found_elem =
+      leaves
+      |> Map.to_list()
+      |> Enum.find(fn
+        {_, ^leaf} ->
+          true
+
+        _ ->
+          false
+      end)
 
-            {acc ++ [{neighbour, true}], hash(leaf <> neighbour)}
-          else
-            neighbour = Enum.at(leaves, leaf_index - 1)
+    if found_elem do
+      {path, root, _index} =
+        for i <- 0..(depth(tree) - 1),
+            reduce: {[], leaf, elem(found_elem, 0)} do
+          {path, node, index} ->
+            # Take the current level of the tree
+            current_nodes = Map.get(tree.nodes, i)
+
+            is_left = (index &&& 1) == 0
+
+            if is_left do
+              # If the node is a left one, take its right sibling
+              sibling = Map.get(current_nodes, index + 1, empty())
+
+              # Hash the node on the left and sibling on the right
+              # The index of its parents is going to be index / 2
+              {path ++ [{sibling, true}], hash(node <> sibling),
+               div(index, 2)}
+            else
+              # If the node is a right one, take its left sibling
+              sibling = Map.get(current_nodes, index - 1, empty())
+
+              # Hash the node on the right and sibling on the left
+              # The index of its parents is going to be (index - 1) / 2
+              {path ++ [{sibling, false}], hash(sibling <> node),
+               div(index - 1, 2)}
+            end
+        end
 
-            {acc ++ [{neighbour, false}], hash(neighbour <> leaf)}
-          end
+      if root == root(tree) do
+        {path, root}
+      else
+        nil
       end
-
-    if root == root(tree) do
-      {frontiers, root}
-    else
-      nil
     end
   end
 
@@ -104,27 +134,69 @@ defmodule Anoma.LocalDomain.MerkleTree do
     calculated_root == root
   end
 
-  defp add_leaf(leaves, index, capacity, leaf) do
-    new_leaves = List.replace_at(leaves, index, leaf)
-    new_index = index + 1
-
-    if new_index == capacity do
-      {new_leaves ++ generate_empty_branch(length(leaves)), new_index,
-       capacity * 2}
+  defp add_leaf(tree, leaf) do
+    index = tree.next_index
+    depth = depth(tree)
+
+    # Add a leaf and recompute needed intermediary nodes
+    new_nodes =
+      compute_nodes(depth, tree.nodes, tree.empty_nodes, index, leaf)
+
+    if index + 1 == tree.capacity do
+      # If the tree is fully filled, we need to recompute a new
+      # root as if by adding an extra empty leaf at next index
+      expanded_nodes =
+        compute_nodes(
+          depth + 1,
+          new_nodes,
+          tree.empty_nodes,
+          index + 1,
+          empty()
+        )
+
+      %MerkleTree{
+        tree
+        | nodes: expanded_nodes,
+          next_index: index + 1,
+          capacity: tree.capacity * 2
+      }
     else
-      {new_leaves, new_index, capacity}
+      %MerkleTree{tree | nodes: new_nodes, next_index: index + 1}
     end
   end
 
-  defp calculate_nodes(leaves, i) do
-    next_level =
-      Enum.chunk_every(leaves, 2, 2, :discard)
-      |> Enum.map(fn [a, b] -> hash(a <> b) end)
+  defp compute_nodes(depth, nodes, empty_nodes, index, leaf) do
+    # Iterate over each level of the tree, updating
+    # only the parent nodes of the leaf
+    {new_nodes, _, _} =
+      for i <- 0..depth, reduce: {nodes, index, leaf} do
+        {acc_nodes, index, node} ->
+          # Fetch the current level of the tree
+          current_nodes = Map.get(acc_nodes, i, %{})
 
-    if length(next_level) > 1 do
-      Map.merge(%{i => next_level}, calculate_nodes(next_level, i + 1))
-    else
-      %{i => next_level}
-    end
+          # Put the updated node at the given index
+          updated_nodes =
+            Map.put(acc_nodes, i, Map.put(current_nodes, index, node))
+
+          is_left = (index &&& 1) == 0
+
+          if is_left do
+            # If the node is a left one, fetch its right sibling
+            sibling =
+              Map.get(current_nodes, index + 1, Map.get(empty_nodes, i))
+
+            # Hash the node on the left and sibling on the right
+            # The index of its parents is going to be index / 2
+            {updated_nodes, div(index, 2), hash(node <> sibling)}
+          else
+            # If the node is a right one, fetch its left sibling
+            sibling = Map.get(current_nodes, index - 1)
+            # Hash the node on the left and sibling on the right
+            # The index of its parents is going to be (index - 1) / 2
+            {updated_nodes, div(index - 1, 2), hash(sibling <> node)}
+          end
+      end
+
+    new_nodes
   end
 end