cw_controllers/
claim.rs

1use cosmwasm_schema::cw_serde;
2use cosmwasm_std::{Addr, BlockInfo, CustomQuery, Deps, StdResult, Storage, Uint128};
3use cw_storage_plus::{Map, Namespace};
4use cw_utils::Expiration;
5
6// TODO: pull into utils?
7#[cw_serde]
8pub struct ClaimsResponse {
9    pub claims: Vec<Claim>,
10}
11
12// TODO: pull into utils?
13#[cw_serde]
14pub struct Claim {
15    pub amount: Uint128,
16    pub release_at: Expiration,
17}
18
19impl Claim {
20    pub fn new(amount: u128, released: Expiration) -> Self {
21        Claim {
22            amount: amount.into(),
23            release_at: released,
24        }
25    }
26}
27
28// TODO: revisit design (split each claim on own key?)
29pub struct Claims(Map<&'static Addr, Vec<Claim>>);
30
31impl Claims {
32    pub const fn new(storage_key: &'static str) -> Self {
33        Claims(Map::new(storage_key))
34    }
35
36    pub fn new_dyn(storage_key: impl Into<Namespace>) -> Self {
37        Claims(Map::new_dyn(storage_key))
38    }
39
40    /// This creates a claim, such that the given address can claim an amount of tokens after
41    /// the release date.
42    pub fn create_claim(
43        &self,
44        storage: &mut dyn Storage,
45        addr: &Addr,
46        amount: Uint128,
47        release_at: Expiration,
48    ) -> StdResult<()> {
49        // add a claim to this user to get their tokens after the unbonding period
50        self.0.update(storage, addr, |old| -> StdResult<_> {
51            let mut claims = old.unwrap_or_default();
52            claims.push(Claim { amount, release_at });
53            Ok(claims)
54        })?;
55        Ok(())
56    }
57
58    /// This iterates over all mature claims for the address, and removes them, up to an optional cap.
59    /// it removes the finished claims and returns the total amount of tokens to be released.
60    pub fn claim_tokens(
61        &self,
62        storage: &mut dyn Storage,
63        addr: &Addr,
64        block: &BlockInfo,
65        cap: Option<Uint128>,
66    ) -> StdResult<Uint128> {
67        let mut to_send = Uint128::zero();
68        self.0.update(storage, addr, |claim| -> StdResult<_> {
69            let (_send, waiting): (Vec<_>, _) =
70                claim.unwrap_or_default().into_iter().partition(|c| {
71                    // if mature and we can pay fully, then include in _send
72                    if c.release_at.is_expired(block) {
73                        if let Some(limit) = cap {
74                            if to_send + c.amount > limit {
75                                return false;
76                            }
77                        }
78                        // TODO: handle partial paying claims?
79                        to_send += c.amount;
80                        true
81                    } else {
82                        // not to send, leave in waiting and save again
83                        false
84                    }
85                });
86            Ok(waiting)
87        })?;
88        Ok(to_send)
89    }
90
91    pub fn query_claims<Q: CustomQuery>(
92        &self,
93        deps: Deps<Q>,
94        address: &Addr,
95    ) -> StdResult<ClaimsResponse> {
96        let claims = self.0.may_load(deps.storage, address)?.unwrap_or_default();
97        Ok(ClaimsResponse { claims })
98    }
99}
100
101#[cfg(test)]
102mod test {
103    use cosmwasm_std::{
104        testing::{mock_dependencies, mock_env},
105        Order,
106    };
107
108    use super::*;
109    const TEST_AMOUNT: u128 = 1000u128;
110    const TEST_EXPIRATION: Expiration = Expiration::AtHeight(10);
111
112    #[test]
113    fn can_create_claim() {
114        let claim = Claim::new(TEST_AMOUNT, TEST_EXPIRATION);
115        assert_eq!(claim.amount, Uint128::from(TEST_AMOUNT));
116        assert_eq!(claim.release_at, TEST_EXPIRATION);
117    }
118
119    #[test]
120    fn can_create_claims() {
121        let deps = mock_dependencies();
122        let claims = Claims::new("claims");
123        // Assert that claims creates a map and there are no keys in the map.
124        assert_eq!(
125            claims
126                .0
127                .range_raw(&deps.storage, None, None, Order::Ascending)
128                .collect::<StdResult<Vec<_>>>()
129                .unwrap()
130                .len(),
131            0
132        );
133    }
134
135    #[test]
136    fn check_create_claim_updates_map() {
137        let mut deps = mock_dependencies();
138        let claims = Claims::new("claims");
139
140        claims
141            .create_claim(
142                deps.as_mut().storage,
143                &Addr::unchecked("addr"),
144                TEST_AMOUNT.into(),
145                TEST_EXPIRATION,
146            )
147            .unwrap();
148
149        // Assert that claims creates a map and there is one claim for the address.
150        let saved_claims = claims
151            .0
152            .load(deps.as_mut().storage, &Addr::unchecked("addr"))
153            .unwrap();
154        assert_eq!(saved_claims.len(), 1);
155        assert_eq!(saved_claims[0].amount, Uint128::from(TEST_AMOUNT));
156        assert_eq!(saved_claims[0].release_at, TEST_EXPIRATION);
157
158        // Adding another claim to same address, make sure that both claims are saved.
159        claims
160            .create_claim(
161                deps.as_mut().storage,
162                &Addr::unchecked("addr"),
163                (TEST_AMOUNT + 100).into(),
164                TEST_EXPIRATION,
165            )
166            .unwrap();
167
168        // Assert that both claims exist for the address.
169        let saved_claims = claims
170            .0
171            .load(deps.as_mut().storage, &Addr::unchecked("addr"))
172            .unwrap();
173        assert_eq!(saved_claims.len(), 2);
174        assert_eq!(saved_claims[0].amount, Uint128::from(TEST_AMOUNT));
175        assert_eq!(saved_claims[0].release_at, TEST_EXPIRATION);
176        assert_eq!(saved_claims[1].amount, Uint128::from(TEST_AMOUNT + 100));
177        assert_eq!(saved_claims[1].release_at, TEST_EXPIRATION);
178
179        // Adding another claim to different address, make sure that other address only has one claim.
180        claims
181            .create_claim(
182                deps.as_mut().storage,
183                &Addr::unchecked("addr2"),
184                (TEST_AMOUNT + 100).into(),
185                TEST_EXPIRATION,
186            )
187            .unwrap();
188
189        // Assert that both claims exist for the address.
190        let saved_claims = claims
191            .0
192            .load(deps.as_mut().storage, &Addr::unchecked("addr"))
193            .unwrap();
194
195        let saved_claims_addr2 = claims
196            .0
197            .load(deps.as_mut().storage, &Addr::unchecked("addr2"))
198            .unwrap();
199        assert_eq!(saved_claims.len(), 2);
200        assert_eq!(saved_claims_addr2.len(), 1);
201    }
202
203    #[test]
204    fn test_claim_tokens_with_no_claims() {
205        let mut deps = mock_dependencies();
206        let claims = Claims::new("claims");
207
208        let amount = claims
209            .claim_tokens(
210                deps.as_mut().storage,
211                &Addr::unchecked("addr"),
212                &mock_env().block,
213                None,
214            )
215            .unwrap();
216        let saved_claims = claims
217            .0
218            .load(deps.as_mut().storage, &Addr::unchecked("addr"))
219            .unwrap();
220
221        assert_eq!(amount, Uint128::zero());
222        assert_eq!(saved_claims.len(), 0);
223    }
224
225    #[test]
226    fn test_claim_tokens_with_no_released_claims() {
227        let mut deps = mock_dependencies();
228        let claims = Claims::new("claims");
229
230        claims
231            .create_claim(
232                deps.as_mut().storage,
233                &Addr::unchecked("addr"),
234                (TEST_AMOUNT + 100).into(),
235                Expiration::AtHeight(10),
236            )
237            .unwrap();
238
239        claims
240            .create_claim(
241                deps.as_mut().storage,
242                &Addr::unchecked("addr"),
243                (TEST_AMOUNT + 100).into(),
244                Expiration::AtHeight(100),
245            )
246            .unwrap();
247
248        let mut env = mock_env();
249        env.block.height = 0;
250        // the address has two claims however they are both not expired
251        let amount = claims
252            .claim_tokens(
253                deps.as_mut().storage,
254                &Addr::unchecked("addr"),
255                &env.block,
256                None,
257            )
258            .unwrap();
259
260        let saved_claims = claims
261            .0
262            .load(deps.as_mut().storage, &Addr::unchecked("addr"))
263            .unwrap();
264
265        assert_eq!(amount, Uint128::zero());
266        assert_eq!(saved_claims.len(), 2);
267        assert_eq!(saved_claims[0].amount, Uint128::from(TEST_AMOUNT + 100));
268        assert_eq!(saved_claims[0].release_at, Expiration::AtHeight(10));
269        assert_eq!(saved_claims[1].amount, Uint128::from(TEST_AMOUNT + 100));
270        assert_eq!(saved_claims[1].release_at, Expiration::AtHeight(100));
271    }
272
273    #[test]
274    fn test_claim_tokens_with_one_released_claim() {
275        let mut deps = mock_dependencies();
276        let claims = Claims::new("claims");
277
278        claims
279            .create_claim(
280                deps.as_mut().storage,
281                &Addr::unchecked("addr"),
282                TEST_AMOUNT.into(),
283                Expiration::AtHeight(10),
284            )
285            .unwrap();
286
287        claims
288            .create_claim(
289                deps.as_mut().storage,
290                &Addr::unchecked("addr"),
291                (TEST_AMOUNT + 100).into(),
292                Expiration::AtHeight(100),
293            )
294            .unwrap();
295
296        let mut env = mock_env();
297        env.block.height = 20;
298        // the address has two claims and the first one can be released
299        let amount = claims
300            .claim_tokens(
301                deps.as_mut().storage,
302                &Addr::unchecked("addr"),
303                &env.block,
304                None,
305            )
306            .unwrap();
307
308        let saved_claims = claims
309            .0
310            .load(deps.as_mut().storage, &Addr::unchecked("addr"))
311            .unwrap();
312
313        assert_eq!(amount, Uint128::from(TEST_AMOUNT));
314        assert_eq!(saved_claims.len(), 1);
315        assert_eq!(saved_claims[0].amount, Uint128::from(TEST_AMOUNT + 100));
316        assert_eq!(saved_claims[0].release_at, Expiration::AtHeight(100));
317    }
318
319    #[test]
320    fn test_claim_tokens_with_all_released_claims() {
321        let mut deps = mock_dependencies();
322        let claims = Claims::new("claims");
323
324        claims
325            .create_claim(
326                deps.as_mut().storage,
327                &Addr::unchecked("addr"),
328                TEST_AMOUNT.into(),
329                Expiration::AtHeight(10),
330            )
331            .unwrap();
332
333        claims
334            .create_claim(
335                deps.as_mut().storage,
336                &Addr::unchecked("addr"),
337                (TEST_AMOUNT + 100).into(),
338                Expiration::AtHeight(100),
339            )
340            .unwrap();
341
342        let mut env = mock_env();
343        env.block.height = 1000;
344        // the address has two claims and both can be released
345        let amount = claims
346            .claim_tokens(
347                deps.as_mut().storage,
348                &Addr::unchecked("addr"),
349                &env.block,
350                None,
351            )
352            .unwrap();
353
354        let saved_claims = claims
355            .0
356            .load(deps.as_mut().storage, &Addr::unchecked("addr"))
357            .unwrap();
358
359        assert_eq!(amount, Uint128::from(TEST_AMOUNT + TEST_AMOUNT + 100));
360        assert_eq!(saved_claims.len(), 0);
361    }
362
363    #[test]
364    fn test_claim_tokens_with_zero_cap() {
365        let mut deps = mock_dependencies();
366        let claims = Claims::new("claims");
367
368        claims
369            .create_claim(
370                deps.as_mut().storage,
371                &Addr::unchecked("addr"),
372                TEST_AMOUNT.into(),
373                Expiration::AtHeight(10),
374            )
375            .unwrap();
376
377        claims
378            .create_claim(
379                deps.as_mut().storage,
380                &Addr::unchecked("addr"),
381                (TEST_AMOUNT + 100).into(),
382                Expiration::AtHeight(100),
383            )
384            .unwrap();
385
386        let mut env = mock_env();
387        env.block.height = 1000;
388
389        let amount = claims
390            .claim_tokens(
391                deps.as_mut().storage,
392                &Addr::unchecked("addr"),
393                &env.block,
394                Some(Uint128::zero()),
395            )
396            .unwrap();
397
398        let saved_claims = claims
399            .0
400            .load(deps.as_mut().storage, &Addr::unchecked("addr"))
401            .unwrap();
402
403        assert_eq!(amount, Uint128::zero());
404        assert_eq!(saved_claims.len(), 2);
405        assert_eq!(saved_claims[0].amount, Uint128::from(TEST_AMOUNT));
406        assert_eq!(saved_claims[0].release_at, Expiration::AtHeight(10));
407        assert_eq!(saved_claims[1].amount, Uint128::from(TEST_AMOUNT + 100));
408        assert_eq!(saved_claims[1].release_at, Expiration::AtHeight(100));
409    }
410
411    #[test]
412    fn test_claim_tokens_with_cap_greater_than_pending_claims() {
413        let mut deps = mock_dependencies();
414        let claims = Claims::new("claims");
415
416        claims
417            .create_claim(
418                deps.as_mut().storage,
419                &Addr::unchecked("addr"),
420                TEST_AMOUNT.into(),
421                Expiration::AtHeight(10),
422            )
423            .unwrap();
424
425        claims
426            .create_claim(
427                deps.as_mut().storage,
428                &Addr::unchecked("addr"),
429                (TEST_AMOUNT + 100).into(),
430                Expiration::AtHeight(100),
431            )
432            .unwrap();
433
434        let mut env = mock_env();
435        env.block.height = 1000;
436
437        let amount = claims
438            .claim_tokens(
439                deps.as_mut().storage,
440                &Addr::unchecked("addr"),
441                &env.block,
442                Some(Uint128::from(2100u128)),
443            )
444            .unwrap();
445
446        let saved_claims = claims
447            .0
448            .load(deps.as_mut().storage, &Addr::unchecked("addr"))
449            .unwrap();
450
451        assert_eq!(amount, Uint128::from(TEST_AMOUNT + TEST_AMOUNT + 100));
452        assert_eq!(saved_claims.len(), 0);
453    }
454
455    #[test]
456    fn test_claim_tokens_with_cap_only_one_claim_released() {
457        let mut deps = mock_dependencies();
458        let claims = Claims::new("claims");
459
460        claims
461            .create_claim(
462                deps.as_mut().storage,
463                &Addr::unchecked("addr"),
464                (TEST_AMOUNT + 100).into(),
465                Expiration::AtHeight(10),
466            )
467            .unwrap();
468
469        claims
470            .create_claim(
471                deps.as_mut().storage,
472                &Addr::unchecked("addr"),
473                TEST_AMOUNT.into(),
474                Expiration::AtHeight(5),
475            )
476            .unwrap();
477
478        let mut env = mock_env();
479        env.block.height = 1000;
480        // the address has two claims and the first one can be released
481        let amount = claims
482            .claim_tokens(
483                deps.as_mut().storage,
484                &Addr::unchecked("addr"),
485                &env.block,
486                Some((TEST_AMOUNT + 50).into()),
487            )
488            .unwrap();
489        assert_eq!(amount, Uint128::from(TEST_AMOUNT));
490
491        let saved_claims = claims
492            .0
493            .load(deps.as_mut().storage, &Addr::unchecked("addr"))
494            .unwrap();
495        assert_eq!(saved_claims.len(), 1);
496        assert_eq!(saved_claims[0].amount, Uint128::from(TEST_AMOUNT + 100));
497        assert_eq!(saved_claims[0].release_at, Expiration::AtHeight(10));
498    }
499
500    #[test]
501    fn test_claim_tokens_with_cap_too_low_no_claims_released() {
502        let mut deps = mock_dependencies();
503        let claims = Claims::new("claims");
504
505        claims
506            .create_claim(
507                deps.as_mut().storage,
508                &Addr::unchecked("addr"),
509                (TEST_AMOUNT + 100).into(),
510                Expiration::AtHeight(10),
511            )
512            .unwrap();
513
514        claims
515            .create_claim(
516                deps.as_mut().storage,
517                &Addr::unchecked("addr"),
518                TEST_AMOUNT.into(),
519                Expiration::AtHeight(5),
520            )
521            .unwrap();
522
523        let mut env = mock_env();
524        env.block.height = 1000;
525        // the address has two claims and the first one can be released
526        let amount = claims
527            .claim_tokens(
528                deps.as_mut().storage,
529                &Addr::unchecked("addr"),
530                &env.block,
531                Some((TEST_AMOUNT - 50).into()),
532            )
533            .unwrap();
534        assert_eq!(amount, Uint128::zero());
535
536        let saved_claims = claims
537            .0
538            .load(deps.as_mut().storage, &Addr::unchecked("addr"))
539            .unwrap();
540        assert_eq!(saved_claims.len(), 2);
541        assert_eq!(saved_claims[0].amount, Uint128::from(TEST_AMOUNT + 100));
542        assert_eq!(saved_claims[0].release_at, Expiration::AtHeight(10));
543        assert_eq!(saved_claims[1].amount, Uint128::from(TEST_AMOUNT));
544        assert_eq!(saved_claims[1].release_at, Expiration::AtHeight(5));
545    }
546
547    #[test]
548    fn test_query_claims_returns_correct_claims() {
549        let mut deps = mock_dependencies();
550        let claims = Claims::new("claims");
551
552        claims
553            .create_claim(
554                deps.as_mut().storage,
555                &Addr::unchecked("addr"),
556                (TEST_AMOUNT + 100).into(),
557                Expiration::AtHeight(10),
558            )
559            .unwrap();
560
561        let queried_claims = claims
562            .query_claims(deps.as_ref(), &Addr::unchecked("addr"))
563            .unwrap();
564        let saved_claims = claims
565            .0
566            .load(deps.as_mut().storage, &Addr::unchecked("addr"))
567            .unwrap();
568        assert_eq!(queried_claims.claims, saved_claims);
569    }
570
571    #[test]
572    fn test_query_claims_returns_empty_for_non_existent_user() {
573        let mut deps = mock_dependencies();
574        let claims = Claims::new("claims");
575
576        claims
577            .create_claim(
578                deps.as_mut().storage,
579                &Addr::unchecked("addr"),
580                (TEST_AMOUNT + 100).into(),
581                Expiration::AtHeight(10),
582            )
583            .unwrap();
584
585        let queried_claims = claims
586            .query_claims(deps.as_ref(), &Addr::unchecked("addr2"))
587            .unwrap();
588
589        assert_eq!(queried_claims.claims.len(), 0);
590    }
591}
cw_controllers/claim.rs

cw_controllers/
claim.rs
