defuse_core/
amounts.rs

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