Fix Archethic.get_balance to prefer UTXO list over aggregated balance

lib/archethic.ex

@@ -13,9 +13,7 @@ defmodule Archethic do
   alias Archethic.P2P.Message
 
   alias Archethic.P2P.Message.{
-    Balance,
     Error,
-    GetBalance,
     NewTransaction,
     Ok,
     StartMining,
@@ -311,38 +309,11 @@ defmodule Archethic do
   @doc """
   Retrieve the balance from an address from the closest nodes
   """
-  @spec get_balance(binary) :: {:ok, UTXO.balance()} | {:error, :network_issue}
+  @spec get_balance(binary) :: UTXO.balance()
   def get_balance(address) when is_binary(address) do
     address
-    |> Election.chain_storage_nodes(P2P.authorized_and_available_nodes())
-    |> get_balance(address)
-  end
-
-  defp get_balance(nodes, address) do
-    case P2P.quorum_read(nodes, %GetBalance{address: address}, &balance_conflict_resolver/1) do
-      {:ok, %Balance{uco: uco, token: token}} -> {:ok, %{uco: uco, token: token}}
-      error -> error
-    end
-  end
-
-  defp balance_conflict_resolver(balances) do
-    %Balance{last_chain_sync_date: highest_date} =
-      Enum.max_by(balances, & &1.last_chain_sync_date, DateTime)
-
-    {max_uco, max_token} =
-      balances
-      |> Enum.filter(&(&1.last_chain_sync_date == highest_date))
-      |> Enum.reduce({0, %{}}, fn
-        %Balance{uco: uco, token: token}, {uco_acc, token_acc} ->
-          token_merger = fn _k, v1, v2 -> max(v1, v2) end
-
-          maximum_token = Map.merge(token, token_acc, token_merger)
-          maximum_uco = max(uco, uco_acc)
-
-          {maximum_uco, maximum_token}
-      end)
-
-    %{uco: max_uco, token: max_token}
+    |> get_unspent_outputs()
+    |> UTXO.get_balance()
   end
 
   @doc """

lib/archethic/contracts/interpreter/library/common/chain_impl.ex

@@ -194,12 +194,9 @@ defmodule Archethic.Contracts.Interpreter.Library.Common.ChainImpl do
   end
 
   defp fetch_balance(address, function) do
-    with {:ok, genesis_address} <- Archethic.fetch_genesis_address(address),
-         {:ok, balance} <- Archethic.get_balance(genesis_address) do
-      balance
-    else
-      _ ->
-        raise Library.Error, message: "Network issue in #{function}"
+    case Archethic.fetch_genesis_address(address) do
+      {:ok, genesis_address} -> Archethic.get_balance(genesis_address)
+      _ -> raise Library.Error, message: "Network issue in #{function}"
     end
   end
 end

lib/archethic/p2p/message.ex

@@ -16,7 +16,6 @@ defmodule Archethic.P2P.Message do
 
   alias __MODULE__.{
     AddressList,
-    Balance,
     BeaconSummaryList,
     BeaconUpdate,
     BootstrappingNodes,
@@ -27,7 +26,6 @@ defmodule Archethic.P2P.Message do
     FirstPublicKey,
     GenesisAddress,
     GetGenesisAddress,
-    GetBalance,
     GetBeaconSummaries,
     GetBeaconSummary,
     GetBeaconSummariesAggregate,
@@ -102,7 +100,6 @@ defmodule Archethic.P2P.Message do
           | CrossValidationDone.t()
           | ReplicateTransaction.t()
           | GetLastTransaction.t()
-          | GetBalance.t()
           | GetTransactionInputs.t()
           | GetTransactionChainLength.t()
           | GetFirstTransactionAddress.t()
@@ -143,7 +140,6 @@ defmodule Archethic.P2P.Message do
           | Transaction.t()
           | NodeList.t()
           | UnspentOutputList.t()
-          | Balance.t()
           | EncryptedStorageNonce.t()
           | BootstrappingNodes.t()
           | TransactionSummaryMessage.t()

lib/archethic/p2p/message/message_id.ex

@@ -19,7 +19,6 @@ defmodule Archethic.P2P.MessageId do
     AcknowledgeStorage,
     NotifyEndOfNodeSync,
     GetLastTransaction,
-    GetBalance,
     GetTransactionInputs,
     GetTransactionChainLength,
     GetFirstPublicKey,
@@ -52,7 +51,6 @@ defmodule Archethic.P2P.MessageId do
     TransactionChainLength,
     BootstrappingNodes,
     EncryptedStorageNonce,
-    Balance,
     NodeList,
     UnspentOutputList,
     TransactionList,
@@ -101,7 +99,7 @@ defmodule Archethic.P2P.MessageId do
     AcknowledgeStorage => 13,
     NotifyEndOfNodeSync => 14,
     GetLastTransaction => 15,
-    GetBalance => 16,
+    # Message number 16 is available
     GetTransactionInputs => 17,
     GetTransactionChainLength => 18,
     RequestChainLock => 19,
@@ -152,7 +150,7 @@ defmodule Archethic.P2P.MessageId do
     TransactionChainLength => 245,
     BootstrappingNodes => 246,
     EncryptedStorageNonce => 247,
-    Balance => 248,
+    # Message number 248 is available
     NodeList => 249,
     UnspentOutputList => 250,
     TransactionList => 251,

lib/archethic/reward.ex

@@ -138,6 +138,8 @@ defmodule Archethic.Reward do
 
     network_pool_balance =
       SharedSecrets.get_network_pool_address()
+      |> UTXO.stream_unspent_outputs()
+      |> Stream.map(& &1.unspent_output)
       |> UTXO.get_balance()
       |> Map.get(:token)
       |> Map.to_list()

lib/archethic/utils/regression/benchmarks/p2p_message.ex

@@ -10,11 +10,11 @@ defmodule Archethic.Utils.Regression.Benchmark.P2PMessage do
 
   alias Archethic.Crypto
 
-  alias Archethic.P2P.Message.Balance
-  alias Archethic.P2P.Message.GetBalance
   alias Archethic.P2P.Message.GetTransaction
   alias Archethic.P2P.Message.GetTransactionChain
   alias Archethic.P2P.Message.TransactionList
+  alias Archethic.P2P.Message.GetUnspentOutputs
+  alias Archethic.P2P.Message.UnspentOutputList
 
   alias Archethic.TransactionChain.Transaction
 
@@ -53,9 +53,9 @@ defmodule Archethic.Utils.Regression.Benchmark.P2PMessage do
              address: get_genesis_address()
            })
        end,
-       "GetBalance" => fn _ ->
-         %Balance{} =
-           __MODULE__.Connection.send_message(conn_pid, %GetBalance{
+       "GetUnspentOutputs" => fn _ ->
+         %UnspentOutputList{} =
+           __MODULE__.Connection.send_message(conn_pid, %GetUnspentOutputs{
              address: get_genesis_address()
            })
        end

lib/archethic/utxo.ex

@@ -272,19 +272,30 @@ defmodule Archethic.UTXO do
 
   @doc """
   Returns the balance for an address using the unspent outputs
+
+  ## Examples
+
+      iex> [
+      ...>   %UnspentOutput{ from: "@Alice10", type: :UCO, amount: 100_000_000},
+      ...>   %UnspentOutput{ from: "@Bob5", type: {:token, "MyToken", 0}, amount: 300_000_000},
+      ...>   %UnspentOutput{ from: "@Charlie5", type: :call},
+      ...>   %UnspentOutput{ from: "@Tom5", type: :state},
+      ...> ]
+      ...> |> UTXO.get_balance()
+      %{
+         uco: 100_000_000,
+         token: %{
+           {"MyToken", 0} => 300_000_000
+         }
+      }
   """
-  @spec get_balance(binary()) :: balance()
-  def get_balance(address) do
-    address
-    |> stream_unspent_outputs()
-    |> Enum.reduce(%{uco: 0, token: %{}}, fn
-      %VersionedUnspentOutput{unspent_output: %UnspentOutput{type: :UCO, amount: amount}}, acc ->
+  @spec get_balance(Enumerable.t() | list(UnspentOutput.t())) :: balance()
+  def get_balance(unspent_outputs) do
+    Enum.reduce(unspent_outputs, %{uco: 0, token: %{}}, fn
+      %UnspentOutput{type: :UCO, amount: amount}, acc ->
         Map.update!(acc, :uco, &(&1 + amount))
 
-      %VersionedUnspentOutput{
-        unspent_output: %UnspentOutput{type: {:token, token_address, token_id}, amount: amount}
-      },
-      acc ->
+      %UnspentOutput{type: {:token, token_address, token_id}, amount: amount}, acc ->
         update_in(acc, [:token, Access.key({token_address, token_id}, 0)], &(&1 + amount))
 
       _, acc ->

lib/archethic_web/api/graphql/schema/resolver.ex

@@ -37,8 +37,9 @@ defmodule ArchethicWeb.API.GraphQL.Schema.Resolver do
   end
 
   def get_balance(address) do
-    with {:ok, genesis_address} <- Archethic.fetch_genesis_address(address),
-         {:ok, %{uco: uco, token: token_balances}} <- Archethic.get_balance(genesis_address) do
+    with {:ok, genesis_address} <- Archethic.fetch_genesis_address(address) do
+      %{uco: uco, token: token_balances} = Archethic.get_balance(genesis_address)
+
       balance = %{
         uco: uco,
         token:

lib/archethic_web/explorer/live/chains/transaction_chain_live.ex

@@ -9,6 +9,7 @@ defmodule ArchethicWeb.Explorer.TransactionChainLive do
 
   alias Archethic.Crypto
   alias Archethic.OracleChain
+  alias Archethic.UTXO
   alias Archethic.TransactionChain.Transaction.ValidationStamp.LedgerOperations.UnspentOutput
   alias ArchethicWeb.Explorer.Components.TransactionsList
   alias ArchethicWeb.Explorer.Components.UnspentOutputList
@@ -136,12 +137,11 @@ defmodule ArchethicWeb.Explorer.TransactionChainLive do
   end
 
   defp get_balance(utxos) do
-    Enum.reduce(utxos, %{}, fn
-      %UnspentOutput{type: type, amount: amount}, acc when amount != nil ->
-        Map.update(acc, type, amount, &(&1 + amount))
+    %{uco: uco, token: tokens} = UTXO.get_balance(utxos)
 
-      _, acc ->
-        acc
+    Enum.reduce(tokens, %{:UCO => uco}, fn
+      {{token_address, token_id}, amount}, acc ->
+        Map.put(acc, {:token, token_address, token_id}, amount)
     end)
   end
 end

lib/archethic_web/explorer/live/settings_live.ex

@@ -160,7 +160,7 @@ defmodule ArchethicWeb.Explorer.SettingsLive do
 
   defp get_token_transfers(previous_reward_address, next_reward_address) do
     {:ok, genesis_address} = Archethic.fetch_genesis_address(previous_reward_address)
-    {:ok, %{token: tokens}} = Archethic.get_balance(genesis_address)
+    %{token: tokens} = Archethic.get_balance(genesis_address)
 
     tokens
     |> Enum.filter(fn {{address, _}, _} -> Reward.is_reward_token?(address) end)

test/archethic/bootstrap/network_init_test.exs

@@ -358,7 +358,13 @@ defmodule Archethic.Bootstrap.NetworkInitTest do
     genesis_pools = Application.get_env(:archethic, NetworkInit)[:genesis_pools]
 
     assert Enum.all?(genesis_pools, fn %{address: address, amount: amount} ->
-             match?(%{uco: ^amount}, UTXO.get_balance(address))
+             balance =
+               address
+               |> UTXO.stream_unspent_outputs()
+               |> Enum.map(& &1.unspent_output)
+               |> UTXO.get_balance()
+
+             match?(%{uco: ^amount}, balance)
            end)
   end
 
@@ -408,7 +414,11 @@ defmodule Archethic.Bootstrap.NetworkInitTest do
 
     network_pool_genesis = Crypto.network_pool_public_key(0) |> Crypto.derive_address()
 
-    assert %{token: %{^key => 3_444_185_300_000_000}} = UTXO.get_balance(network_pool_genesis)
+    assert %{token: %{^key => 3_444_185_300_000_000}} =
+             network_pool_genesis
+             |> UTXO.stream_unspent_outputs()
+             |> Enum.map(& &1.unspent_output)
+             |> UTXO.get_balance()
   end
 
   test "init_software_origin_shared_secrets_chain/1 should create first origin shared secret transaction" do