Skip to content

Commit

Permalink
feat: Add in-mem UniqueMultimap for testing (#84)
Browse files Browse the repository at this point in the history
* feat: Add in-mem UniqueMultimap for testing

* test: Add test for UnifyAccount signed extension

* chore: Fix clippy errors

* test: Add unit tests for in-mem UniqueMultimap
  • Loading branch information
conr2d authored Oct 25, 2024
1 parent 38f28dc commit 3765dbc
Show file tree
Hide file tree
Showing 8 changed files with 268 additions and 16 deletions.
2 changes: 1 addition & 1 deletion frame/babel/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ sp-io = { git = "https://github.com/paritytech/polkadot-sdk", branch = "stable24
sp-runtime = { git = "https://github.com/paritytech/polkadot-sdk", branch = "stable2409", default-features = false }

[dev-dependencies]
hex = "0.4.3"
const-hex = "1.13"
# substrate
pallet-assets = { git = "https://github.com/paritytech/polkadot-sdk", branch = "stable2409" }
pallet-balances = { git = "https://github.com/paritytech/polkadot-sdk", branch = "stable2409" }
Expand Down
6 changes: 3 additions & 3 deletions frame/babel/src/cosmos/address.rs
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ where
pub struct AccountToAddr<T>(PhantomData<T>);
impl<T> Convert<AccountIdOf<T>, String> for AccountToAddr<T>
where
T: pallet_cosmwasm::Config + unify_account::Config,
T: pallet_cosmwasm::Config + unify_account::Config<AccountId = AccountIdOf<T>>,
{
fn convert(account: AccountIdOf<T>) -> String {
let addresses = T::AddressMap::get(&account);
Expand All @@ -61,7 +61,7 @@ where

impl<T> Convert<String, Result<AccountIdOf<T>, ()>> for AccountToAddr<T>
where
T: pallet_cosmwasm::Config + unify_account::Config,
T: pallet_cosmwasm::Config + unify_account::Config<AccountId = AccountIdOf<T>>,
{
fn convert(address: String) -> Result<AccountIdOf<T>, ()> {
let (_hrp, address_raw) = acc_address_from_bech32(&address).map_err(|_| ())?;
Expand All @@ -71,7 +71,7 @@ where

impl<T> Convert<Vec<u8>, Result<AccountIdOf<T>, ()>> for AccountToAddr<T>
where
T: pallet_cosmwasm::Config + unify_account::Config,
T: pallet_cosmwasm::Config + unify_account::Config<AccountId = AccountIdOf<T>>,
{
fn convert(address: Vec<u8>) -> Result<AccountIdOf<T>, ()> {
match address.len() {
Expand Down
2 changes: 1 addition & 1 deletion frame/babel/src/cosmos/precompile.rs
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,6 @@ mod tests {
let ExecuteMsg::Dispatch { input } =
serde_json_wasm::from_slice(message.as_bytes()).unwrap();

assert_eq!(input, hex::decode("0a030090b5ab205c6974c9ea841be688864633dc9ca8a357843eeacf2314649965fe220f0000c16ff28623").unwrap());
assert_eq!(input, const_hex::decode("0a030090b5ab205c6974c9ea841be688864633dc9ca8a357843eeacf2314649965fe220f0000c16ff28623").unwrap());
}
}
61 changes: 52 additions & 9 deletions frame/babel/src/extensions/unify_account.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,10 @@
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

use crate::traits::AccountIdProvider;
use core::marker::PhantomData;
use frame_support::traits::tokens::{fungible, Fortitude, Preservation};
use np_babel::{CosmosAddress, EthereumAddress, VarAddress};
use np_babel::VarAddress;
use pallet_multimap::traits::UniqueMultimap;
use parity_scale_codec::{Decode, Encode, MaxEncodedLen};
use scale_info::TypeInfo;
Expand All @@ -28,12 +29,14 @@ use sp_runtime::{
transaction_validity::{TransactionValidityError, ValidTransaction},
};

type AccountIdOf<T> = <T as AccountIdProvider>::AccountId;

/// A configuration for UnifyAccount signed extension.
pub trait Config: frame_system::Config<AccountId: From<H256> + TryInto<ecdsa::Public>> {
pub trait Config: AccountIdProvider<AccountId: From<H256> + TryInto<ecdsa::Public>> {
/// A map from account to addresses.
type AddressMap: UniqueMultimap<Self::AccountId, VarAddress>;
type AddressMap: UniqueMultimap<AccountIdOf<Self>, VarAddress>;
/// Drain account balance when unifying accounts.
type DrainBalance: DrainBalance<Self::AccountId>;
type DrainBalance: DrainBalance<AccountIdOf<Self>>;
}

/// Unifies the accounts associated with the same public key.
Expand All @@ -56,19 +59,19 @@ impl<T> UnifyAccount<T> {
}

impl<T: Config> UnifyAccount<T> {
pub fn unify_ecdsa(who: &T::AccountId) -> Result<(), &'static str> {
pub fn unify_ecdsa(who: &AccountIdOf<T>) -> Result<(), &'static str> {
if let Ok(public) = who.clone().try_into() {
#[cfg(feature = "ethereum")]
{
let address = EthereumAddress::from(public);
let address = np_babel::EthereumAddress::from(public);
let interim = address.clone().into_account_truncating();
T::DrainBalance::drain_balance(&interim, who)?;
T::AddressMap::try_insert(who, VarAddress::Ethereum(address))
.map_err(|_| "account unification failed: ethereum")?;
}
#[cfg(feature = "cosmos")]
{
let address = CosmosAddress::from(public);
let address = np_babel::CosmosAddress::from(public);
let interim = address.clone().into_account_truncating();
T::DrainBalance::drain_balance(&interim, who)?;
T::AddressMap::try_insert(who, VarAddress::Cosmos(address))
Expand Down Expand Up @@ -105,8 +108,11 @@ impl<T> core::fmt::Debug for UnifyAccount<T> {
}
}

impl<T: Config> SignedExtension for UnifyAccount<T> {
type AccountId = T::AccountId;
impl<T> SignedExtension for UnifyAccount<T>
where
T: Config + frame_system::Config<AccountId = AccountIdOf<T>>,
{
type AccountId = AccountIdOf<T>;
type Call = T::RuntimeCall;
type AdditionalSigned = ();
type Pre = ();
Expand Down Expand Up @@ -167,3 +173,40 @@ where
.map(|_| amount)
}
}

#[cfg(test)]
mod tests {
use super::*;
use np_runtime::{AccountId32, MultiSigner};

type AccountId = AccountId32<MultiSigner>;

struct MockConfig;

impl Config for MockConfig {
type AddressMap = pallet_multimap::traits::in_mem::UniqueMultimap<AccountId, VarAddress>;
type DrainBalance = ();
}

impl AccountIdProvider for MockConfig {
type AccountId = AccountId;
}

fn dev_public() -> ecdsa::Public {
const_hex::decode_to_array(
b"02509540919faacf9ab52146c9aa40db68172d83777250b28e4679176e49ccdd9f",
)
.unwrap()
.into()
}

#[test]
fn unify_ecdsa_works() {
let who = AccountId::from(dev_public());
let _ = UnifyAccount::<MockConfig>::unify_ecdsa(&who);
let cosmos = VarAddress::Cosmos(dev_public().into());
assert_eq!(<MockConfig as Config>::AddressMap::find_key(cosmos), Some(who.clone()));
let ethereum = VarAddress::Ethereum(dev_public().into());
assert_eq!(<MockConfig as Config>::AddressMap::find_key(ethereum), Some(who));
}
}
6 changes: 4 additions & 2 deletions frame/babel/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,13 +21,15 @@

extern crate alloc;

#[cfg(test)]
mod mock;

#[cfg(feature = "cosmos")]
pub mod cosmos;
#[cfg(feature = "ethereum")]
pub mod ethereum;
pub mod extensions;
#[cfg(test)]
mod mock;
pub mod traits;

pub use extensions::unify_account::UnifyAccount;
pub use np_babel::VarAddress;
Expand Down
39 changes: 39 additions & 0 deletions frame/babel/src/traits.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
// This file is part of Noir.

// Copyright (c) Haderech Pte. Ltd.
// SPDX-License-Identifier: GPL-3.0-or-later

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

use core::fmt::Debug;
use frame_support::{
sp_runtime::traits::{MaybeDisplay, MaybeSerializeDeserialize, Member},
Parameter,
};
use parity_scale_codec::MaxEncodedLen;

/// Trait for providing the account id type.
pub trait AccountIdProvider {
type AccountId: Parameter
+ Member
+ MaybeSerializeDeserialize
+ Debug
+ MaybeDisplay
+ Ord
+ MaxEncodedLen;
}

impl<T: frame_system::Config> AccountIdProvider for T {
type AccountId = T::AccountId;
}
3 changes: 3 additions & 0 deletions frame/multimap/src/traits.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,9 @@

use crate::*;

#[cfg(feature = "std")]
pub mod in_mem;

use alloc::collections::BTreeSet;
use frame_support::ensure;
use parity_scale_codec::{Codec, EncodeLike, FullCodec};
Expand Down
165 changes: 165 additions & 0 deletions frame/multimap/src/traits/in_mem.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,165 @@
// This file is part of Noir.

// Copyright (c) Haderech Pte. Ltd.
// SPDX-License-Identifier: Apache-2.0

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use parity_scale_codec::{EncodeLike, FullCodec};
use std::{
cell::RefCell,
collections::{BTreeMap, BTreeSet},
marker::PhantomData,
};

thread_local! {
static UNIQUE_MULTIMAP: RefCell<BTreeMap<Vec<u8>, BTreeSet<Vec<u8>>>> = const { RefCell::new(BTreeMap::new()) };
static UNIQUE_MULTIMAP_INDEX: RefCell<BTreeMap<Vec<u8>, Vec<u8>>> = const { RefCell::new(BTreeMap::new()) };
}

pub struct UniqueMultimap<K, V>(PhantomData<(K, V)>);

impl<K: FullCodec, V: FullCodec + Ord> super::UniqueMultimap<K, V> for UniqueMultimap<K, V> {
type Error = &'static str;

fn try_insert<KeyArg: EncodeLike<K>, ValArg: EncodeLike<V>>(
key: KeyArg,
value: ValArg,
) -> Result<bool, Self::Error> {
let key = key.encode();
let value = value.encode();
match UNIQUE_MULTIMAP_INDEX.with(|index| {
let mut index = index.borrow_mut();
if let Some(existing_key) = index.get(&value) {
if existing_key != &key {
Err("Duplicate value")
} else {
Ok(false)
}
} else {
index.insert(value.clone(), key.clone());
Ok(true)
}
}) {
Ok(true) => {
UNIQUE_MULTIMAP.with(|map| {
let mut map = map.borrow_mut();
map.entry(key).or_default().insert(value);
});
Ok(true)
},
Ok(false) => Ok(false),
Err(e) => Err(e),
}
}

fn get<KeyArg: EncodeLike<K>>(key: KeyArg) -> BTreeSet<V> {
let key = key.encode();
UNIQUE_MULTIMAP.with(|map| {
map.borrow()
.get(&key)
.map(|values| {
values
.iter()
.map(|value| V::decode(&mut &value[..]).expect("Decoding failed"))
.collect()
})
.unwrap_or_default()
})
}

fn find_key<ValArg: EncodeLike<V>>(value: ValArg) -> Option<K> {
let value = value.encode();
UNIQUE_MULTIMAP_INDEX.with(|index| {
index
.borrow()
.get(&value)
.map(|key| K::decode(&mut &key[..]).expect("Decoding failed"))
})
}

fn remove<KeyArg: EncodeLike<K>, ValArg: EncodeLike<V>>(key: KeyArg, value: ValArg) -> bool {
let key = key.encode();
let value = value.encode();
if UNIQUE_MULTIMAP_INDEX.with(|index| {
let mut index = index.borrow_mut();
index.remove(&value).is_some()
}) {
UNIQUE_MULTIMAP.with(|map| {
let mut map = map.borrow_mut();
if let Some(values) = map.get_mut(&key) {
values.remove(&value);
}
});
true
} else {
false
}
}

fn remove_all<KeyArg: EncodeLike<K>>(key: KeyArg) -> bool {
let key = key.encode();
if let Some(values) = UNIQUE_MULTIMAP.with(|map| map.borrow_mut().remove(&key)) {
UNIQUE_MULTIMAP_INDEX.with(|index| {
let mut index = index.borrow_mut();
for value in values {
index.remove(&value);
}
});
true
} else {
false
}
}
}

#[cfg(test)]
mod tests {
use crate::traits::{in_mem::UniqueMultimap, UniqueMultimap as _};
use std::collections::BTreeSet;

type Multimap = UniqueMultimap<String, u32>;

#[test]
fn unique_multimap_insert() {
assert_eq!(Multimap::try_insert("alice".to_string(), 1), Ok(true));
assert_eq!(Multimap::try_insert("alice".to_string(), 1), Ok(false));
assert_eq!(Multimap::try_insert("bob".to_string(), 1), Err("Duplicate value"));
assert_eq!(Multimap::try_insert("alice".to_string(), 2), Ok(true));
assert_eq!(Multimap::try_insert("bob".to_string(), 3), Ok(true));
assert_eq!(Multimap::try_insert("bob".to_string(), 1), Err("Duplicate value"));
}

#[test]
fn unique_multimap_get() {
let _ = Multimap::try_insert("alice".to_string(), 1);
let _ = Multimap::try_insert("alice".to_string(), 2);
let _ = Multimap::try_insert("bob".to_string(), 3);
assert_eq!(Multimap::get("alice"), BTreeSet::from([1, 2]));
assert_eq!(Multimap::find_key(1), Some("alice".to_string()));
assert_eq!(Multimap::find_key(3), Some("bob".to_string()));
}

#[test]
fn unique_multimap_remove() {
let _ = Multimap::try_insert("alice".to_string(), 1);
let _ = Multimap::try_insert("alice".to_string(), 2);
let _ = Multimap::try_insert("alice".to_string(), 3);
assert!(Multimap::remove("alice", 2));
assert_eq!(Multimap::get("alice"), BTreeSet::from([1, 3]));
assert!(!Multimap::remove("alice", 2));
assert!(Multimap::remove_all("alice"));
assert!(Multimap::get("alice").is_empty());
assert!(!Multimap::remove_all("alice"));
}
}

0 comments on commit 3765dbc

Please sign in to comment.