pallet_torus0/
stake.rs

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

pallet_torus0/
stake.rs
