cw20_stake/
math.rs

1use std::{convert::TryInto, ops::Div};
2
3use cosmwasm_std::{Uint128, Uint256};
4
5/// Computes the amount to add to an address' staked balance when
6/// staking.
7///
8/// # Arguments
9///
10/// * `staked_total` - The number of tokens that have been staked.
11/// * `balance` - The number of tokens the contract has (staked_total + rewards).
12/// * `sent` - The number of tokens the user has sent to be staked.
13pub(crate) fn amount_to_stake(staked_total: Uint128, balance: Uint128, sent: Uint128) -> Uint128 {
14    if staked_total.is_zero() || balance.is_zero() {
15        sent
16    } else {
17        staked_total
18            .full_mul(sent)
19            .div(Uint256::from(balance))
20            .try_into()
21            .unwrap() // balance := staked_total + rewards
22                      // => balance >= staked_total
23                      // => staked_total / balance <= 1
24                      // => staked_total * sent / balance <= sent
25                      // => we can safely unwrap here as sent fits into a u128 by construction.
26    }
27}
28
29/// Computes the number of tokens to return to an address when
30/// claiming.
31///
32/// # Arguments
33///
34/// * `staked_total` - The number of tokens that have been staked.
35/// * `balance` - The number of tokens the contract has (staked_total + rewards).
36/// * `ask` - The number of tokens being claimed.
37///
38/// # Invariants
39///
40/// These must be checked by the caller. If checked, this function is
41/// guarenteed not to panic.
42///
43/// 1. staked_total != 0.
44/// 2. ask + balance <= 2^128
45/// 3. ask <= staked_total
46///
47/// For information on the panic conditions for math, see:
48/// <https://rust-lang.github.io/rfcs/0560-integer-overflow.html>
49pub(crate) fn amount_to_claim(staked_total: Uint128, balance: Uint128, ask: Uint128) -> Uint128 {
50    // we know that:
51    //
52    // 1. cw20's max supply is 2^128
53    // 2. balance := staked_total + rewards
54    //
55    // for non-malicious inputs:
56    //
57    // 3. 1 => ask + balance <= 2^128
58    // 4. ask <= staked_total
59    // 5. staked_total != 0
60    // 6. 4 => ask / staked_total <= 1
61    // 7. 3 => balance <= 2^128
62    // 8. 6 + 7 => ask / staked_total * balance <= 2^128
63    //
64    // which, as addition and division are communative, proves that
65    // ask * balance / staked_total will fit into a 128 bit integer.
66    ask.full_mul(balance)
67        .div(Uint256::from(staked_total))
68        .try_into()
69        .unwrap()
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75
76    #[test]
77    fn test_amount_to_stake_no_overflow() {
78        let sent = Uint128::new(2);
79        let balance = Uint128::MAX - sent;
80
81        let overflows_naively = sent.checked_mul(balance).is_err();
82        assert!(overflows_naively);
83
84        // will panic and fail the test if we've done this wrong.
85        amount_to_stake(balance, balance, sent);
86    }
87
88    #[test]
89    fn test_amount_to_stake_with_zeros() {
90        let sent = Uint128::new(42);
91        let balance = Uint128::zero();
92        let amount = amount_to_stake(balance, balance, sent);
93        assert_eq!(amount, sent);
94    }
95
96    #[test]
97    fn test_amount_to_claim_no_overflow() {
98        let ask = Uint128::new(2);
99        let balance = Uint128::MAX - ask;
100
101        let overflows_naively = ask.checked_mul(balance).is_err();
102        assert!(overflows_naively);
103
104        amount_to_claim(balance, balance, ask);
105    }
106
107    // check that our invariants are indeed invariants.
108
109    #[test]
110    #[should_panic(expected = "attempt to divide by zero")]
111    fn test_amount_to_claim_invariant_one() {
112        let ask = Uint128::new(2);
113        let balance = Uint128::zero();
114
115        amount_to_claim(balance, balance, ask);
116    }
117
118    #[test]
119    #[should_panic(expected = "ConversionOverflowError")]
120    fn test_amount_to_claim_invariant_two() {
121        // Could end up in a situation like this if there are a lot of
122        // rewards, but very few staked tokens.
123        let ask = Uint128::new(2);
124        let balance = Uint128::MAX;
125        let staked_total = Uint128::new(1);
126
127        amount_to_claim(staked_total, balance, ask);
128    }
129
130    #[test]
131    #[should_panic(expected = "ConversionOverflowError")]
132    fn test_amount_to_claim_invariant_three() {
133        let ask = Uint128::new(2);
134        let balance = Uint128::MAX;
135        let staked_total = Uint128::new(1);
136
137        amount_to_claim(staked_total, balance, ask);
138    }
139}
cw20_stake/math.rs

cw20_stake/
math.rs
