Final implementation of auto returning stake on endblock

contracts/tg4-stake/Cargo.toml

@@ -31,6 +31,7 @@ cosmwasm-std = { version = "0.16.0" }
 schemars = "0.8.1"
 serde = { version = "1.0.103", default-features = false, features = ["derive"] }
 thiserror = { version = "1.0.21" }
+itertools = "0.10"
 
 [dev-dependencies]
 cosmwasm-schema = { version = "0.16.0" }

contracts/tg4-stake/src/claim.rs

@@ -1,6 +1,7 @@
 // Copied from cw-plus repository: https://github.com/CosmWasm/cw-plus/tree/main/packages/controllers
 // Original file distributed on Apache license
 
+use itertools::Itertools;
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
@@ -118,21 +119,23 @@ impl<'a> Claims<'a> {
                 Err(_) => true,
             });
 
-        let (claims, amount) = self.filter_claims(claims, cap.map(u128::from), None)?;
+        let claims = self.filter_claims(claims, cap.map(u128::from), None)?;
+        let amount = claims.iter().map(|claim| claim.amount).sum();
+
         self.release_claims(storage, claims)?;
 
-        Ok(amount.into())
+        Ok(amount)
     }
 
     /// This iterates over all mature claims of any addresses, and removes them. Up to `limit`
     /// claims would be processed, starting from the oldest. It removes the finished claims and
-    /// returns the total amount of tokens to be released.
+    /// returns vector of pairs: `(addr, amount)`, representing amount of tokens to be released to particular addresses
     pub fn claim_expired(
         &self,
         storage: &mut dyn Storage,
         block: &BlockInfo,
         limit: impl Into<Option<u64>>,
-    ) -> StdResult<Uint128> {
+    ) -> StdResult<Vec<(Addr, Uint128)>> {
         let claims = self
             .claims
             .idx
@@ -148,46 +151,49 @@ impl<'a> Claims<'a> {
                 Order::Ascending,
             );
 
-        let (claims, amount) = self.filter_claims(claims, None, limit.into())?;
+        let mut claims = self.filter_claims(claims, None, limit.into())?;
+        claims.sort_by_key(|claim| claim.addr.clone());
+
+        let releases = claims
+            .iter()
+            // TODO: use `slice::group_by` in place of `Itertools::group_by` when `slice_group_by`
+            // is stabilized [https://github.com/rust-lang/rust/issues/80552]
+            .group_by(|claim| &claim.addr)
+            .into_iter()
+            .map(|(addr, group)| (addr.clone(), group.map(|claim| claim.amount).sum()))
+            .collect();
+
         self.release_claims(storage, claims)?;
 
-        Ok(amount.into())
+        Ok(releases)
     }
 
-    /// Processes claims filtering those which are to be released. Returns vector of keys of claims
-    /// to be released, and accumulated amount of tokens to be released
+    /// Processes claims filtering those which are to be released. Returns vector of claims to be
+    /// released
     fn filter_claims(
         &self,
         claims: impl IntoIterator<Item = StdResult<(Vec<u8>, Claim)>>,
         cap: Option<u128>,
         limit: Option<u64>,
-    ) -> StdResult<(Vec<(Addr, ExpirationKey)>, u128)> {
-        // will be filled at the final step of processing
-        let mut amount = 0;
-
+    ) -> StdResult<Vec<Claim>> {
         let claims = claims
             .into_iter()
-            // calculate sum for claims up to this one and pair it with index
+            // calculate sum for claims up to this one for cap filtering
             .scan(0u128, |sum, claim| match claim {
                 Ok((_, claim)) => {
                     *sum += u128::from(claim.amount);
-                    let idx = (claim.addr, claim.release_at.into());
-                    Some(Ok((idx, *sum)))
+                    Some(Ok((*sum, claim)))
                 }
                 Err(err) => Some(Err(err)),
             })
             // stop when sum exceeds limit
             .take_while(|claim| match (cap, claim) {
-                (Some(cap), Ok((_, sum))) => cap <= *sum,
+                (Some(cap), Ok((sum, _))) => cap <= *sum,
                 _ => true,
             })
-            // now only proper claims as in iterator, so just map and store amount - the last one
-            // would be the one stored
+            // now only proper claims as in iterator, so just map them back to claim
             .map(|claim| match claim {
-                Ok((idx, claim)) => {
-                    amount = claim;
-                    Ok(idx)
-                }
+                Ok((_, claim)) => Ok(claim),
                 Err(err) => Err(err),
             });
 
@@ -201,17 +207,18 @@ impl<'a> Claims<'a> {
             claims.collect::<StdResult<_>>()
         }?;
 
-        Ok((claims, amount))
+        Ok(claims)
     }
 
     /// Releases given claims by removing them from storage
     fn release_claims(
         &self,
         storage: &mut dyn Storage,
-        claims: impl IntoIterator<Item = (Addr, ExpirationKey)>,
+        claims: impl IntoIterator<Item = Claim>,
     ) -> StdResult<()> {
         for claim in claims {
-            self.claims.remove(storage, claim)?;
+            self.claims
+                .remove(storage, (claim.addr, claim.release_at.into()))?;
         }
 
         Ok(())

contracts/tg4-stake/src/contract.rs

@@ -268,13 +268,12 @@ pub fn execute_claim(
     }
 
     let config = CONFIG.load(deps.storage)?;
-    let amount;
-    match &config.denom {
-        Denom::Native(denom) => amount = coins(release.u128(), denom),
+    let amount = match &config.denom {
+        Denom::Native(denom) => coins(release.into(), denom),
         Denom::Cw20(_addr) => {
-            unimplemented!("The CW20 coins release functionality is in progress")
+            return Err(ContractError::Cw20CoinsRelease {});
         }
-    }
+    };
 
     let res = Response::new()
         .add_attribute("action", "claim")
@@ -284,6 +283,7 @@ pub fn execute_claim(
             to_address: info.sender.into(),
             amount,
         });
+
     Ok(res)
 }
 
@@ -297,9 +297,10 @@ fn coins_to_string(coins: &[Coin]) -> String {
 }
 
 #[cfg_attr(not(feature = "library"), entry_point)]
-pub fn sudo(deps: DepsMut, _env: Env, msg: TgradeSudoMsg) -> Result<Response, ContractError> {
+pub fn sudo(deps: DepsMut, env: Env, msg: TgradeSudoMsg) -> Result<Response, ContractError> {
     match msg {
         TgradeSudoMsg::PrivilegeChange(PrivilegeChangeMsg::Promoted {}) => privilege_promote(deps),
+        TgradeSudoMsg::EndBlock {} => end_block(deps, env),
         _ => Err(ContractError::UnknownSudoMsg {}),
     }
 }
@@ -315,9 +316,41 @@ fn privilege_promote(deps: DepsMut) -> Result<Response, ContractError> {
     }
 }
 
-fn _end_block(_deps: DepsMut) -> Result<Response, ContractError> {
-    // TODO: Auto return claims
-    todo!()
+fn end_block(deps: DepsMut, env: Env) -> Result<Response, ContractError> {
+    let mut resp = Response::new();
+
+    let config = CONFIG.load(deps.storage)?;
+    if config.auto_return_limit > 0 {
+        let sub_msgs = release_expired_claims(deps, env, config)?;
+        resp = resp.add_submessages(sub_msgs);
+    }
+
+    Ok(resp)
+}
+
+fn release_expired_claims(
+    deps: DepsMut,
+    env: Env,
+    config: Config,
+) -> Result<Vec<SubMsg>, ContractError> {
+    let releases = claims().claim_expired(deps.storage, &env.block, config.auto_return_limit)?;
+
+    releases
+        .into_iter()
+        .map(|(addr, amount)| {
+            let amount = match &config.denom {
+                Denom::Native(denom) => coins(amount.into(), denom),
+                Denom::Cw20(_) => {
+                    return Err(ContractError::Cw20CoinsRelease {});
+                }
+            };
+
+            Ok(SubMsg::new(BankMsg::Send {
+                to_address: addr.into(),
+                amount,
+            }))
+        })
+        .collect()
 }
 
 #[cfg_attr(not(feature = "library"), entry_point)]

contracts/tg4-stake/src/error.rs

@@ -40,4 +40,7 @@ pub enum ContractError {
 
     #[error("Unrecognized sudo message")]
     UnknownSudoMsg {},
+
+    #[error("Cw20 coins release functionality is in progress")]
+    Cw20CoinsRelease {},
 }