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}