whisky_common/data/
value.rs

1use serde::{Deserialize, Serialize};
2
3use crate::{
4    data::{ByteString, Int, Map, PlutusDataJson},
5    models::Asset,
6};
7use std::collections::BTreeMap;
8
9#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
10pub struct Value(pub BTreeMap<String, u64>);
11
12impl Value {
13    pub fn new() -> Self {
14        Value(BTreeMap::new())
15    }
16
17    pub fn from_asset(asset: &Asset) -> Self {
18        let mut asset_map = BTreeMap::new();
19        asset_map.insert(
20            Value::sanitize_unit(&asset.unit()),
21            asset.quantity().parse::<u64>().unwrap(),
22        );
23        Value(asset_map)
24    }
25
26    pub fn from_asset_vec(assets: &[Asset]) -> Self {
27        let mut asset_map = BTreeMap::new();
28        for asset in assets {
29            let current_value = asset_map
30                .entry(Value::sanitize_unit(&asset.unit()))
31                .or_insert(0);
32            *current_value += asset.quantity().parse::<u64>().unwrap();
33        }
34        Value(asset_map)
35    }
36
37    pub fn add_asset(&mut self, unit: &str, quantity: u64) -> &mut Self {
38        let current_value = self.0.entry(Value::sanitize_unit(unit)).or_insert(0);
39        *current_value += quantity;
40        self
41    }
42
43    pub fn add_assets(&mut self, assets: &[Asset]) -> &mut Self {
44        for asset in assets {
45            self.add_asset(&asset.unit(), asset.quantity().parse::<u64>().unwrap());
46        }
47        self
48    }
49
50    pub fn merge(&mut self, other: &Value) -> &mut Self {
51        for (key, value) in other.0.clone() {
52            let current_value = self.0.entry(Value::sanitize_unit(&key)).or_insert(0);
53            *current_value += value;
54        }
55        self
56    }
57
58    pub fn negate_asset(&mut self, unit: &str, quantity: u64) -> &mut Self {
59        let current_value = self.0.entry(Value::sanitize_unit(unit)).or_insert(0);
60        if *current_value <= quantity {
61            self.0.remove(unit);
62        } else {
63            *current_value -= quantity;
64        };
65        self
66    }
67
68    pub fn negate_assets(&mut self, other: &[Asset]) -> &mut Self {
69        for asset in other {
70            self.negate_asset(&asset.unit(), asset.quantity().parse::<u64>().unwrap());
71        }
72        self
73    }
74
75    pub fn negate_value(&mut self, other: &Value) -> &mut Self {
76        for (key, value) in other.0.clone() {
77            let unit = Value::sanitize_unit(&key);
78            let current_value = self.0.entry(unit.clone()).or_insert(0);
79            if *current_value <= value {
80                self.0.remove(&unit);
81            } else {
82                *current_value -= value;
83            }
84        }
85        self
86    }
87
88    pub fn to_asset_vec(&self) -> Vec<Asset> {
89        let mut assets = vec![];
90        for (unit, quantity) in &self.0 {
91            assets.push(Asset::new(Value::sanitize_unit(unit), quantity.to_string()));
92        }
93        assets
94    }
95
96    // Accessor
97    pub fn get(&self, key: &str) -> u64 {
98        let key = if key.is_empty() { "lovelace" } else { key };
99        match self.0.get(key) {
100            Some(value) => *value,
101            None => 0,
102        }
103    }
104
105    pub fn get_policy_assets(&self, policy_id: &str) -> Vec<Asset> {
106        let mut assets = vec![];
107        for (unit, quantity) in &self.0 {
108            if unit.starts_with(policy_id) {
109                assets.push(Asset::new(unit.clone(), quantity.to_string()));
110            }
111        }
112        assets
113    }
114
115    pub fn keys(&self) -> Vec<String> {
116        self.0.keys().cloned().collect()
117    }
118
119    pub fn len(&self) -> usize {
120        self.0.len()
121    }
122
123    // Comparison function
124    pub fn geq(&self, other: &Value) -> bool {
125        for (key, value) in &other.0 {
126            if self
127                .0
128                .get(&Value::sanitize_unit(key))
129                .is_some_and(|v| v < value)
130            {
131                return false;
132            }
133        }
134        true
135    }
136
137    pub fn leq(&self, other: &Value) -> bool {
138        for (key, value) in &other.0 {
139            if self
140                .0
141                .get(&Value::sanitize_unit(key))
142                .is_some_and(|v| v > value)
143            {
144                return false;
145            }
146        }
147        true
148    }
149
150    pub fn is_empty(&self) -> bool {
151        self.0.is_empty()
152    }
153
154    pub fn sanitize_unit(unit: &str) -> String {
155        if unit.is_empty() {
156            "lovelace".to_string()
157        } else {
158            unit.to_string()
159        }
160    }
161}
162
163impl PlutusDataJson for Value {
164    fn to_json(&self) -> serde_json::Value {
165        let mut value_map: BTreeMap<String, BTreeMap<String, u64>> = BTreeMap::new();
166
167        self.0.iter().for_each(|(unit, quantity)| {
168            let sanitized_name = unit.replace("lovelace", "");
169            let policy = &sanitized_name[..56.min(sanitized_name.len())];
170            let token = &sanitized_name[56.min(sanitized_name.len())..];
171
172            value_map
173                .entry(policy.to_string())
174                .or_insert_with(BTreeMap::new)
175                .entry(token.to_string())
176                .and_modify(|q| *q += quantity)
177                .or_insert(*quantity);
178        });
179
180        let json_map = value_map
181            .into_iter() // Keys will already be sorted
182            .map(|(policy, tokens)| {
183                (
184                    ByteString::new(&policy),
185                    tokens
186                        .into_iter() // Token keys will already be sorted
187                        .map(|(token, quantity)| {
188                            (ByteString::new(&token), Int::new(quantity as i128))
189                        })
190                        .collect(),
191                )
192            })
193            .collect::<Map<ByteString, Map<ByteString, Int>>>();
194        json_map.to_json()
195    }
196}
197
198impl Default for Value {
199    fn default() -> Self {
200        Value::new()
201    }
202}