make balance search use async

Author: SachinMeier

lib/proof_of_reserves.ex

@@ -119,6 +119,98 @@ defmodule ProofOfReserves do
     end)
   end
 
+  @doc """
+  async_find_balances_for_accounts is the async version of find_balances_for_accounts.
+  It is a recursive function that splits the leaves into two lists, finds the balances
+  for each account in each list, and then merges the results of the two lists.
+  """
+  @spec async_find_balances_for_accounts(
+          list(MerkleSumTree.Node.t()),
+          non_neg_integer(),
+          list(%{
+            account_id: non_neg_integer(),
+            account_subkey: binary()
+          })
+        ) ::
+          list(%{
+            account_id: non_neg_integer(),
+            balance: non_neg_integer(),
+            attestation_key: binary()
+          })
+  def async_find_balances_for_accounts(leaves, block_height, accounts) do
+    account_balances =
+      Enum.map(accounts, fn %{account_id: account_id, account_subkey: subkey} ->
+        %{
+          account_id: account_id,
+          balance: 0,
+          attestation_key: Util.calculate_attestation_key(subkey, block_height, account_id)
+        }
+      end)
+
+    async_find_balances(leaves, 0, account_balances)
+  end
+
+  # helper function for async_find_balances_for_accounts
+  @spec async_find_balances(
+          list(MerkleSumTree.Node.t()),
+          non_neg_integer(),
+          list(%{
+            account_id: non_neg_integer(),
+            balance: non_neg_integer(),
+            attestation_key: binary()
+          })
+        ) ::
+          list(%{
+            account_id: non_neg_integer(),
+            balance: non_neg_integer(),
+            attestation_key: binary()
+          })
+  defp async_find_balances([], _leaf_idx, account_balances), do: account_balances
+
+  defp async_find_balances([leaf], leaf_idx, account_balances) do
+    # when there is only one leaf, we can simply check if the leaf hash matches the leaf hash
+    # for each account in the list. If it does, we add the value to the balance.
+    Enum.map(account_balances, fn %{balance: balance, attestation_key: attestation_key} =
+                                    account_balance ->
+      # TODO: only a single account should match the leaf hash at most, so we can return early
+      if Util.leaf_hash(leaf.value, attestation_key, leaf_idx) == leaf.hash do
+        # if the leaf hash matches, we add the value to the balance
+        Map.put(account_balance, :balance, balance + leaf.value)
+      else
+        account_balance
+      end
+    end)
+  end
+
+  defp async_find_balances(leaves, leaf_idx, account_balances) do
+    # Split the leaves into two lists
+    leaf_ct = length(leaves)
+    mid_idx = div(leaf_ct, 2)
+
+    {left_leaves, right_leaves} = Enum.split(leaves, mid_idx)
+
+    # Recursively & asynchronously find the balances for each account in both lists
+    left_task = Task.async(fn -> async_find_balances(left_leaves, leaf_idx, account_balances) end)
+
+    right_task =
+      Task.async(fn -> async_find_balances(right_leaves, leaf_idx + mid_idx, account_balances) end)
+
+    left_balances = Task.await(left_task, :infinity)
+    right_balances = Task.await(right_task, :infinity)
+
+    # Merge the results of the two lists
+    Enum.zip_with(left_balances, right_balances, fn
+      # the matching on account_id and attestation_key ensures that the balances are added to the correct account
+      %{account_id: account_id, balance: left_balance, attestation_key: attestation_key},
+      %{account_id: account_id, balance: right_balance, attestation_key: attestation_key} ->
+        %{
+          account_id: account_id,
+          balance: left_balance + right_balance,
+          attestation_key: attestation_key
+        }
+    end)
+  end
+
   @doc """
   get_tree_root returns the root node of a Merkle Sum Tree.
   """

verify_liabilities.exs

@@ -157,12 +157,10 @@ defmodule VerifyLiabilities do
     account_uids =
       Map.new(accounts, fn %{account_id: id, account_uid: account_uid} -> {id, account_uid} end)
 
-    IO.puts("Verifying balances in Merkle Tree...")
-
     balances =
       tree
       |> ProofOfReserves.MerkleSumTree.get_leaves()
-      |> ProofOfReserves.find_balances_for_accounts(block_height, accounts)
+      |> ProofOfReserves.async_find_balances_for_accounts(block_height, accounts)
       |> Enum.map(fn %{balance: balance, account_id: account_id} ->
         %{
           account_uid: Map.fetch!(account_uids, account_id),