Full fix

Author: Dinonard

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

@@ -627,6 +627,7 @@ pub(crate) fn assert_unstake(
         .expect("Entry must exist since 'unstake' is being called.");
     let pre_era_info = pre_snapshot.current_era_info;
 
+    let unstake_era = pre_snapshot.active_protocol_state.era;
     let unstake_period = pre_snapshot.active_protocol_state.period_number();
     let unstake_subperiod = pre_snapshot.active_protocol_state.subperiod();
 
@@ -679,9 +680,8 @@ pub(crate) fn assert_unstake(
     // 2. verify staker info
     // =====================
     // =====================
-    let staker_delta_stake =
-        pre_staker_info.staked.total() - pre_staker_info.previous_staked.total();
 
+    // Verify that expected unstake amounts are applied.
     if is_full_unstake {
         assert!(
             !StakerInfo::<Test>::contains_key(&account, smart_contract),
@@ -722,19 +722,19 @@ pub(crate) fn assert_unstake(
             is_loyal,
             "If 'Voting' stake amount is reduced in B&E period, loyalty flag must be set to false."
         );
+    }
 
-        if staker_delta_stake.is_zero() {
-            assert!(post_staker_info.previous_staked.is_empty(),);
-        } else {
-            assert_eq!(
-                post_staker_info.previous_staked.total(),
-                pre_staker_info.staked.total()
-            );
-            assert_eq!(
-                post_staker_info.previous_staked.era,
-                pre_staker_info.staked.era - 1
-            );
-        }
+    let unstaked_amount_era_pairs =
+        pre_staker_info
+            .clone()
+            .unstake(amount, unstake_period, 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, amount);
     }
 
     // 3. verify contract stake
@@ -753,30 +753,21 @@ pub(crate) fn assert_unstake(
         "Staked amount must decreased by the 'amount'"
     );
 
-    // Ensure that former era stake is updated correctly.
-    // If this is true, it means that the staker must have had a stake in some previous era.
-    // It's not possible to unstake more than user has staked, so this has to always hold true.
-    if staker_delta_stake < amount {
-        let expected_adjustment = amount - staker_delta_stake;
-
-        let latest_contract_stake_era = post_contract_stake
-            .latest_stake_era()
-            .expect("Has to be defined");
-
-        let previous_former_era_stake_amount = pre_contract_stake
-            .get(latest_contract_stake_era - 1, unstake_period)
-            .expect("Has to be defined")
-            .total();
-        let current_former_era_stake_amount = post_contract_stake
-            .get(latest_contract_stake_era - 1, unstake_period)
-            .expect("Has to be defined")
-            .total();
-
-        assert_eq!(
-            current_former_era_stake_amount,
-            previous_former_era_stake_amount - expected_adjustment,
-            "Former era stake can only be adjusted by the amount that was actually unstaked from it."
-        );
+    // 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
+            );
+        }
     }
 
     // 4. verify era info

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

@@ -2049,6 +2049,7 @@ fn singular_staking_info_basics_are_ok() {
         era_1 + 1,
         "Stake era should remain valid."
     );
+    assert!(staking_info.previous_staked.is_empty());
 
     // Add some staked amount during `BuildAndEarn` period
     let era_2 = 9;
@@ -2067,6 +2068,12 @@ fn singular_staking_info_basics_are_ok() {
         bep_stake_amount_1
     );
     assert_eq!(staking_info.era(), era_2 + 1);
+
+    assert_eq!(staking_info.previous_staked.total(), vote_stake_amount_1);
+    assert_eq!(
+        staking_info.previous_staked.era, era_2,
+        "Must be equal to the previous staked era."
+    );
 }
 
 #[test]
@@ -2102,7 +2109,7 @@ 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 - 1, remaining_stake), (era_2, remaining_stake)]
+        vec![(era_2, remaining_stake)]
     );
     assert!(staking_info.total_staked_amount().is_zero());
     assert!(
@@ -2174,8 +2181,7 @@ fn singular_staking_info_unstake_during_bep_is_ok() {
 
     assert_eq!(
         staking_info.unstake(unstake_2, era_2, Subperiod::BuildAndEarn),
-        // Both amounts are the same since previous staked era is less than last staked era - 1
-        vec![(era_2 - 1, unstake_2), (era_2, unstake_2)]
+        vec![(era_2, unstake_2)]
     );
     assert_eq!(
         staking_info.total_staked_amount(),

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

@@ -1002,6 +1002,9 @@ impl SingularStakingInfo {
         // Keep the previous stake amount for future reference
         self.previous_staked = self.staked;
         self.previous_staked.era = current_era;
+        if self.previous_staked.total().is_zero() {
+            self.previous_staked = Default::default();
+        }
 
         // Stake is only valid from the next era so we keep it consistent here
         self.staked.add(amount, subperiod);
@@ -1021,22 +1024,6 @@ impl SingularStakingInfo {
     /// This means that the returned value will at most cover two eras - the last staked era, and the one before it.
     ///
     /// Last staked era can be the current era, or the era after.
-    /// If after the unstake operation, delta between previous & current era is larger than 1, previous staked entry is ignored.
-    ///
-    /// #### Example (simplified)
-    ///
-    /// 1. previous_staked: (era: 4, amount: 50), staked: (era: 5, amount: 100)
-    /// 2. unstake 30 in era 6 is called
-    /// 3. Return value is: (era: 5, amount: 30), (era: 6, amount: 30)
-    /// 4. previous_staked: (era: 5, amount: 70), (era: 6, amount: 70)
-    ///
-    /// In case same example is repeated, but unstake is done in era 5:
-    ///
-    /// 1. previous_staked: (era: 4, amount: 50), staked: (era: 5, amount: 100)
-    /// 2. unstake 30 in era 5 is called
-    /// 3. Return value is: (era: 5, amount: 30)
-    /// 4. previous_staked: (era: 4, amount: 50), (era: 5, amount: 70)
-    ///
     pub fn unstake(
         &mut self,
         amount: Balance,
@@ -1046,51 +1033,64 @@ impl SingularStakingInfo {
         let mut result = Vec::new();
         let staked_snapshot = self.staked;
 
-        // 1. Modify current staked amount
+        // 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);
         result.push((self.staked.era, unstaked_amount));
 
-        // 2. Calculate previous stake delta.
-        // In case new stake era is exactly one era after the previous stake era, we can calculate this as a delta.
-        // Otherwise, if it's further far away in the past, the previous era stake is equal to the snapshot.
-        let stake_delta = if self.staked.era == self.previous_staked.era.saturating_add(1) {
-            staked_snapshot
-                .total()
-                .saturating_sub(self.previous_staked.total())
-        } else {
-            Balance::zero()
-        };
-
-        // 2. Update loyal staker flag accordingly
+        // 2. Update loyal staker flag accordingly.
         self.loyal_staker = self.loyal_staker
             && match subperiod {
                 Subperiod::Voting => !self.staked.voting.is_zero(),
                 Subperiod::BuildAndEarn => self.staked.voting == staked_snapshot.voting,
             };
 
-        // 3. Move over the snapshot to the previous snapshot field and make sure
-        // to update era in case something was staked before.
-        if stake_delta.is_zero() {
-            self.previous_staked = Default::default();
+        // 3. Determine what was the previous staked amount.
+        // This is done by simply comparing where does the _previous era_ fit in the current context.
+        let previous_era = self.staked.era.saturating_sub(1);
+
+        self.previous_staked = if staked_snapshot.era <= previous_era {
+            let mut previous_staked = staked_snapshot;
+            previous_staked.era = previous_era;
+            previous_staked
+        } else if !self.previous_staked.is_empty() && self.previous_staked.era <= previous_era {
+            let mut previous_staked = self.previous_staked;
+            previous_staked.era = previous_era;
+            previous_staked
         } else {
-            self.previous_staked = staked_snapshot;
-            self.previous_staked.era = self.staked.era.saturating_sub(1);
-        }
+            Default::default()
+        };
 
-        // 4. If something was unstaked from the previous era, add it to the result
-        if unstaked_amount > stake_delta {
-            // Doesn't need to be at 0-th index, but we still add it for clarity
-            result.insert(
-                0,
-                (
-                    self.staked.era.saturating_sub(1),
-                    // The diff between unstake amount & stake delta tells us how much was
-                    // actually unstaked from the previous era.
-                    unstaked_amount.saturating_sub(stake_delta),
-                ),
-            )
+        // 4. Calculate how much is being unstaked from the previous staked era entry, in case its era equals the current era.
+        //
+        // Simples way to explain this is via an example.
+        // Let's assume a simplification where stake amount entries are in `(era, amount)` format.
+        //
+        // 1. Values: previous_staked: **(2, 10)**, staked: **(3, 15)**
+        // 2. User calls unstake during **era 2**, and unstakes amount **6**.
+        //    Clearly some amount was staked during era 2, which resulted in era 3 stake being increased by 5.
+        //    Calling unstake immediately in the same era should not necessarily reduce current era stake amount.
+        //    This should be allowed to happen only if the unstaked amount is larger than the difference between the staked amount of two eras.
+        // 3. Values: previous_staked: **(2, 9)**, staked: **(3, 9)**
+        //
+        // An alternative scenario, where user calls unstake during **era 2**, and unstakes amount **4**.
+        // 3. Values: previous_staked: **(2, 10)**, staked: **(3, 11)**
+        //
+        // Note that the unstake operation didn't chip away from the current era, only the next one.
+        if self.previous_staked.era == current_era {
+            let maybe_stake_delta = staked_snapshot
+                .total()
+                .checked_sub(self.previous_staked.total());
+            match maybe_stake_delta {
+                Some(stake_delta) if unstaked_amount > stake_delta => {
+                    let overflow_amount = unstaked_amount - stake_delta;
+                    self.previous_staked.subtract(overflow_amount);
+
+                    result.insert(0, (self.previous_staked.era, overflow_amount));
+                }
+                _ => {}
+            }
         }
 
         result