verify_proof Path Naming Change

lib/examples/e_merkle_tree.ex

@@ -1,6 +1,7 @@
 defmodule Examples.EMerkleTree do
   require ExUnit.Assertions
   import ExUnit.Assertions
+  alias Anoma.LocalDomain.BatchMerkleTree
   alias Anoma.LocalDomain.MerkleTree
 
   def empty_tree() do
@@ -12,7 +13,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 +48,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 +75,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
@@ -78,4 +94,33 @@ defmodule Examples.EMerkleTree do
 
     new_tree
   end
+
+  def merkle_tree_differential_root_comparisson(leaves_number \\ 10_000) do
+    tree1 = MerkleTree.new()
+    tree2 = BatchMerkleTree.new()
+
+    for i <- 1..leaves_number, reduce: {tree1, tree2} do
+      {tree1, tree2} ->
+        leaf = :crypto.hash(:sha256, :erlang.term_to_binary(i))
+        upd_tree1 = MerkleTree.add(tree1, [leaf])
+        upd_tree2 = BatchMerkleTree.add(tree2, [leaf])
+
+        assert MerkleTree.root(upd_tree1) ==
+                 BatchMerkleTree.root(upd_tree2)
+
+        {upd_tree1, upd_tree2}
+    end
+  end
+
+  def merkle_tree_differential_proof_comparisson(leaves_number \\ 100) do
+    {tree1, tree2} =
+      merkle_tree_differential_root_comparisson(leaves_number)
+
+    leaves = Map.get(tree1.nodes, 0)
+
+    for leaf <- leaves do
+      assert MerkleTree.generate_proof(tree1, leaf) ==
+               BatchMerkleTree.generate_proof(tree2, leaf)
+    end
+  end
 end

lib/local_domain/structures/batch_merkle_tree.ex

@@ -0,0 +1,212 @@
+defmodule Anoma.LocalDomain.BatchMerkleTree do
+  @moduledoc """
+  I implement merkle tree behaviour for use within local domain applications.
+
+  Has a variable depth.
+  """
+
+  import Bitwise
+  use TypedStruct
+
+  alias __MODULE__
+
+  typedstruct enforce: true do
+    # map from levels to the map from index to the node
+    field(:nodes, %{integer() => %{integer() => binary()}},
+      default: %{
+        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)
+
+    # leaf map with leaves as keys
+    # for efficient path computation
+    field(:leaf_map, %{binary() => integer()}, default: %{})
+  end
+
+  def hash(bytes) do
+    :crypto.hash(:sha256, bytes)
+  end
+
+  def empty() do
+    hash("EMPTY")
+  end
+
+  def new() do
+    # 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
+
+    %BatchMerkleTree{empty_nodes: empty_nodes}
+  end
+
+  def depth(tree) do
+    tree.capacity |> :math.log2() |> trunc()
+  end
+
+  def root(tree) do
+    Map.get(Map.get(tree.nodes, depth(tree)), 0)
+  end
+
+  def add(tree, leaves) do
+    index = tree.next_index
+
+    # Compute the new leaf map and new commitment length
+    # in the same loop
+    {new_leaf_map, new_commitment_length} =
+      for leaf <- leaves, reduce: {tree.leaf_map, index} do
+        {map_acc, index_acc} ->
+          {Map.put(map_acc, index_acc, leaf), index + 1}
+      end
+
+    {depth, capacity} =
+      if new_commitment_length >= tree.capacity do
+        # If the tree capacity is exceeded, calculate the new minimal depth
+        new_depth =
+          (new_commitment_length |> :math.log2() |> trunc()) + 1
+
+        {new_depth, Integer.pow(2, new_depth)}
+      else
+        {depth(tree), tree.capacity}
+      end
+
+    # Add leaves and recompute needed intermediary nodes
+    new_nodes =
+      compute_nodes(depth, tree.nodes, tree.empty_nodes, index, leaves)
+
+    %BatchMerkleTree{
+      tree
+      | nodes: new_nodes,
+        next_index: new_commitment_length,
+        capacity: capacity,
+        leaf_map: new_leaf_map
+    }
+  end
+
+  def generate_proof(tree, leaf) do
+    # Get the index of the leaf
+    index_found = Map.get(tree.leaf_map, leaf)
+
+    if index_found do
+      {path, root, _index} =
+        for i <- 0..(depth(tree) - 1),
+            reduce: {[], leaf, index_found} 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
+
+      if root == root(tree) do
+        {path, root}
+      else
+        nil
+      end
+    end
+  end
+
+  def verify_proof(leaf, path, root) do
+    calculated_root =
+      Enum.reduce(path, leaf, fn {neighbour, is_left}, acc ->
+        if is_left do
+          hash(acc <> neighbour)
+        else
+          hash(neighbour <> acc)
+        end
+      end)
+
+    calculated_root == root
+  end
+
+  defp compute_nodes(depth, nodes, empty_nodes, index, leaves) do
+    # Iterate over each level of the tree, updating
+    # only the parent nodes of the given leaves
+    {new_nodes, _, _} =
+      for i <- 0..depth, reduce: {nodes, index, leaves} do
+        {acc_nodes, index, nodes} ->
+          # Fetch the current level of the tree
+          current_nodes = Map.get(acc_nodes, i, %{})
+
+          # If the first node is the right one, fetch its sibling
+          initial_left_sibling =
+            if is_left(index) do
+              nil
+            else
+              Map.get(current_nodes, index - 1)
+            end
+
+          # Iterate over all the nodes
+          # Populate the current level with them
+          # Also calculate the list of parent nodes
+          {updated_current_nodes, parents, _, final_left_sibling} =
+            for node <- nodes,
+                reduce: {current_nodes, [], index, initial_left_sibling} do
+              {acc_current_nodes, parents, j, left_sibling} ->
+                if left_sibling do
+                  # Hash the left sibling with the current node
+                  {Map.put(acc_current_nodes, j, node),
+                   [hash(left_sibling <> node) | parents], j + 1, nil}
+                else
+                  # record the current node as a left sibling
+                  {Map.put(acc_current_nodes, j, node), parents, j + 1,
+                   node}
+                end
+            end
+
+          # If the final node was a left one, we have to compute one more parent
+          # by hashing with an empty node of the appropriate level
+          final_parents =
+            if final_left_sibling do
+              [
+                hash(final_left_sibling <> Map.get(empty_nodes, i))
+                | parents
+              ]
+            else
+              parents
+            end
+
+          {Map.put(acc_nodes, i, updated_current_nodes), div(index, 2),
+           Enum.reverse(final_parents)}
+      end
+
+    new_nodes
+  end
+
+  defp is_left(index) do
+    (index &&& 1) == 0
+  end
+end