use sp_runtime::{traits::Saturating, RuntimeDebug};
use sp_std::{
collections::{btree_map::BTreeMap, btree_set::BTreeSet},
mem,
prelude::*,
};
use xcm::latest::{
AssetId, AssetInstance,
Fungibility::{Fungible, NonFungible},
MultiAsset, MultiAssetFilter, MultiAssets, MultiLocation,
WildFungibility::{Fungible as WildFungible, NonFungible as WildNonFungible},
WildMultiAsset::{All, AllOf},
};
#[derive(Default, Clone, RuntimeDebug, Eq, PartialEq)]
pub struct Assets {
pub fungible: BTreeMap<AssetId, u128>,
pub non_fungible: BTreeSet<(AssetId, AssetInstance)>,
}
impl From<MultiAsset> for Assets {
fn from(asset: MultiAsset) -> Assets {
let mut result = Self::default();
result.subsume(asset);
result
}
}
impl From<Vec<MultiAsset>> for Assets {
fn from(assets: Vec<MultiAsset>) -> Assets {
let mut result = Self::default();
for asset in assets.into_iter() {
result.subsume(asset)
}
result
}
}
impl From<MultiAssets> for Assets {
fn from(assets: MultiAssets) -> Assets {
assets.drain().into()
}
}
impl From<Assets> for Vec<MultiAsset> {
fn from(a: Assets) -> Self {
a.into_assets_iter().collect()
}
}
impl From<Assets> for MultiAssets {
fn from(a: Assets) -> Self {
a.into_assets_iter().collect::<Vec<MultiAsset>>().into()
}
}
#[derive(Debug)]
pub enum TakeError {
AssetUnderflow(MultiAsset),
}
impl Assets {
pub fn new() -> Self {
Self::default()
}
pub fn len(&self) -> usize {
self.fungible.len() + self.non_fungible.len()
}
pub fn is_empty(&self) -> bool {
self.fungible.is_empty() && self.non_fungible.is_empty()
}
pub fn fungible_assets_iter(&self) -> impl Iterator<Item = MultiAsset> + '_ {
self.fungible
.iter()
.map(|(id, &amount)| MultiAsset { fun: Fungible(amount), id: id.clone() })
}
pub fn non_fungible_assets_iter(&self) -> impl Iterator<Item = MultiAsset> + '_ {
self.non_fungible
.iter()
.map(|(id, instance)| MultiAsset { fun: NonFungible(instance.clone()), id: id.clone() })
}
pub fn into_assets_iter(self) -> impl Iterator<Item = MultiAsset> {
self.fungible
.into_iter()
.map(|(id, amount)| MultiAsset { fun: Fungible(amount), id })
.chain(
self.non_fungible
.into_iter()
.map(|(id, instance)| MultiAsset { fun: NonFungible(instance), id }),
)
}
pub fn assets_iter(&self) -> impl Iterator<Item = MultiAsset> + '_ {
self.fungible_assets_iter().chain(self.non_fungible_assets_iter())
}
pub fn subsume_assets(&mut self, mut assets: Assets) {
let mut f_iter = assets.fungible.iter_mut();
let mut g_iter = self.fungible.iter_mut();
if let (Some(mut f), Some(mut g)) = (f_iter.next(), g_iter.next()) {
loop {
if f.0 == g.0 {
(*f.1).saturating_accrue(*g.1);
}
if f.0 <= g.0 {
f = match f_iter.next() {
Some(x) => x,
None => break,
};
}
if f.0 >= g.0 {
g = match g_iter.next() {
Some(x) => x,
None => break,
};
}
}
}
self.fungible.append(&mut assets.fungible);
self.non_fungible.append(&mut assets.non_fungible);
}
pub fn subsume(&mut self, asset: MultiAsset) {
match asset.fun {
Fungible(amount) => {
self.fungible
.entry(asset.id)
.and_modify(|e| *e = e.saturating_add(amount))
.or_insert(amount);
},
NonFungible(instance) => {
self.non_fungible.insert((asset.id, instance));
},
}
}
pub fn swapped(&mut self, mut with: Assets) -> Self {
mem::swap(&mut *self, &mut with);
with
}
pub fn prepend_location(&mut self, prepend: &MultiLocation) {
let mut fungible = Default::default();
mem::swap(&mut self.fungible, &mut fungible);
self.fungible = fungible
.into_iter()
.map(|(mut id, amount)| {
let _ = id.prepend_with(prepend);
(id, amount)
})
.collect();
let mut non_fungible = Default::default();
mem::swap(&mut self.non_fungible, &mut non_fungible);
self.non_fungible = non_fungible
.into_iter()
.map(|(mut class, inst)| {
let _ = class.prepend_with(prepend);
(class, inst)
})
.collect();
}
pub fn reanchor(
&mut self,
target: &MultiLocation,
ancestry: &MultiLocation,
mut maybe_failed_bin: Option<&mut Self>,
) {
let mut fungible = Default::default();
mem::swap(&mut self.fungible, &mut fungible);
self.fungible = fungible
.into_iter()
.filter_map(|(mut id, amount)| match id.reanchor(target, ancestry) {
Ok(()) => Some((id, amount)),
Err(()) => {
maybe_failed_bin.as_mut().map(|f| f.fungible.insert(id, amount));
None
},
})
.collect();
let mut non_fungible = Default::default();
mem::swap(&mut self.non_fungible, &mut non_fungible);
self.non_fungible = non_fungible
.into_iter()
.filter_map(|(mut class, inst)| match class.reanchor(target, ancestry) {
Ok(()) => Some((class, inst)),
Err(()) => {
maybe_failed_bin.as_mut().map(|f| f.non_fungible.insert((class, inst)));
None
},
})
.collect();
}
pub fn ensure_contains(&self, assets: &MultiAssets) -> Result<(), TakeError> {
for asset in assets.inner().iter() {
match asset {
MultiAsset { fun: Fungible(ref amount), ref id } => {
if self.fungible.get(id).map_or(true, |a| a < amount) {
return Err(TakeError::AssetUnderflow((id.clone(), *amount).into()))
}
},
MultiAsset { fun: NonFungible(ref instance), ref id } => {
let id_instance = (id.clone(), instance.clone());
if !self.non_fungible.contains(&id_instance) {
return Err(TakeError::AssetUnderflow(id_instance.into()))
}
},
}
}
return Ok(())
}
fn general_take(
&mut self,
mask: MultiAssetFilter,
saturate: bool,
limit: usize,
) -> Result<Assets, TakeError> {
let mut taken = Assets::new();
match mask {
MultiAssetFilter::Wild(All) =>
if self.fungible.len() + self.non_fungible.len() <= limit {
return Ok(self.swapped(Assets::new()))
} else {
let fungible = mem::replace(&mut self.fungible, Default::default());
fungible.into_iter().for_each(|(c, amount)| {
if taken.len() < limit {
taken.fungible.insert(c, amount);
} else {
self.fungible.insert(c, amount);
}
});
let non_fungible = mem::replace(&mut self.non_fungible, Default::default());
non_fungible.into_iter().for_each(|(c, instance)| {
if taken.len() < limit {
taken.non_fungible.insert((c, instance));
} else {
self.non_fungible.insert((c, instance));
}
});
},
MultiAssetFilter::Wild(AllOf { fun: WildFungible, id }) => {
if let Some((id, amount)) = self.fungible.remove_entry(&id) {
taken.fungible.insert(id, amount);
}
},
MultiAssetFilter::Wild(AllOf { fun: WildNonFungible, id }) => {
let non_fungible = mem::replace(&mut self.non_fungible, Default::default());
non_fungible.into_iter().for_each(|(c, instance)| {
if c == id && taken.len() < limit {
taken.non_fungible.insert((c, instance));
} else {
self.non_fungible.insert((c, instance));
}
});
},
MultiAssetFilter::Definite(assets) => {
if !saturate {
self.ensure_contains(&assets)?;
}
for asset in assets.drain().into_iter() {
match asset {
MultiAsset { fun: Fungible(amount), id } => {
let (remove, amount) = match self.fungible.get_mut(&id) {
Some(self_amount) => {
let amount = amount.min(*self_amount);
*self_amount -= amount;
(*self_amount == 0, amount)
},
None => (false, 0),
};
if remove {
self.fungible.remove(&id);
}
if amount > 0 {
taken.subsume(MultiAsset::from((id, amount)).into());
}
},
MultiAsset { fun: NonFungible(instance), id } => {
let id_instance = (id, instance);
if self.non_fungible.remove(&id_instance) {
taken.subsume(id_instance.into())
}
},
}
if taken.len() == limit {
break
}
}
},
}
Ok(taken)
}
pub fn saturating_take(&mut self, asset: MultiAssetFilter) -> Assets {
self.general_take(asset, true, usize::max_value())
.expect("general_take never results in error when saturating")
}
pub fn limited_saturating_take(&mut self, asset: MultiAssetFilter, limit: usize) -> Assets {
self.general_take(asset, true, limit)
.expect("general_take never results in error when saturating")
}
pub fn try_take(&mut self, mask: MultiAssetFilter) -> Result<Assets, TakeError> {
self.general_take(mask, false, usize::max_value())
}
pub fn checked_sub(mut self, asset: MultiAsset) -> Result<Assets, Assets> {
match asset.fun {
Fungible(amount) => {
let remove = if let Some(balance) = self.fungible.get_mut(&asset.id) {
if *balance >= amount {
*balance -= amount;
*balance == 0
} else {
return Err(self)
}
} else {
return Err(self)
};
if remove {
self.fungible.remove(&asset.id);
}
Ok(self)
},
NonFungible(instance) =>
if self.non_fungible.remove(&(asset.id, instance)) {
Ok(self)
} else {
Err(self)
},
}
}
pub fn min(&self, mask: &MultiAssetFilter) -> Assets {
let mut masked = Assets::new();
match mask {
MultiAssetFilter::Wild(All) => return self.clone(),
MultiAssetFilter::Wild(AllOf { fun: WildFungible, id }) => {
if let Some(&amount) = self.fungible.get(&id) {
masked.fungible.insert(id.clone(), amount);
}
},
MultiAssetFilter::Wild(AllOf { fun: WildNonFungible, id }) => {
self.non_fungible.iter().for_each(|(ref c, ref instance)| {
if c == id {
masked.non_fungible.insert((c.clone(), instance.clone()));
}
});
},
MultiAssetFilter::Definite(assets) =>
for asset in assets.inner().iter() {
match asset {
MultiAsset { fun: Fungible(ref amount), ref id } => {
if let Some(m) = self.fungible.get(id) {
masked.subsume((id.clone(), Fungible(*amount.min(m))).into());
}
},
MultiAsset { fun: NonFungible(ref instance), ref id } => {
let id_instance = (id.clone(), instance.clone());
if self.non_fungible.contains(&id_instance) {
masked.subsume(id_instance.into());
}
},
}
},
}
masked
}
}
#[cfg(test)]
mod tests {
use super::*;
use xcm::latest::prelude::*;
#[allow(non_snake_case)]
fn AF(id: u8, amount: u128) -> MultiAsset {
(vec![id], amount).into()
}
#[allow(non_snake_case)]
fn ANF(class: u8, instance_id: u8) -> MultiAsset {
(vec![class], vec![instance_id]).into()
}
#[allow(non_snake_case)]
fn CF(amount: u128) -> MultiAsset {
(Here, amount).into()
}
#[allow(non_snake_case)]
fn CNF(instance_id: u8) -> MultiAsset {
(Here, [instance_id; 4]).into()
}
fn test_assets() -> Assets {
let mut assets = Assets::new();
assets.subsume(AF(1, 100));
assets.subsume(ANF(2, 20));
assets.subsume(CF(300));
assets.subsume(CNF(40));
assets
}
#[test]
fn subsume_assets_works() {
let t1 = test_assets();
let mut t2 = Assets::new();
t2.subsume(AF(1, 50));
t2.subsume(ANF(2, 10));
t2.subsume(CF(300));
t2.subsume(CNF(50));
let mut r1 = t1.clone();
r1.subsume_assets(t2.clone());
let mut r2 = t1.clone();
for a in t2.assets_iter() {
r2.subsume(a)
}
assert_eq!(r1, r2);
}
#[test]
fn checked_sub_works() {
let t = test_assets();
let t = t.checked_sub(AF(1, 50)).unwrap();
let t = t.checked_sub(AF(1, 51)).unwrap_err();
let t = t.checked_sub(AF(1, 50)).unwrap();
let t = t.checked_sub(AF(1, 1)).unwrap_err();
let t = t.checked_sub(CF(150)).unwrap();
let t = t.checked_sub(CF(151)).unwrap_err();
let t = t.checked_sub(CF(150)).unwrap();
let t = t.checked_sub(CF(1)).unwrap_err();
let t = t.checked_sub(ANF(2, 21)).unwrap_err();
let t = t.checked_sub(ANF(2, 20)).unwrap();
let t = t.checked_sub(ANF(2, 20)).unwrap_err();
let t = t.checked_sub(CNF(41)).unwrap_err();
let t = t.checked_sub(CNF(40)).unwrap();
let t = t.checked_sub(CNF(40)).unwrap_err();
assert_eq!(t, Assets::new());
}
#[test]
fn into_assets_iter_works() {
let assets = test_assets();
let mut iter = assets.into_assets_iter();
assert_eq!(Some(CF(300)), iter.next());
assert_eq!(Some(AF(1, 100)), iter.next());
assert_eq!(Some(CNF(40)), iter.next());
assert_eq!(Some(ANF(2, 20)), iter.next());
assert_eq!(None, iter.next());
}
#[test]
fn assets_into_works() {
let mut assets_vec: Vec<MultiAsset> = Vec::new();
assets_vec.push(AF(1, 100));
assets_vec.push(ANF(2, 20));
assets_vec.push(CF(300));
assets_vec.push(CNF(40));
assets_vec.push(AF(1, 100));
assets_vec.push(ANF(2, 20));
assets_vec.push(CF(300));
assets_vec.push(CNF(40));
let assets: Assets = assets_vec.into();
let mut iter = assets.into_assets_iter();
assert_eq!(Some(CF(600)), iter.next());
assert_eq!(Some(AF(1, 200)), iter.next());
assert_eq!(Some(CNF(40)), iter.next());
assert_eq!(Some(ANF(2, 20)), iter.next());
assert_eq!(None, iter.next());
}
#[test]
fn min_all_and_none_works() {
let assets = test_assets();
let none = MultiAssets::new().into();
let all = All.into();
let none_min = assets.min(&none);
assert_eq!(None, none_min.assets_iter().next());
let all_min = assets.min(&all);
assert!(all_min.assets_iter().eq(assets.assets_iter()));
}
#[test]
fn min_all_abstract_works() {
let assets = test_assets();
let fungible = Wild((vec![1], WildFungible).into());
let non_fungible = Wild((vec![2], WildNonFungible).into());
let fungible = assets.min(&fungible);
let fungible = fungible.assets_iter().collect::<Vec<_>>();
assert_eq!(fungible, vec![AF(1, 100)]);
let non_fungible = assets.min(&non_fungible);
let non_fungible = non_fungible.assets_iter().collect::<Vec<_>>();
assert_eq!(non_fungible, vec![ANF(2, 20)]);
}
#[test]
fn min_all_concrete_works() {
let assets = test_assets();
let fungible = Wild((Here, WildFungible).into());
let non_fungible = Wild((Here, WildNonFungible).into());
let fungible = assets.min(&fungible);
let fungible = fungible.assets_iter().collect::<Vec<_>>();
assert_eq!(fungible, vec![CF(300)]);
let non_fungible = assets.min(&non_fungible);
let non_fungible = non_fungible.assets_iter().collect::<Vec<_>>();
assert_eq!(non_fungible, vec![CNF(40)]);
}
#[test]
fn min_basic_works() {
let assets1 = test_assets();
let mut assets2 = Assets::new();
assets2.subsume(AF(1, 50));
assets2.subsume(ANF(2, 40));
assets2.subsume(CF(600));
assets2.subsume(CNF(40));
let assets2: MultiAssets = assets2.into();
let assets_min = assets1.min(&assets2.into());
let assets_min = assets_min.into_assets_iter().collect::<Vec<_>>();
assert_eq!(assets_min, vec![CF(300), AF(1, 50), CNF(40)]);
}
#[test]
fn saturating_take_all_and_none_works() {
let mut assets = test_assets();
let taken_none = assets.saturating_take(vec![].into());
assert_eq!(None, taken_none.assets_iter().next());
let taken_all = assets.saturating_take(All.into());
assert_eq!(None, assets.assets_iter().next());
let all_iter = taken_all.assets_iter();
assert!(all_iter.eq(test_assets().assets_iter()));
}
#[test]
fn saturating_take_all_abstract_works() {
let mut assets = test_assets();
let fungible = Wild((vec![1], WildFungible).into());
let non_fungible = Wild((vec![2], WildNonFungible).into());
let fungible = assets.saturating_take(fungible);
let fungible = fungible.assets_iter().collect::<Vec<_>>();
assert_eq!(fungible, vec![AF(1, 100)]);
let non_fungible = assets.saturating_take(non_fungible);
let non_fungible = non_fungible.assets_iter().collect::<Vec<_>>();
assert_eq!(non_fungible, vec![ANF(2, 20)]);
let final_assets = assets.assets_iter().collect::<Vec<_>>();
assert_eq!(final_assets, vec![CF(300), CNF(40)]);
}
#[test]
fn saturating_take_all_concrete_works() {
let mut assets = test_assets();
let fungible = Wild((Here, WildFungible).into());
let non_fungible = Wild((Here, WildNonFungible).into());
let fungible = assets.saturating_take(fungible);
let fungible = fungible.assets_iter().collect::<Vec<_>>();
assert_eq!(fungible, vec![CF(300)]);
let non_fungible = assets.saturating_take(non_fungible);
let non_fungible = non_fungible.assets_iter().collect::<Vec<_>>();
assert_eq!(non_fungible, vec![CNF(40)]);
let assets = assets.assets_iter().collect::<Vec<_>>();
assert_eq!(assets, vec![AF(1, 100), ANF(2, 20)]);
}
#[test]
fn saturating_take_basic_works() {
let mut assets1 = test_assets();
let mut assets2 = Assets::new();
assets2.subsume(AF(1, 50));
assets2.subsume(ANF(2, 40));
assets2.subsume(CF(600));
assets2.subsume(CNF(40));
let assets2: MultiAssets = assets2.into();
let taken = assets1.saturating_take(assets2.into());
let taken = taken.into_assets_iter().collect::<Vec<_>>();
assert_eq!(taken, vec![CF(300), AF(1, 50), CNF(40)]);
let assets = assets1.into_assets_iter().collect::<Vec<_>>();
assert_eq!(assets, vec![AF(1, 50), ANF(2, 20)]);
}
}