Skip to main content

tx3_tir/model/
assets.rs

1use serde::{Deserialize, Serialize};
2use std::collections::{HashMap, HashSet};
3
4pub type AssetPolicy = Vec<u8>;
5pub type AssetName = Vec<u8>;
6
7#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash, Ord, PartialOrd)]
8pub enum AssetClass {
9    Naked,
10    Named(AssetName),
11    Defined(AssetPolicy, AssetName),
12}
13
14impl AssetClass {
15    pub fn is_defined(&self) -> bool {
16        matches!(self, AssetClass::Defined(_, _))
17    }
18
19    pub fn is_named(&self) -> bool {
20        matches!(self, AssetClass::Named(_))
21    }
22
23    pub fn is_naked(&self) -> bool {
24        matches!(self, AssetClass::Naked)
25    }
26
27    pub fn policy(&self) -> Option<&[u8]> {
28        match self {
29            AssetClass::Defined(policy, _) => Some(policy),
30            _ => None,
31        }
32    }
33
34    pub fn name(&self) -> Option<&[u8]> {
35        match self {
36            AssetClass::Defined(_, name) => Some(name),
37            AssetClass::Named(name) => Some(name),
38            _ => None,
39        }
40    }
41}
42
43impl std::fmt::Display for AssetClass {
44    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45        match self {
46            AssetClass::Naked => write!(f, "naked")?,
47            AssetClass::Named(name) => write!(f, "{}", hex::encode(name))?,
48            AssetClass::Defined(policy, name) => {
49                write!(f, "{}.{}", hex::encode(policy), hex::encode(name))?
50            }
51        }
52
53        Ok(())
54    }
55}
56
57#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
58pub struct CanonicalAssets(HashMap<AssetClass, i128>);
59
60impl std::fmt::Display for CanonicalAssets {
61    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
62        write!(f, "CanonicalAssets {{")?;
63
64        for (class, amount) in self.iter() {
65            write!(f, "{}:{}", class, amount)?;
66        }
67
68        write!(f, "}}")?;
69
70        Ok(())
71    }
72}
73
74impl Default for CanonicalAssets {
75    fn default() -> Self {
76        Self::empty()
77    }
78}
79
80impl std::ops::Deref for CanonicalAssets {
81    type Target = HashMap<AssetClass, i128>;
82
83    fn deref(&self) -> &Self::Target {
84        &self.0
85    }
86}
87
88impl CanonicalAssets {
89    pub fn empty() -> Self {
90        Self(HashMap::new())
91    }
92
93    pub fn from_class_and_amount(class: AssetClass, amount: i128) -> Self {
94        Self(HashMap::from([(class, amount)]))
95    }
96
97    pub fn from_naked_amount(amount: i128) -> Self {
98        Self(HashMap::from([(AssetClass::Naked, amount)]))
99    }
100
101    pub fn from_named_asset(asset_name: &[u8], amount: i128) -> Self {
102        if asset_name.is_empty() {
103            return Self::from_naked_amount(amount);
104        }
105
106        Self(HashMap::from([(
107            AssetClass::Named(asset_name.to_vec()),
108            amount,
109        )]))
110    }
111
112    pub fn from_defined_asset(policy: &[u8], asset_name: &[u8], amount: i128) -> Self {
113        if policy.is_empty() {
114            return Self::from_named_asset(asset_name, amount);
115        }
116
117        Self(HashMap::from([(
118            AssetClass::Defined(policy.to_vec(), asset_name.to_vec()),
119            amount,
120        )]))
121    }
122
123    pub fn from_asset(policy: Option<&[u8]>, name: Option<&[u8]>, amount: i128) -> Self {
124        match (policy, name) {
125            (Some(policy), Some(name)) => Self::from_defined_asset(policy, name, amount),
126            (Some(policy), None) => Self::from_defined_asset(policy, &[], amount),
127            (None, Some(name)) => Self::from_named_asset(name, amount),
128            (None, None) => Self::from_naked_amount(amount),
129        }
130    }
131
132    pub fn classes(&self) -> HashSet<AssetClass> {
133        self.iter().map(|(class, _)| class.clone()).collect()
134    }
135
136    pub fn naked_amount(&self) -> Option<i128> {
137        self.get(&AssetClass::Naked).cloned()
138    }
139
140    pub fn asset_amount2(&self, policy: &[u8], name: &[u8]) -> Option<i128> {
141        self.get(&AssetClass::Defined(policy.to_vec(), name.to_vec()))
142            .cloned()
143    }
144
145    pub fn asset_amount(&self, asset: &AssetClass) -> Option<i128> {
146        self.get(asset).cloned()
147    }
148
149    pub fn contains_total(&self, other: &Self) -> bool {
150        for (class, other_amount) in other.iter() {
151            if *other_amount == 0 {
152                continue;
153            }
154
155            if *other_amount < 0 {
156                return false;
157            }
158
159            let Some(self_amount) = self.get(class) else {
160                return false;
161            };
162
163            if *self_amount < 0 {
164                return false;
165            }
166
167            if self_amount < other_amount {
168                return false;
169            }
170        }
171
172        true
173    }
174
175    pub fn contains_some(&self, other: &Self) -> bool {
176        if other.is_empty() {
177            return true;
178        }
179
180        if self.is_empty() {
181            return false;
182        }
183
184        for (class, other_amount) in other.iter() {
185            if *other_amount == 0 {
186                continue;
187            }
188
189            let Some(self_amount) = self.get(class) else {
190                continue;
191            };
192
193            if *self_amount > 0 {
194                return true;
195            }
196        }
197
198        false
199    }
200
201    pub fn is_empty(&self) -> bool {
202        self.iter().all(|(_, value)| *value == 0)
203    }
204
205    pub fn is_empty_or_negative(&self) -> bool {
206        for (_, value) in self.iter() {
207            if *value > 0 {
208                return false;
209            }
210        }
211
212        true
213    }
214
215    pub fn is_only_naked(&self) -> bool {
216        self.iter().all(|(x, _)| x.is_naked())
217    }
218
219    pub fn as_homogenous_asset(&self) -> Option<(AssetClass, i128)> {
220        if self.0.len() != 1 {
221            return None;
222        }
223
224        let (class, amount) = self.0.iter().next().unwrap();
225        Some((class.clone(), *amount))
226    }
227}
228
229impl From<CanonicalAssets> for HashMap<AssetClass, i128> {
230    fn from(assets: CanonicalAssets) -> Self {
231        assets.0
232    }
233}
234
235impl IntoIterator for CanonicalAssets {
236    type Item = (AssetClass, i128);
237    type IntoIter = std::collections::hash_map::IntoIter<AssetClass, i128>;
238
239    fn into_iter(self) -> Self::IntoIter {
240        self.0.into_iter()
241    }
242}
243
244impl std::ops::Neg for CanonicalAssets {
245    type Output = Self;
246
247    fn neg(self) -> Self {
248        let mut negated = self.0;
249
250        for (_, value) in negated.iter_mut() {
251            *value = -*value;
252        }
253
254        Self(negated)
255    }
256}
257
258impl std::ops::Add for CanonicalAssets {
259    type Output = Self;
260
261    fn add(self, other: Self) -> Self {
262        let mut aggregated = self.0;
263
264        for (key, value) in other.0 {
265            *aggregated.entry(key).or_default() += value;
266        }
267
268        aggregated.retain(|_, &mut value| value != 0);
269
270        Self(aggregated)
271    }
272}
273
274impl std::ops::Sub for CanonicalAssets {
275    type Output = Self;
276
277    fn sub(self, other: Self) -> Self {
278        let mut aggregated = self.0;
279
280        for (key, value) in other.0 {
281            *aggregated.entry(key).or_default() -= value;
282        }
283
284        aggregated.retain(|_, &mut value| value != 0);
285
286        Self(aggregated)
287    }
288}
289
290#[cfg(test)]
291mod tests {
292    use super::*;
293    use proptest::prelude::*;
294
295    prop_compose! {
296      fn any_asset() (
297        policy in any::<Vec<u8>>(),
298        name in any::<Vec<u8>>(),
299        amount in any::<i128>(),
300      ) -> CanonicalAssets {
301        CanonicalAssets::from_defined_asset(&policy, &name, amount)
302      }
303    }
304
305    prop_compose! {
306      fn any_positive_asset() (
307        policy in any::<Vec<u8>>(),
308        name in any::<Vec<u8>>(),
309        amount in 1..i128::MAX,
310      ) -> CanonicalAssets {
311        CanonicalAssets::from_defined_asset(&policy, &name, amount)
312      }
313    }
314
315    prop_compose! {
316      fn any_positive_composite_asset() (
317        naked_amount in 0..i128::MAX,
318        defined1 in any_positive_asset(),
319        defined2 in any_positive_asset(),
320      ) -> CanonicalAssets {
321        let naked = CanonicalAssets::from_naked_amount(naked_amount);
322        let composite = naked + defined1 + defined2;
323        composite
324      }
325    }
326
327    proptest! {
328        #[test]
329        fn empty_doesnt_contain_anything(asset in any_asset()) {
330            let x = CanonicalAssets::empty();
331            assert!(!x.contains_total(&asset));
332            assert!(!x.contains_some(&asset));
333        }
334    }
335
336    proptest! {
337        #[test]
338        fn empty_is_contained_in_everything(asset in any_asset()) {
339            let x = CanonicalAssets::empty();
340            assert!(asset.contains_total(&x));
341            assert!(asset.contains_some(&x));
342        }
343    }
344
345    proptest! {
346        #[test]
347        fn add_positive_makes_it_present(asset in any_positive_asset()) {
348            let x = CanonicalAssets::empty();
349            let x = x + asset.clone();
350            assert!(x.contains_total(&asset));
351            assert!(x.contains_some(&asset));
352            assert!(!x.is_empty_or_negative());
353        }
354    }
355
356    proptest! {
357        #[test]
358        fn sub_on_empty_makes_it_negative(asset in any_positive_asset()) {
359            let x = CanonicalAssets::empty();
360            let x = x - asset.clone();
361            assert!(!x.contains_total(&asset));
362            assert!(!x.contains_some(&asset));
363            assert!(x.is_empty_or_negative());
364        }
365    }
366
367    proptest! {
368        #[test]
369        fn add_is_inverse_of_sub(original in any_asset(), subtracted in any_asset()) {
370            let x = original.clone();
371            let x = x - subtracted.clone();
372            let x = x + subtracted.clone().clone();
373            assert_eq!(x, original);
374        }
375    }
376
377    proptest! {
378        #[test]
379        fn composite_contains_some_naked(composite in any_positive_composite_asset()) {
380            assert!(composite.contains_some(&CanonicalAssets::from_naked_amount(1)));
381        }
382    }
383
384    proptest! {
385        #[test]
386        fn composite_contains_some_composite(composite1 in any_positive_composite_asset(), composite2 in any_positive_composite_asset()) {
387            assert!(composite1.contains_some(&composite2));
388        }
389    }
390}