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")))]
231const _: () = {
232 use near_sdk::schemars::{
233 JsonSchema,
234 r#gen::SchemaGenerator,
235 schema::{InstanceType, Schema, SchemaObject},
236 };
237 use serde_with::schemars_0_8::JsonSchemaAs;
238
239 impl<T, As> JsonSchemaAs<Amounts<T>> for Amounts<As>
240 where
241 As: JsonSchemaAs<T>,
242 {
243 fn schema_name() -> String {
244 As::schema_name()
245 }
246
247 fn is_referenceable() -> bool {
248 false
249 }
250
251 fn json_schema(generator: &mut SchemaGenerator) -> Schema {
252 As::json_schema(generator)
253 }
254 }
255};
256
257#[cfg(test)]
258mod tests {
259
260 use crate::token_id::nep141::Nep141TokenId;
261
262 use super::*;
263
264 #[test]
265 fn invariant() {
266 let [t1, t2] =
267 ["t1.near", "t2.near"].map(|t| TokenId::Nep141(Nep141TokenId::new(t.parse().unwrap())));
268
269 assert!(Amounts::<BTreeMap<TokenId, i128>>::default().is_empty());
270 assert!(
271 Amounts::<BTreeMap<_, i128>>::default()
272 .with_apply_deltas([(t1.clone(), 0)])
273 .unwrap()
274 .is_empty()
275 );
276
277 assert!(
278 !Amounts::<BTreeMap<_, i128>>::default()
279 .with_apply_deltas([(t1.clone(), 1)])
280 .unwrap()
281 .is_empty()
282 );
283
284 assert!(
285 !Amounts::<BTreeMap<_, i128>>::default()
286 .with_apply_deltas([(t1.clone(), -1)])
287 .unwrap()
288 .is_empty()
289 );
290
291 assert!(
292 Amounts::<BTreeMap<_, i128>>::default()
293 .with_apply_deltas([(t1.clone(), 1), (t1.clone(), -1)])
294 .unwrap()
295 .is_empty()
296 );
297
298 assert!(
299 !Amounts::<BTreeMap<_, i128>>::default()
300 .with_apply_deltas([(t1.clone(), 1), (t1.clone(), -1), (t2.clone(), -1)])
301 .unwrap()
302 .is_empty()
303 );
304
305 assert!(
306 Amounts::<BTreeMap<_, i128>>::default()
307 .with_apply_deltas([(t1.clone(), 1), (t1, -1), (t2.clone(), -1), (t2, 1)])
308 .unwrap()
309 .is_empty()
310 );
311 }
312}