SC new Chain functions

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

@@ -2,37 +2,82 @@ defmodule Archethic.Contracts.Interpreter.Library.Common.Chain do
   @moduledoc false
   @behaviour Archethic.Contracts.Interpreter.Library
 
-  alias Archethic.Crypto
-
   alias Archethic.Contracts.Interpreter.ASTHelper, as: AST
   alias Archethic.Contracts.Interpreter.Library.Common.ChainImpl
 
   use Knigge, otp_app: :archethic, default: ChainImpl, delegate_at_runtime?: true
 
-  @callback get_genesis_address(binary()) :: Crypto.prepended_hash()
-  @callback get_first_transaction_address(binary()) :: Crypto.prepended_hash() | nil
-  @callback get_genesis_public_key(binary()) :: Crypto.key() | nil
+  @callback get_genesis_address(binary()) :: binary()
+  @callback get_first_transaction_address(binary()) :: binary() | nil
+  @callback get_genesis_public_key(binary()) :: binary() | nil
   @callback get_transaction(binary()) :: map()
-  @callback get_burn_address() :: Crypto.prepended_hash()
+  @callback get_burn_address() :: binary()
+  @callback get_previous_address(binary() | map()) :: binary()
+  @callback get_balance(binary()) :: map()
+  @callback get_uco_balance(binary()) :: float()
+  @callback get_token_balance(binary(), binary()) :: float()
+  @callback get_token_balance(binary(), binary(), non_neg_integer()) :: float()
+  @callback get_tokens_balance(binary()) :: list()
+  @callback get_tokens_balance(binary(), list()) :: list()
 
   @spec check_types(atom(), list()) :: boolean()
   def check_types(:get_genesis_address, [first]) do
-    AST.is_binary?(first) || AST.is_variable_or_function_call?(first)
+    binary_or_variable_or_function?(first)
   end
 
   def check_types(:get_first_transaction_address, [first]) do
-    AST.is_binary?(first) || AST.is_variable_or_function_call?(first)
+    binary_or_variable_or_function?(first)
   end
 
   def check_types(:get_genesis_public_key, [first]) do
-    AST.is_binary?(first) || AST.is_variable_or_function_call?(first)
+    binary_or_variable_or_function?(first)
   end
 
   def check_types(:get_transaction, [first]) do
-    AST.is_binary?(first) || AST.is_variable_or_function_call?(first)
+    binary_or_variable_or_function?(first)
   end
 
   def check_types(:get_burn_address, []), do: true
 
+  def check_types(:get_previous_address, [first]) do
+    AST.is_map?(first) || binary_or_variable_or_function?(first)
+  end
+
+  def check_types(:get_balance, [first]) do
+    binary_or_variable_or_function?(first)
+  end
+
+  def check_types(:get_uco_balance, [first]) do
+    binary_or_variable_or_function?(first)
+  end
+
+  def check_types(:get_token_balance, [first, second]) do
+    binary_or_variable_or_function?(first) && binary_or_variable_or_function?(second)
+  end
+
+  def check_types(:get_token_balance, [first, second, third]) do
+    check_types(:get_token_balance, [first, second]) && number_or_variable_or_function?(third)
+  end
+
+  def check_types(:get_tokens_balance, [first]) do
+    binary_or_variable_or_function?(first)
+  end
+
+  def check_types(:get_tokens_balance, [first, second]) do
+    check_types(:get_tokens_balance, [first]) && list_or_variable_or_function?(second)
+  end
+
   def check_types(_, _), do: false
+
+  defp binary_or_variable_or_function?(arg) do
+    AST.is_binary?(arg) || AST.is_variable_or_function_call?(arg)
+  end
+
+  defp number_or_variable_or_function?(arg) do
+    AST.is_number?(arg) || AST.is_variable_or_function_call?(arg)
+  end
+
+  defp list_or_variable_or_function?(arg) do
+    AST.is_list?(arg) || AST.is_variable_or_function_call?(arg)
+  end
 end

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

@@ -2,12 +2,19 @@ defmodule Archethic.Contracts.Interpreter.Library.Common.ChainImpl do
   @moduledoc false
 
   alias Archethic.Contracts.Interpreter.Legacy
+  alias Archethic.Contracts.Interpreter.Library
   alias Archethic.Contracts.Interpreter.Library.Common.Chain
   alias Archethic.Contracts.Interpreter.Legacy.UtilsInterpreter
   alias Archethic.Contracts.ContractConstants, as: Constants
+
+  alias Archethic.Crypto
+
   alias Archethic.Tag
+
   alias Archethic.TransactionChain.Transaction.ValidationStamp.LedgerOperations
 
+  alias Archethic.Utils
+
   @behaviour Chain
   use Tag
 
@@ -44,14 +51,127 @@ defmodule Archethic.Contracts.Interpreter.Library.Common.ChainImpl do
     |> UtilsInterpreter.get_address("Chain.get_transaction/1")
     |> Archethic.search_transaction()
     |> then(fn
-      {:ok, tx} ->
-        Constants.from_transaction(tx)
-
-      {:error, _} ->
-        nil
+      {:ok, tx} -> Constants.from_transaction(tx)
+      {:error, _} -> nil
     end)
   end
 
   @impl Chain
   def get_burn_address(), do: LedgerOperations.burning_address() |> Base.encode16()
+
+  @impl Chain
+  def get_previous_address(previous_public_key) when is_binary(previous_public_key),
+    do: previous_address(previous_public_key)
+
+  def get_previous_address(%{"previous_public_key" => previous_public_key}),
+    do: previous_address(previous_public_key)
+
+  def get_previous_address(arg),
+    do:
+      raise(Library.Error,
+        message:
+          "Invalid arg for Chain.get_previous_address(), expected string or map, got #{inspect(arg)}"
+      )
+
+  defp previous_address(previous_public_key) do
+    with {:ok, pub_key} <- Base.decode16(previous_public_key),
+         true <- Crypto.valid_public_key?(pub_key) do
+      pub_key |> Crypto.derive_address() |> Base.encode16()
+    else
+      _ ->
+        raise Library.Error,
+          message:
+            "Invalid previous public key in Chain.get_previous_address(), got #{inspect(previous_public_key)}"
+    end
+  end
+
+  @impl Chain
+  @tag [:io]
+  def get_balance(address_hex) do
+    function = "Chain.get_balance"
+
+    %{uco: uco_amount, token: tokens} =
+      address_hex |> get_binary_address(function) |> fetch_balance(function)
+
+    tokens =
+      Enum.reduce(tokens, %{}, fn {{token_address, token_id}, amount}, acc ->
+        key = %{"token_address" => Base.encode16(token_address), "token_id" => token_id}
+        Map.put(acc, key, Utils.from_bigint(amount))
+      end)
+
+    %{"uco" => Utils.from_bigint(uco_amount), "tokens" => tokens}
+  end
+
+  @impl Chain
+  @tag [:io]
+  def get_uco_balance(address_hex) do
+    function = "Chain.get_balance"
+    %{uco: uco_amount} = address_hex |> get_binary_address(function) |> fetch_balance(function)
+
+    Utils.from_bigint(uco_amount)
+  end
+
+  @impl Chain
+  @tag [:io]
+  def get_token_balance(address_hex, token_address_hex, token_id \\ 0)
+
+  def get_token_balance(address_hex, token_address_hex, token_id) do
+    function = "Chain.get_token_balance"
+    token_address = get_binary_address(token_address_hex, function)
+    %{token: tokens} = address_hex |> get_binary_address(function) |> fetch_balance(function)
+
+    tokens |> Map.get({token_address, token_id}, 0) |> Utils.from_bigint()
+  end
+
+  @impl Chain
+  @tag [:io]
+  def get_tokens_balance(address_hex) do
+    function = "Chain.get_tokens_balance"
+
+    %{token: tokens} = address_hex |> get_binary_address(function) |> fetch_balance(function)
+
+    Enum.reduce(tokens, %{}, fn {{token_address, token_id}, amount}, acc ->
+      key = %{"token_address" => Base.encode16(token_address), "token_id" => token_id}
+      Map.put(acc, key, Utils.from_bigint(amount))
+    end)
+  end
+
+  @impl Chain
+  @tag [:io]
+  def get_tokens_balance(address_hex, requested_tokens) do
+    function = "Chain.get_tokens_balance"
+
+    %{token: tokens} = address_hex |> get_binary_address(function) |> fetch_balance(function)
+
+    Enum.reduce(
+      requested_tokens,
+      %{},
+      fn token = %{"token_address" => token_address_hex, "token_id" => token_id}, acc ->
+        key = {get_binary_address(token_address_hex, function), token_id}
+        amount = Map.get(tokens, key, 0) |> Utils.from_bigint()
+        Map.put(acc, token, amount)
+      end
+    )
+  end
+
+  defp get_binary_address(address_hex, function) do
+    with {:ok, address} <- Base.decode16(address_hex),
+         true <- Crypto.valid_address?(address) do
+      address
+    else
+      _ ->
+        raise Library.Error,
+          message: "Invalid address in #{function}, got #{inspect(address_hex)}"
+    end
+  end
+
+  defp fetch_balance(address, function) do
+    with {:ok, last_address} <- Archethic.get_last_transaction_address(address),
+         {:ok, balance} <- Archethic.get_balance(last_address) do
+      balance
+    else
+      _ ->
+        raise Library.Error, message: "Network issue in #{function}"
+    end
+  end
 end