tw_asset_plus/
impl_cw.rs

1use cosmwasm_std::{
2    from_slice, to_binary, to_vec, Api, BankMsg, Coin, CosmosMsg, QuerierWrapper, StdError,
3    StdResult, Uint128, WasmMsg,
4};
5use cw20::{BalanceResponse, Cw20ExecuteMsg, Cw20QueryMsg};
6use cw_storage_plus::PrimaryKey;
7
8use crate::{asset::AssetInfo, Asset};
9
10impl PrimaryKey<'_> for AssetInfo {
11    type Prefix = ();
12
13    type SubPrefix = ();
14
15    fn key(&self) -> Vec<&[u8]> {
16        vec![self.as_bytes()]
17    }
18}
19
20impl<'a> PrimaryKey<'a> for &'a AssetInfo {
21    type Prefix = ();
22
23    type SubPrefix = ();
24
25    fn key(&self) -> Vec<&[u8]> {
26        vec![self.as_bytes()]
27    }
28}
29
30impl Asset {
31    pub fn new<A: Into<Uint128>>(info: AssetInfo, amount: A) -> Self {
32        Asset {
33            info,
34            amount: amount.into(),
35        }
36    }
37
38    pub fn is_empty(&self) -> bool {
39        self.amount.is_zero()
40    }
41
42    pub fn transfer_all_msg<T: Into<String>>(&self, to_address: T) -> StdResult<CosmosMsg> {
43        self.info.transfer_msg(to_address, self.amount)
44    }
45
46    pub fn assert_sent_token(&self, coins: &[Coin]) -> StdResult<()> {
47        if let AssetInfo::NativeToken { denom } = &self.info {
48            match coins.iter().find(|c| &c.denom == denom) {
49                Some(c) => (c.amount == self.amount).then(|| ()).ok_or_else(|| {
50                    StdError::generic_err(format!(
51                        "Expected amount {} but found {} for denom {}",
52                        self.amount, c.amount, denom
53                    ))
54                }),
55                None => Err(StdError::generic_err(format!(
56                    "Denom {} not found in sent coins",
57                    denom
58                ))),
59            }?;
60        }
61
62        Ok(())
63    }
64}
65
66impl AssetInfo {
67    pub fn assert_eq(&self, other: &Self) -> StdResult<()> {
68        (self == other)
69            .then(|| ())
70            .ok_or_else(|| StdError::generic_err("Asset info mismatch"))
71    }
72
73    pub fn as_bytes(&self) -> &[u8] {
74        match &self {
75            AssetInfo::Token { contract_addr } => contract_addr.as_bytes(),
76            AssetInfo::NativeToken { denom } => denom.as_bytes(),
77        }
78    }
79
80    pub fn from_bytes(b: &[u8], api: &dyn Api) -> StdResult<Self> {
81        let s = String::from_utf8(b.to_vec())
82            .map_err(|_| StdError::invalid_utf8("String parsing error"))?;
83        Ok(match api.addr_validate(&s) {
84            Ok(addr) => AssetInfo::Token {
85                contract_addr: addr,
86            },
87            Err(_) => AssetInfo::NativeToken { denom: s },
88        })
89    }
90
91    pub fn to_serde_vec(&self) -> StdResult<Vec<u8>> {
92        to_vec(self)
93    }
94
95    pub fn from_serde_slice(b: &[u8]) -> StdResult<Self> {
96        from_slice(b)
97    }
98
99    pub fn transfer_msg<T: Into<String>, A: Into<Uint128>>(
100        &self,
101        to_address: T,
102        amount: A,
103    ) -> StdResult<CosmosMsg> {
104        let msg = match self {
105            AssetInfo::Token { contract_addr } => CosmosMsg::Wasm(WasmMsg::Execute {
106                contract_addr: contract_addr.into(),
107                msg: to_binary(&Cw20ExecuteMsg::Transfer {
108                    recipient: to_address.into(),
109                    amount: amount.into(),
110                })?,
111                funds: vec![],
112            }),
113            AssetInfo::NativeToken { denom } => CosmosMsg::Bank(BankMsg::Send {
114                to_address: to_address.into(),
115                amount: vec![Coin {
116                    denom: denom.clone(),
117                    amount: amount.into(),
118                }],
119            }),
120        };
121
122        Ok(msg)
123    }
124
125    pub fn query_balance<T: Into<String>>(
126        &self,
127        querier: &QuerierWrapper,
128        address: T,
129    ) -> StdResult<Uint128> {
130        match self {
131            AssetInfo::Token { contract_addr } => {
132                let bal: BalanceResponse = querier.query_wasm_smart(
133                    contract_addr,
134                    &Cw20QueryMsg::Balance {
135                        address: address.into(),
136                    },
137                )?;
138                Ok(bal.balance)
139            }
140            AssetInfo::NativeToken { denom } => Ok(querier.query_balance(address, denom)?.amount),
141        }
142    }
143}
tw_asset_plus/impl_cw.rs

tw_asset_plus/
impl_cw.rs
