refactor: Use EncodeLike args for Multimap traits (#77)

NoirHQ · Oct 14, 2024 · f261c23 · f261c23
1 parent b4b513a
commit f261c23
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 35 deletions.
diff --git a/frame/babel/src/cosmos/address.rs b/frame/babel/src/cosmos/address.rs
@@ -45,7 +45,7 @@ where
 	T: pallet_cosmwasm::Config + unify_account::Config,
 {
 	fn convert(account: AccountIdOf<T>) -> String {
-		let addresses = T::AddressMap::get(account.clone());
+		let addresses = T::AddressMap::get(&account);
 		let address: Option<&CosmosAddress> = addresses.iter().find_map(|address| match address {
 			VarAddress::Cosmos(address) => Some(address),
 			_ => None,

diff --git a/frame/babel/src/extensions/unify_account.rs b/frame/babel/src/extensions/unify_account.rs
@@ -63,15 +63,15 @@ impl<T: Config> UnifyAccount<T> {
 				let address = EthereumAddress::from(public);
 				let interim = address.clone().into_account_truncating();
 				T::DrainBalance::drain_balance(&interim, who)?;
-				T::AddressMap::try_insert(who.clone(), VarAddress::Ethereum(address))
+				T::AddressMap::try_insert(who, VarAddress::Ethereum(address))
 					.map_err(|_| "account unification failed: ethereum")?;
 			}
 			#[cfg(feature = "cosmos")]
 			{
 				let address = CosmosAddress::from(public);
 				let interim = address.clone().into_account_truncating();
 				T::DrainBalance::drain_balance(&interim, who)?;
-				T::AddressMap::try_insert(who.clone(), VarAddress::Cosmos(address))
+				T::AddressMap::try_insert(who, VarAddress::Cosmos(address))
 					.map_err(|_| "account unification failed: cosmos")?;
 			}
 		}

diff --git a/frame/babel/src/lib.rs b/frame/babel/src/lib.rs
@@ -126,7 +126,7 @@ pub mod pallet {
 			transaction: Vec<u8>,
 		) -> DispatchResultWithPostInfo {
 			let who = ensure_signed(origin)?;
-			let address = T::AddressMap::get(who.clone())
+			let address = T::AddressMap::get(&who)
 				.iter()
 				.find_map(|address| match address {
 					VarAddress::Ethereum(address) => Some(address.clone()),

diff --git a/frame/multimap/src/traits.rs b/frame/multimap/src/traits.rs
@@ -19,31 +19,44 @@ use crate::*;
 
 use alloc::collections::BTreeSet;
 use frame_support::ensure;
+use parity_scale_codec::{Codec, EncodeLike, FullCodec};
+
+fn transmute<T: EncodeLike<R>, R: Codec>(value: T) -> R {
+	value.using_encoded(|encoded| R::decode(&mut &encoded[..]).expect("Decoding failed"))
+}
 
 /// Unique multimap whose values are unique across all keys.
-pub trait UniqueMultimap<K, V> {
+pub trait UniqueMultimap<K: FullCodec, V: FullCodec> {
 	type Error;
 
 	/// Tries to insert a value into the multimap.
-	fn try_insert(key: K, value: V) -> Result<bool, Self::Error>;
+	fn try_insert<KeyArg: EncodeLike<K>, ValArg: EncodeLike<V>>(
+		key: KeyArg,
+		value: ValArg,
+	) -> Result<bool, Self::Error>;
 
 	/// Gets all values for a key.
-	fn get(key: K) -> BTreeSet<V>;
+	fn get<KeyArg: EncodeLike<K>>(key: KeyArg) -> BTreeSet<V>;
 
 	/// Finds the key for a value.
-	fn find_key(value: V) -> Option<K>;
+	fn find_key<ValArg: EncodeLike<V>>(value: ValArg) -> Option<K>;
 
 	/// Removes a value from the multimap.
-	fn remove(key: K, value: V) -> bool;
+	fn remove<KeyArg: EncodeLike<K>, ValArg: EncodeLike<V>>(key: KeyArg, value: ValArg) -> bool;
 
 	/// Removes all values for a key.
-	fn remove_all(key: K) -> bool;
+	fn remove_all<KeyArg: EncodeLike<K>>(key: KeyArg) -> bool;
 }
 
 impl<T: Config<I>, I: 'static> UniqueMultimap<T::Key, T::Value> for Pallet<T, I> {
 	type Error = Error<T, I>;
 
-	fn try_insert(key: T::Key, value: T::Value) -> Result<bool, Error<T, I>> {
+	fn try_insert<K: EncodeLike<T::Key>, V: EncodeLike<T::Value>>(
+		key: K,
+		value: V,
+	) -> Result<bool, Error<T, I>> {
+		let key = transmute(key);
+		let value = transmute(value);
 		Map::<T, I>::try_mutate(&key, |values| {
 			ensure!(
 				Index::<T, I>::get(&value).filter(|k| *k != key).is_none(),
@@ -53,57 +66,66 @@ impl<T: Config<I>, I: 'static> UniqueMultimap<T::Key, T::Value> for Pallet<T, I>
 				.try_insert(value.clone())
 				.inspect(|ok| {
 					if *ok {
-						Index::<T, I>::insert(&value, &key);
+						Index::<T, I>::insert(value, key.clone());
 					}
 				})
 				.map_err(|_| Error::<T, I>::CapacityOverflow)
 		})
 	}
 
-	fn get(key: T::Key) -> BTreeSet<T::Value> {
-		Map::<T, I>::get(&key).into()
+	fn get<K: EncodeLike<T::Key>>(key: K) -> BTreeSet<T::Value> {
+		Map::<T, I>::get(key).into()
 	}
 
-	fn find_key(value: T::Value) -> Option<T::Key> {
-		Index::<T, I>::get(&value)
+	fn find_key<V: EncodeLike<T::Value>>(value: V) -> Option<T::Key> {
+		Index::<T, I>::get(value)
 	}
 
-	fn remove(key: T::Key, value: T::Value) -> bool {
-		Map::<T, I>::try_mutate(&key, |values| -> Result<bool, ()> {
-			Ok(values.remove(&value).then(|| Index::<T, I>::remove(&value)).is_some())
+	fn remove<K: EncodeLike<T::Key>, V: EncodeLike<T::Value>>(key: K, value: V) -> bool {
+		let value = transmute(value);
+		Map::<T, I>::try_mutate(key, |values| -> Result<bool, ()> {
+			Ok(values.remove(&value).then(|| Index::<T, I>::remove(value)).is_some())
 		})
 		.unwrap_or(false)
 	}
 
-	fn remove_all(key: T::Key) -> bool {
-		Map::<T, I>::take(&key).into_iter().fold(false, |_, value| {
-			Index::<T, I>::remove(&value);
+	fn remove_all<K: EncodeLike<T::Key>>(key: K) -> bool {
+		Map::<T, I>::take(key).into_iter().fold(false, |_, value| {
+			Index::<T, I>::remove(value);
 			true
 		})
 	}
 }
 
 /// Unique map whose the value is unique across all keys.
-pub trait UniqueMap<K, V> {
+pub trait UniqueMap<K: FullCodec, V: FullCodec> {
 	type Error;
 
 	/// Tries to insert a value into the map.
-	fn try_insert(key: K, value: V) -> Result<bool, Self::Error>;
+	fn try_insert<KeyArg: EncodeLike<K>, ValArg: EncodeLike<V>>(
+		key: KeyArg,
+		value: ValArg,
+	) -> Result<bool, Self::Error>;
 
 	/// Gets the value for a key.
-	fn get(key: K) -> Option<V>;
+	fn get<KeyArg: EncodeLike<K>>(key: KeyArg) -> Option<V>;
 
 	/// Finds the key for a value.
-	fn find_key(value: V) -> Option<K>;
+	fn find_key<ValArg: EncodeLike<V>>(value: ValArg) -> Option<K>;
 
 	/// Removes a value from the map.
-	fn remove(key: K);
+	fn remove<KeyArg: EncodeLike<K>>(key: KeyArg);
 }
 
 impl<T: Config<I>, I: 'static> UniqueMap<T::Key, T::Value> for Pallet<T, I> {
 	type Error = Error<T, I>;
 
-	fn try_insert(key: T::Key, value: T::Value) -> Result<bool, Error<T, I>> {
+	fn try_insert<K: EncodeLike<T::Key>, V: EncodeLike<T::Value>>(
+		key: K,
+		value: V,
+	) -> Result<bool, Error<T, I>> {
+		let key = transmute(key);
+		let value = transmute(value);
 		Map::<T, I>::try_mutate(&key, |values| {
 			ensure!(
 				Index::<T, I>::get(&value).filter(|k| *k != key).is_none(),
@@ -113,21 +135,21 @@ impl<T: Config<I>, I: 'static> UniqueMap<T::Key, T::Value> for Pallet<T, I> {
 			*values = BTreeSet::from([value.clone()])
 				.try_into()
 				.map_err(|_| Error::<T, I>::CapacityOverflow)?;
-			Index::<T, I>::insert(&value, &key);
+			Index::<T, I>::insert(value, key.clone());
 
 			Ok(true)
 		})
 	}
 
-	fn get(key: T::Key) -> Option<T::Value> {
-		Map::<T, I>::get(&key).first().cloned()
+	fn get<K: EncodeLike<T::Key>>(key: K) -> Option<T::Value> {
+		Map::<T, I>::get(key).first().cloned()
 	}
 
-	fn find_key(value: T::Value) -> Option<T::Key> {
-		Index::<T, I>::get(&value)
+	fn find_key<V: EncodeLike<T::Value>>(value: V) -> Option<T::Key> {
+		Index::<T, I>::get(value)
 	}
 
-	fn remove(key: T::Key) {
-		Map::<T, I>::remove(&key);
+	fn remove<K: EncodeLike<T::Key>>(key: K) {
+		Map::<T, I>::remove(key);
 	}
 }