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