pallet_torus0/
stake.rs

1use polkadot_sdk::{
2    frame_support::{dispatch::DispatchResult, ensure, traits::NamedReservableCurrency},
3    sp_std::{collections::btree_map::BTreeMap, vec::Vec},
4};
5
6use crate::{AccountIdOf, BalanceOf};
7use crate::{StakedBy, StakingTo, TotalStake, agent};
8
9pub const STAKE_IDENTIFIER: &[u8; 8] = b"torstake";
10
11/// Stakes `amount` tokens from `staker` to `staked` by withdrawing the tokens
12/// and adding them to the [`crate::StakingTo`] and [`crate::StakedBy`] maps.
13pub fn add_stake<T: crate::Config>(
14    staker: AccountIdOf<T>,
15    staked: AccountIdOf<T>,
16    amount: BalanceOf<T>,
17) -> DispatchResult {
18    ensure!(
19        agent::exists::<T>(&staked),
20        crate::Error::<T>::AgentDoesNotExist
21    );
22
23    T::Currency::reserve_named(STAKE_IDENTIFIER, &staker, amount)
24        .map_err(|_| crate::Error::<T>::NotEnoughBalanceToStake)?;
25
26    StakedBy::<T>::mutate(&staked, &staker, |stake| {
27        *stake = Some(stake.unwrap_or(0).saturating_add(amount))
28    });
29    StakingTo::<T>::mutate(&staker, &staked, |stake| {
30        *stake = Some(stake.unwrap_or(0).saturating_add(amount))
31    });
32
33    TotalStake::<T>::mutate(|total_stake| *total_stake = total_stake.saturating_add(amount));
34
35    crate::Pallet::<T>::deposit_event(crate::Event::<T>::StakeAdded(staker, staked, amount));
36
37    Ok(())
38}
39
40/// Withdraws stake from an agent and gives it back to the staker.
41pub fn remove_stake<T: crate::Config>(
42    staker: AccountIdOf<T>,
43    staked: AccountIdOf<T>,
44    amount: BalanceOf<T>,
45) -> DispatchResult {
46    ensure!(
47        agent::exists::<T>(&staked),
48        crate::Error::<T>::AgentDoesNotExist
49    );
50
51    ensure!(
52        StakingTo::<T>::get(&staker, &staked).unwrap_or(0) >= amount,
53        crate::Error::<T>::NotEnoughStakeToWithdraw
54    );
55
56    remove_stake0::<T>(staker, staked, amount, true);
57
58    Ok(())
59}
60
61fn remove_stake0<T: crate::Config>(
62    staker: AccountIdOf<T>,
63    staked: AccountIdOf<T>,
64    amount: BalanceOf<T>,
65    keep: bool,
66) {
67    let Some(stake) = StakingTo::<T>::get(&staker, &staked) else {
68        return;
69    };
70
71    let amount = stake.min(amount);
72    let new_stake = stake.saturating_sub(amount);
73    let new_stake = if keep || new_stake > 0 {
74        Some(new_stake)
75    } else {
76        None
77    };
78
79    StakingTo::<T>::set(&staker, &staked, new_stake);
80    StakedBy::<T>::set(&staked, &staker, new_stake);
81    TotalStake::<T>::mutate(|total_stake| *total_stake = total_stake.saturating_sub(amount));
82
83    T::Currency::unreserve_named(STAKE_IDENTIFIER, &staker, amount);
84
85    crate::Pallet::<T>::deposit_event(crate::Event::<T>::StakeRemoved(staker, staked, amount));
86}
87
88/// Transfers stake from an account to another (see [`remove_stake`],
89/// [`add_stake`]).
90pub fn transfer_stake<T: crate::Config>(
91    staker: AccountIdOf<T>,
92    old_staked: AccountIdOf<T>,
93    new_staked: AccountIdOf<T>,
94    amount: BalanceOf<T>,
95) -> DispatchResult {
96    remove_stake::<T>(staker.clone(), old_staked, amount)?;
97    add_stake::<T>(staker, new_staked, amount)?;
98    Ok(())
99}
100
101/// Usually called when de-registering an agent, removes all stakes on a given
102/// key.
103pub(crate) fn clear_key<T: crate::Config>(key: &AccountIdOf<T>) -> DispatchResult {
104    let stakes: Vec<_> = StakingTo::<T>::iter().collect();
105    for (staker, staked, amount) in stakes {
106        if &staker == key || &staked == key {
107            remove_stake0::<T>(staker, staked, amount, false);
108        }
109    }
110
111    Ok(())
112}
113
114#[inline]
115pub fn sum_staking_to<T: crate::Config>(staker: &AccountIdOf<T>) -> BalanceOf<T> {
116    StakingTo::<T>::iter_prefix_values(staker).sum()
117}
118
119#[inline]
120pub fn get_staking_to_vector<T: crate::Config>(
121    staker: &AccountIdOf<T>,
122) -> BTreeMap<T::AccountId, BalanceOf<T>> {
123    StakingTo::<T>::iter_prefix(staker).collect()
124}
125
126#[inline]
127pub fn get_staked_by_vector<T: crate::Config>(
128    staked: &AccountIdOf<T>,
129) -> Vec<(T::AccountId, BalanceOf<T>)> {
130    StakedBy::<T>::iter_prefix(staked).collect()
131}
132
133#[inline]
134pub fn sum_staked_by<T: crate::Config>(staked: &AccountIdOf<T>) -> BalanceOf<T> {
135    StakedBy::<T>::iter_prefix_values(staked).sum()
136}
pallet_torus0/stake.rs

pallet_torus0/
stake.rs
