Support arbitrary addresses in balances.toml

crates/apps_lib/src/config/genesis.rs

@@ -30,6 +30,110 @@ use namada_sdk::token::Denomination;
 use namada_sdk::{storage, token};
 use serde::{Deserialize, Serialize};
 
+#[derive(
+    Clone,
+    Debug,
+    BorshSerialize,
+    BorshDeserialize,
+    BorshDeserializer,
+    PartialEq,
+    Eq,
+    Ord,
+    PartialOrd,
+    Hash,
+)]
+#[allow(missing_docs)]
+pub enum GenesisBalanceAddress {
+    PublicKey(StringEncoded<common::PublicKey>),
+    Address(Address),
+}
+
+impl GenesisBalanceAddress {
+    /// Return an [`Address`] from this [`GenesisBalanceAddress`].
+    #[inline]
+    pub fn address(&self) -> Address {
+        match self {
+            Self::Address(addr) => addr.clone(),
+            Self::PublicKey(pk) => (&pk.raw).into(),
+        }
+    }
+}
+
+impl Serialize for GenesisBalanceAddress {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        match self {
+            GenesisBalanceAddress::Address(address) => {
+                Serialize::serialize(&address, serializer)
+            }
+            GenesisBalanceAddress::PublicKey(pk) => {
+                Serialize::serialize(pk, serializer)
+            }
+        }
+    }
+}
+
+impl<'de> Deserialize<'de> for GenesisBalanceAddress {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        struct FieldVisitor;
+
+        impl<'de> serde::de::Visitor<'de> for FieldVisitor {
+            type Value = GenesisBalanceAddress;
+
+            fn expecting(
+                &self,
+                formatter: &mut Formatter<'_>,
+            ) -> std::fmt::Result {
+                formatter
+                    .write_str("a bech32m encoded public key or an address")
+            }
+
+            fn visit_str<E>(self, value: &str) -> Result<Self::Value, E>
+            where
+                E: serde::de::Error,
+            {
+                GenesisBalanceAddress::from_str(value)
+                    .map_err(serde::de::Error::custom)
+            }
+        }
+
+        deserializer.deserialize_str(FieldVisitor)
+    }
+}
+
+impl Display for GenesisBalanceAddress {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            GenesisBalanceAddress::Address(address) => write!(f, "{address}"),
+            GenesisBalanceAddress::PublicKey(pk) => write!(f, "{}", pk),
+        }
+    }
+}
+
+impl FromStr for GenesisBalanceAddress {
+    type Err = String;
+
+    fn from_str(value: &str) -> Result<Self, Self::Err> {
+        // Try to deserialize a PK first
+        let maybe_pk = StringEncoded::<common::PublicKey>::from_str(value);
+        match maybe_pk {
+            Ok(pk) => Ok(GenesisBalanceAddress::PublicKey(pk)),
+            Err(_) => {
+                // If that doesn't work, attempt to retrieve
+                // an address
+                let address =
+                    Address::from_str(value).map_err(|err| err.to_string())?;
+                Ok(GenesisBalanceAddress::Address(address))
+            }
+        }
+    }
+}
+
 #[derive(
     Clone,
     Debug,
@@ -49,6 +153,18 @@ pub enum GenesisAddress {
     EstablishedAddress(EstablishedAddress),
 }
 
+impl From<GenesisAddress> for GenesisBalanceAddress {
+    #[inline]
+    fn from(genesis_addr: GenesisAddress) -> Self {
+        match genesis_addr {
+            GenesisAddress::PublicKey(pk) => Self::PublicKey(pk),
+            GenesisAddress::EstablishedAddress(addr) => {
+                Self::Address(Address::Established(addr))
+            }
+        }
+    }
+}
+
 impl GenesisAddress {
     /// Return an [`Address`] from this [`GenesisAddress`].
     #[inline]
@@ -427,7 +543,8 @@ pub fn make_dev_genesis(
             .first()
             .unwrap();
         let genesis_addr =
-            GenesisAddress::EstablishedAddress(tx.tx.data.address.raw.clone());
+            GenesisAddress::EstablishedAddress(tx.tx.data.address.raw.clone())
+                .into();
 
         let balance = *nam_balances.0.get(&genesis_addr).unwrap();
         let bonded = {
@@ -544,10 +661,12 @@ pub fn make_dev_genesis(
                 .unwrap();
 
             let validator_addr =
-                GenesisAddress::EstablishedAddress(validator_address.clone());
+                GenesisAddress::EstablishedAddress(validator_address.clone())
+                    .into();
             let account_pk = GenesisAddress::PublicKey(StringEncoded::new(
                 consensus_keypair.ref_to(),
-            ));
+            ))
+            .into();
 
             nam_balances.0.insert(validator_addr, first_val_balance);
             nam_balances.0.insert(account_pk, first_val_balance);

crates/apps_lib/src/config/genesis/templates.rs

@@ -24,7 +24,7 @@ use super::transactions::{self, Transactions};
 use super::utils::{read_toml, write_toml};
 use crate::config::genesis::chain::DeriveEstablishedAddress;
 use crate::config::genesis::transactions::{BondTx, SignedBondTx};
-use crate::config::genesis::GenesisAddress;
+use crate::config::genesis::GenesisBalanceAddress;
 use crate::wallet::Alias;
 
 pub const BALANCES_FILE_NAME: &str = "balances.toml";
@@ -141,7 +141,7 @@ pub struct DenominatedBalances {
     Eq,
 )]
 pub struct RawTokenBalances(
-    pub BTreeMap<GenesisAddress, token::DenominatedAmount>,
+    pub BTreeMap<GenesisBalanceAddress, token::DenominatedAmount>,
 );
 
 /// Genesis balances for a given token
@@ -157,7 +157,7 @@ pub struct RawTokenBalances(
     Eq,
 )]
 pub struct TokenBalances(
-    pub BTreeMap<GenesisAddress, token::DenominatedAmount>,
+    pub BTreeMap<GenesisBalanceAddress, token::DenominatedAmount>,
 );
 
 /// Genesis validity predicates
@@ -522,7 +522,7 @@ pub struct IbcParams {
 }
 
 impl TokenBalances {
-    pub fn get(&self, addr: &GenesisAddress) -> Option<token::Amount> {
+    pub fn get(&self, addr: &GenesisBalanceAddress) -> Option<token::Amount> {
         self.0.get(addr).map(|amt| amt.amount())
     }
 }
@@ -1048,7 +1048,7 @@ mod tests {
         let sk = key::testing::keypair_1();
         let pk = sk.ref_to();
         let address =
-            GenesisAddress::PublicKey(StringEncoded { raw: pk.clone() });
+            GenesisBalanceAddress::PublicKey(StringEncoded { raw: pk.clone() });
         let balance = token::Amount::from(101_000_001);
         let token_alias = Alias::from("Some_token".to_string());
         let contents = format!(

crates/apps_lib/src/config/genesis/transactions.rs

@@ -42,7 +42,7 @@ use crate::config::genesis::chain::DeriveEstablishedAddress;
 use crate::config::genesis::templates::{
     TemplateValidation, Unvalidated, Validated,
 };
-use crate::config::genesis::{utils, GenesisAddress};
+use crate::config::genesis::{utils, GenesisAddress, GenesisBalanceAddress};
 use crate::wallet::{CliWalletUtils, WalletTransport};
 
 /// Dummy chain id used to sign [`Tx`] objects at pre-genesis.
@@ -1197,9 +1197,10 @@ fn validate_bond(
 
     // Check and update token balance of the source
     let native_token = &parameters.parameters.native_token;
+    let source = GenesisBalanceAddress::from(source.clone());
     match balances.get_mut(native_token) {
         Some(balances) => {
-            let balance = balances.amounts.get_mut(source);
+            let balance = balances.amounts.get_mut(&source);
             match balance {
                 Some(balance) => {
                     if *balance < *amount {
@@ -1213,7 +1214,7 @@ fn validate_bond(
                     } else {
                         // Deduct the amount from source
                         if amount == balance {
-                            balances.amounts.remove(source);
+                            balances.amounts.remove(&source);
                         } else if let Some(new_balance) =
                             balance.checked_sub(*amount)
                         {
@@ -1254,7 +1255,7 @@ fn validate_bond(
 #[derive(Clone, Debug)]
 pub struct TokenBalancesForValidation {
     /// Accumulator for tokens transferred to accounts
-    pub amounts: BTreeMap<GenesisAddress, DenominatedAmount>,
+    pub amounts: BTreeMap<GenesisBalanceAddress, DenominatedAmount>,
 }
 
 pub fn validate_established_account(

crates/node/src/shell/init_chain.rs

@@ -512,7 +512,7 @@ where
             };
 
             for (owner, balance) in balances {
-                if let genesis::GenesisAddress::PublicKey(pk) = owner {
+                if let genesis::GenesisBalanceAddress::PublicKey(pk) = owner {
                     namada_sdk::account::init_account_storage(
                         &mut self.state,
                         &owner.address(),

crates/tests/src/e2e/setup.rs

@@ -251,15 +251,16 @@ where
             .get_mut(&Alias::from_str("nam").expect("Infallible"))
             .expect("NAM balances should exist in pre-genesis wallet already");
         nam_balances.0.insert(
-            GenesisAddress::PublicKey(StringEncoded::new(sk.ref_to())),
+            GenesisAddress::PublicKey(StringEncoded::new(sk.ref_to())).into(),
             token::DenominatedAmount::new(
                 token::Amount::from_uint(1000000, NATIVE_MAX_DECIMAL_PLACES)
                     .unwrap(),
                 NATIVE_MAX_DECIMAL_PLACES.into(),
             ),
         );
         nam_balances.0.insert(
-            GenesisAddress::EstablishedAddress(validator_address.clone()),
+            GenesisAddress::EstablishedAddress(validator_address.clone())
+                .into(),
             token::DenominatedAmount::new(
                 token::Amount::from_uint(2000000, NATIVE_MAX_DECIMAL_PLACES)
                     .unwrap(),