1use core::fmt::Debug;
2use std::{borrow::Cow, collections::BTreeMap};
3
4use borsh::{BorshDeserialize, BorshSerialize};
5use defuse_map_utils::{IterableMap, cleanup::DefaultMap};
6use defuse_num_utils::{CheckedAdd, CheckedSub};
7use impl_tools::autoimpl;
8use serde::{Deserialize, Deserializer, Serialize, Serializer};
9use serde_with::{DeserializeAs, SerializeAs};
10
11use crate::token_id::TokenId;
12
13#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
14#[cfg_attr(feature = "borsh-schema", derive(::borsh::BorshSchema))]
15#[cfg_attr(feature = "schemars-v0_8", derive(::schemars::JsonSchema))]
16#[derive(
17 Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize,
18)]
19#[autoimpl(Deref using self.0)]
20#[repr(transparent)]
21pub struct Amounts<T = BTreeMap<TokenId, u128>>(T);
22
23impl<T> Amounts<T> {
24 #[inline]
25 pub const fn new(map: T) -> Self {
26 Self(map)
27 }
28
29 #[inline]
30 pub fn into_inner(self) -> T {
31 self.0
32 }
33}
34
35impl<T> Amounts<T>
36where
37 T: DefaultMap,
38 T::V: Copy,
39{
40 #[inline]
41 pub fn amount_for(&self, k: &T::K) -> T::V {
42 self.0.get(k).copied().unwrap_or_default()
43 }
44
45 #[must_use]
46 #[inline]
47 pub fn add(&mut self, k: T::K, amount: u128) -> Option<T::V>
48 where
49 T::V: CheckedAdd<u128>,
50 {
51 self.checked_apply(k, |a| a.checked_add(amount))
52 }
53
54 #[must_use]
55 #[inline]
56 pub fn with_add(mut self, k: T::K, amount: u128) -> Option<Self>
57 where
58 T::V: CheckedAdd<u128>,
59 {
60 self.add(k, amount)?;
61 Some(self)
62 }
63
64 #[must_use]
65 #[inline]
66 pub fn with_add_many(self, amounts: impl IntoIterator<Item = (T::K, u128)>) -> Option<Self>
67 where
68 T::V: CheckedAdd<u128>,
69 {
70 amounts
71 .into_iter()
72 .try_fold(self, |amounts, (k, amount)| amounts.with_add(k, amount))
73 }
74
75 #[must_use]
76 #[inline]
77 pub fn sub(&mut self, k: T::K, amount: u128) -> Option<T::V>
78 where
79 T::V: CheckedSub<u128>,
80 {
81 self.checked_apply(k, |a| a.checked_sub(amount))
82 }
83
84 #[must_use]
85 #[inline]
86 pub fn with_sub(mut self, k: T::K, amount: u128) -> Option<Self>
87 where
88 T::V: CheckedSub<u128>,
89 {
90 self.sub(k, amount)?;
91 Some(self)
92 }
93
94 #[must_use]
95 #[inline]
96 pub fn with_sub_many(self, amounts: impl IntoIterator<Item = (T::K, u128)>) -> Option<Self>
97 where
98 T::V: CheckedSub<u128>,
99 {
100 amounts
101 .into_iter()
102 .try_fold(self, |amounts, (k, amount)| amounts.with_sub(k, amount))
103 }
104
105 #[must_use]
106 #[inline]
107 pub fn apply_delta(&mut self, k: T::K, delta: i128) -> Option<T::V>
108 where
109 T::V: CheckedAdd<i128>,
110 {
111 self.checked_apply(k, |a| a.checked_add(delta))
112 }
113
114 #[must_use]
115 #[inline]
116 pub fn with_apply_delta(mut self, k: T::K, delta: i128) -> Option<Self>
117 where
118 T::V: CheckedAdd<i128>,
119 {
120 self.apply_delta(k, delta)?;
121 Some(self)
122 }
123
124 #[must_use]
125 #[inline]
126 pub fn with_apply_deltas(self, amounts: impl IntoIterator<Item = (T::K, i128)>) -> Option<Self>
127 where
128 T::V: CheckedAdd<i128>,
129 {
130 amounts.into_iter().try_fold(self, |amounts, (k, delta)| {
131 amounts.with_apply_delta(k, delta)
132 })
133 }
134
135 #[must_use]
136 #[inline]
137 fn checked_apply(&mut self, k: T::K, f: impl FnOnce(T::V) -> Option<T::V>) -> Option<T::V> {
138 let mut a = self.0.entry_or_default(k);
139 *a = f(*a)?;
140 Some(*a)
141 }
142}
143
144#[allow(clippy::iter_without_into_iter)]
145impl<T> Amounts<T>
146where
147 T: IterableMap,
148{
149 pub fn iter(&self) -> T::Iter<'_> {
150 self.0.iter()
151 }
152}
153
154impl<T> IntoIterator for Amounts<T>
155where
156 T: IntoIterator,
157{
158 type Item = T::Item;
159
160 type IntoIter = T::IntoIter;
161
162 #[inline]
163 fn into_iter(self) -> Self::IntoIter {
164 self.into_inner().into_iter()
165 }
166}
167
168impl<'a, T> IntoIterator for &'a Amounts<T>
169where
170 &'a T: IntoIterator,
171{
172 type Item = <&'a T as IntoIterator>::Item;
173
174 type IntoIter = <&'a T as IntoIterator>::IntoIter;
175
176 #[inline]
177 fn into_iter(self) -> Self::IntoIter {
178 self.0.into_iter()
179 }
180}
181
182impl<T> Amounts<T>
183where
184 T: IterableMap,
185{
186 #[inline]
187 pub fn len(&self) -> usize {
188 self.0.len()
189 }
190
191 #[inline]
192 pub fn is_empty(&self) -> bool {
193 self.0.is_empty()
194 }
195}
196
197impl<T> From<Amounts<T>> for Cow<'_, Amounts<T>>
198where
199 T: Clone,
200{
201 fn from(value: Amounts<T>) -> Self {
202 Self::Owned(value)
203 }
204}
205
206impl<T, As> SerializeAs<Amounts<T>> for Amounts<As>
207where
208 As: SerializeAs<T>,
209{
210 #[inline]
211 fn serialize_as<S>(source: &Amounts<T>, serializer: S) -> Result<S::Ok, S::Error>
212 where
213 S: Serializer,
214 {
215 As::serialize_as(&source.0, serializer)
216 }
217}
218
219impl<'de, T, As> DeserializeAs<'de, Amounts<T>> for Amounts<As>
220where
221 As: DeserializeAs<'de, T>,
222{
223 #[inline]
224 fn deserialize_as<D>(deserializer: D) -> Result<Amounts<T>, D::Error>
225 where
226 D: Deserializer<'de>,
227 {
228 As::deserialize_as(deserializer).map(Amounts)
229 }
230}
231
232#[cfg(feature = "schemars-v0_8")]
233const _: () = {
234 use schemars::{r#gen::SchemaGenerator, schema::Schema};
235 use serde_with::schemars_0_8::JsonSchemaAs;
236
237 impl<T, As> JsonSchemaAs<Amounts<T>> for Amounts<As>
238 where
239 As: JsonSchemaAs<T>,
240 {
241 fn schema_name() -> String {
242 As::schema_name()
243 }
244
245 fn is_referenceable() -> bool {
246 false
247 }
248
249 fn json_schema(generator: &mut SchemaGenerator) -> Schema {
250 As::json_schema(generator)
251 }
252 }
253};
254
255#[cfg(test)]
256mod tests {
257
258 use crate::{AccountId, token_id::nep141::Nep141TokenId};
259
260 use super::*;
261
262 #[test]
263 fn invariant() {
264 let [t1, t2] = ["t1.near", "t2.near"]
265 .map(|t| TokenId::Nep141(Nep141TokenId::new(t.parse::<AccountId>().unwrap())));
266
267 assert!(Amounts::<BTreeMap<TokenId, i128>>::default().is_empty());
268 assert!(
269 Amounts::<BTreeMap<_, i128>>::default()
270 .with_apply_deltas([(t1.clone(), 0)])
271 .unwrap()
272 .is_empty()
273 );
274
275 assert!(
276 !Amounts::<BTreeMap<_, i128>>::default()
277 .with_apply_deltas([(t1.clone(), 1)])
278 .unwrap()
279 .is_empty()
280 );
281
282 assert!(
283 !Amounts::<BTreeMap<_, i128>>::default()
284 .with_apply_deltas([(t1.clone(), -1)])
285 .unwrap()
286 .is_empty()
287 );
288
289 assert!(
290 Amounts::<BTreeMap<_, i128>>::default()
291 .with_apply_deltas([(t1.clone(), 1), (t1.clone(), -1)])
292 .unwrap()
293 .is_empty()
294 );
295
296 assert!(
297 !Amounts::<BTreeMap<_, i128>>::default()
298 .with_apply_deltas([(t1.clone(), 1), (t1.clone(), -1), (t2.clone(), -1)])
299 .unwrap()
300 .is_empty()
301 );
302
303 assert!(
304 Amounts::<BTreeMap<_, i128>>::default()
305 .with_apply_deltas([(t1.clone(), 1), (t1, -1), (t2.clone(), -1), (t2, 1)])
306 .unwrap()
307 .is_empty()
308 );
309 }
310}