starfleit/
asset.rs

1use schemars::JsonSchema;
2use serde::{Deserialize, Serialize};
3use std::fmt;
4
5use crate::querier::{query_balance, query_native_decimals, query_token_balance, query_token_info};
6use cosmwasm_std::{
7    to_binary, Addr, Api, BankMsg, CanonicalAddr, Coin, CosmosMsg, MessageInfo, QuerierWrapper,
8    StdError, StdResult, SubMsg, Uint128, WasmMsg,
9};
10use cw20::Cw20ExecuteMsg;
11
12#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
13pub struct Asset {
14    pub info: AssetInfo,
15    pub amount: Uint128,
16}
17
18impl fmt::Display for Asset {
19    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
20        write!(f, "{}{}", self.amount, self.info)
21    }
22}
23
24impl Asset {
25    pub fn is_native_token(&self) -> bool {
26        self.info.is_native_token()
27    }
28
29    pub fn into_msg(self, recipient: Addr) -> StdResult<CosmosMsg> {
30        let amount = self.amount;
31
32        match &self.info {
33            AssetInfo::Token { contract_addr } => Ok(CosmosMsg::Wasm(WasmMsg::Execute {
34                contract_addr: contract_addr.to_string(),
35                msg: to_binary(&Cw20ExecuteMsg::Transfer {
36                    recipient: recipient.to_string(),
37                    amount,
38                })?,
39                funds: vec![],
40            })),
41            AssetInfo::NativeToken { denom } => Ok(CosmosMsg::Bank(BankMsg::Send {
42                to_address: recipient.to_string(),
43                amount: vec![Coin {
44                    amount: self.amount,
45                    denom: denom.to_string(),
46                }],
47            })),
48        }
49    }
50
51    pub fn into_submsg(self, recipient: Addr) -> StdResult<SubMsg> {
52        Ok(SubMsg::new(self.into_msg(recipient)?))
53    }
54
55    pub fn assert_sent_native_token_balance(&self, message_info: &MessageInfo) -> StdResult<()> {
56        if let AssetInfo::NativeToken { denom } = &self.info {
57            match message_info.funds.iter().find(|x| x.denom == *denom) {
58                Some(coin) => {
59                    if self.amount == coin.amount {
60                        Ok(())
61                    } else {
62                        Err(StdError::generic_err("Native token balance mismatch between the argument and the transferred"))
63                    }
64                }
65                None => {
66                    if self.amount.is_zero() {
67                        Ok(())
68                    } else {
69                        Err(StdError::generic_err("Native token balance mismatch between the argument and the transferred"))
70                    }
71                }
72            }
73        } else {
74            Ok(())
75        }
76    }
77
78    pub fn to_raw(&self, api: &dyn Api) -> StdResult<AssetRaw> {
79        Ok(AssetRaw {
80            info: match &self.info {
81                AssetInfo::NativeToken { denom } => AssetInfoRaw::NativeToken {
82                    denom: denom.to_string(),
83                },
84                AssetInfo::Token { contract_addr } => AssetInfoRaw::Token {
85                    contract_addr: api.addr_canonicalize(contract_addr.as_str())?,
86                },
87            },
88            amount: self.amount,
89        })
90    }
91}
92
93/// AssetInfo contract_addr is usually passed from the cw20 hook
94/// so we can trust the contract_addr is properly validated.
95#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
96#[serde(rename_all = "snake_case")]
97pub enum AssetInfo {
98    Token { contract_addr: String },
99    NativeToken { denom: String },
100}
101
102impl fmt::Display for AssetInfo {
103    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
104        match self {
105            AssetInfo::NativeToken { denom } => write!(f, "{denom}"),
106            AssetInfo::Token { contract_addr } => write!(f, "{contract_addr}"),
107        }
108    }
109}
110
111impl AssetInfo {
112    pub fn to_raw(&self, api: &dyn Api) -> StdResult<AssetInfoRaw> {
113        match self {
114            AssetInfo::NativeToken { denom } => Ok(AssetInfoRaw::NativeToken {
115                denom: denom.to_string(),
116            }),
117            AssetInfo::Token { contract_addr } => Ok(AssetInfoRaw::Token {
118                contract_addr: api.addr_canonicalize(contract_addr.as_str())?,
119            }),
120        }
121    }
122
123    pub fn is_native_token(&self) -> bool {
124        match self {
125            AssetInfo::NativeToken { .. } => true,
126            AssetInfo::Token { .. } => false,
127        }
128    }
129    pub fn query_pool(
130        &self,
131        querier: &QuerierWrapper,
132        api: &dyn Api,
133        pool_addr: Addr,
134    ) -> StdResult<Uint128> {
135        match self {
136            AssetInfo::Token { contract_addr, .. } => query_token_balance(
137                querier,
138                api.addr_validate(contract_addr.as_str())?,
139                pool_addr,
140            ),
141            AssetInfo::NativeToken { denom, .. } => {
142                query_balance(querier, pool_addr, denom.to_string())
143            }
144        }
145    }
146
147    pub fn equal(&self, asset: &AssetInfo) -> bool {
148        match self {
149            AssetInfo::Token { contract_addr, .. } => {
150                let self_contract_addr = contract_addr;
151                match asset {
152                    AssetInfo::Token { contract_addr, .. } => self_contract_addr == contract_addr,
153                    AssetInfo::NativeToken { .. } => false,
154                }
155            }
156            AssetInfo::NativeToken { denom, .. } => {
157                let self_denom = denom;
158                match asset {
159                    AssetInfo::Token { .. } => false,
160                    AssetInfo::NativeToken { denom, .. } => self_denom == denom,
161                }
162            }
163        }
164    }
165
166    pub fn query_decimals(&self, account_addr: Addr, querier: &QuerierWrapper) -> StdResult<u8> {
167        match self {
168            AssetInfo::NativeToken { denom } => {
169                query_native_decimals(querier, account_addr, denom.to_string())
170            }
171            AssetInfo::Token { contract_addr } => {
172                let token_info = query_token_info(querier, Addr::unchecked(contract_addr))?;
173                Ok(token_info.decimals)
174            }
175        }
176    }
177}
178
179#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
180pub struct AssetRaw {
181    pub info: AssetInfoRaw,
182    pub amount: Uint128,
183}
184
185impl AssetRaw {
186    pub fn to_normal(&self, api: &dyn Api) -> StdResult<Asset> {
187        Ok(Asset {
188            info: match &self.info {
189                AssetInfoRaw::NativeToken { denom } => AssetInfo::NativeToken {
190                    denom: denom.to_string(),
191                },
192                AssetInfoRaw::Token { contract_addr } => AssetInfo::Token {
193                    contract_addr: api.addr_humanize(contract_addr)?.to_string(),
194                },
195            },
196            amount: self.amount,
197        })
198    }
199}
200
201#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
202pub enum AssetInfoRaw {
203    Token { contract_addr: CanonicalAddr },
204    NativeToken { denom: String },
205}
206
207impl AssetInfoRaw {
208    pub fn to_normal(&self, api: &dyn Api) -> StdResult<AssetInfo> {
209        match self {
210            AssetInfoRaw::NativeToken { denom } => Ok(AssetInfo::NativeToken {
211                denom: denom.to_string(),
212            }),
213            AssetInfoRaw::Token { contract_addr } => Ok(AssetInfo::Token {
214                contract_addr: api.addr_humanize(contract_addr)?.to_string(),
215            }),
216        }
217    }
218
219    pub fn as_bytes(&self) -> &[u8] {
220        match self {
221            AssetInfoRaw::NativeToken { denom } => denom.as_bytes(),
222            AssetInfoRaw::Token { contract_addr } => contract_addr.as_slice(),
223        }
224    }
225
226    pub fn equal(&self, asset: &AssetInfoRaw) -> bool {
227        match self {
228            AssetInfoRaw::Token { contract_addr, .. } => {
229                let self_contract_addr = contract_addr;
230                match asset {
231                    AssetInfoRaw::Token { contract_addr, .. } => {
232                        self_contract_addr == contract_addr
233                    }
234                    AssetInfoRaw::NativeToken { .. } => false,
235                }
236            }
237            AssetInfoRaw::NativeToken { denom, .. } => {
238                let self_denom = denom;
239                match asset {
240                    AssetInfoRaw::Token { .. } => false,
241                    AssetInfoRaw::NativeToken { denom, .. } => self_denom == denom,
242                }
243            }
244        }
245    }
246}
247
248// We define a custom struct for each query response
249#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
250pub struct PairInfo {
251    pub asset_infos: [AssetInfo; 2],
252    pub contract_addr: String,
253    pub liquidity_token: String,
254    pub asset_decimals: [u8; 2],
255}
256
257#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
258pub struct PairInfoRaw {
259    pub asset_infos: [AssetInfoRaw; 2],
260    pub contract_addr: CanonicalAddr,
261    pub liquidity_token: CanonicalAddr,
262    pub asset_decimals: [u8; 2],
263}
264
265impl PairInfoRaw {
266    pub fn to_normal(&self, api: &dyn Api) -> StdResult<PairInfo> {
267        Ok(PairInfo {
268            liquidity_token: api.addr_humanize(&self.liquidity_token)?.to_string(),
269            contract_addr: api.addr_humanize(&self.contract_addr)?.to_string(),
270            asset_infos: [
271                self.asset_infos[0].to_normal(api)?,
272                self.asset_infos[1].to_normal(api)?,
273            ],
274            asset_decimals: self.asset_decimals,
275        })
276    }
277
278    pub fn query_pools(
279        &self,
280        querier: &QuerierWrapper,
281        api: &dyn Api,
282        contract_addr: Addr,
283    ) -> StdResult<[Asset; 2]> {
284        let info_0: AssetInfo = self.asset_infos[0].to_normal(api)?;
285        let info_1: AssetInfo = self.asset_infos[1].to_normal(api)?;
286        Ok([
287            Asset {
288                amount: info_0.query_pool(querier, api, contract_addr.clone())?,
289                info: info_0,
290            },
291            Asset {
292                amount: info_1.query_pool(querier, api, contract_addr)?,
293                info: info_1,
294            },
295        ])
296    }
297}