feat: use enum in uma trait

chain-extensions/types/unified-accounts/src/lib.rs

@@ -39,8 +39,10 @@ pub enum Command {
 pub enum UnifiedAddress<T: Encode + Decode> {
     /// The address fetched from the mappings and the account
     /// is unified
+    #[codec(index = 0)]
     Mapped(T),
     /// The default address associated with account as there
     /// is no mapping found and accounts are not unified
+    #[codec(index = 1)]
     Default(T),
 }

chain-extensions/unified-accounts/src/lib.rs

@@ -18,24 +18,17 @@
 
 #![cfg_attr(not(feature = "std"), no_std)]
 
-use astar_primitives::{
-    ethereum_checked::AccountMapping,
-    evm::{EvmAddress, UnifiedAddressMapper},
-};
+use astar_primitives::evm::{EvmAddress, UnifiedAddressMapper};
 use core::marker::PhantomData;
 use sp_runtime::DispatchError;
 
 use frame_support::DefaultNoBound;
 use pallet_contracts::chain_extension::{
     ChainExtension, Environment, Ext, InitState, Result as DispatchResult, RetVal,
 };
-use pallet_evm::AddressMapping;
 use pallet_unified_accounts::WeightInfo;
 use parity_scale_codec::Encode;
-pub use unified_accounts_chain_extension_types::{
-    Command::{self, *},
-    UnifiedAddress,
-};
+pub use unified_accounts_chain_extension_types::Command::{self, *};
 
 type UAWeight<T> = <T as pallet_unified_accounts::Config>::WeightInfo;
 
@@ -67,13 +60,8 @@ where
                 // charge weight
                 env.charge_weight(UAWeight::<T>::to_h160_or_default())?;
 
-                let evm_address = if let Some(h160) = UA::to_h160(&account_id) {
-                    UnifiedAddress::Mapped(h160)
-                } else {
-                    UnifiedAddress::Default(T::DefaultNativeToEvm::into_h160(account_id))
-                };
                 // write to buffer
-                evm_address.using_encoded(|r| env.write(r, false, None))?;
+                UA::to_h160_or_default(&account_id).using_encoded(|r| env.write(r, false, None))?;
             }
             GetNativeAddress => {
                 let evm_address: EvmAddress = env.read_as()?;
@@ -87,15 +75,9 @@ where
                 // charge weight
                 env.charge_weight(UAWeight::<T>::to_account_id_or_default())?;
 
-                // read the storage item
-                let native_address = if let Some(native) = UA::to_account_id(&evm_address) {
-                    UnifiedAddress::Mapped(native)
-                } else {
-                    UnifiedAddress::Default(T::DefaultEvmToNative::into_account_id(evm_address))
-                };
-
                 // write to buffer
-                native_address.using_encoded(|r| env.write(r, false, None))?;
+                UA::to_account_id_or_default(&evm_address)
+                    .using_encoded(|r| env.write(r, false, None))?;
             }
         };
         Ok(RetVal::Converging(0))

pallets/unified-accounts/src/lib.rs

@@ -67,7 +67,7 @@
 
 use astar_primitives::{
     ethereum_checked::AccountMapping,
-    evm::{EvmAddress, UnifiedAddressMapper},
+    evm::{EvmAddress, UnifiedAddress, UnifiedAddressMapper},
     Balance,
 };
 use frame_support::{
@@ -351,11 +351,12 @@ impl<T: Config> UnifiedAddressMapper<T::AccountId> for Pallet<T> {
         EvmToNative::<T>::get(evm_address)
     }
 
-    fn to_account_id_or_default(evm_address: &EvmAddress) -> T::AccountId {
-        EvmToNative::<T>::get(evm_address).unwrap_or_else(|| {
-            // fallback to default account_id
-            T::DefaultEvmToNative::into_account_id(evm_address.clone())
-        })
+    fn to_account_id_or_default(evm_address: &EvmAddress) -> UnifiedAddress<T::AccountId> {
+        if let Some(native) = Self::to_account_id(&evm_address) {
+            UnifiedAddress::Mapped(native)
+        } else {
+            UnifiedAddress::Default(T::DefaultEvmToNative::into_account_id(evm_address.clone()))
+        }
     }
 
     fn to_default_account_id(evm_address: &EvmAddress) -> T::AccountId {
@@ -366,11 +367,12 @@ impl<T: Config> UnifiedAddressMapper<T::AccountId> for Pallet<T> {
         NativeToEvm::<T>::get(account_id)
     }
 
-    fn to_h160_or_default(account_id: &T::AccountId) -> EvmAddress {
-        NativeToEvm::<T>::get(account_id).unwrap_or_else(|| {
-            // fallback to default account_id
-            T::DefaultNativeToEvm::into_h160(account_id.clone())
-        })
+    fn to_h160_or_default(account_id: &T::AccountId) -> UnifiedAddress<EvmAddress> {
+        if let Some(h160) = Self::to_h160(&account_id) {
+            UnifiedAddress::Mapped(h160)
+        } else {
+            UnifiedAddress::Default(T::DefaultNativeToEvm::into_h160(account_id.clone()))
+        }
     }
 
     fn to_default_h160(account_id: &T::AccountId) -> EvmAddress {
@@ -381,14 +383,15 @@ impl<T: Config> UnifiedAddressMapper<T::AccountId> for Pallet<T> {
 /// AccountMapping wrapper implementation
 impl<T: Config> AccountMapping<T::AccountId> for Pallet<T> {
     fn into_h160(account: T::AccountId) -> H160 {
-        <Self as UnifiedAddressMapper<T::AccountId>>::to_h160_or_default(&account)
+        <Self as UnifiedAddressMapper<T::AccountId>>::to_h160_or_default(&account).into_address()
     }
 }
 
 /// AddressMapping wrapper implementation
 impl<T: Config> AddressMapping<T::AccountId> for Pallet<T> {
     fn into_account_id(evm_address: H160) -> T::AccountId {
         <Self as UnifiedAddressMapper<T::AccountId>>::to_account_id_or_default(&evm_address)
+            .into_address()
     }
 }
 
@@ -415,7 +418,8 @@ impl<T: Config> StaticLookup for Pallet<T> {
             MultiAddress::Address20(i) => Ok(
                 <Self as UnifiedAddressMapper<T::AccountId>>::to_account_id_or_default(
                     &EvmAddress::from_slice(&i),
-                ),
+                )
+                .into_address(),
             ),
             _ => Err(LookupError),
         }

primitives/src/evm.rs

@@ -19,6 +19,8 @@
 use crate::{AccountId, AssetId};
 
 use frame_support::ensure;
+use parity_scale_codec::{Decode, Encode, MaxEncodedLen};
+use scale_info::TypeInfo;
 use sp_core::H160;
 use sp_std::marker::PhantomData;
 
@@ -64,15 +66,37 @@ pub trait UnifiedAddressMapper<AccountId> {
     fn to_account_id(evm_address: &EvmAddress) -> Option<AccountId>;
     /// Gets the account id associated with given evm address.
     /// If no mapping exists, then return the default evm address.
-    fn to_account_id_or_default(evm_address: &EvmAddress) -> AccountId;
+    fn to_account_id_or_default(evm_address: &EvmAddress) -> UnifiedAddress<AccountId>;
     /// Gets the default account id which is associated with given evm address.
     fn to_default_account_id(evm_address: &EvmAddress) -> AccountId;
 
     /// Gets the evm address associated with given account id, if mapped else None.
     fn to_h160(account_id: &AccountId) -> Option<EvmAddress>;
     /// Gets the evm address associated with given account id.
     /// If no mapping exists, then return the default account id.
-    fn to_h160_or_default(account_id: &AccountId) -> EvmAddress;
+    fn to_h160_or_default(account_id: &AccountId) -> UnifiedAddress<EvmAddress>;
     /// Gets the default evm address which is associated with given account id.
     fn to_default_h160(account_id: &AccountId) -> EvmAddress;
 }
+
+#[derive(Debug, Clone, PartialEq, Eq, Encode, Decode, MaxEncodedLen, TypeInfo)]
+pub enum UnifiedAddress<Address> {
+    /// The address fetched from the mappings and the account
+    /// is unified
+    #[codec(index = 0)]
+    Mapped(Address),
+    /// The default address associated with account as there
+    /// is no mapping found and accounts are not unified
+    #[codec(index = 1)]
+    Default(Address),
+}
+
+impl<Address> UnifiedAddress<Address> {
+    /// Get the underlying address
+    pub fn into_address(self) -> Address {
+        match self {
+            Self::Default(a) => a,
+            Self::Mapped(a) => a,
+        }
+    }
+}

tests/integration/src/unified_accounts.rs

@@ -17,9 +17,9 @@
 // along with Astar. If not, see <http://www.gnu.org/licenses/>.
 
 use crate::setup::*;
+use astar_primitives::evm::UnifiedAddress;
 use parity_scale_codec::Encode;
 use sp_io::hashing::keccak_256;
-use unified_accounts_chain_extension_types::UnifiedAddress;
 
 const AU_CE_GETTER: &'static str = "au_ce_getters";
 
@@ -70,7 +70,7 @@ fn unified_accounts_chain_extension_works() {
                 contract_id.clone(),
                 [GET_H160_OR_DEFAULT.to_vec(), ALICE.encode()].concat()
             ),
-            UnifiedAddress::Default(UnifiedAccounts::to_h160_or_default(&ALICE))
+            UnifiedAccounts::to_h160_or_default(&ALICE)
         );
         // mapped native address should be None
         assert_eq!(
@@ -88,7 +88,7 @@ fn unified_accounts_chain_extension_works() {
                 contract_id.clone(),
                 [GET_NATIVE_OR_DEFAULT.to_vec(), alith().encode()].concat()
             ),
-            UnifiedAddress::Default(UnifiedAccounts::to_account_id_or_default(&alith()))
+            UnifiedAccounts::to_account_id_or_default(&alith())
         );
 
         //