Fix staked amount too high after unstake issue (#1270)

* Fix staked amount too high after unstake issue

* Additonal UT

pallets/dapp-staking-v3/src/test/testing_utils.rs

@@ -727,15 +727,16 @@ pub(crate) fn assert_unstake(
     let unstaked_amount_era_pairs =
         pre_staker_info
             .clone()
-            .unstake(expected_amount, unstake_period, unstake_subperiod);
+            .unstake(expected_amount, unstake_era, unstake_subperiod);
     assert!(unstaked_amount_era_pairs.len() <= 2 && unstaked_amount_era_pairs.len() > 0);
-    {
-        let (last_unstake_era, last_unstake_amount) = unstaked_amount_era_pairs
-            .last()
-            .expect("Has to exist due to success of previous check");
-        assert_eq!(*last_unstake_era, unstake_era.max(pre_staker_info.era()));
-        assert_eq!(*last_unstake_amount, expected_amount);
-    }
+
+    // If unstake from next era exists, it must exactly match the expected unstake amount.
+    unstaked_amount_era_pairs
+        .iter()
+        .filter(|(era, _)| *era > unstake_era)
+        .for_each(|(_, amount)| {
+            assert_eq!(*amount, expected_amount);
+        });
 
     // 3. verify contract stake
     // =========================
@@ -753,21 +754,31 @@ pub(crate) fn assert_unstake(
         "Staked amount must decreased by the 'amount'"
     );
 
-    // Ensure staked amounts are updated as expected, unless it's full unstake.
-    if !is_full_unstake {
-        for (unstake_era_iter, unstake_amount_iter) in unstaked_amount_era_pairs {
-            assert_eq!(
-                post_contract_stake
-                    .get(unstake_era_iter, unstake_period)
-                    .expect("Must exist.")
-                    .total(),
-                pre_contract_stake
-                    .get(unstake_era_iter, unstake_period)
-                    .expect("Must exist")
-                    .total()
-                    - unstake_amount_iter
-            );
-        }
+    // A generic check, comparing what was received in the (era, amount) pairs and the impact it had on the contract stake.
+    for (unstake_era_iter, unstake_amount_iter) in unstaked_amount_era_pairs {
+        assert_eq!(
+            post_contract_stake
+                .get(unstake_era_iter, unstake_period)
+                .unwrap_or_default() // it's possible that full unstake cleared the entry
+                .total(),
+            pre_contract_stake
+                .get(unstake_era_iter, unstake_period)
+                .expect("Must exist")
+                .total()
+                - unstake_amount_iter
+        );
+    }
+
+    // More precise check, independent of the generic check above.
+    // If next era entry exists, it must be reduced by the unstake amount, nothing less.
+    if let Some(entry) = pre_contract_stake.get(unstake_era + 1, unstake_period) {
+        assert_eq!(
+            post_contract_stake
+                .get(unstake_era + 1, unstake_period)
+                .unwrap_or_default()
+                .total(),
+            entry.total() - expected_amount
+        );
     }
 
     // 4. verify era info

pallets/dapp-staking-v3/src/test/tests.rs

@@ -3242,3 +3242,47 @@ fn claim_dapp_reward_with_rank() {
         }));
     })
 }
+
+#[test]
+fn unstake_correctly_reduces_future_contract_stake() {
+    ExtBuilder::build().execute_with(|| {
+        // 0. Register smart contract, lock&stake some amount with staker 1 during the voting subperiod
+        let smart_contract = MockSmartContract::wasm(1 as AccountId);
+        assert_register(1, &smart_contract);
+
+        let (staker_1, amount_1) = (1, 29);
+        assert_lock(staker_1, amount_1);
+        assert_stake(staker_1, &smart_contract, amount_1);
+
+        // 1. Advance to the build&earn subperiod, stake some amount with staker 2
+        advance_to_next_era();
+        let (staker_2, amount_2) = (2, 11);
+        assert_lock(staker_2, amount_2);
+        assert_stake(staker_2, &smart_contract, amount_2);
+
+        // 2. Advance a few eras, creating a gap but remaining within the same period.
+        //    Claim all rewards for staker 1.
+        //    Lock & stake some amount with staker 3.
+        advance_to_era(ActiveProtocolState::<Test>::get().era + 3);
+        assert_eq!(
+            ActiveProtocolState::<Test>::get().period_number(),
+            1,
+            "Sanity check."
+        );
+        for _ in 0..required_number_of_reward_claims(staker_1) {
+            assert_claim_staker_rewards(staker_1);
+        }
+
+        // This ensures contract stake entry is aligned to the current era, and future entry refers to the era after this one.
+        //
+        // This is important to reproduce an issue where the (era, amount) pairs returned by the `unstake` function don't correctly
+        // cover the next era.
+        let (staker_3, amount_3) = (3, 13);
+        assert_lock(staker_3, amount_3);
+        assert_stake(staker_3, &smart_contract, amount_3);
+
+        // 3. Unstake from staker 1, and ensure the future stake is reduced.
+        //    Unstake amount should be slightly higher than the 2nd stake amount to ensure whole b&e stake amount is removed.
+        assert_unstake(staker_1, &smart_contract, amount_2 + 3);
+    })
+}

pallets/dapp-staking-v3/src/test/tests_types.rs

@@ -2115,7 +2115,8 @@ fn singular_staking_info_unstake_during_voting_is_ok() {
     let remaining_stake = staking_info.total_staked_amount();
     assert_eq!(
         staking_info.unstake(remaining_stake + 1, era_2, Subperiod::Voting),
-        vec![(era_2, remaining_stake)]
+        vec![(era_2, remaining_stake), (era_2 + 1, remaining_stake)],
+        "Also chipping away from the next era since the unstake is relevant to the ongoing era."
     );
     assert!(staking_info.total_staked_amount().is_zero());
     assert!(
@@ -2217,7 +2218,8 @@ fn singular_staking_info_unstake_during_bep_is_ok() {
 
     assert_eq!(
         staking_info.unstake(unstake_2, era_2, Subperiod::BuildAndEarn),
-        vec![(era_2, unstake_2)]
+        vec![(era_2, unstake_2), (era_2 + 1, unstake_2)],
+        "Also chipping away from the next era since the unstake is relevant to the ongoing era."
     );
     assert_eq!(
         staking_info.total_staked_amount(),
@@ -2240,6 +2242,59 @@ fn singular_staking_info_unstake_during_bep_is_ok() {
     assert_eq!(staking_info.previous_staked.era, era_2 - 1);
 }
 
+#[test]
+fn singular_staking_info_unstake_era_amount_pairs_are_ok() {
+    let period_number = 1;
+    let subperiod = Subperiod::BuildAndEarn;
+
+    // 1. Unstake only reduces the amount from a the future era
+    {
+        let era = 3;
+        let stake_amount = 13;
+        let unstake_amount = 3;
+        let mut staking_info = SingularStakingInfo::new(period_number, subperiod);
+        staking_info.stake(stake_amount, era, Subperiod::BuildAndEarn);
+
+        assert_eq!(
+            staking_info.unstake(unstake_amount, era, Subperiod::BuildAndEarn),
+            vec![(era + 1, unstake_amount)]
+        );
+    }
+
+    // 2. Unstake reduces the amount from the current & next era.
+    {
+        let era = 3;
+        let stake_amount = 17;
+        let unstake_amount = 5;
+        let mut staking_info = SingularStakingInfo::new(period_number, subperiod);
+        staking_info.stake(stake_amount, era, Subperiod::BuildAndEarn);
+
+        assert_eq!(
+            staking_info
+                .clone()
+                .unstake(unstake_amount, era + 1, Subperiod::BuildAndEarn),
+            vec![(era + 1, unstake_amount), (era + 2, unstake_amount)]
+        );
+    }
+
+    // 3. Unstake reduces the amount from the current & next era.
+    //    Unlike the previous example, entries are not aligned with the current era
+    {
+        let era = 3;
+        let stake_amount = 17;
+        let unstake_amount = 5;
+        let mut staking_info = SingularStakingInfo::new(period_number, subperiod);
+        staking_info.stake(stake_amount, era, Subperiod::BuildAndEarn);
+
+        assert_eq!(
+            staking_info
+                .clone()
+                .unstake(unstake_amount, era + 2, Subperiod::BuildAndEarn),
+            vec![(era + 2, unstake_amount), (era + 3, unstake_amount)]
+        );
+    }
+}
+
 #[test]
 fn contract_stake_amount_basic_get_checks_work() {
     // Sanity checks for empty struct
@@ -2413,9 +2468,18 @@ fn contract_stake_amount_stake_is_ok() {
     let stake_era_2 = era_2 + 1;
     let amount_2 = 37;
     contract_stake.stake(amount_2, period_info_1, era_2);
-    let entry_2_1 = contract_stake.get(stake_era_1, period_1).unwrap();
+    let entry_2_1 = contract_stake
+        .get(era_2, period_1)
+        .expect("Since stake will change next era, entries should be aligned.");
     let entry_2_2 = contract_stake.get(stake_era_2, period_1).unwrap();
-    assert_eq!(entry_2_1, entry_1_2, "Old entry must remain unchanged.");
+    assert_eq!(
+        entry_2_1.for_type(Subperiod::Voting),
+        entry_1_2.for_type(Subperiod::Voting)
+    );
+    assert_eq!(
+        entry_2_1.for_type(Subperiod::BuildAndEarn),
+        entry_1_2.for_type(Subperiod::BuildAndEarn)
+    );
     assert_eq!(entry_2_2.era, stake_era_2);
     assert_eq!(entry_2_2.period, period_1);
     assert_eq!(
@@ -2442,7 +2506,7 @@ fn contract_stake_amount_stake_is_ok() {
 
     contract_stake.stake(amount_3, period_info_2, era_3);
     assert!(
-        contract_stake.get(stake_era_1, period_1).is_none(),
+        contract_stake.get(era_2, period_1).is_none(),
         "Old period must be removed."
     );
     assert!(

pallets/dapp-staking-v3/src/types.rs

@@ -1047,7 +1047,7 @@ impl SingularStakingInfo {
         let mut result = Vec::new();
         let staked_snapshot = self.staked;
 
-        // 1. Modify current staked amount, and update the result.
+        // 1. Modify 'current' staked amount, and update the result.
         self.staked.subtract(amount);
         let unstaked_amount = staked_snapshot.total().saturating_sub(self.staked.total());
         self.staked.era = self.staked.era.max(current_era);
@@ -1105,6 +1105,9 @@ impl SingularStakingInfo {
                 }
                 _ => {}
             }
+        } else if self.staked.era == current_era {
+            // In case the `staked` era was already the current era, it also means we're chipping away from the future era.
+            result.push((self.staked.era.saturating_add(1), unstaked_amount));
         }
 
         // 5. Convenience cleanup
@@ -1251,6 +1254,8 @@ impl ContractStakeAmount {
             // Future entry has an older era, but periods match so overwrite the 'current' entry with it
             Some(stake_amount) if stake_amount.period == period_info.number => {
                 self.staked = *stake_amount;
+                // Align the eras to keep it simple
+                self.staked.era = current_era;
             }
             // Otherwise do nothing
             _ => (),
@@ -1303,7 +1308,6 @@ impl ContractStakeAmount {
         }
 
         // 2. Value updates - only after alignment
-
         for (era, amount) in era_and_amount_pairs {
             if self.staked.era == era {
                 self.staked.subtract(amount);