Add reconciliation mechanism

contracts/hub/Cargo.toml

@@ -20,5 +20,9 @@ cw-storage-plus = "0.9"
 steak = { path = "../../packages/steak" }
 terra-cosmwasm = "2.2"
 
+# These two are only used in legacy code; to be removed later
+schemars = "0.8.1"
+serde = { version = "1.0.103", default-features = false, features = ["derive"] }
+
 [dev-dependencies]
 serde = { version = "1.0.103", default-features = false, features = ["derive"] }

contracts/hub/src/contract.rs

@@ -8,6 +8,7 @@ use terra_cosmwasm::TerraMsgWrapper;
 use steak::hub::{CallbackMsg, ExecuteMsg, InstantiateMsg, MigrateMsg, QueryMsg, ReceiveMsg};
 
 use crate::helpers::{parse_received_fund, unwrap_reply};
+use crate::legacy::migrate_batches;
 use crate::state::State;
 use crate::{execute, queries};
 
@@ -59,6 +60,7 @@ pub fn execute(
         ExecuteMsg::AcceptOwnership {} => execute::accept_ownership(deps, info.sender),
         ExecuteMsg::Harvest {} => execute::harvest(deps, env),
         ExecuteMsg::Rebalance {} => execute::rebalance(deps, env),
+        ExecuteMsg::Reconcile {} => execute::reconcile(deps, env),
         ExecuteMsg::SubmitBatch {} => execute::submit_batch(deps, env),
         ExecuteMsg::Callback(callback_msg) => callback(deps, env, info, callback_msg),
     }
@@ -159,6 +161,10 @@ pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> {
 }
 
 #[entry_point]
-pub fn migrate(_deps: DepsMut, _env: Env, _msg: MigrateMsg) -> StdResult<Response> {
-    Ok(Response::new())
+pub fn migrate(deps: DepsMut, env: Env, _msg: MigrateMsg) -> StdResult<Response<TerraMsgWrapper>> {
+    let event = migrate_batches(deps.storage)?;
+
+    let res = execute::reconcile(deps, env)?;
+
+    Ok(res.add_event(event))
 }

contracts/hub/src/execute.rs

@@ -14,7 +14,7 @@ use steak::hub::{Batch, CallbackMsg, ExecuteMsg, InstantiateMsg, PendingBatch, U
 use crate::helpers::{query_cw20_total_supply, query_delegations, query_delegation};
 use crate::math::{
     compute_mint_amount, compute_redelegations_for_rebalancing, compute_redelegations_for_removal,
-    compute_unbond_amount, compute_undelegations,
+    compute_unbond_amount, compute_undelegations, reconcile_batches,
 };
 use crate::state::State;
 use crate::types::{Coins, Delegation};
@@ -410,6 +410,7 @@ pub fn submit_batch(deps: DepsMut, env: Env) -> StdResult<Response<TerraMsgWrapp
         pending_batch.id.into(),
         &Batch {
             id: pending_batch.id,
+            reconciled: false,
             total_shares: pending_batch.usteak_to_burn,
             uluna_unclaimed: uluna_to_unbond,
             est_unbond_end_time: current_time + unbond_period,
@@ -453,6 +454,70 @@ pub fn submit_batch(deps: DepsMut, env: Env) -> StdResult<Response<TerraMsgWrapp
         .add_attribute("action", "steakhub/unbond"))
 }
 
+pub fn reconcile(deps: DepsMut, env: Env) -> StdResult<Response<TerraMsgWrapper>> {
+    let state = State::default();
+    let current_time = env.block.time.seconds();
+
+    // Load batches that have not been reconciled
+    let all_batches = state
+        .previous_batches
+        .idx
+        .reconciled
+        .prefix(false.into())
+        .range(deps.storage, None, None, Order::Ascending)
+        .map(|item| {
+            let (_, v) = item?;
+            Ok(v)
+        })
+        .collect::<StdResult<Vec<_>>>()?;
+
+    let mut batches = all_batches
+        .into_iter()
+        .filter(|b| current_time > b.est_unbond_end_time)
+        .collect::<Vec<_>>();
+
+    let uluna_expected_received: Uint128 = batches
+        .iter()
+        .map(|b| b.uluna_unclaimed)
+        .sum();
+
+    if uluna_expected_received.is_zero() {
+        return Ok(Response::new());
+    }
+
+    let unlocked_coins = state.unlocked_coins.load(deps.storage)?;
+    let uluna_expected_unlocked = Coins(unlocked_coins).find("uluna").amount;
+
+    let uluna_expected = uluna_expected_received + uluna_expected_unlocked;
+    let uluna_actual = deps.querier.query_balance(&env.contract.address, "uluna")?.amount;
+
+    if uluna_actual >= uluna_expected {
+        return Ok(Response::new());
+    }
+
+    let uluna_to_deduct = uluna_expected - uluna_actual;
+
+    reconcile_batches(&mut batches, uluna_to_deduct);
+
+    for batch in &batches {
+        state.previous_batches.save(deps.storage, batch.id.into(), batch)?;
+    }
+
+    let ids = batches
+        .iter()
+        .map(|b| b.id.to_string())
+        .collect::<Vec<_>>()
+        .join(",");
+
+    let event = Event::new("steakhub/reconciled")
+        .add_attribute("ids", ids)
+        .add_attribute("uluna_deducted", uluna_to_deduct.to_string());
+
+    Ok(Response::new()
+        .add_event(event)
+        .add_attribute("action", "steakhub/reconcile"))
+}
+
 pub fn withdraw_unbonded(
     deps: DepsMut,
     env: Env,
@@ -477,11 +542,17 @@ pub fn withdraw_unbonded(
         })
         .collect::<StdResult<Vec<_>>>()?;
 
+    // NOTE: Luna in the following batches are withdrawn it the batch:
+    // - is a _previous_ batch, not a _pending_ batch
+    // - is reconciled
+    // - has finished unbonding
+    // If not sure whether the batches have been reconciled, the user should first invoke `ExecuteMsg::Reconcile`
+    // before withdrawing.
     let mut total_uluna_to_refund = Uint128::zero();
     let mut ids: Vec<String> = vec![];
     for request in &requests {
         if let Ok(mut batch) = state.previous_batches.load(deps.storage, request.id.into()) {
-            if batch.est_unbond_end_time < current_time {
+            if batch.reconciled && batch.est_unbond_end_time < current_time {
                 let uluna_to_refund = batch
                     .uluna_unclaimed
                     .multiply_ratio(request.shares, batch.total_shares);
@@ -493,7 +564,7 @@ pub fn withdraw_unbonded(
                 batch.uluna_unclaimed -= uluna_to_refund;
 
                 if batch.total_shares.is_zero() {
-                    state.previous_batches.remove(deps.storage, request.id.into());
+                    state.previous_batches.remove(deps.storage, request.id.into())?;
                 } else {
                     state.previous_batches.save(deps.storage, batch.id.into(), &batch)?;
                 }
@@ -543,8 +614,14 @@ pub fn rebalance(deps: DepsMut, env: Env) -> StdResult<Response<TerraMsgWrapper>
         .map(|rd| SubMsg::reply_on_success(rd.to_cosmos_msg(), 2))
         .collect::<Vec<_>>();
 
+    let amount: u128 = new_redelegations.iter().map(|rd| rd.amount).sum();
+
+    let event = Event::new("steakhub/rebalanced")
+        .add_attribute("uluna_moved", amount.to_string());
+
     Ok(Response::new()
         .add_submessages(redelegate_submsgs)
+        .add_event(event)
         .add_attribute("action", "steakhub/rebalance"))
 }
 

contracts/hub/src/legacy.rs

@@ -0,0 +1,129 @@
+use cosmwasm_std::{Storage, Order, StdResult, Uint128, Event};
+use cw_storage_plus::{Map, U64Key};
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
+
+use steak::hub::Batch;
+
+use crate::state::State;
+
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
+pub struct LegacyBatch {
+    pub id: u64,
+    pub total_shares: Uint128,
+    pub uluna_unclaimed: Uint128,
+    pub est_unbond_end_time: u64,
+}
+
+const LEGACY_BATCHES: Map<U64Key, LegacyBatch> = Map::new("previous_batches");
+
+pub(crate) fn migrate_batches(storage: &mut dyn Storage) -> StdResult<Event> {
+    let state = State::default();
+
+    // Find all previous batches
+    let legacy_batches = LEGACY_BATCHES
+        .range(storage, None, None, Order::Ascending)
+        .map(|item| {
+            let (_, v) = item?;
+            Ok(v)
+        })
+        .collect::<StdResult<Vec<_>>>()?;
+
+    // Cast legacy data to the new type
+    let batches = legacy_batches
+        .iter()
+        .map(|lb| Batch {
+            id: lb.id,
+            reconciled: false,
+            total_shares: lb.total_shares,
+            uluna_unclaimed: lb.uluna_unclaimed,
+            est_unbond_end_time: lb.est_unbond_end_time,
+        })
+        .collect::<Vec<_>>();
+
+    // Delete the legacy data
+    legacy_batches
+        .iter()
+        .for_each(|lb| {
+            LEGACY_BATCHES.remove(storage, lb.id.into())
+        });
+
+    // Save the new type data
+    // We use unwrap here, which is undesired, but it's ok with me since this code will only be in
+    // the contract temporarily
+    batches
+        .iter()
+        .for_each(|b| {
+            state.previous_batches.save(storage, b.id.into(), b).unwrap()
+        });
+
+    let ids = batches.iter().map(|b| b.id.to_string()).collect::<Vec<_>>();
+
+    Ok(Event::new("steakhub/batches_migrated")
+        .add_attribute("ids", ids.join(",")))
+}
+
+#[cfg(test)]
+mod tests {
+    use cosmwasm_std::testing::mock_dependencies;
+
+    use super::*;
+    use crate::queries;
+
+    #[test]
+    fn migrating_batches() {
+        let mut deps = mock_dependencies(&[]);
+
+        let legacy_batches = vec![
+            LegacyBatch {
+                id: 1,
+                total_shares: Uint128::new(123),
+                uluna_unclaimed: Uint128::new(678),
+                est_unbond_end_time: 10000,
+            },
+            LegacyBatch {
+                id: 2,
+                total_shares: Uint128::new(234),
+                uluna_unclaimed: Uint128::new(789),
+                est_unbond_end_time: 15000,
+            },
+            LegacyBatch {
+                id: 3,
+                total_shares: Uint128::new(345),
+                uluna_unclaimed: Uint128::new(890),
+                est_unbond_end_time: 20000,
+            },
+            LegacyBatch {
+                id: 4,
+                total_shares: Uint128::new(456),
+                uluna_unclaimed: Uint128::new(999),
+                est_unbond_end_time: 25000,
+            },
+        ];
+
+        for legacy_batch in &legacy_batches {
+            LEGACY_BATCHES.save(deps.as_mut().storage, legacy_batch.id.into(), legacy_batch).unwrap();
+        }
+
+        let event = migrate_batches(deps.as_mut().storage).unwrap();
+        assert_eq!(
+            event,
+            Event::new("steakhub/batches_migrated").add_attribute("ids", "1,2,3,4")
+        );
+
+        let batches = queries::previous_batches(deps.as_ref(), None, None).unwrap();
+
+        let expected = legacy_batches
+            .iter()
+            .map(|lb| Batch {
+                id: lb.id,
+                reconciled: false,
+                total_shares: lb.total_shares,
+                uluna_unclaimed: lb.uluna_unclaimed,
+                est_unbond_end_time: lb.est_unbond_end_time,
+            })
+            .collect::<Vec<_>>();
+
+        assert_eq!(batches, expected);
+    }
+}

contracts/hub/src/lib.rs

@@ -10,3 +10,6 @@ pub mod types;
 
 #[cfg(test)]
 mod testing;
+
+// Legacy code; only used in migrations
+mod legacy;

contracts/hub/src/math.rs

@@ -2,6 +2,8 @@ use std::{cmp, cmp::Ordering};
 
 use cosmwasm_std::Uint128;
 
+use steak::hub::Batch;
+
 use crate::types::{Delegation, Redelegation, Undelegation};
 
 //--------------------------------------------------------------------------------------------------
@@ -192,3 +194,27 @@ pub(crate) fn compute_redelegations_for_rebalancing(
 
     new_redelegations
 }
+
+//--------------------------------------------------------------------------------------------------
+// Batch logics
+//--------------------------------------------------------------------------------------------------
+
+/// If the received uluna amount after the unbonding period is less than expected, e.g. due to rounding
+/// error or the validator(s) being slashed, then deduct the difference in amount evenly from each
+/// unreconciled batch.
+///
+/// The idea of "reconciling" is based on Stader's implementation:
+/// https://github.com/stader-labs/stader-liquid-token/blob/v0.2.1/contracts/staking/src/contract.rs#L968-L1048
+pub(crate) fn reconcile_batches(batches: &mut [Batch], uluna_to_deduct: Uint128) {
+    let batch_count = batches.len() as u128;
+    let uluna_per_batch = uluna_to_deduct.u128() / batch_count;
+    let remainder = uluna_to_deduct.u128() % batch_count;
+
+    for (i, batch) in batches.iter_mut().enumerate() {
+        let remainder_for_batch: u128 = if (i + 1) as u128 <= remainder { 1 } else { 0 };
+        let uluna_for_batch = uluna_per_batch + remainder_for_batch;
+
+        batch.uluna_unclaimed -= Uint128::new(uluna_for_batch);
+        batch.reconciled = true;
+    }
+}