From 14a90bc99fd3b906f2e3061de3ad64df4828af5b Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Thu, 29 Feb 2024 10:13:14 +0100
Subject: [PATCH 01/28] Reproduce issue

---
 .../dapp-staking-v3/src/test/testing_utils.rs | 45 ++++++++++++++
 pallets/dapp-staking-v3/src/test/tests.rs     | 58 ++++++++++++++++++-
 2 files changed, 100 insertions(+), 3 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/test/testing_utils.rs b/pallets/dapp-staking-v3/src/test/testing_utils.rs
index 5ddb48bf4f..35ef13353b 100644
--- a/pallets/dapp-staking-v3/src/test/testing_utils.rs
+++ b/pallets/dapp-staking-v3/src/test/testing_utils.rs
@@ -627,6 +627,17 @@ pub(crate) fn assert_unstake(
         .expect("Entry must exist since 'unstake' is being called.");
     let pre_era_info = pre_snapshot.current_era_info;
 
+    // Calculate expected unstake amounts from the voting and build&earn subperiods.
+    let (expected_voting_unstake, expected_build_and_earn_unstake) =
+        if pre_staker_info.staked_amount(Subperiod::BuildAndEarn) >= amount {
+            (Balance::zero(), amount)
+        } else {
+            (
+                amount - pre_staker_info.staked_amount(Subperiod::BuildAndEarn),
+                pre_staker_info.staked_amount(Subperiod::BuildAndEarn),
+            )
+        };
+
     let unstake_period = pre_snapshot.active_protocol_state.period_number();
     let unstake_subperiod = pre_snapshot.active_protocol_state.subperiod();
 
@@ -736,6 +747,40 @@ pub(crate) fn assert_unstake(
             .saturating_sub(amount),
         "Staked amount must decreased by the 'amount'"
     );
+    assert_eq!(
+        post_contract_stake.staked.voting,
+        pre_contract_stake.staked.voting.saturating_sub(expected_voting_unstake),
+        "Voting stake amount must be decreased exactly by how much staker's voting amount was reduced."
+    );
+
+    assert_eq!(
+        post_contract_stake.staked.build_and_earn,
+        pre_contract_stake
+            .staked
+            .build_and_earn
+            .saturating_sub(expected_build_and_earn_unstake),
+        "B&E stake amount must be decreased exactly by how much staker's B&E amount was reduced."
+    );
+    match (
+        pre_contract_stake.staked_future,
+        post_contract_stake.staked_future,
+    ) {
+        (Some(pre_future), Some(post_future)) => {
+            assert_eq!(
+                post_future.voting,
+                pre_future.voting.saturating_sub(expected_voting_unstake),
+                "Voting stake amount must be decreased exactly by how much staker's voting amount was reduced."
+            );
+            assert_eq!(
+                post_future.build_and_earn,
+                pre_future
+                    .build_and_earn
+                    .saturating_sub(expected_build_and_earn_unstake),
+                "B&E stake amount must be decreased exactly by how much staker's B&E amount was reduced."
+            );
+        }
+        (_, _) => {}
+    }
 
     // 4. verify era info
     // =========================
diff --git a/pallets/dapp-staking-v3/src/test/tests.rs b/pallets/dapp-staking-v3/src/test/tests.rs
index f196dab81a..62375948ac 100644
--- a/pallets/dapp-staking-v3/src/test/tests.rs
+++ b/pallets/dapp-staking-v3/src/test/tests.rs
@@ -18,9 +18,9 @@
 
 use crate::test::{mock::*, testing_utils::*};
 use crate::{
-    pallet::Config, ActiveProtocolState, DAppId, EraRewards, Error, Event, ForcingType,
-    IntegratedDApps, Ledger, NextDAppId, PeriodNumber, Safeguard, StakerInfo, Subperiod,
-    TierConfig,
+    pallet::Config, ActiveProtocolState, ContractStake, DAppId, EraRewards, Error, Event,
+    ForcingType, IntegratedDApps, Ledger, NextDAppId, PeriodNumber, Safeguard, StakerInfo,
+    Subperiod, TierConfig,
 };
 
 use frame_support::{
@@ -2594,6 +2594,58 @@ fn stake_and_unstake_after_reward_claim_is_ok() {
     })
 }
 
+#[test]
+fn stake_and_unstake_correctly_update_staked_amounts() {
+    ExtBuilder::build().execute_with(|| {
+        // Register smart contract
+        let dev_account = 1;
+        let smart_contract = MockSmartContract::wasm(1 as AccountId);
+        assert_register(dev_account, &smart_contract);
+        let smart_contract_id = IntegratedDApps::<Test>::get(&smart_contract).unwrap().id;
+
+        // Lock & stake some amount by the first staker, and lock some amount by the second staker
+        let account_1 = 2;
+        let amount_1 = 50;
+        assert_lock(account_1, amount_1);
+        assert_stake(account_1, &smart_contract, amount_1);
+
+        let account_2 = 3;
+        let amount_2 = 10;
+        assert_lock(account_2, amount_2);
+
+        // Stake & unstake repeatedly by the second staker, in the current era
+        let contract_stake_snapshot = ContractStake::<Test>::get(&smart_contract_id);
+
+        for _ in 0..20 {
+            assert_stake(account_2, &smart_contract, amount_2);
+            assert_unstake(account_2, &smart_contract, amount_2);
+        }
+
+        // Check that the contract stake snapshot is the same as before
+        assert_eq!(
+            contract_stake_snapshot,
+            ContractStake::<Test>::get(&smart_contract_id),
+            "Contract stake snapshot should not change."
+        );
+
+        // Stake & unstake repeatedly in the next era
+        advance_to_next_era();
+        let contract_stake_snapshot = ContractStake::<Test>::get(&smart_contract_id);
+
+        for _ in 0..20 {
+            assert_stake(account_2, &smart_contract, amount_2);
+            assert_unstake(account_2, &smart_contract, amount_2);
+        }
+
+        // Check that the contract stake snapshot is the same as before
+        assert_eq!(
+            contract_stake_snapshot,
+            ContractStake::<Test>::get(&smart_contract_id),
+            "Contract stake snapshot should not change."
+        );
+    })
+}
+
 #[test]
 fn stake_after_period_ends_with_max_staked_contracts() {
     ExtBuilder::build().execute_with(|| {

From 8f02d9527af066e541002fe3e2631cfb493b8bc7 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Thu, 29 Feb 2024 10:24:10 +0100
Subject: [PATCH 02/28] Update unstake logic

---
 pallets/dapp-staking-v3/src/lib.rs | 70 ++++++++++++++++++------------
 1 file changed, 43 insertions(+), 27 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/lib.rs b/pallets/dapp-staking-v3/src/lib.rs
index 7e9d8d7d7a..0b0a6fd1df 100644
--- a/pallets/dapp-staking-v3/src/lib.rs
+++ b/pallets/dapp-staking-v3/src/lib.rs
@@ -1106,34 +1106,50 @@ pub mod pallet {
 
             // 1.
             // Update `StakerInfo` storage with the reduced stake amount on the specified contract.
-            let (new_staking_info, amount) = match StakerInfo::<T>::get(&account, &smart_contract) {
-                Some(mut staking_info) => {
-                    ensure!(
-                        staking_info.period_number() == protocol_state.period_number(),
-                        Error::<T>::UnstakeFromPastPeriod
-                    );
-                    ensure!(
-                        staking_info.total_staked_amount() >= amount,
-                        Error::<T>::UnstakeAmountTooLarge
-                    );
-
-                    // If unstaking would take the total staked amount below the minimum required value,
-                    // unstake everything.
-                    let amount = if staking_info.total_staked_amount().saturating_sub(amount)
-                        < T::MinimumStakeAmount::get()
-                    {
-                        staking_info.total_staked_amount()
-                    } else {
-                        amount
-                    };
+            let (new_staking_info, voting_unstake_amount, bep_unstake_amount) =
+                match StakerInfo::<T>::get(&account, &smart_contract) {
+                    Some(mut staking_info) => {
+                        ensure!(
+                            staking_info.period_number() == protocol_state.period_number(),
+                            Error::<T>::UnstakeFromPastPeriod
+                        );
+                        ensure!(
+                            staking_info.total_staked_amount() >= amount,
+                            Error::<T>::UnstakeAmountTooLarge
+                        );
 
-                    staking_info.unstake(amount, current_era, protocol_state.subperiod());
-                    (staking_info, amount)
-                }
-                None => {
-                    return Err(Error::<T>::NoStakingInfo.into());
-                }
-            };
+                        // If unstaking would take the total staked amount below the minimum required value,
+                        // unstake everything.
+                        let amount = if staking_info.total_staked_amount().saturating_sub(amount)
+                            < T::MinimumStakeAmount::get()
+                        {
+                            staking_info.total_staked_amount()
+                        } else {
+                            amount
+                        };
+
+                        // We need to know the exact amounts to unstake from each subperiod.
+                        let (voting_unstake_amount, bep_unstake_amount) =
+                            if staking_info.staked_amount(Subperiod::BuildAndEarn) >= amount {
+                                (Balance::zero(), amount)
+                            } else {
+                                (
+                                    amount.saturating_sub(
+                                        staking_info.staked_amount(Subperiod::BuildAndEarn),
+                                    ),
+                                    staking_info.staked_amount(Subperiod::BuildAndEarn),
+                                )
+                            };
+
+                        staking_info.unstake(amount, current_era, protocol_state.subperiod());
+
+                        (staking_info, voting_unstake_amount, bep_unstake_amount)
+                    }
+                    None => {
+                        return Err(Error::<T>::NoStakingInfo.into());
+                    }
+                };
+            let amount = voting_unstake_amount.saturating_add(bep_unstake_amount);
 
             // 2.
             // Reduce stake amount

From a689b2406f2d58a63013792762ba78835d3f09e5 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Thu, 29 Feb 2024 11:06:50 +0100
Subject: [PATCH 03/28] Fix

---
 pallets/dapp-staking-v3/src/lib.rs              |  7 ++++++-
 pallets/dapp-staking-v3/src/test/tests.rs       | 16 +++++++++++-----
 pallets/dapp-staking-v3/src/test/tests_types.rs | 10 ++++++----
 pallets/dapp-staking-v3/src/types.rs            | 14 +++++++++++---
 4 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/lib.rs b/pallets/dapp-staking-v3/src/lib.rs
index 0b0a6fd1df..f232ec2569 100644
--- a/pallets/dapp-staking-v3/src/lib.rs
+++ b/pallets/dapp-staking-v3/src/lib.rs
@@ -1169,7 +1169,12 @@ pub mod pallet {
             // 3.
             // Update `ContractStake` storage with the reduced stake amount on the specified contract.
             let mut contract_stake_info = ContractStake::<T>::get(&dapp_info.id);
-            contract_stake_info.unstake(amount, protocol_state.period_info, current_era);
+            contract_stake_info.unstake(
+                voting_unstake_amount,
+                bep_unstake_amount,
+                protocol_state.period_info,
+                current_era,
+            );
 
             // 4.
             // Update total staked amount for the next era.
diff --git a/pallets/dapp-staking-v3/src/test/tests.rs b/pallets/dapp-staking-v3/src/test/tests.rs
index 62375948ac..068b2b197a 100644
--- a/pallets/dapp-staking-v3/src/test/tests.rs
+++ b/pallets/dapp-staking-v3/src/test/tests.rs
@@ -2595,7 +2595,7 @@ fn stake_and_unstake_after_reward_claim_is_ok() {
 }
 
 #[test]
-fn stake_and_unstake_correctly_update_staked_amounts() {
+fn stake_and_unstake_correctly_updates_staked_amounts() {
     ExtBuilder::build().execute_with(|| {
         // Register smart contract
         let dev_account = 1;
@@ -2637,11 +2637,17 @@ fn stake_and_unstake_correctly_update_staked_amounts() {
             assert_unstake(account_2, &smart_contract, amount_2);
         }
 
-        // Check that the contract stake snapshot is the same as before
+        // Check that the contract stake snapshot staked amount is the same as before
         assert_eq!(
-            contract_stake_snapshot,
-            ContractStake::<Test>::get(&smart_contract_id),
-            "Contract stake snapshot should not change."
+            contract_stake_snapshot.staked_amount(1, Subperiod::Voting),
+            ContractStake::<Test>::get(&smart_contract_id).staked_amount(1, Subperiod::Voting),
+            "Voting staked amount should not change."
+        );
+        assert_eq!(
+            contract_stake_snapshot.staked_amount(1, Subperiod::BuildAndEarn),
+            ContractStake::<Test>::get(&smart_contract_id)
+                .staked_amount(1, Subperiod::BuildAndEarn),
+            "Build&Earn staked amount should not change."
         );
     })
 }
diff --git a/pallets/dapp-staking-v3/src/test/tests_types.rs b/pallets/dapp-staking-v3/src/test/tests_types.rs
index 64da02a23e..9f3ef0722e 100644
--- a/pallets/dapp-staking-v3/src/test/tests_types.rs
+++ b/pallets/dapp-staking-v3/src/test/tests_types.rs
@@ -2430,6 +2430,7 @@ fn contract_stake_amount_stake_is_ok() {
     assert!(contract_stake.staked_future.is_some());
 }
 
+// TODO: improve this a bit
 #[test]
 fn contract_stake_amount_unstake_is_ok() {
     let mut contract_stake = ContractStakeAmount::default();
@@ -2447,7 +2448,7 @@ fn contract_stake_amount_unstake_is_ok() {
 
     // 1st scenario - unstake in the same era
     let amount_1 = 5;
-    contract_stake.unstake(amount_1, period_info, era_1);
+    contract_stake.unstake(amount_1, 0, period_info, era_1);
     assert_eq!(
         contract_stake.total_staked_amount(period),
         stake_amount - amount_1
@@ -2467,7 +2468,7 @@ fn contract_stake_amount_unstake_is_ok() {
     };
     let era_2 = era_1 + 1;
 
-    contract_stake.unstake(amount_1, period_info, era_2);
+    contract_stake.unstake(0, amount_1, period_info, era_2);
     assert_eq!(
         contract_stake.total_staked_amount(period),
         stake_amount - amount_1 * 2
@@ -2488,7 +2489,7 @@ fn contract_stake_amount_unstake_is_ok() {
     // 3rd scenario - bump up unstake eras by more than 1, entries should be aligned to the current era
     let era_3 = era_2 + 3;
     let amount_2 = 7;
-    contract_stake.unstake(amount_2, period_info, era_3);
+    contract_stake.unstake(0, amount_2, period_info, era_3);
     assert_eq!(
         contract_stake.total_staked_amount(period),
         stake_amount - amount_1 * 2 - amount_2
@@ -2509,7 +2510,8 @@ fn contract_stake_amount_unstake_is_ok() {
     // 4th scenario - do a full unstake with existing future entry, expect a cleanup
     contract_stake.stake(stake_amount, period_info, era_3);
     contract_stake.unstake(
-        contract_stake.total_staked_amount(period),
+        contract_stake.staked_amount(period, Subperiod::Voting),
+        contract_stake.staked_amount(period, Subperiod::BuildAndEarn),
         period_info,
         era_3,
     );
diff --git a/pallets/dapp-staking-v3/src/types.rs b/pallets/dapp-staking-v3/src/types.rs
index 4101fa1990..17a89bd5f8 100644
--- a/pallets/dapp-staking-v3/src/types.rs
+++ b/pallets/dapp-staking-v3/src/types.rs
@@ -1199,7 +1199,13 @@ impl ContractStakeAmount {
     }
 
     /// Unstake the specified `amount` from the contract, for the specified `subperiod` and `era`.
-    pub fn unstake(&mut self, amount: Balance, period_info: PeriodInfo, current_era: EraNumber) {
+    pub fn unstake(
+        &mut self,
+        voting_amount: Balance,
+        bep_amount: Balance,
+        period_info: PeriodInfo,
+        current_era: EraNumber,
+    ) {
         // First align entries - we only need to keep track of the current era, and the next one
         match self.staked_future {
             // Future entry exists, but it covers current or older era.
@@ -1219,9 +1225,11 @@ impl ContractStakeAmount {
         }
 
         // Subtract both amounts
-        self.staked.subtract(amount, period_info.subperiod);
+        self.staked.voting.saturating_reduce(voting_amount);
+        self.staked.build_and_earn.saturating_reduce(bep_amount);
         if let Some(stake_amount) = self.staked_future.as_mut() {
-            stake_amount.subtract(amount, period_info.subperiod);
+            stake_amount.voting.saturating_reduce(voting_amount);
+            stake_amount.build_and_earn.saturating_reduce(bep_amount);
         }
 
         // Convenience cleanup

From f32ab472ddff95e569291d2f690ab8bf20c972f1 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Thu, 29 Feb 2024 11:46:58 +0100
Subject: [PATCH 04/28] Fix cont.

---
 .../dapp-staking-v3/src/test/testing_utils.rs | 56 +++++------------
 .../dapp-staking-v3/src/test/tests_types.rs   | 63 ++++++++++++++-----
 pallets/dapp-staking-v3/src/types.rs          | 15 +++--
 3 files changed, 75 insertions(+), 59 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/test/testing_utils.rs b/pallets/dapp-staking-v3/src/test/testing_utils.rs
index 35ef13353b..def8dfa05d 100644
--- a/pallets/dapp-staking-v3/src/test/testing_utils.rs
+++ b/pallets/dapp-staking-v3/src/test/testing_utils.rs
@@ -627,17 +627,6 @@ pub(crate) fn assert_unstake(
         .expect("Entry must exist since 'unstake' is being called.");
     let pre_era_info = pre_snapshot.current_era_info;
 
-    // Calculate expected unstake amounts from the voting and build&earn subperiods.
-    let (expected_voting_unstake, expected_build_and_earn_unstake) =
-        if pre_staker_info.staked_amount(Subperiod::BuildAndEarn) >= amount {
-            (Balance::zero(), amount)
-        } else {
-            (
-                amount - pre_staker_info.staked_amount(Subperiod::BuildAndEarn),
-                pre_staker_info.staked_amount(Subperiod::BuildAndEarn),
-            )
-        };
-
     let unstake_period = pre_snapshot.active_protocol_state.period_number();
     let unstake_subperiod = pre_snapshot.active_protocol_state.subperiod();
 
@@ -652,6 +641,17 @@ pub(crate) fn assert_unstake(
         amount
     };
 
+    // Calculate expected unstake amounts from the voting and build&earn subperiods.
+    let (expected_voting_unstake, expected_build_and_earn_unstake) =
+        if pre_staker_info.staked_amount(Subperiod::BuildAndEarn) >= amount {
+            (Balance::zero(), amount)
+        } else {
+            (
+                amount - pre_staker_info.staked_amount(Subperiod::BuildAndEarn),
+                pre_staker_info.staked_amount(Subperiod::BuildAndEarn),
+            )
+        };
+
     // Unstake from smart contract & verify event
     assert_ok!(DappStaking::unstake(
         RuntimeOrigin::signed(account),
@@ -747,40 +747,18 @@ pub(crate) fn assert_unstake(
             .saturating_sub(amount),
         "Staked amount must decreased by the 'amount'"
     );
+
     assert_eq!(
-        post_contract_stake.staked.voting,
-        pre_contract_stake.staked.voting.saturating_sub(expected_voting_unstake),
+        post_contract_stake.staked_amount(unstake_period, Subperiod::Voting),
+        pre_contract_stake.staked_amount(unstake_period, Subperiod::Voting) - expected_voting_unstake,
         "Voting stake amount must be decreased exactly by how much staker's voting amount was reduced."
     );
-
     assert_eq!(
-        post_contract_stake.staked.build_and_earn,
-        pre_contract_stake
-            .staked
-            .build_and_earn
-            .saturating_sub(expected_build_and_earn_unstake),
+        post_contract_stake.staked_amount(unstake_period, Subperiod::BuildAndEarn),
+        pre_contract_stake.staked_amount(unstake_period, Subperiod::BuildAndEarn)
+            - expected_build_and_earn_unstake,
         "B&E stake amount must be decreased exactly by how much staker's B&E amount was reduced."
     );
-    match (
-        pre_contract_stake.staked_future,
-        post_contract_stake.staked_future,
-    ) {
-        (Some(pre_future), Some(post_future)) => {
-            assert_eq!(
-                post_future.voting,
-                pre_future.voting.saturating_sub(expected_voting_unstake),
-                "Voting stake amount must be decreased exactly by how much staker's voting amount was reduced."
-            );
-            assert_eq!(
-                post_future.build_and_earn,
-                pre_future
-                    .build_and_earn
-                    .saturating_sub(expected_build_and_earn_unstake),
-                "B&E stake amount must be decreased exactly by how much staker's B&E amount was reduced."
-            );
-        }
-        (_, _) => {}
-    }
 
     // 4. verify era info
     // =========================
diff --git a/pallets/dapp-staking-v3/src/test/tests_types.rs b/pallets/dapp-staking-v3/src/test/tests_types.rs
index 9f3ef0722e..1946da7cbb 100644
--- a/pallets/dapp-staking-v3/src/test/tests_types.rs
+++ b/pallets/dapp-staking-v3/src/test/tests_types.rs
@@ -2430,32 +2430,42 @@ fn contract_stake_amount_stake_is_ok() {
     assert!(contract_stake.staked_future.is_some());
 }
 
-// TODO: improve this a bit
 #[test]
 fn contract_stake_amount_unstake_is_ok() {
     let mut contract_stake = ContractStakeAmount::default();
 
     // Prep action - create a stake entry
     let era_1 = 2;
+    let era_2 = era_1 + 1;
     let period = 3;
     let period_info = PeriodInfo {
         number: period,
         subperiod: Subperiod::Voting,
         next_subperiod_start_era: 20,
     };
-    let stake_amount = 100;
-    contract_stake.stake(stake_amount, period_info, era_1);
+    let vp_stake_amount = 47;
+    let bep_stake_amount = 53;
+    contract_stake.stake(vp_stake_amount, period_info, era_1);
+    contract_stake.stake(
+        bep_stake_amount,
+        PeriodInfo {
+            subperiod: Subperiod::BuildAndEarn,
+            ..period_info
+        },
+        era_1,
+    );
+    let total_stake_amount = vp_stake_amount + bep_stake_amount;
 
-    // 1st scenario - unstake in the same era
+    // 1st scenario - unstake in the same era, from `Voting` subperiod
     let amount_1 = 5;
     contract_stake.unstake(amount_1, 0, period_info, era_1);
     assert_eq!(
         contract_stake.total_staked_amount(period),
-        stake_amount - amount_1
+        total_stake_amount - amount_1
     );
     assert_eq!(
         contract_stake.staked_amount(period, Subperiod::Voting),
-        stake_amount - amount_1
+        vp_stake_amount - amount_1
     );
     assert!(contract_stake.staked.is_empty());
     assert!(contract_stake.staked_future.is_some());
@@ -2466,16 +2476,20 @@ fn contract_stake_amount_unstake_is_ok() {
         subperiod: Subperiod::BuildAndEarn,
         next_subperiod_start_era: 40,
     };
-    let era_2 = era_1 + 1;
 
-    contract_stake.unstake(0, amount_1, period_info, era_2);
+    contract_stake.unstake(amount_1, 0, period_info, era_2);
     assert_eq!(
         contract_stake.total_staked_amount(period),
-        stake_amount - amount_1 * 2
+        total_stake_amount - amount_1 * 2
     );
     assert_eq!(
         contract_stake.staked_amount(period, Subperiod::Voting),
-        stake_amount - amount_1 * 2
+        vp_stake_amount - amount_1 * 2
+    );
+    assert_eq!(
+        contract_stake.staked_amount(period, Subperiod::BuildAndEarn),
+        bep_stake_amount,
+        "Must remain unchanged."
     );
     assert!(
         !contract_stake.staked.is_empty(),
@@ -2486,17 +2500,36 @@ fn contract_stake_amount_unstake_is_ok() {
         "future entry should be cleaned up since it refers to the current era"
     );
 
-    // 3rd scenario - bump up unstake eras by more than 1, entries should be aligned to the current era
+    // 3rd scenario - same as previous scenario, but for `BuildAndEarn` subperiod
+    contract_stake.unstake(0, amount_1, period_info, era_2);
+    assert_eq!(
+        contract_stake.total_staked_amount(period),
+        total_stake_amount - amount_1 * 3
+    );
+    assert_eq!(
+        contract_stake.staked_amount(period, Subperiod::Voting),
+        vp_stake_amount - amount_1 * 2
+    );
+    assert_eq!(
+        contract_stake.staked_amount(period, Subperiod::BuildAndEarn),
+        bep_stake_amount - amount_1
+    );
+
+    // 4th scenario - bump up unstake eras by more than 1, entries should be aligned to the current era
     let era_3 = era_2 + 3;
     let amount_2 = 7;
     contract_stake.unstake(0, amount_2, period_info, era_3);
     assert_eq!(
         contract_stake.total_staked_amount(period),
-        stake_amount - amount_1 * 2 - amount_2
+        total_stake_amount - amount_1 * 3 - amount_2
     );
     assert_eq!(
         contract_stake.staked_amount(period, Subperiod::Voting),
-        stake_amount - amount_1 * 2 - amount_2
+        vp_stake_amount - amount_1 * 2
+    );
+    assert_eq!(
+        contract_stake.staked_amount(period, Subperiod::BuildAndEarn),
+        bep_stake_amount - amount_1 - amount_2
     );
     assert_eq!(
         contract_stake.staked.era, era_3,
@@ -2507,8 +2540,8 @@ fn contract_stake_amount_unstake_is_ok() {
         "future entry should remain 'None'"
     );
 
-    // 4th scenario - do a full unstake with existing future entry, expect a cleanup
-    contract_stake.stake(stake_amount, period_info, era_3);
+    // 5th scenario - do a full unstake with existing future entry, expect a cleanup
+    contract_stake.stake(total_stake_amount, period_info, era_3);
     contract_stake.unstake(
         contract_stake.staked_amount(period, Subperiod::Voting),
         contract_stake.staked_amount(period, Subperiod::BuildAndEarn),
diff --git a/pallets/dapp-staking-v3/src/types.rs b/pallets/dapp-staking-v3/src/types.rs
index 17a89bd5f8..948661e122 100644
--- a/pallets/dapp-staking-v3/src/types.rs
+++ b/pallets/dapp-staking-v3/src/types.rs
@@ -882,6 +882,13 @@ impl StakeAmount {
             }
         }
     }
+
+    /// Saturating reduce the specified `amount` from the appropriate subperiod.
+    /// If the amount is larger than the subperiod amount, the subperiod amount is set to zero.
+    pub fn subtract_subperiods(&mut self, voting_amount: Balance, bep_amount: Balance) {
+        self.voting.saturating_reduce(voting_amount);
+        self.build_and_earn.saturating_reduce(bep_amount);
+    }
 }
 
 /// Info about an era, including the rewards, how much is locked, unlocking, etc.
@@ -1198,7 +1205,7 @@ impl ContractStakeAmount {
         }
     }
 
-    /// Unstake the specified `amount` from the contract, for the specified `subperiod` and `era`.
+    /// Unstake the specified `Voting` and `Build&Earn` subperiod amounts from the contract, for the specified `subperiod` and `era`.
     pub fn unstake(
         &mut self,
         voting_amount: Balance,
@@ -1225,11 +1232,9 @@ impl ContractStakeAmount {
         }
 
         // Subtract both amounts
-        self.staked.voting.saturating_reduce(voting_amount);
-        self.staked.build_and_earn.saturating_reduce(bep_amount);
+        self.staked.subtract_subperiods(voting_amount, bep_amount);
         if let Some(stake_amount) = self.staked_future.as_mut() {
-            stake_amount.voting.saturating_reduce(voting_amount);
-            stake_amount.build_and_earn.saturating_reduce(bep_amount);
+            stake_amount.subtract_subperiods(voting_amount, bep_amount);
         }
 
         // Convenience cleanup

From edfb62da110437375625f34a5e0c3c493576c692 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Fri, 1 Mar 2024 10:24:40 +0100
Subject: [PATCH 05/28] Improve reproduction test

---
 pallets/dapp-staking-v3/src/test/tests.rs | 65 ++++++++++++++++++-----
 1 file changed, 53 insertions(+), 12 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/test/tests.rs b/pallets/dapp-staking-v3/src/test/tests.rs
index 068b2b197a..e1dd357a4d 100644
--- a/pallets/dapp-staking-v3/src/test/tests.rs
+++ b/pallets/dapp-staking-v3/src/test/tests.rs
@@ -2613,7 +2613,7 @@ fn stake_and_unstake_correctly_updates_staked_amounts() {
         let amount_2 = 10;
         assert_lock(account_2, amount_2);
 
-        // Stake & unstake repeatedly by the second staker, in the current era
+        // 1st scenario: repeated stake & unstake in the `Voting` subperiod
         let contract_stake_snapshot = ContractStake::<Test>::get(&smart_contract_id);
 
         for _ in 0..20 {
@@ -2621,14 +2621,20 @@ fn stake_and_unstake_correctly_updates_staked_amounts() {
             assert_unstake(account_2, &smart_contract, amount_2);
         }
 
-        // Check that the contract stake snapshot is the same as before
+        // Check that the staked amount for the upcoming era is same as before
+        let current_era = ActiveProtocolState::<Test>::get().era;
+        let period_number = ActiveProtocolState::<Test>::get().period_number();
         assert_eq!(
-            contract_stake_snapshot,
-            ContractStake::<Test>::get(&smart_contract_id),
-            "Contract stake snapshot should not change."
+            contract_stake_snapshot
+                .get(current_era + 1, period_number)
+                .expect("Entry must exist."),
+            ContractStake::<Test>::get(&smart_contract_id)
+                .get(current_era + 1, period_number)
+                .expect("Entry must exist."),
+            "Ongoing era staked amount must not change."
         );
 
-        // Stake & unstake repeatedly in the next era
+        // 2nd scenario: repeated stake & unstake in the first era of the `Build&Earn` subperiod
         advance_to_next_era();
         let contract_stake_snapshot = ContractStake::<Test>::get(&smart_contract_id);
 
@@ -2638,16 +2644,51 @@ fn stake_and_unstake_correctly_updates_staked_amounts() {
         }
 
         // Check that the contract stake snapshot staked amount is the same as before
+        let current_era = ActiveProtocolState::<Test>::get().era;
         assert_eq!(
-            contract_stake_snapshot.staked_amount(1, Subperiod::Voting),
-            ContractStake::<Test>::get(&smart_contract_id).staked_amount(1, Subperiod::Voting),
-            "Voting staked amount should not change."
+            contract_stake_snapshot
+                .get(current_era, period_number)
+                .expect("Entry must exist."),
+            ContractStake::<Test>::get(&smart_contract_id)
+                .get(current_era, period_number)
+                .expect("Entry must exist."),
+            "Ongoing era staked amount must not change."
         );
+
+        assert_eq!(
+            contract_stake_snapshot
+                .get(current_era, period_number)
+                .expect("Entry must exist.")
+                .total(),
+            ContractStake::<Test>::get(&smart_contract_id)
+                .get(current_era + 1, period_number)
+                .expect("Entry must exist.")
+                .total(),
+            "Ongoing era staked amount must be equal to the upcoming era stake."
+        );
+
+        // 3rd scenario: repeated stake & unstake in the second era of the `Build&Earn` subperiod
+        assert_stake(account_2, &smart_contract, amount_2);
+        assert_lock(account_2, amount_2);
+        advance_to_next_era();
+
+        let contract_stake_snapshot = ContractStake::<Test>::get(&smart_contract_id);
+
+        for _ in 0..20 {
+            assert_stake(account_2, &smart_contract, amount_2);
+            assert_unstake(account_2, &smart_contract, amount_2);
+        }
+
+        // Check that the contract stake snapshot staked amount is the same as before
+        let current_era = ActiveProtocolState::<Test>::get().era;
         assert_eq!(
-            contract_stake_snapshot.staked_amount(1, Subperiod::BuildAndEarn),
+            contract_stake_snapshot
+                .get(current_era, period_number)
+                .expect("Entry must exist."),
             ContractStake::<Test>::get(&smart_contract_id)
-                .staked_amount(1, Subperiod::BuildAndEarn),
-            "Build&Earn staked amount should not change."
+                .get(current_era, period_number)
+                .expect("Entry must exist."),
+            "Ongoing era staked amount must not change."
         );
     })
 }

From c246c2910cf8faa0373003eeb23e933c41faaac8 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Fri, 1 Mar 2024 17:12:37 +0100
Subject: [PATCH 06/28] Improvements start

---
 pallets/dapp-staking-v3/src/lib.rs            | 21 +++----------------
 .../dapp-staking-v3/src/test/testing_utils.rs | 12 +----------
 pallets/dapp-staking-v3/src/types.rs          | 14 +++----------
 3 files changed, 7 insertions(+), 40 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/lib.rs b/pallets/dapp-staking-v3/src/lib.rs
index f232ec2569..7adcde7e4b 100644
--- a/pallets/dapp-staking-v3/src/lib.rs
+++ b/pallets/dapp-staking-v3/src/lib.rs
@@ -1106,7 +1106,7 @@ pub mod pallet {
 
             // 1.
             // Update `StakerInfo` storage with the reduced stake amount on the specified contract.
-            let (new_staking_info, voting_unstake_amount, bep_unstake_amount) =
+            let (new_staking_info, amount) =
                 match StakerInfo::<T>::get(&account, &smart_contract) {
                     Some(mut staking_info) => {
                         ensure!(
@@ -1128,28 +1128,14 @@ pub mod pallet {
                             amount
                         };
 
-                        // We need to know the exact amounts to unstake from each subperiod.
-                        let (voting_unstake_amount, bep_unstake_amount) =
-                            if staking_info.staked_amount(Subperiod::BuildAndEarn) >= amount {
-                                (Balance::zero(), amount)
-                            } else {
-                                (
-                                    amount.saturating_sub(
-                                        staking_info.staked_amount(Subperiod::BuildAndEarn),
-                                    ),
-                                    staking_info.staked_amount(Subperiod::BuildAndEarn),
-                                )
-                            };
-
                         staking_info.unstake(amount, current_era, protocol_state.subperiod());
 
-                        (staking_info, voting_unstake_amount, bep_unstake_amount)
+                        (staking_info, amount)
                     }
                     None => {
                         return Err(Error::<T>::NoStakingInfo.into());
                     }
                 };
-            let amount = voting_unstake_amount.saturating_add(bep_unstake_amount);
 
             // 2.
             // Reduce stake amount
@@ -1170,8 +1156,7 @@ pub mod pallet {
             // Update `ContractStake` storage with the reduced stake amount on the specified contract.
             let mut contract_stake_info = ContractStake::<T>::get(&dapp_info.id);
             contract_stake_info.unstake(
-                voting_unstake_amount,
-                bep_unstake_amount,
+                amount,
                 protocol_state.period_info,
                 current_era,
             );
diff --git a/pallets/dapp-staking-v3/src/test/testing_utils.rs b/pallets/dapp-staking-v3/src/test/testing_utils.rs
index def8dfa05d..59e9bb153d 100644
--- a/pallets/dapp-staking-v3/src/test/testing_utils.rs
+++ b/pallets/dapp-staking-v3/src/test/testing_utils.rs
@@ -748,17 +748,7 @@ pub(crate) fn assert_unstake(
         "Staked amount must decreased by the 'amount'"
     );
 
-    assert_eq!(
-        post_contract_stake.staked_amount(unstake_period, Subperiod::Voting),
-        pre_contract_stake.staked_amount(unstake_period, Subperiod::Voting) - expected_voting_unstake,
-        "Voting stake amount must be decreased exactly by how much staker's voting amount was reduced."
-    );
-    assert_eq!(
-        post_contract_stake.staked_amount(unstake_period, Subperiod::BuildAndEarn),
-        pre_contract_stake.staked_amount(unstake_period, Subperiod::BuildAndEarn)
-            - expected_build_and_earn_unstake,
-        "B&E stake amount must be decreased exactly by how much staker's B&E amount was reduced."
-    );
+    // TODO: expand this check!
 
     // 4. verify era info
     // =========================
diff --git a/pallets/dapp-staking-v3/src/types.rs b/pallets/dapp-staking-v3/src/types.rs
index 948661e122..da3e62bdc8 100644
--- a/pallets/dapp-staking-v3/src/types.rs
+++ b/pallets/dapp-staking-v3/src/types.rs
@@ -882,13 +882,6 @@ impl StakeAmount {
             }
         }
     }
-
-    /// Saturating reduce the specified `amount` from the appropriate subperiod.
-    /// If the amount is larger than the subperiod amount, the subperiod amount is set to zero.
-    pub fn subtract_subperiods(&mut self, voting_amount: Balance, bep_amount: Balance) {
-        self.voting.saturating_reduce(voting_amount);
-        self.build_and_earn.saturating_reduce(bep_amount);
-    }
 }
 
 /// Info about an era, including the rewards, how much is locked, unlocking, etc.
@@ -1208,8 +1201,7 @@ impl ContractStakeAmount {
     /// Unstake the specified `Voting` and `Build&Earn` subperiod amounts from the contract, for the specified `subperiod` and `era`.
     pub fn unstake(
         &mut self,
-        voting_amount: Balance,
-        bep_amount: Balance,
+        amount: Balance,
         period_info: PeriodInfo,
         current_era: EraNumber,
     ) {
@@ -1232,9 +1224,9 @@ impl ContractStakeAmount {
         }
 
         // Subtract both amounts
-        self.staked.subtract_subperiods(voting_amount, bep_amount);
+        self.staked.subtract(amount, period_info.subperiod);
         if let Some(stake_amount) = self.staked_future.as_mut() {
-            stake_amount.subtract_subperiods(voting_amount, bep_amount);
+            stake_amount.subtract(amount, period_info.subperiod);
         }
 
         // Convenience cleanup

From c7afa2bbfcc10a1b398564e6fdaa217471460669 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Fri, 1 Mar 2024 19:15:47 +0100
Subject: [PATCH 07/28] Fix for the issue cont.

---
 pallets/dapp-staking-v3/src/lib.rs      | 10 ++-
 pallets/dapp-staking-v3/src/test/mod.rs |  3 +-
 pallets/dapp-staking-v3/src/types.rs    | 99 +++++++++++++++++++++----
 3 files changed, 93 insertions(+), 19 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/lib.rs b/pallets/dapp-staking-v3/src/lib.rs
index 7adcde7e4b..94fc364f46 100644
--- a/pallets/dapp-staking-v3/src/lib.rs
+++ b/pallets/dapp-staking-v3/src/lib.rs
@@ -1106,7 +1106,7 @@ pub mod pallet {
 
             // 1.
             // Update `StakerInfo` storage with the reduced stake amount on the specified contract.
-            let (new_staking_info, amount) =
+            let (new_staking_info, amount, unstake_info) =
                 match StakerInfo::<T>::get(&account, &smart_contract) {
                     Some(mut staking_info) => {
                         ensure!(
@@ -1128,9 +1128,10 @@ pub mod pallet {
                             amount
                         };
 
-                        staking_info.unstake(amount, current_era, protocol_state.subperiod());
+                        let unstake_info =
+                            staking_info.unstake(amount, current_era, protocol_state.subperiod());
 
-                        (staking_info, amount)
+                        (staking_info, amount, unstake_info)
                     }
                     None => {
                         return Err(Error::<T>::NoStakingInfo.into());
@@ -1155,10 +1156,11 @@ pub mod pallet {
             // 3.
             // Update `ContractStake` storage with the reduced stake amount on the specified contract.
             let mut contract_stake_info = ContractStake::<T>::get(&dapp_info.id);
-            contract_stake_info.unstake(
+            contract_stake_info.new_unstake(
                 amount,
                 protocol_state.period_info,
                 current_era,
+                unstake_info,
             );
 
             // 4.
diff --git a/pallets/dapp-staking-v3/src/test/mod.rs b/pallets/dapp-staking-v3/src/test/mod.rs
index 0774935213..4bec1fab58 100644
--- a/pallets/dapp-staking-v3/src/test/mod.rs
+++ b/pallets/dapp-staking-v3/src/test/mod.rs
@@ -19,4 +19,5 @@
 pub(crate) mod mock;
 mod testing_utils;
 mod tests;
-mod tests_types;
+// TODO: uncomment this later, once tests are fixed
+// mod tests_types;
diff --git a/pallets/dapp-staking-v3/src/types.rs b/pallets/dapp-staking-v3/src/types.rs
index da3e62bdc8..f496dd585c 100644
--- a/pallets/dapp-staking-v3/src/types.rs
+++ b/pallets/dapp-staking-v3/src/types.rs
@@ -857,6 +857,8 @@ impl StakeAmount {
         }
     }
 
+    // TODO: remove subperiod argument since it's pointless. If BEP is > 0, then it's build&earn, otherwise it's voting or build&earn but it doesn't matter.
+
     /// Unstake the specified `amount` for the specified `subperiod`.
     ///
     /// In case subperiod is `Voting`, the amount is subtracted from the voting subperiod.
@@ -977,6 +979,8 @@ impl EraInfo {
 /// Keeps track of amount staked in the 'voting subperiod', as well as 'build&earn subperiod'.
 #[derive(Encode, Decode, MaxEncodedLen, Copy, Clone, Debug, PartialEq, Eq, TypeInfo, Default)]
 pub struct SingularStakingInfo {
+    /// Amount staked before, if anything.
+    previous_staked: StakeAmount,
     /// Staked amount
     staked: StakeAmount,
     /// Indicates whether a staker is a loyal staker or not.
@@ -992,11 +996,10 @@ impl SingularStakingInfo {
     /// `subperiod` - subperiod during which this entry is created.
     pub fn new(period: PeriodNumber, subperiod: Subperiod) -> Self {
         Self {
+            previous_staked: Default::default(),
             staked: StakeAmount {
-                voting: Balance::zero(),
-                build_and_earn: Balance::zero(),
-                era: 0,
                 period,
+                ..Default::default()
             },
             // Loyalty staking is only possible if stake is first made during the voting subperiod.
             loyal_staker: subperiod == Subperiod::Voting,
@@ -1005,11 +1008,17 @@ impl SingularStakingInfo {
 
     /// Stake the specified amount on the contract, for the specified subperiod.
     pub fn stake(&mut self, amount: Balance, current_era: EraNumber, subperiod: Subperiod) {
-        self.staked.add(amount, subperiod);
+        // Keep the previous stake amount for future reference
+        self.previous_staked = self.staked;
+        self.previous_staked.era = current_era;
+
         // Stake is only valid from the next era so we keep it consistent here
+        self.staked.add(amount, subperiod);
         self.staked.era = current_era.saturating_add(1);
     }
 
+    // TODO: add better docs below and in the function!
+
     /// Unstakes some of the specified amount from the contract.
     ///
     /// In case the `amount` being unstaked is larger than the amount staked in the `Voting` subperiod,
@@ -1021,10 +1030,18 @@ impl SingularStakingInfo {
         amount: Balance,
         current_era: EraNumber,
         subperiod: Subperiod,
-    ) -> (Balance, Balance) {
+    ) -> (EraNumber, Balance) {
+        // Keep the previous stake amounts for final calculation
         let snapshot = self.staked;
 
+        let stake_delta = self
+            .staked
+            .total()
+            .saturating_sub(self.previous_staked.total());
+
+        // Modify inner fields
         self.staked.subtract(amount, subperiod);
+        let unstaked_amount = snapshot.total().saturating_sub(self.staked.total());
         // Keep the latest era for which the entry is valid
         self.staked.era = self.staked.era.max(current_era);
 
@@ -1034,13 +1051,23 @@ impl SingularStakingInfo {
                 Subperiod::BuildAndEarn => self.staked.voting == snapshot.voting,
             };
 
-        // Amount that was unstaked
-        (
-            snapshot.voting.saturating_sub(self.staked.voting),
-            snapshot
-                .build_and_earn
-                .saturating_sub(self.staked.build_and_earn),
-        )
+        // 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();
+        } else {
+            self.previous_staked = snapshot;
+            self.previous_staked.era = self.staked.era.saturating_sub(1);
+        }
+
+        if stake_delta.is_zero() {
+            (0, 0)
+        } else {
+            (
+                self.staked.era.saturating_sub(1),
+                unstaked_amount.saturating_sub(stake_delta),
+            )
+        }
     }
 
     /// Total staked on the contract by the user. Both subperiod stakes are included.
@@ -1199,11 +1226,49 @@ impl ContractStakeAmount {
     }
 
     /// Unstake the specified `Voting` and `Build&Earn` subperiod amounts from the contract, for the specified `subperiod` and `era`.
-    pub fn unstake(
+    pub fn unstake(&mut self, amount: Balance, period_info: PeriodInfo, current_era: EraNumber) {
+        // First align entries - we only need to keep track of the current era, and the next one
+        match self.staked_future {
+            // Future entry exists, but it covers current or older era.
+            Some(stake_amount)
+                if stake_amount.era <= current_era && stake_amount.period == period_info.number =>
+            {
+                self.staked = stake_amount;
+                self.staked.era = current_era;
+                self.staked_future = None;
+            }
+            _ => (),
+        }
+
+        // Current entry is from the right period, but older era. Shift it to the current era.
+        if self.staked.era < current_era && self.staked.period == period_info.number {
+            self.staked.era = current_era;
+        }
+
+        // Subtract both amounts
+        self.staked.subtract(amount, period_info.subperiod);
+        if let Some(stake_amount) = self.staked_future.as_mut() {
+            stake_amount.subtract(amount, period_info.subperiod);
+        }
+
+        // Convenience cleanup
+        if self.staked.is_empty() {
+            self.staked = Default::default();
+        }
+        if let Some(stake_amount) = self.staked_future {
+            if stake_amount.is_empty() {
+                self.staked_future = None;
+            }
+        }
+    }
+
+    // TODO
+    pub fn new_unstake(
         &mut self,
         amount: Balance,
         period_info: PeriodInfo,
         current_era: EraNumber,
+        additional_unstake_info: (EraNumber, Balance),
     ) {
         // First align entries - we only need to keep track of the current era, and the next one
         match self.staked_future {
@@ -1224,9 +1289,15 @@ impl ContractStakeAmount {
         }
 
         // Subtract both amounts
-        self.staked.subtract(amount, period_info.subperiod);
         if let Some(stake_amount) = self.staked_future.as_mut() {
             stake_amount.subtract(amount, period_info.subperiod);
+
+            let (old_era, old_amount) = additional_unstake_info;
+            if self.staked.era == old_era {
+                self.staked.subtract(old_amount, period_info.subperiod);
+            }
+        } else {
+            self.staked.subtract(amount, period_info.subperiod);
         }
 
         // Convenience cleanup

From 9c63ab40253ac8c84c5169c837ab17b32ce4d514 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Mon, 4 Mar 2024 12:09:03 +0100
Subject: [PATCH 08/28] Expanded test, cleanup

---
 .../dapp-staking-v3/src/test/testing_utils.rs | 11 --------
 pallets/dapp-staking-v3/src/test/tests.rs     | 26 +++++++++++++++++++
 pallets/dapp-staking-v3/src/types.rs          |  4 +--
 3 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/test/testing_utils.rs b/pallets/dapp-staking-v3/src/test/testing_utils.rs
index 59e9bb153d..6372a0d5a7 100644
--- a/pallets/dapp-staking-v3/src/test/testing_utils.rs
+++ b/pallets/dapp-staking-v3/src/test/testing_utils.rs
@@ -641,17 +641,6 @@ pub(crate) fn assert_unstake(
         amount
     };
 
-    // Calculate expected unstake amounts from the voting and build&earn subperiods.
-    let (expected_voting_unstake, expected_build_and_earn_unstake) =
-        if pre_staker_info.staked_amount(Subperiod::BuildAndEarn) >= amount {
-            (Balance::zero(), amount)
-        } else {
-            (
-                amount - pre_staker_info.staked_amount(Subperiod::BuildAndEarn),
-                pre_staker_info.staked_amount(Subperiod::BuildAndEarn),
-            )
-        };
-
     // Unstake from smart contract & verify event
     assert_ok!(DappStaking::unstake(
         RuntimeOrigin::signed(account),
diff --git a/pallets/dapp-staking-v3/src/test/tests.rs b/pallets/dapp-staking-v3/src/test/tests.rs
index e1dd357a4d..129e074da8 100644
--- a/pallets/dapp-staking-v3/src/test/tests.rs
+++ b/pallets/dapp-staking-v3/src/test/tests.rs
@@ -2690,6 +2690,32 @@ fn stake_and_unstake_correctly_updates_staked_amounts() {
                 .expect("Entry must exist."),
             "Ongoing era staked amount must not change."
         );
+
+        // 4th scenario: Unstake with more than was staked for the next era
+        let delta = 5;
+        let amount_3 = amount_2 + delta;
+        assert_stake(account_2, &smart_contract, amount_2);
+
+        let contract_stake_snapshot = ContractStake::<Test>::get(&smart_contract_id);
+        for _ in 0..20 {
+            assert_unstake(account_2, &smart_contract, amount_3);
+            assert_stake(account_2, &smart_contract, amount_3);
+        }
+
+        // Check that the contract stake snapshot staked amount is the same as before
+        let current_era = ActiveProtocolState::<Test>::get().era;
+        assert_eq!(
+            contract_stake_snapshot
+                .get(current_era, period_number)
+                .expect("Entry must exist.")
+                .total(),
+            ContractStake::<Test>::get(&smart_contract_id)
+                .get(current_era, period_number)
+                .expect("Entry must exist.")
+                .total()
+                + delta,
+            "Ongoing era stake must be reduced by the `delta` amount."
+        );
     })
 }
 
diff --git a/pallets/dapp-staking-v3/src/types.rs b/pallets/dapp-staking-v3/src/types.rs
index f496dd585c..7f1ef3a317 100644
--- a/pallets/dapp-staking-v3/src/types.rs
+++ b/pallets/dapp-staking-v3/src/types.rs
@@ -1024,14 +1024,14 @@ impl SingularStakingInfo {
     /// In case the `amount` being unstaked is larger than the amount staked in the `Voting` subperiod,
     /// and `Voting` subperiod has passed, this will remove the _loyalty_ flag from the staker.
     ///
-    /// Returns the amount that was unstaked from the `Voting` subperiod stake, and from the `Build&Earn` subperiod stake.
+    /// Returns how much was unstaked from the _previous_ staked era.
     pub fn unstake(
         &mut self,
         amount: Balance,
         current_era: EraNumber,
         subperiod: Subperiod,
     ) -> (EraNumber, Balance) {
-        // Keep the previous stake amounts for final calculation
+        // Keep the previous stake amounts for the final calculation
         let snapshot = self.staked;
 
         let stake_delta = self

From 6a71a85b16aadbb74cfa1266d5f52a2b5319fa4d Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Mon, 4 Mar 2024 12:52:45 +0100
Subject: [PATCH 09/28] Additional checks

---
 .../dapp-staking-v3/src/test/testing_utils.rs | 31 ++++++++++++++++++-
 pallets/dapp-staking-v3/src/types.rs          |  6 ++--
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/test/testing_utils.rs b/pallets/dapp-staking-v3/src/test/testing_utils.rs
index 6372a0d5a7..4fa425aa1d 100644
--- a/pallets/dapp-staking-v3/src/test/testing_utils.rs
+++ b/pallets/dapp-staking-v3/src/test/testing_utils.rs
@@ -721,6 +721,8 @@ pub(crate) fn assert_unstake(
         );
     }
 
+    // TODO: expand checks for previous_staked entry!
+
     // 3. verify contract stake
     // =========================
     // =========================
@@ -737,7 +739,34 @@ pub(crate) fn assert_unstake(
         "Staked amount must decreased by the 'amount'"
     );
 
-    // TODO: expand this check!
+    let staker_delta_stake =
+        pre_staker_info.staked.total() - pre_staker_info.previous_staked.total();
+
+    // 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 current_era = pre_snapshot.active_protocol_state.era;
+        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."
+        );
+    }
 
     // 4. verify era info
     // =========================
diff --git a/pallets/dapp-staking-v3/src/types.rs b/pallets/dapp-staking-v3/src/types.rs
index 7f1ef3a317..0069c73415 100644
--- a/pallets/dapp-staking-v3/src/types.rs
+++ b/pallets/dapp-staking-v3/src/types.rs
@@ -980,11 +980,11 @@ impl EraInfo {
 #[derive(Encode, Decode, MaxEncodedLen, Copy, Clone, Debug, PartialEq, Eq, TypeInfo, Default)]
 pub struct SingularStakingInfo {
     /// Amount staked before, if anything.
-    previous_staked: StakeAmount,
+    pub(crate) previous_staked: StakeAmount,
     /// Staked amount
-    staked: StakeAmount,
+    pub(crate) staked: StakeAmount,
     /// Indicates whether a staker is a loyal staker or not.
-    loyal_staker: bool,
+    pub(crate) loyal_staker: bool,
 }
 
 impl SingularStakingInfo {

From d97b93665c0d4d349eff192bcf78f90f7510669e Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Mon, 4 Mar 2024 13:11:52 +0100
Subject: [PATCH 10/28] More checks, resolve TODOs

---
 pallets/dapp-staking-v3/src/lib.rs            |  2 +-
 .../dapp-staking-v3/src/test/testing_utils.rs | 23 +++++++---
 pallets/dapp-staking-v3/src/types.rs          | 45 +++----------------
 3 files changed, 23 insertions(+), 47 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/lib.rs b/pallets/dapp-staking-v3/src/lib.rs
index 94fc364f46..49812d04f3 100644
--- a/pallets/dapp-staking-v3/src/lib.rs
+++ b/pallets/dapp-staking-v3/src/lib.rs
@@ -1156,7 +1156,7 @@ pub mod pallet {
             // 3.
             // Update `ContractStake` storage with the reduced stake amount on the specified contract.
             let mut contract_stake_info = ContractStake::<T>::get(&dapp_info.id);
-            contract_stake_info.new_unstake(
+            contract_stake_info.unstake(
                 amount,
                 protocol_state.period_info,
                 current_era,
diff --git a/pallets/dapp-staking-v3/src/test/testing_utils.rs b/pallets/dapp-staking-v3/src/test/testing_utils.rs
index 4fa425aa1d..744dc7f8b1 100644
--- a/pallets/dapp-staking-v3/src/test/testing_utils.rs
+++ b/pallets/dapp-staking-v3/src/test/testing_utils.rs
@@ -679,6 +679,9 @@ pub(crate) fn assert_unstake(
     // 2. verify staker info
     // =====================
     // =====================
+    let staker_delta_stake =
+        pre_staker_info.staked.total() - pre_staker_info.previous_staked.total();
+
     if is_full_unstake {
         assert!(
             !StakerInfo::<Test>::contains_key(&account, smart_contract),
@@ -719,9 +722,20 @@ pub(crate) fn assert_unstake(
             is_loyal,
             "If 'Voting' stake amount is reduced in B&E period, loyalty flag must be set to false."
         );
-    }
 
-    // TODO: expand checks for previous_staked entry!
+        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
+            );
+        }
+    }
 
     // 3. verify contract stake
     // =========================
@@ -739,15 +753,12 @@ pub(crate) fn assert_unstake(
         "Staked amount must decreased by the 'amount'"
     );
 
-    let staker_delta_stake =
-        pre_staker_info.staked.total() - pre_staker_info.previous_staked.total();
-
     // 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 current_era = pre_snapshot.active_protocol_state.era;
+
         let latest_contract_stake_era = post_contract_stake
             .latest_stake_era()
             .expect("Has to be defined");
diff --git a/pallets/dapp-staking-v3/src/types.rs b/pallets/dapp-staking-v3/src/types.rs
index 0069c73415..9807f530e1 100644
--- a/pallets/dapp-staking-v3/src/types.rs
+++ b/pallets/dapp-staking-v3/src/types.rs
@@ -1017,8 +1017,6 @@ impl SingularStakingInfo {
         self.staked.era = current_era.saturating_add(1);
     }
 
-    // TODO: add better docs below and in the function!
-
     /// Unstakes some of the specified amount from the contract.
     ///
     /// In case the `amount` being unstaked is larger than the amount staked in the `Voting` subperiod,
@@ -1030,6 +1028,7 @@ impl SingularStakingInfo {
         amount: Balance,
         current_era: EraNumber,
         subperiod: Subperiod,
+        // TODO: change the entire approach to return vector of unstaked amounts per era (more generic that way?) Item for the future?
     ) -> (EraNumber, Balance) {
         // Keep the previous stake amounts for the final calculation
         let snapshot = self.staked;
@@ -1065,6 +1064,8 @@ impl SingularStakingInfo {
         } else {
             (
                 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),
             )
         }
@@ -1226,44 +1227,7 @@ impl ContractStakeAmount {
     }
 
     /// Unstake the specified `Voting` and `Build&Earn` subperiod amounts from the contract, for the specified `subperiod` and `era`.
-    pub fn unstake(&mut self, amount: Balance, period_info: PeriodInfo, current_era: EraNumber) {
-        // First align entries - we only need to keep track of the current era, and the next one
-        match self.staked_future {
-            // Future entry exists, but it covers current or older era.
-            Some(stake_amount)
-                if stake_amount.era <= current_era && stake_amount.period == period_info.number =>
-            {
-                self.staked = stake_amount;
-                self.staked.era = current_era;
-                self.staked_future = None;
-            }
-            _ => (),
-        }
-
-        // Current entry is from the right period, but older era. Shift it to the current era.
-        if self.staked.era < current_era && self.staked.period == period_info.number {
-            self.staked.era = current_era;
-        }
-
-        // Subtract both amounts
-        self.staked.subtract(amount, period_info.subperiod);
-        if let Some(stake_amount) = self.staked_future.as_mut() {
-            stake_amount.subtract(amount, period_info.subperiod);
-        }
-
-        // Convenience cleanup
-        if self.staked.is_empty() {
-            self.staked = Default::default();
-        }
-        if let Some(stake_amount) = self.staked_future {
-            if stake_amount.is_empty() {
-                self.staked_future = None;
-            }
-        }
-    }
-
-    // TODO
-    pub fn new_unstake(
+    pub fn unstake(
         &mut self,
         amount: Balance,
         period_info: PeriodInfo,
@@ -1297,6 +1261,7 @@ impl ContractStakeAmount {
                 self.staked.subtract(old_amount, period_info.subperiod);
             }
         } else {
+            // No future entry exists, so only current one needs to be updated.
             self.staked.subtract(amount, period_info.subperiod);
         }
 

From 9717b9d9b80b7d80eee1f3ee95217e8884d631d2 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Mon, 4 Mar 2024 14:09:56 +0100
Subject: [PATCH 11/28] Working solution

---
 pallets/dapp-staking-v3/src/test/mod.rs       |  3 +-
 .../dapp-staking-v3/src/test/tests_types.rs   | 56 +++++++++++++------
 pallets/dapp-staking-v3/src/types.rs          |  4 ++
 3 files changed, 44 insertions(+), 19 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/test/mod.rs b/pallets/dapp-staking-v3/src/test/mod.rs
index 4bec1fab58..0774935213 100644
--- a/pallets/dapp-staking-v3/src/test/mod.rs
+++ b/pallets/dapp-staking-v3/src/test/mod.rs
@@ -19,5 +19,4 @@
 pub(crate) mod mock;
 mod testing_utils;
 mod tests;
-// TODO: uncomment this later, once tests are fixed
-// mod tests_types;
+mod tests_types;
diff --git a/pallets/dapp-staking-v3/src/test/tests_types.rs b/pallets/dapp-staking-v3/src/test/tests_types.rs
index 1946da7cbb..408350f70a 100644
--- a/pallets/dapp-staking-v3/src/test/tests_types.rs
+++ b/pallets/dapp-staking-v3/src/test/tests_types.rs
@@ -2093,7 +2093,7 @@ fn singular_staking_info_unstake_during_voting_is_ok() {
     let unstake_amount_1 = 5;
     assert_eq!(
         staking_info.unstake(unstake_amount_1, era_1, Subperiod::Voting),
-        (unstake_amount_1, Balance::zero())
+        (era_1, Balance::zero())
     );
     assert_eq!(
         staking_info.total_staked_amount(),
@@ -2111,7 +2111,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),
-        (remaining_stake, Balance::zero())
+        (EraNumber::zero(), Balance::zero())
     );
     assert!(staking_info.total_staked_amount().is_zero());
     assert!(
@@ -2138,7 +2138,7 @@ fn singular_staking_info_unstake_during_bep_is_ok() {
     let unstake_1 = 5;
     assert_eq!(
         staking_info.unstake(5, era_1, Subperiod::BuildAndEarn),
-        (Balance::zero(), unstake_1)
+        (era_1, Balance::zero())
     );
     assert_eq!(
         staking_info.total_staked_amount(),
@@ -2169,7 +2169,7 @@ fn singular_staking_info_unstake_during_bep_is_ok() {
 
     assert_eq!(
         staking_info.unstake(unstake_2, era_2, Subperiod::BuildAndEarn),
-        (voting_stake_overflow, current_bep_stake)
+        (EraNumber::zero(), Balance::zero())
     );
     assert_eq!(
         staking_info.total_staked_amount(),
@@ -2458,7 +2458,12 @@ fn contract_stake_amount_unstake_is_ok() {
 
     // 1st scenario - unstake in the same era, from `Voting` subperiod
     let amount_1 = 5;
-    contract_stake.unstake(amount_1, 0, period_info, era_1);
+    contract_stake.unstake(
+        amount_1,
+        period_info,
+        era_1,
+        (EraNumber::zero(), Balance::zero()),
+    );
     assert_eq!(
         contract_stake.total_staked_amount(period),
         total_stake_amount - amount_1
@@ -2477,19 +2482,23 @@ fn contract_stake_amount_unstake_is_ok() {
         next_subperiod_start_era: 40,
     };
 
-    contract_stake.unstake(amount_1, 0, period_info, era_2);
+    contract_stake.unstake(
+        amount_1,
+        period_info,
+        era_2,
+        (EraNumber::zero(), Balance::zero()),
+    );
     assert_eq!(
         contract_stake.total_staked_amount(period),
         total_stake_amount - amount_1 * 2
     );
     assert_eq!(
         contract_stake.staked_amount(period, Subperiod::Voting),
-        vp_stake_amount - amount_1 * 2
+        vp_stake_amount - amount_1
     );
     assert_eq!(
         contract_stake.staked_amount(period, Subperiod::BuildAndEarn),
-        bep_stake_amount,
-        "Must remain unchanged."
+        bep_stake_amount - amount_1,
     );
     assert!(
         !contract_stake.staked.is_empty(),
@@ -2501,35 +2510,45 @@ fn contract_stake_amount_unstake_is_ok() {
     );
 
     // 3rd scenario - same as previous scenario, but for `BuildAndEarn` subperiod
-    contract_stake.unstake(0, amount_1, period_info, era_2);
+    contract_stake.unstake(
+        amount_1,
+        period_info,
+        era_2,
+        (EraNumber::zero(), Balance::zero()),
+    );
     assert_eq!(
         contract_stake.total_staked_amount(period),
         total_stake_amount - amount_1 * 3
     );
     assert_eq!(
         contract_stake.staked_amount(period, Subperiod::Voting),
-        vp_stake_amount - amount_1 * 2
+        vp_stake_amount - amount_1
     );
     assert_eq!(
         contract_stake.staked_amount(period, Subperiod::BuildAndEarn),
-        bep_stake_amount - amount_1
+        bep_stake_amount - amount_1 * 2
     );
 
     // 4th scenario - bump up unstake eras by more than 1, entries should be aligned to the current era
     let era_3 = era_2 + 3;
     let amount_2 = 7;
-    contract_stake.unstake(0, amount_2, period_info, era_3);
+    contract_stake.unstake(
+        amount_2,
+        period_info,
+        era_3,
+        (EraNumber::zero(), Balance::zero()),
+    );
     assert_eq!(
         contract_stake.total_staked_amount(period),
         total_stake_amount - amount_1 * 3 - amount_2
     );
     assert_eq!(
         contract_stake.staked_amount(period, Subperiod::Voting),
-        vp_stake_amount - amount_1 * 2
+        vp_stake_amount - amount_1
     );
     assert_eq!(
         contract_stake.staked_amount(period, Subperiod::BuildAndEarn),
-        bep_stake_amount - amount_1 - amount_2
+        bep_stake_amount - amount_1 * 2 - amount_2
     );
     assert_eq!(
         contract_stake.staked.era, era_3,
@@ -2542,12 +2561,15 @@ fn contract_stake_amount_unstake_is_ok() {
 
     // 5th scenario - do a full unstake with existing future entry, expect a cleanup
     contract_stake.stake(total_stake_amount, period_info, era_3);
+
+    println!("{:?}", contract_stake);
     contract_stake.unstake(
-        contract_stake.staked_amount(period, Subperiod::Voting),
-        contract_stake.staked_amount(period, Subperiod::BuildAndEarn),
+        contract_stake.total_staked_amount(period),
         period_info,
         era_3,
+        (EraNumber::zero(), Balance::zero()),
     );
+    println!("Contract stake after: {:?}", contract_stake);
     assert!(contract_stake.staked.is_empty());
     assert!(contract_stake.staked_future.is_none());
 }
diff --git a/pallets/dapp-staking-v3/src/types.rs b/pallets/dapp-staking-v3/src/types.rs
index 9807f530e1..ea773c6eaa 100644
--- a/pallets/dapp-staking-v3/src/types.rs
+++ b/pallets/dapp-staking-v3/src/types.rs
@@ -1252,6 +1252,8 @@ impl ContractStakeAmount {
             self.staked.era = current_era;
         }
 
+        // TODO: there should be a cleaner way to do this
+
         // Subtract both amounts
         if let Some(stake_amount) = self.staked_future.as_mut() {
             stake_amount.subtract(amount, period_info.subperiod);
@@ -1259,6 +1261,8 @@ impl ContractStakeAmount {
             let (old_era, old_amount) = additional_unstake_info;
             if self.staked.era == old_era {
                 self.staked.subtract(old_amount, period_info.subperiod);
+            } else if stake_amount.total().is_zero() {
+                self.staked = Default::default();
             }
         } else {
             // No future entry exists, so only current one needs to be updated.

From 6d13b54aaf1b7292f0a1c9a7ea1268054acf7bdb Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Mon, 4 Mar 2024 16:19:16 +0100
Subject: [PATCH 12/28] Remove println

---
 pallets/dapp-staking-v3/src/test/tests_types.rs | 2 --
 1 file changed, 2 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/test/tests_types.rs b/pallets/dapp-staking-v3/src/test/tests_types.rs
index 408350f70a..cf4af28323 100644
--- a/pallets/dapp-staking-v3/src/test/tests_types.rs
+++ b/pallets/dapp-staking-v3/src/test/tests_types.rs
@@ -2562,14 +2562,12 @@ fn contract_stake_amount_unstake_is_ok() {
     // 5th scenario - do a full unstake with existing future entry, expect a cleanup
     contract_stake.stake(total_stake_amount, period_info, era_3);
 
-    println!("{:?}", contract_stake);
     contract_stake.unstake(
         contract_stake.total_staked_amount(period),
         period_info,
         era_3,
         (EraNumber::zero(), Balance::zero()),
     );
-    println!("Contract stake after: {:?}", contract_stake);
     assert!(contract_stake.staked.is_empty());
     assert!(contract_stake.staked_future.is_none());
 }

From 72d6f05658a92bafdb4015f9df1db02e35f39498 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Mon, 4 Mar 2024 17:33:49 +0100
Subject: [PATCH 13/28] Refactoring

---
 pallets/dapp-staking-v3/src/lib.rs      |  9 ++--
 pallets/dapp-staking-v3/src/test/mod.rs |  3 +-
 pallets/dapp-staking-v3/src/types.rs    | 69 ++++++++++++++-----------
 3 files changed, 44 insertions(+), 37 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/lib.rs b/pallets/dapp-staking-v3/src/lib.rs
index 49812d04f3..26a94cb33d 100644
--- a/pallets/dapp-staking-v3/src/lib.rs
+++ b/pallets/dapp-staking-v3/src/lib.rs
@@ -1106,7 +1106,7 @@ pub mod pallet {
 
             // 1.
             // Update `StakerInfo` storage with the reduced stake amount on the specified contract.
-            let (new_staking_info, amount, unstake_info) =
+            let (new_staking_info, era_and_amount_pairs) =
                 match StakerInfo::<T>::get(&account, &smart_contract) {
                     Some(mut staking_info) => {
                         ensure!(
@@ -1128,10 +1128,10 @@ pub mod pallet {
                             amount
                         };
 
-                        let unstake_info =
+                        let era_and_amount_pairs =
                             staking_info.unstake(amount, current_era, protocol_state.subperiod());
 
-                        (staking_info, amount, unstake_info)
+                        (staking_info, era_and_amount_pairs)
                     }
                     None => {
                         return Err(Error::<T>::NoStakingInfo.into());
@@ -1157,10 +1157,9 @@ pub mod pallet {
             // Update `ContractStake` storage with the reduced stake amount on the specified contract.
             let mut contract_stake_info = ContractStake::<T>::get(&dapp_info.id);
             contract_stake_info.unstake(
-                amount,
+                era_and_amount_pairs,
                 protocol_state.period_info,
                 current_era,
-                unstake_info,
             );
 
             // 4.
diff --git a/pallets/dapp-staking-v3/src/test/mod.rs b/pallets/dapp-staking-v3/src/test/mod.rs
index 0774935213..d98484de4e 100644
--- a/pallets/dapp-staking-v3/src/test/mod.rs
+++ b/pallets/dapp-staking-v3/src/test/mod.rs
@@ -19,4 +19,5 @@
 pub(crate) mod mock;
 mod testing_utils;
 mod tests;
-mod tests_types;
+// TODO: return this later, and fix the type tests
+// mod tests_types;
diff --git a/pallets/dapp-staking-v3/src/types.rs b/pallets/dapp-staking-v3/src/types.rs
index ea773c6eaa..d7166deac9 100644
--- a/pallets/dapp-staking-v3/src/types.rs
+++ b/pallets/dapp-staking-v3/src/types.rs
@@ -1022,28 +1022,32 @@ impl SingularStakingInfo {
     /// In case the `amount` being unstaked is larger than the amount staked in the `Voting` subperiod,
     /// and `Voting` subperiod has passed, this will remove the _loyalty_ flag from the staker.
     ///
-    /// Returns how much was unstaked from the _previous_ staked era.
+    /// Returns a vector of `(era, amount)` pairs, where `era` is the era in which the unstake happened,
+    /// and the amount is the corresponding amount.
     pub fn unstake(
         &mut self,
         amount: Balance,
         current_era: EraNumber,
         subperiod: Subperiod,
-        // TODO: change the entire approach to return vector of unstaked amounts per era (more generic that way?) Item for the future?
-    ) -> (EraNumber, Balance) {
-        // Keep the previous stake amounts for the final calculation
-        let snapshot = self.staked;
+    ) -> Vec<(EraNumber, Balance)> {
+        let mut result = Vec::new();
 
+        // Prepare values for further calculations
+        let snapshot = self.staked;
         let stake_delta = self
             .staked
             .total()
             .saturating_sub(self.previous_staked.total());
 
-        // Modify inner fields
+        // Modify current staked amount
         self.staked.subtract(amount, subperiod);
         let unstaked_amount = snapshot.total().saturating_sub(self.staked.total());
-        // Keep the latest era for which the entry is valid
         self.staked.era = self.staked.era.max(current_era);
 
+        // Push the result into result vector
+        result.push((self.staked.era, unstaked_amount));
+
+        // Update loyal staker flag accordingly
         self.loyal_staker = self.loyal_staker
             && match subperiod {
                 Subperiod::Voting => !self.staked.voting.is_zero(),
@@ -1059,16 +1063,19 @@ impl SingularStakingInfo {
             self.previous_staked.era = self.staked.era.saturating_sub(1);
         }
 
-        if stake_delta.is_zero() {
-            (0, 0)
-        } else {
-            (
-                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),
+        if unstaked_amount > stake_delta {
+            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),
+                ),
             )
         }
+
+        result
     }
 
     /// Total staked on the contract by the user. Both subperiod stakes are included.
@@ -1229,12 +1236,12 @@ impl ContractStakeAmount {
     /// Unstake the specified `Voting` and `Build&Earn` subperiod amounts from the contract, for the specified `subperiod` and `era`.
     pub fn unstake(
         &mut self,
-        amount: Balance,
+        era_and_amount_pairs: Vec<(EraNumber, Balance)>,
         period_info: PeriodInfo,
         current_era: EraNumber,
-        additional_unstake_info: (EraNumber, Balance),
     ) {
-        // First align entries - we only need to keep track of the current era, and the next one
+        // 1. Entry alignment
+        // e only need to keep track of the current era, and the next one
         match self.staked_future {
             // Future entry exists, but it covers current or older era.
             Some(stake_amount)
@@ -1252,24 +1259,24 @@ impl ContractStakeAmount {
             self.staked.era = current_era;
         }
 
-        // TODO: there should be a cleaner way to do this
+        // 2. Value updates - only after alignment
 
-        // Subtract both amounts
-        if let Some(stake_amount) = self.staked_future.as_mut() {
-            stake_amount.subtract(amount, period_info.subperiod);
+        for (era, amount) in era_and_amount_pairs {
+            if self.staked.era == era {
+                self.staked.subtract(amount, period_info.subperiod);
+                continue;
+            }
 
-            let (old_era, old_amount) = additional_unstake_info;
-            if self.staked.era == old_era {
-                self.staked.subtract(old_amount, period_info.subperiod);
-            } else if stake_amount.total().is_zero() {
-                self.staked = Default::default();
+            match self.staked_future.as_mut() {
+                Some(future_stake_amount) if future_stake_amount.era == era => {
+                    future_stake_amount.subtract(amount, period_info.subperiod);
+                }
+                // Otherwise do nothing
+                _ => (),
             }
-        } else {
-            // No future entry exists, so only current one needs to be updated.
-            self.staked.subtract(amount, period_info.subperiod);
         }
 
-        // Convenience cleanup
+        // 3. Convenience cleanup
         if self.staked.is_empty() {
             self.staked = Default::default();
         }

From fcaa8d5127a6d84760da4f00f3639520cdbd70cb Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Mon, 4 Mar 2024 17:53:10 +0100
Subject: [PATCH 14/28] Remove redundant args

---
 pallets/dapp-staking-v3/src/lib.rs   |  4 +-
 pallets/dapp-staking-v3/src/types.rs | 57 ++++++++++++----------------
 2 files changed, 26 insertions(+), 35 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/lib.rs b/pallets/dapp-staking-v3/src/lib.rs
index 26a94cb33d..bbc2d835ed 100644
--- a/pallets/dapp-staking-v3/src/lib.rs
+++ b/pallets/dapp-staking-v3/src/lib.rs
@@ -1165,7 +1165,7 @@ pub mod pallet {
             // 4.
             // Update total staked amount for the next era.
             CurrentEraInfo::<T>::mutate(|era_info| {
-                era_info.unstake_amount(amount, protocol_state.subperiod());
+                era_info.unstake_amount(amount);
             });
 
             // 5.
@@ -1454,7 +1454,7 @@ pub mod pallet {
             // This means 'fake' stake total amount has been kept until now, even though contract was unregistered.
             // Although strange, it's been requested to keep it like this from the team.
             CurrentEraInfo::<T>::mutate(|era_info| {
-                era_info.unstake_amount(amount, protocol_state.subperiod());
+                era_info.unstake_amount(amount);
             });
 
             // Update remaining storage entries
diff --git a/pallets/dapp-staking-v3/src/types.rs b/pallets/dapp-staking-v3/src/types.rs
index d7166deac9..cfee04adbb 100644
--- a/pallets/dapp-staking-v3/src/types.rs
+++ b/pallets/dapp-staking-v3/src/types.rs
@@ -629,7 +629,7 @@ where
             return Err(AccountLedgerError::UnstakeAmountLargerThanStake);
         }
 
-        self.staked.subtract(amount, current_period_info.subperiod);
+        self.staked.subtract(amount);
 
         // Convenience cleanup
         if self.staked.is_empty() {
@@ -637,7 +637,7 @@ where
         }
 
         if let Some(mut stake_amount) = self.staked_future {
-            stake_amount.subtract(amount, current_period_info.subperiod);
+            stake_amount.subtract(amount);
 
             self.staked_future = if stake_amount.is_empty() {
                 None
@@ -857,31 +857,22 @@ impl StakeAmount {
         }
     }
 
-    // TODO: remove subperiod argument since it's pointless. If BEP is > 0, then it's build&earn, otherwise it's voting or build&earn but it doesn't matter.
-
-    /// Unstake the specified `amount` for the specified `subperiod`.
-    ///
-    /// In case subperiod is `Voting`, the amount is subtracted from the voting subperiod.
+    /// Unstake the specified `amount`.
     ///
-    /// In case subperiod is `Build&Earn`, the amount is first subtracted from the
-    /// build&earn amount, and any rollover is subtracted from the voting subperiod.
-    pub fn subtract(&mut self, amount: Balance, subperiod: Subperiod) {
-        match subperiod {
-            Subperiod::Voting => self.voting.saturating_reduce(amount),
-            Subperiod::BuildAndEarn => {
-                if self.build_and_earn >= amount {
-                    self.build_and_earn.saturating_reduce(amount);
-                } else {
-                    // Rollover from build&earn to voting, is guaranteed to be larger than zero due to previous check
-                    // E.g. voting = 10, build&earn = 5, amount = 7
-                    // underflow = build&earn - amount = 5 - 7 = -2
-                    // voting = 10 - 2 = 8
-                    // build&earn = 0
-                    let remainder = amount.saturating_sub(self.build_and_earn);
-                    self.build_and_earn = Balance::zero();
-                    self.voting.saturating_reduce(remainder);
-                }
-            }
+    /// Attempt to subtract from `Build&Earn` subperiod amount is done first. Any rollover is subtracted from
+    /// the `Voting` subperiod amount.
+    pub fn subtract(&mut self, amount: Balance) {
+        if self.build_and_earn >= amount {
+            self.build_and_earn.saturating_reduce(amount);
+        } else {
+            // Rollover from build&earn to voting, is guaranteed to be larger than zero due to previous check
+            // E.g. voting = 10, build&earn = 5, amount = 7
+            // underflow = build&earn - amount = 5 - 7 = -2
+            // voting = 10 - 2 = 8
+            // build&earn = 0
+            let remainder = amount.saturating_sub(self.build_and_earn);
+            self.build_and_earn = Balance::zero();
+            self.voting.saturating_reduce(remainder);
         }
     }
 }
@@ -925,10 +916,10 @@ impl EraInfo {
         self.next_stake_amount.add(amount, subperiod);
     }
 
-    /// Subtract the specified `amount` from the appropriate stake amount, based on the `Subperiod`.
-    pub fn unstake_amount(&mut self, amount: Balance, subperiod: Subperiod) {
-        self.current_stake_amount.subtract(amount, subperiod);
-        self.next_stake_amount.subtract(amount, subperiod);
+    /// Subtract the specified `amount` from the appropriate stake amount.
+    pub fn unstake_amount(&mut self, amount: Balance) {
+        self.current_stake_amount.subtract(amount);
+        self.next_stake_amount.subtract(amount);
     }
 
     /// Total staked amount in this era.
@@ -1040,7 +1031,7 @@ impl SingularStakingInfo {
             .saturating_sub(self.previous_staked.total());
 
         // Modify current staked amount
-        self.staked.subtract(amount, subperiod);
+        self.staked.subtract(amount);
         let unstaked_amount = snapshot.total().saturating_sub(self.staked.total());
         self.staked.era = self.staked.era.max(current_era);
 
@@ -1263,13 +1254,13 @@ impl ContractStakeAmount {
 
         for (era, amount) in era_and_amount_pairs {
             if self.staked.era == era {
-                self.staked.subtract(amount, period_info.subperiod);
+                self.staked.subtract(amount);
                 continue;
             }
 
             match self.staked_future.as_mut() {
                 Some(future_stake_amount) if future_stake_amount.era == era => {
-                    future_stake_amount.subtract(amount, period_info.subperiod);
+                    future_stake_amount.subtract(amount);
                 }
                 // Otherwise do nothing
                 _ => (),

From 55377303173a2b4a00b965ba3216ccae4bd452ed Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Mon, 4 Mar 2024 18:40:46 +0100
Subject: [PATCH 15/28] Tests, comments, improvements

---
 pallets/dapp-staking-v3/src/test/mod.rs       |  3 +-
 .../dapp-staking-v3/src/test/tests_types.rs   | 97 +++++++------------
 pallets/dapp-staking-v3/src/types.rs          | 12 +--
 3 files changed, 42 insertions(+), 70 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/test/mod.rs b/pallets/dapp-staking-v3/src/test/mod.rs
index d98484de4e..0774935213 100644
--- a/pallets/dapp-staking-v3/src/test/mod.rs
+++ b/pallets/dapp-staking-v3/src/test/mod.rs
@@ -19,5 +19,4 @@
 pub(crate) mod mock;
 mod testing_utils;
 mod tests;
-// TODO: return this later, and fix the type tests
-// mod tests_types;
+mod tests_types;
diff --git a/pallets/dapp-staking-v3/src/test/tests_types.rs b/pallets/dapp-staking-v3/src/test/tests_types.rs
index cf4af28323..deaecc61e9 100644
--- a/pallets/dapp-staking-v3/src/test/tests_types.rs
+++ b/pallets/dapp-staking-v3/src/test/tests_types.rs
@@ -1832,7 +1832,7 @@ fn era_info_unstake_works() {
 
     // 1st scenario - unstake some amount, no overflow
     let unstake_amount_1 = bep_stake_amount_1;
-    era_info.unstake_amount(unstake_amount_1, Subperiod::BuildAndEarn);
+    era_info.unstake_amount(unstake_amount_1);
 
     // Current era
     assert_eq!(
@@ -1859,7 +1859,7 @@ fn era_info_unstake_works() {
     // 2nd scenario - unstake some more, but with overflow
     let overflow = 2;
     let unstake_amount_2 = bep_stake_amount_2 - unstake_amount_1 + overflow;
-    era_info.unstake_amount(unstake_amount_2, Subperiod::BuildAndEarn);
+    era_info.unstake_amount(unstake_amount_2);
 
     // Current era
     assert_eq!(
@@ -1972,58 +1972,49 @@ fn stake_amount_works() {
     assert!(stake_amount.for_type(Subperiod::Voting).is_zero());
     assert!(stake_amount.for_type(Subperiod::BuildAndEarn).is_zero());
 
-    // Stake some amount in voting period
+    // Stake some amount in voting subperiod
     let vp_stake_1 = 11;
     stake_amount.add(vp_stake_1, Subperiod::Voting);
     assert_eq!(stake_amount.total(), vp_stake_1);
     assert_eq!(stake_amount.for_type(Subperiod::Voting), vp_stake_1);
     assert!(stake_amount.for_type(Subperiod::BuildAndEarn).is_zero());
 
-    // Stake some amount in build&earn period
+    // Stake some amount in build&earn subperiod
     let bep_stake_1 = 13;
     stake_amount.add(bep_stake_1, Subperiod::BuildAndEarn);
     assert_eq!(stake_amount.total(), vp_stake_1 + bep_stake_1);
     assert_eq!(stake_amount.for_type(Subperiod::Voting), vp_stake_1);
     assert_eq!(stake_amount.for_type(Subperiod::BuildAndEarn), bep_stake_1);
 
-    // Unstake some amount from voting period
-    let vp_unstake_1 = 5;
-    stake_amount.subtract(5, Subperiod::Voting);
-    assert_eq!(
-        stake_amount.total(),
-        vp_stake_1 + bep_stake_1 - vp_unstake_1
-    );
+    // Unstake some amount, expect build&earn subperiod to be reduced
+    let unstake_1 = 5;
+    stake_amount.subtract(5);
+    assert_eq!(stake_amount.total(), vp_stake_1 + bep_stake_1 - unstake_1);
+    assert_eq!(stake_amount.for_type(Subperiod::Voting), vp_stake_1);
     assert_eq!(
-        stake_amount.for_type(Subperiod::Voting),
-        vp_stake_1 - vp_unstake_1
+        stake_amount.for_type(Subperiod::BuildAndEarn),
+        bep_stake_1 - unstake_1
     );
-    assert_eq!(stake_amount.for_type(Subperiod::BuildAndEarn), bep_stake_1);
 
-    // Unstake some amount from build&earn period
-    let bep_unstake_1 = 2;
-    stake_amount.subtract(bep_unstake_1, Subperiod::BuildAndEarn);
+    // Unstake some amount, once again expect build&earn subperiod to be reduced
+    let unstake_2 = 2;
+    stake_amount.subtract(unstake_2);
     assert_eq!(
         stake_amount.total(),
-        vp_stake_1 + bep_stake_1 - vp_unstake_1 - bep_unstake_1
-    );
-    assert_eq!(
-        stake_amount.for_type(Subperiod::Voting),
-        vp_stake_1 - vp_unstake_1
+        vp_stake_1 + bep_stake_1 - unstake_1 - unstake_2
     );
+    assert_eq!(stake_amount.for_type(Subperiod::Voting), vp_stake_1);
     assert_eq!(
         stake_amount.for_type(Subperiod::BuildAndEarn),
-        bep_stake_1 - bep_unstake_1
+        bep_stake_1 - unstake_1 - unstake_2
     );
 
-    // Unstake some more from build&earn period, and chip away from the voting period
-    let total_stake = vp_stake_1 + bep_stake_1 - vp_unstake_1 - bep_unstake_1;
-    let bep_unstake_2 = bep_stake_1 - bep_unstake_1 + 1;
-    stake_amount.subtract(bep_unstake_2, Subperiod::BuildAndEarn);
-    assert_eq!(stake_amount.total(), total_stake - bep_unstake_2);
-    assert_eq!(
-        stake_amount.for_type(Subperiod::Voting),
-        vp_stake_1 - vp_unstake_1 - 1
-    );
+    // Unstake even more, but this time expect voting subperiod amount to be reduced
+    let total_stake = vp_stake_1 + bep_stake_1 - unstake_1 - unstake_2;
+    let unstake_3 = bep_stake_1 - unstake_1 - unstake_2 + 1;
+    stake_amount.subtract(unstake_3);
+    assert_eq!(stake_amount.total(), total_stake - unstake_3);
+    assert_eq!(stake_amount.for_type(Subperiod::Voting), vp_stake_1 - 1);
     assert!(stake_amount.for_type(Subperiod::BuildAndEarn).is_zero());
 }
 
@@ -2093,7 +2084,7 @@ fn singular_staking_info_unstake_during_voting_is_ok() {
     let unstake_amount_1 = 5;
     assert_eq!(
         staking_info.unstake(unstake_amount_1, era_1, Subperiod::Voting),
-        (era_1, Balance::zero())
+        vec![(era_1 + 1, vote_stake_amount_1 - unstake_amount_1)]
     );
     assert_eq!(
         staking_info.total_staked_amount(),
@@ -2111,7 +2102,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),
-        (EraNumber::zero(), Balance::zero())
+        vec![(era_2, remaining_stake)]
     );
     assert!(staking_info.total_staked_amount().is_zero());
     assert!(
@@ -2138,7 +2129,7 @@ fn singular_staking_info_unstake_during_bep_is_ok() {
     let unstake_1 = 5;
     assert_eq!(
         staking_info.unstake(5, era_1, Subperiod::BuildAndEarn),
-        (era_1, Balance::zero())
+        vec![(era_1 + 1, bep_stake_amount_1 - unstake_1)]
     );
     assert_eq!(
         staking_info.total_staked_amount(),
@@ -2169,7 +2160,10 @@ fn singular_staking_info_unstake_during_bep_is_ok() {
 
     assert_eq!(
         staking_info.unstake(unstake_2, era_2, Subperiod::BuildAndEarn),
-        (EraNumber::zero(), Balance::zero())
+        vec![
+            (era_2 - 1, voting_stake_overflow),
+            (era_2, current_bep_stake)
+        ]
     );
     assert_eq!(
         staking_info.total_staked_amount(),
@@ -2458,12 +2452,7 @@ fn contract_stake_amount_unstake_is_ok() {
 
     // 1st scenario - unstake in the same era, from `Voting` subperiod
     let amount_1 = 5;
-    contract_stake.unstake(
-        amount_1,
-        period_info,
-        era_1,
-        (EraNumber::zero(), Balance::zero()),
-    );
+    contract_stake.unstake(vec![(era_1, amount_1)], period_info, era_1);
     assert_eq!(
         contract_stake.total_staked_amount(period),
         total_stake_amount - amount_1
@@ -2482,12 +2471,7 @@ fn contract_stake_amount_unstake_is_ok() {
         next_subperiod_start_era: 40,
     };
 
-    contract_stake.unstake(
-        amount_1,
-        period_info,
-        era_2,
-        (EraNumber::zero(), Balance::zero()),
-    );
+    contract_stake.unstake(vec![(era_2, amount_1)], period_info, era_2);
     assert_eq!(
         contract_stake.total_staked_amount(period),
         total_stake_amount - amount_1 * 2
@@ -2510,12 +2494,7 @@ fn contract_stake_amount_unstake_is_ok() {
     );
 
     // 3rd scenario - same as previous scenario, but for `BuildAndEarn` subperiod
-    contract_stake.unstake(
-        amount_1,
-        period_info,
-        era_2,
-        (EraNumber::zero(), Balance::zero()),
-    );
+    contract_stake.unstake(vec![(era_2, amount_1)], period_info, era_2);
     assert_eq!(
         contract_stake.total_staked_amount(period),
         total_stake_amount - amount_1 * 3
@@ -2532,12 +2511,7 @@ fn contract_stake_amount_unstake_is_ok() {
     // 4th scenario - bump up unstake eras by more than 1, entries should be aligned to the current era
     let era_3 = era_2 + 3;
     let amount_2 = 7;
-    contract_stake.unstake(
-        amount_2,
-        period_info,
-        era_3,
-        (EraNumber::zero(), Balance::zero()),
-    );
+    contract_stake.unstake(vec![(era_3, amount_2)], period_info, era_3);
     assert_eq!(
         contract_stake.total_staked_amount(period),
         total_stake_amount - amount_1 * 3 - amount_2
@@ -2563,10 +2537,9 @@ fn contract_stake_amount_unstake_is_ok() {
     contract_stake.stake(total_stake_amount, period_info, era_3);
 
     contract_stake.unstake(
-        contract_stake.total_staked_amount(period),
+        vec![(era_3, contract_stake.total_staked_amount(period))], // TODO: revisit this
         period_info,
         era_3,
-        (EraNumber::zero(), Balance::zero()),
     );
     assert!(contract_stake.staked.is_empty());
     assert!(contract_stake.staked_future.is_none());
diff --git a/pallets/dapp-staking-v3/src/types.rs b/pallets/dapp-staking-v3/src/types.rs
index cfee04adbb..f80e3bd8f5 100644
--- a/pallets/dapp-staking-v3/src/types.rs
+++ b/pallets/dapp-staking-v3/src/types.rs
@@ -1023,29 +1023,27 @@ impl SingularStakingInfo {
     ) -> Vec<(EraNumber, Balance)> {
         let mut result = Vec::new();
 
-        // Prepare values for further calculations
+        // 0. Prepare values for further calculations
         let snapshot = self.staked;
         let stake_delta = self
             .staked
             .total()
             .saturating_sub(self.previous_staked.total());
 
-        // Modify current staked amount
+        // 1. Modify current staked amount
         self.staked.subtract(amount);
         let unstaked_amount = snapshot.total().saturating_sub(self.staked.total());
         self.staked.era = self.staked.era.max(current_era);
-
-        // Push the result into result vector
         result.push((self.staked.era, unstaked_amount));
 
-        // 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 == snapshot.voting,
             };
 
-        // Move over the snapshot to the previous snapshot field and make sure
+        // 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();
@@ -1054,7 +1052,9 @@ impl SingularStakingInfo {
             self.previous_staked.era = self.staked.era.saturating_sub(1);
         }
 
+        // 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,
                 (

From f51febefa58c3c2abef738d524466c374d154e74 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Tue, 5 Mar 2024 09:42:52 +0100
Subject: [PATCH 16/28] Passing type tests

---
 .../dapp-staking-v3/src/test/tests_types.rs   | 68 ++++++++++---------
 1 file changed, 35 insertions(+), 33 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/test/tests_types.rs b/pallets/dapp-staking-v3/src/test/tests_types.rs
index deaecc61e9..55fdc1d3fb 100644
--- a/pallets/dapp-staking-v3/src/test/tests_types.rs
+++ b/pallets/dapp-staking-v3/src/test/tests_types.rs
@@ -2084,7 +2084,7 @@ fn singular_staking_info_unstake_during_voting_is_ok() {
     let unstake_amount_1 = 5;
     assert_eq!(
         staking_info.unstake(unstake_amount_1, era_1, Subperiod::Voting),
-        vec![(era_1 + 1, vote_stake_amount_1 - unstake_amount_1)]
+        vec![(era_1 + 1, unstake_amount_1)]
     );
     assert_eq!(
         staking_info.total_staked_amount(),
@@ -2102,7 +2102,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, remaining_stake)]
+        vec![(era_2 - 1, remaining_stake), (era_2, remaining_stake)]
     );
     assert!(staking_info.total_staked_amount().is_zero());
     assert!(
@@ -2129,7 +2129,8 @@ fn singular_staking_info_unstake_during_bep_is_ok() {
     let unstake_1 = 5;
     assert_eq!(
         staking_info.unstake(5, era_1, Subperiod::BuildAndEarn),
-        vec![(era_1 + 1, bep_stake_amount_1 - unstake_1)]
+        // We're unstaking from the `era_1 + 1` because stake was made for that era
+        vec![(era_1 + 1, unstake_1)]
     );
     assert_eq!(
         staking_info.total_staked_amount(),
@@ -2160,10 +2161,8 @@ fn singular_staking_info_unstake_during_bep_is_ok() {
 
     assert_eq!(
         staking_info.unstake(unstake_2, era_2, Subperiod::BuildAndEarn),
-        vec![
-            (era_2 - 1, voting_stake_overflow),
-            (era_2, current_bep_stake)
-        ]
+        // Both amounts are the same since current stake is less than the previous one
+        vec![(era_2 - 1, unstake_2), (era_2, unstake_2)]
     );
     assert_eq!(
         staking_info.total_staked_amount(),
@@ -2450,21 +2449,25 @@ fn contract_stake_amount_unstake_is_ok() {
     );
     let total_stake_amount = vp_stake_amount + bep_stake_amount;
 
-    // 1st scenario - unstake in the same era, from `Voting` subperiod
+    // 1st scenario - unstake some amount from the next era, B&E subperiod
     let amount_1 = 5;
-    contract_stake.unstake(vec![(era_1, amount_1)], period_info, era_1);
+    contract_stake.unstake(vec![(era_1 + 1, amount_1)], period_info, era_1);
     assert_eq!(
         contract_stake.total_staked_amount(period),
         total_stake_amount - amount_1
     );
     assert_eq!(
         contract_stake.staked_amount(period, Subperiod::Voting),
-        vp_stake_amount - amount_1
+        vp_stake_amount
+    );
+    assert_eq!(
+        contract_stake.staked_amount(period, Subperiod::BuildAndEarn),
+        bep_stake_amount - amount_1
     );
     assert!(contract_stake.staked.is_empty());
     assert!(contract_stake.staked_future.is_some());
 
-    // 2nd scenario - unstake in the next era
+    // 2nd scenario - unstake in the next era, expect entry alignment
     let period_info = PeriodInfo {
         number: period,
         subperiod: Subperiod::BuildAndEarn,
@@ -2478,11 +2481,11 @@ fn contract_stake_amount_unstake_is_ok() {
     );
     assert_eq!(
         contract_stake.staked_amount(period, Subperiod::Voting),
-        vp_stake_amount - amount_1
+        vp_stake_amount
     );
     assert_eq!(
         contract_stake.staked_amount(period, Subperiod::BuildAndEarn),
-        bep_stake_amount - amount_1,
+        bep_stake_amount - amount_1 * 2,
     );
     assert!(
         !contract_stake.staked.is_empty(),
@@ -2493,37 +2496,38 @@ fn contract_stake_amount_unstake_is_ok() {
         "future entry should be cleaned up since it refers to the current era"
     );
 
-    // 3rd scenario - same as previous scenario, but for `BuildAndEarn` subperiod
-    contract_stake.unstake(vec![(era_2, amount_1)], period_info, era_2);
+    // 3rd scenario - unstake such amount we chip away from the Voting subperiod stake amount
+    let voting_unstake_amount = 2;
+    let amount_2 =
+        contract_stake.staked_amount(period, Subperiod::BuildAndEarn) + voting_unstake_amount;
+    contract_stake.unstake(vec![(era_2, amount_2)], period_info, era_2);
     assert_eq!(
         contract_stake.total_staked_amount(period),
-        total_stake_amount - amount_1 * 3
+        total_stake_amount - amount_1 * 2 - amount_2
     );
     assert_eq!(
         contract_stake.staked_amount(period, Subperiod::Voting),
-        vp_stake_amount - amount_1
-    );
-    assert_eq!(
-        contract_stake.staked_amount(period, Subperiod::BuildAndEarn),
-        bep_stake_amount - amount_1 * 2
+        vp_stake_amount - voting_unstake_amount
     );
+    assert!(contract_stake
+        .staked_amount(period, Subperiod::BuildAndEarn)
+        .is_zero(),);
 
     // 4th scenario - bump up unstake eras by more than 1, entries should be aligned to the current era
     let era_3 = era_2 + 3;
-    let amount_2 = 7;
-    contract_stake.unstake(vec![(era_3, amount_2)], period_info, era_3);
+    let amount_3 = 7;
+    contract_stake.unstake(vec![(era_3, amount_3)], period_info, era_3);
     assert_eq!(
         contract_stake.total_staked_amount(period),
-        total_stake_amount - amount_1 * 3 - amount_2
+        total_stake_amount - amount_1 * 2 - amount_2 - amount_3
     );
     assert_eq!(
         contract_stake.staked_amount(period, Subperiod::Voting),
-        vp_stake_amount - amount_1
-    );
-    assert_eq!(
-        contract_stake.staked_amount(period, Subperiod::BuildAndEarn),
-        bep_stake_amount - amount_1 * 2 - amount_2
+        vp_stake_amount - voting_unstake_amount - amount_3
     );
+    assert!(contract_stake
+        .staked_amount(period, Subperiod::BuildAndEarn)
+        .is_zero());
     assert_eq!(
         contract_stake.staked.era, era_3,
         "Should be aligned to the current era."
@@ -2533,11 +2537,9 @@ fn contract_stake_amount_unstake_is_ok() {
         "future entry should remain 'None'"
     );
 
-    // 5th scenario - do a full unstake with existing future entry, expect a cleanup
-    contract_stake.stake(total_stake_amount, period_info, era_3);
-
+    // 5th scenario - do a full unstake, even with overflow, with existing future entry, expect a cleanup
     contract_stake.unstake(
-        vec![(era_3, contract_stake.total_staked_amount(period))], // TODO: revisit this
+        vec![(era_3, contract_stake.total_staked_amount(period) + 1)],
         period_info,
         era_3,
     );

From 37c5e9823b80d261c5a0e7f22c3a438fd47930bc Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Tue, 5 Mar 2024 10:35:26 +0100
Subject: [PATCH 17/28] More tests & docs

---
 .../dapp-staking-v3/src/test/tests_types.rs   | 33 +++++++++----
 pallets/dapp-staking-v3/src/types.rs          | 47 +++++++++++++++----
 2 files changed, 60 insertions(+), 20 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/test/tests_types.rs b/pallets/dapp-staking-v3/src/test/tests_types.rs
index 55fdc1d3fb..989be9cd5a 100644
--- a/pallets/dapp-staking-v3/src/test/tests_types.rs
+++ b/pallets/dapp-staking-v3/src/test/tests_types.rs
@@ -57,7 +57,7 @@ fn period_info_basic_checks() {
     assert_eq!(info.subperiod, Subperiod::Voting);
     assert_eq!(info.next_subperiod_start_era, next_subperiod_start_era);
 
-    // Voting period checks
+    // Voting subperiod checks
     assert!(!info.is_next_period(next_subperiod_start_era - 1));
     assert!(!info.is_next_period(next_subperiod_start_era));
     assert!(!info.is_next_period(next_subperiod_start_era + 1));
@@ -68,7 +68,7 @@ fn period_info_basic_checks() {
     ] {
         assert!(
             !info.is_next_period(era),
-            "Cannot trigger 'true' in the Voting period type."
+            "Cannot trigger 'true' in the Voting subperiod type."
         );
     }
 
@@ -520,7 +520,7 @@ fn account_ledger_add_stake_amount_basic_example_with_different_subperiods_works
     assert!(acc_ledger.staked.is_empty());
     assert!(acc_ledger.staked_future.is_none());
 
-    // 1st scenario - stake some amount in Voting period, and ensure values are as expected.
+    // 1st scenario - stake some amount in Voting subperiod, and ensure values are as expected.
     let era_1 = 1;
     let period_1 = 1;
     let period_info_1 = PeriodInfo {
@@ -1775,7 +1775,7 @@ fn era_info_stake_works() {
     // Sanity check
     assert!(era_info.total_locked.is_zero());
 
-    // Add some voting period stake
+    // Add some voting subperiod stake
     let vp_stake_amount = 7;
     era_info.add_stake_amount(vp_stake_amount, Subperiod::Voting);
     assert_eq!(era_info.total_staked_amount_next_era(), vp_stake_amount);
@@ -2080,7 +2080,7 @@ fn singular_staking_info_unstake_during_voting_is_ok() {
     let vote_stake_amount_1 = 11;
     staking_info.stake(vote_stake_amount_1, era_1, Subperiod::Voting);
 
-    // Unstake some amount during `Voting` period, loyalty should remain as expected.
+    // 1. Unstake some amount during `Voting` period, loyalty should remain as expected.
     let unstake_amount_1 = 5;
     assert_eq!(
         staking_info.unstake(unstake_amount_1, era_1, Subperiod::Voting),
@@ -2097,7 +2097,7 @@ fn singular_staking_info_unstake_during_voting_is_ok() {
         "Stake era should remain valid."
     );
 
-    // Fully unstake, attempting to underflow, and ensure loyalty flag has been removed.
+    // 2. Fully unstake, attempting to underflow, and ensure loyalty flag has been removed.
     let era_2 = era_1 + 2;
     let remaining_stake = staking_info.total_staked_amount();
     assert_eq!(
@@ -2151,8 +2151,21 @@ fn singular_staking_info_unstake_during_bep_is_ok() {
         "Stake era should remain valid."
     );
 
-    // 2nd scenario - unstake all of the amount staked during B&E period, and then some more.
-    // The point is to take a chunk from the voting period stake too.
+    // 2nd scenario - Ensure that staked amount is larger than the previous stake amount, and then
+    // unstake enough to result in some overflow of the stake delta.
+    let bep_stake_amount_2 = 13;
+    staking_info.stake(bep_stake_amount_2, era_1, Subperiod::BuildAndEarn);
+    let delta = staking_info.staked.total() - staking_info.previous_staked.total();
+    let overflow = 1;
+    let unstake_2 = delta + overflow;
+
+    assert_eq!(
+        staking_info.unstake(unstake_2, era_1, Subperiod::BuildAndEarn),
+        vec![(era_1, overflow), (era_1 + 1, unstake_2)]
+    );
+
+    // 3rd scenario - unstake all of the amount staked during B&E period, and then some more.
+    // The point is to take a chunk from the voting subperiod stake too.
     let current_total_stake = staking_info.total_staked_amount();
     let current_bep_stake = staking_info.staked_amount(Subperiod::BuildAndEarn);
     let voting_stake_overflow = 2;
@@ -2161,7 +2174,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 current stake is less than the previous one
+        // 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)]
     );
     assert_eq!(
@@ -2177,7 +2190,7 @@ fn singular_staking_info_unstake_during_bep_is_ok() {
         .is_zero());
     assert!(
         !staking_info.is_loyal(),
-        "Loyalty flag should have been removed due to non-zero voting period unstake"
+        "Loyalty flag should have been removed due to non-zero voting subperiod unstake"
     );
     assert_eq!(staking_info.era(), era_2);
 }
diff --git a/pallets/dapp-staking-v3/src/types.rs b/pallets/dapp-staking-v3/src/types.rs
index f80e3bd8f5..7e0a51bfdb 100644
--- a/pallets/dapp-staking-v3/src/types.rs
+++ b/pallets/dapp-staking-v3/src/types.rs
@@ -1015,6 +1015,28 @@ impl SingularStakingInfo {
     ///
     /// Returns a vector of `(era, amount)` pairs, where `era` is the era in which the unstake happened,
     /// and the amount is the corresponding amount.
+    ///
+    /// ### NOTE
+    /// `SingularStakingInfo` always aims to keep track of the staked amount between two consecutive eras.
+    /// 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,
@@ -1022,25 +1044,30 @@ impl SingularStakingInfo {
         subperiod: Subperiod,
     ) -> Vec<(EraNumber, Balance)> {
         let mut result = Vec::new();
-
-        // 0. Prepare values for further calculations
-        let snapshot = self.staked;
-        let stake_delta = self
-            .staked
-            .total()
-            .saturating_sub(self.previous_staked.total());
+        let staked_snapshot = self.staked;
 
         // 1. Modify current staked amount
         self.staked.subtract(amount);
-        let unstaked_amount = snapshot.total().saturating_sub(self.staked.total());
+        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
         self.loyal_staker = self.loyal_staker
             && match subperiod {
                 Subperiod::Voting => !self.staked.voting.is_zero(),
-                Subperiod::BuildAndEarn => self.staked.voting == snapshot.voting,
+                Subperiod::BuildAndEarn => self.staked.voting == staked_snapshot.voting,
             };
 
         // 3. Move over the snapshot to the previous snapshot field and make sure
@@ -1048,7 +1075,7 @@ impl SingularStakingInfo {
         if stake_delta.is_zero() {
             self.previous_staked = Default::default();
         } else {
-            self.previous_staked = snapshot;
+            self.previous_staked = staked_snapshot;
             self.previous_staked.era = self.staked.era.saturating_sub(1);
         }
 

From 40d527a0d7224d1ead1e97e801075e4e3e11ea55 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Tue, 5 Mar 2024 16:04:27 +0100
Subject: [PATCH 18/28] Full fix

---
 .../dapp-staking-v3/src/test/testing_utils.rs | 67 ++++++-------
 .../dapp-staking-v3/src/test/tests_types.rs   | 12 ++-
 pallets/dapp-staking-v3/src/types.rs          | 96 +++++++++----------
 3 files changed, 86 insertions(+), 89 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/test/testing_utils.rs b/pallets/dapp-staking-v3/src/test/testing_utils.rs
index 744dc7f8b1..df6ee25430 100644
--- a/pallets/dapp-staking-v3/src/test/testing_utils.rs
+++ b/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
diff --git a/pallets/dapp-staking-v3/src/test/tests_types.rs b/pallets/dapp-staking-v3/src/test/tests_types.rs
index 989be9cd5a..6682def895 100644
--- a/pallets/dapp-staking-v3/src/test/tests_types.rs
+++ b/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(),
diff --git a/pallets/dapp-staking-v3/src/types.rs b/pallets/dapp-staking-v3/src/types.rs
index 7e0a51bfdb..af204bdeae 100644
--- a/pallets/dapp-staking-v3/src/types.rs
+++ b/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

From ebde1459f5ad8877b42bed22e9d48a4d65970cd8 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Tue, 5 Mar 2024 16:35:19 +0100
Subject: [PATCH 19/28] Improved tests & docs

---
 .../dapp-staking-v3/src/test/tests_types.rs   | 38 ++++++++++++++++++-
 pallets/dapp-staking-v3/src/types.rs          | 18 +++++++--
 2 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/test/tests_types.rs b/pallets/dapp-staking-v3/src/test/tests_types.rs
index 6682def895..f5438333c9 100644
--- a/pallets/dapp-staking-v3/src/test/tests_types.rs
+++ b/pallets/dapp-staking-v3/src/test/tests_types.rs
@@ -2104,6 +2104,9 @@ fn singular_staking_info_unstake_during_voting_is_ok() {
         "Stake era should remain valid."
     );
 
+    assert!(staking_info.previous_staked.is_empty());
+    assert!(staking_info.previous_staked.era.is_zero());
+
     // 2. Fully unstake, attempting to underflow, and ensure loyalty flag has been removed.
     let era_2 = era_1 + 2;
     let remaining_stake = staking_info.total_staked_amount();
@@ -2116,7 +2119,10 @@ fn singular_staking_info_unstake_during_voting_is_ok() {
         !staking_info.is_loyal(),
         "Loyalty flag should have been removed since it was full unstake."
     );
-    assert_eq!(staking_info.era(), era_2);
+    assert!(staking_info.era().is_zero());
+
+    assert!(staking_info.previous_staked.is_empty());
+    assert!(staking_info.previous_staked.era.is_zero());
 }
 
 #[test]
@@ -2132,6 +2138,9 @@ fn singular_staking_info_unstake_during_bep_is_ok() {
     let bep_stake_amount_1 = 23;
     staking_info.stake(bep_stake_amount_1, era_1, Subperiod::BuildAndEarn);
 
+    assert_eq!(staking_info.previous_staked.total(), vote_stake_amount_1);
+    assert_eq!(staking_info.previous_staked.era, era_1);
+
     // 1st scenario - Unstake some of the amount staked during B&E period
     let unstake_1 = 5;
     assert_eq!(
@@ -2158,10 +2167,15 @@ fn singular_staking_info_unstake_during_bep_is_ok() {
         "Stake era should remain valid."
     );
 
+    // No changes to the previous staked amount
+    assert_eq!(staking_info.previous_staked.total(), vote_stake_amount_1);
+    assert_eq!(staking_info.previous_staked.era, era_1);
+
     // 2nd scenario - Ensure that staked amount is larger than the previous stake amount, and then
     // unstake enough to result in some overflow of the stake delta.
     let bep_stake_amount_2 = 13;
     staking_info.stake(bep_stake_amount_2, era_1, Subperiod::BuildAndEarn);
+    let previous_total_stake = staking_info.previous_staked.total();
     let delta = staking_info.staked.total() - staking_info.previous_staked.total();
     let overflow = 1;
     let unstake_2 = delta + overflow;
@@ -2171,6 +2185,25 @@ fn singular_staking_info_unstake_during_bep_is_ok() {
         vec![(era_1, overflow), (era_1 + 1, unstake_2)]
     );
 
+    assert_eq!(
+        staking_info.total_staked_amount(),
+        vote_stake_amount_1 + bep_stake_amount_1 + bep_stake_amount_2 - unstake_1 - unstake_2
+    );
+    assert_eq!(
+        staking_info.staked_amount(Subperiod::Voting),
+        vote_stake_amount_1
+    );
+    assert_eq!(
+        staking_info.staked_amount(Subperiod::BuildAndEarn),
+        bep_stake_amount_1 + bep_stake_amount_2 - unstake_1 - unstake_2
+    );
+
+    assert_eq!(
+        staking_info.previous_staked.total(),
+        previous_total_stake - overflow
+    );
+    assert_eq!(staking_info.previous_staked.era, era_1);
+
     // 3rd scenario - unstake all of the amount staked during B&E period, and then some more.
     // The point is to take a chunk from the voting subperiod stake too.
     let current_total_stake = staking_info.total_staked_amount();
@@ -2199,6 +2232,9 @@ fn singular_staking_info_unstake_during_bep_is_ok() {
         "Loyalty flag should have been removed due to non-zero voting subperiod unstake"
     );
     assert_eq!(staking_info.era(), era_2);
+
+    assert_eq!(staking_info.previous_staked.total(), current_total_stake);
+    assert_eq!(staking_info.previous_staked.era, era_2 - 1);
 }
 
 #[test]
diff --git a/pallets/dapp-staking-v3/src/types.rs b/pallets/dapp-staking-v3/src/types.rs
index af204bdeae..df4b1df274 100644
--- a/pallets/dapp-staking-v3/src/types.rs
+++ b/pallets/dapp-staking-v3/src/types.rs
@@ -1067,15 +1067,15 @@ impl SingularStakingInfo {
         // 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**.
+        // a. Values: previous_staked: **(2, 10)**, staked: **(3, 15)**
+        // b. 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)**
+        // c. 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)**
+        // c. 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 {
@@ -1093,6 +1093,16 @@ impl SingularStakingInfo {
             }
         }
 
+        // 5. Convenience cleanup
+        if self.previous_staked.is_empty() {
+            self.previous_staked = Default::default();
+        }
+        if self.staked.is_empty() {
+            self.staked = Default::default();
+            // No longer relevant.
+            self.previous_staked = Default::default();
+        }
+
         result
     }
 

From fd8380d3e4f924688def0bf782e82d445666340b Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Tue, 5 Mar 2024 16:58:50 +0100
Subject: [PATCH 20/28] Additional test

---
 .../dapp-staking-v3/src/test/tests_types.rs   | 62 ++++++++++++++++++-
 pallets/dapp-staking-v3/src/types.rs          |  2 +-
 2 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/test/tests_types.rs b/pallets/dapp-staking-v3/src/test/tests_types.rs
index f5438333c9..9fbd1e7ac1 100644
--- a/pallets/dapp-staking-v3/src/test/tests_types.rs
+++ b/pallets/dapp-staking-v3/src/test/tests_types.rs
@@ -2479,7 +2479,7 @@ fn contract_stake_amount_stake_is_ok() {
 }
 
 #[test]
-fn contract_stake_amount_unstake_is_ok() {
+fn contract_stake_amount_basic_unstake_is_ok() {
     let mut contract_stake = ContractStakeAmount::default();
 
     // Prep action - create a stake entry
@@ -2602,6 +2602,66 @@ fn contract_stake_amount_unstake_is_ok() {
     assert!(contract_stake.staked_future.is_none());
 }
 
+#[test]
+fn contract_stake_amount_advanced_unstake_is_ok() {
+    let mut contract_stake = ContractStakeAmount::default();
+
+    // Prep action - create staked & staked_future fields
+    let era_1 = 3;
+    let era_2 = era_1 + 1;
+    let period = 1;
+    let period_info = PeriodInfo {
+        number: period,
+        subperiod: Subperiod::Voting,
+        next_subperiod_start_era: 20,
+    };
+    let vp_stake_amount = 31;
+    let bep_stake_amount = 19;
+
+    // Stake in two consecutive eras. Entries will be aligned.
+    contract_stake.stake(vp_stake_amount, period_info, era_1);
+    contract_stake.stake(
+        bep_stake_amount,
+        PeriodInfo {
+            subperiod: Subperiod::BuildAndEarn,
+            ..period_info
+        },
+        era_2,
+    );
+    let total_stake_amount = vp_stake_amount + bep_stake_amount;
+
+    // Unstake some amount from both staked & staked_future fields
+    let amount_1 = 2;
+    let amount_2 = 3;
+    contract_stake.unstake(
+        vec![(era_2, amount_1), (era_2 + 1, amount_2)],
+        period_info,
+        era_2,
+    );
+
+    // Verify future era staked values
+    assert_eq!(
+        contract_stake.staked_future.expect("Must exist").total(),
+        total_stake_amount - amount_2
+    );
+    assert_eq!(
+        contract_stake.staked_future.expect("Must exist").voting,
+        vp_stake_amount
+    );
+    assert_eq!(
+        contract_stake
+            .staked_future
+            .expect("Must exist")
+            .build_and_earn,
+        bep_stake_amount - amount_2
+    );
+
+    // Verify current era stake values
+    assert_eq!(contract_stake.staked.total(), vp_stake_amount - amount_1);
+    assert_eq!(contract_stake.staked.voting, vp_stake_amount - amount_1);
+    assert!(contract_stake.staked.build_and_earn.is_zero());
+}
+
 #[test]
 fn era_reward_span_push_and_get_works() {
     get_u32_type!(SpanLength, 8);
diff --git a/pallets/dapp-staking-v3/src/types.rs b/pallets/dapp-staking-v3/src/types.rs
index df4b1df274..4869649e19 100644
--- a/pallets/dapp-staking-v3/src/types.rs
+++ b/pallets/dapp-staking-v3/src/types.rs
@@ -1269,7 +1269,7 @@ impl ContractStakeAmount {
         current_era: EraNumber,
     ) {
         // 1. Entry alignment
-        // e only need to keep track of the current era, and the next one
+        // We only need to keep track of the current era, and the next one.
         match self.staked_future {
             // Future entry exists, but it covers current or older era.
             Some(stake_amount)

From 34546b91cf49fd38f68cb61a0e11d9f684bad4d2 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Tue, 5 Mar 2024 17:11:09 +0100
Subject: [PATCH 21/28] Comments

---
 pallets/dapp-staking-v3/src/types.rs | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/pallets/dapp-staking-v3/src/types.rs b/pallets/dapp-staking-v3/src/types.rs
index 4869649e19..dc3f4fbf39 100644
--- a/pallets/dapp-staking-v3/src/types.rs
+++ b/pallets/dapp-staking-v3/src/types.rs
@@ -1261,7 +1261,8 @@ impl ContractStakeAmount {
         }
     }
 
-    /// Unstake the specified `Voting` and `Build&Earn` subperiod amounts from the contract, for the specified `subperiod` and `era`.
+    /// Unstake the specified `(era, amount)` pairs from the contract.
+    // Important to account for the ongoing specified `subperiod` and `era` in order to align the entries.
     pub fn unstake(
         &mut self,
         era_and_amount_pairs: Vec<(EraNumber, Balance)>,

From 1d14cdad54df30a2a39a3a31d266827a100ed3b4 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Tue, 12 Mar 2024 14:43:59 +0100
Subject: [PATCH 22/28] Implemented migration

---
 Cargo.lock                                    |   2 +-
 pallets/dapp-staking-migration/Cargo.toml     |   3 -
 .../src/benchmarking.rs                       | 119 +---
 pallets/dapp-staking-migration/src/lib.rs     | 574 +++++-------------
 pallets/dapp-staking-migration/src/mock.rs    | 293 ---------
 pallets/dapp-staking-migration/src/tests.rs   | 198 ------
 pallets/dapp-staking-migration/src/weights.rs | 214 ++-----
 pallets/dapp-staking-v3/src/lib.rs            |   2 +-
 pallets/dapp-staking-v3/src/types.rs          |  13 +
 runtime/astar/src/lib.rs                      |  52 +-
 runtime/shibuya/Cargo.toml                    |   4 +
 runtime/shibuya/src/lib.rs                    |  11 +-
 runtime/shiden/src/lib.rs                     |   9 +
 13 files changed, 263 insertions(+), 1231 deletions(-)
 delete mode 100644 pallets/dapp-staking-migration/src/mock.rs
 delete mode 100644 pallets/dapp-staking-migration/src/tests.rs

diff --git a/Cargo.lock b/Cargo.lock
index aa7af7f9d1..07d9ab1129 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -7705,7 +7705,6 @@ dependencies = [
  "frame-system",
  "pallet-balances",
  "pallet-dapp-staking-v3",
- "pallet-dapps-staking",
  "parity-scale-codec",
  "scale-info",
  "sp-arithmetic",
@@ -13262,6 +13261,7 @@ dependencies = [
  "pallet-collective",
  "pallet-contracts",
  "pallet-contracts-primitives",
+ "pallet-dapp-staking-migration",
  "pallet-dapp-staking-v3",
  "pallet-democracy",
  "pallet-dynamic-evm-base-fee",
diff --git a/pallets/dapp-staking-migration/Cargo.toml b/pallets/dapp-staking-migration/Cargo.toml
index ab503dbe17..8ba22f9ea0 100644
--- a/pallets/dapp-staking-migration/Cargo.toml
+++ b/pallets/dapp-staking-migration/Cargo.toml
@@ -20,7 +20,6 @@ sp-std = { workspace = true }
 
 astar-primitives = { workspace = true, optional = true }
 pallet-dapp-staking-v3 = { workspace = true }
-pallet-dapps-staking = { workspace = true }
 
 [dev-dependencies]
 astar-primitives = { workspace = true }
@@ -38,7 +37,6 @@ std = [
 	"frame-support/std",
 	"frame-system/std",
 	"pallet-dapp-staking-v3/std",
-	"pallet-dapps-staking/std",
 	"frame-benchmarking/std",
 	"astar-primitives?/std",
 	"sp-core/std",
@@ -50,7 +48,6 @@ runtime-benchmarks = [
 	"frame-system/runtime-benchmarks",
 	"sp-runtime/runtime-benchmarks",
 	"pallet-dapp-staking-v3/runtime-benchmarks",
-	"pallet-dapps-staking/runtime-benchmarks",
 	"astar-primitives/runtime-benchmarks",
 ]
 try-runtime = ["frame-support/try-runtime", "astar-primitives"]
diff --git a/pallets/dapp-staking-migration/src/benchmarking.rs b/pallets/dapp-staking-migration/src/benchmarking.rs
index 047e272d1b..2b21acf3a1 100644
--- a/pallets/dapp-staking-migration/src/benchmarking.rs
+++ b/pallets/dapp-staking-migration/src/benchmarking.rs
@@ -19,49 +19,33 @@
 use super::{Pallet as Migration, *};
 
 use frame_benchmarking::{account as benchmark_account, v2::*};
-use frame_support::{assert_ok, traits::Currency};
 
-/// Generate an unique smart contract using the provided index as a sort-of indetifier
-fn smart_contract<T: pallet_dapps_staking::Config>(index: u8) -> T::SmartContract {
-    // This is a hacky approach to provide different smart contracts without touching the smart contract trait.
-    let mut encoded_smart_contract = T::SmartContract::default().encode();
-    *encoded_smart_contract.last_mut().unwrap() = index;
+use astar_primitives::{dapp_staking::SmartContractHandle, Balance};
+use pallet_dapp_staking_v3::StakeAmount;
 
-    Decode::decode(&mut TrailingZeroInput::new(encoded_smart_contract.as_ref()))
-        .expect("Shouldn't occur as long as EVM is the default type.")
+fn smart_contract<T: Config>(idx: u8) -> T::SmartContract {
+    let address: T::AccountId = benchmark_account("smart_contract", idx.into(), 456);
+    T::SmartContract::wasm(address)
 }
 
-/// Initialize the old dApp staking pallet with some storage.
 pub(super) fn initial_config<T: Config>() {
-    let dapps_number = <T as pallet_dapp_staking_v3::Config>::MaxNumberOfContracts::get();
-    let dapps_number = (dapps_number as u8).min(100);
-
-    // Add some dummy dApps to the old pallet.
-    for idx in 0..dapps_number {
-        let developer: T::AccountId = benchmark_account("developer", idx.into(), 123);
-        <T as pallet_dapps_staking::Config>::Currency::make_free_balance_be(
-            &developer,
-            <T as pallet_dapps_staking::Config>::RegisterDeposit::get() * 2,
-        );
+    for idx in 0..10 {
+        let account: T::AccountId = benchmark_account("developer", idx.into(), 123);
         let smart_contract = smart_contract::<T>(idx);
-        assert_ok!(pallet_dapps_staking::Pallet::<T>::register(
-            RawOrigin::Root.into(),
-            developer,
-            smart_contract.clone(),
-        ));
 
-        let staker: T::AccountId = benchmark_account("staker", idx.into(), 123);
-        let lock_amount = <T as pallet_dapps_staking::Config>::MinimumStakingAmount::get()
-            .max(<T as pallet_dapp_staking_v3::Config>::MinimumLockedAmount::get());
-        <T as pallet_dapps_staking::Config>::Currency::make_free_balance_be(
-            &staker,
-            lock_amount * 100,
+        v5::StakerInfo::<T>::insert(
+            &account,
+            &smart_contract,
+            v5::SingularStakingInfo {
+                staked: StakeAmount {
+                    voting: 123 * (idx as Balance + 1),
+                    build_and_earn: 345 * (idx as Balance + 1),
+                    era: 1,
+                    period: 2,
+                },
+                loyal_staker: true,
+            },
         );
-        assert_ok!(pallet_dapps_staking::Pallet::<T>::bond_and_stake(
-            RawOrigin::Signed(staker.clone()).into(),
-            smart_contract,
-            lock_amount,
-        ));
     }
 }
 
@@ -70,77 +54,20 @@ mod benchmarks {
     use super::*;
 
     #[benchmark]
-    fn migrate_dapps_success() {
-        initial_config::<T>();
-
-        #[block]
-        {
-            assert!(Migration::<T>::migrate_dapps().is_ok());
-        }
-    }
-
-    #[benchmark]
-    fn migrate_dapps_noop() {
-        #[block]
-        {
-            assert!(Migration::<T>::migrate_dapps().is_err());
-        }
-    }
-
-    #[benchmark]
-    fn migrate_ledger_success() {
-        initial_config::<T>();
-
-        #[block]
-        {
-            assert!(Migration::<T>::migrate_ledger().is_ok());
-        }
-    }
-
-    #[benchmark]
-    fn migrate_ledger_noop() {
-        #[block]
-        {
-            assert!(Migration::<T>::migrate_ledger().is_err());
-        }
-    }
-
-    #[benchmark]
-    fn cleanup_old_storage_success(x: Linear<1, 5>) {
+    fn translate_staking_info_success() {
         initial_config::<T>();
 
         #[block]
         {
-            // TODO: for some reason, tests always fail here, nothing gets removed from storage.
-            // When tested against real runtime, it works just fine.
-            let _ = Migration::<T>::cleanup_old_storage(x.into());
+            assert!(Migration::<T>::translate_staking_info(None).is_ok());
         }
     }
 
     #[benchmark]
-    fn cleanup_old_storage_noop() {
-        let hashed_prefix = twox_128(pallet_dapps_staking::Pallet::<T>::name().as_bytes());
-        let _ = clear_prefix(&hashed_prefix, None);
-
+    fn translate_staking_info_success_noop() {
         #[block]
         {
-            assert!(Migration::<T>::cleanup_old_storage(1).is_err());
+            assert!(Migration::<T>::translate_staking_info(None).is_err());
         }
     }
-
-    impl_benchmark_test_suite!(
-        Pallet,
-        crate::benchmarking::tests::new_test_ext(),
-        crate::mock::Test,
-    );
-}
-
-#[cfg(test)]
-mod tests {
-    use crate::mock;
-    use sp_io::TestExternalities;
-
-    pub fn new_test_ext() -> TestExternalities {
-        mock::ExtBuilder::build()
-    }
 }
diff --git a/pallets/dapp-staking-migration/src/lib.rs b/pallets/dapp-staking-migration/src/lib.rs
index ef13155484..68460c5d4c 100644
--- a/pallets/dapp-staking-migration/src/lib.rs
+++ b/pallets/dapp-staking-migration/src/lib.rs
@@ -18,67 +18,27 @@
 
 #![cfg_attr(not(feature = "std"), no_std)]
 
-//! ## Summary
-//!
-//! Purpose of this pallet is to provide multi-stage migration for moving
-//! from the old _dapps_staking_v2_ over to the new _dapp_staking_v3_.
-//!
-//! ## Approach
-//!
-//! ### Multi-Stage Migration
-//!
-//! Since a lot of data has to be cleaned up & migrated, it is necessary to do this in multiple steps.
-//! To reduce the risk of something going wrong, nothing is done in _mandatory hooks_, like `on_initialize` or `on_idle`.
-//! Instead, a dedicated extrinsic call is introduced, which can be called to move the migration forward.
-//! As long as this call moves the migration forward, its cost is refunded to the user.
-//! Once migration finishes, the extrinsic call will no longer do anything but won't refund the call cost either.
-//!
-//! ### Migration Steps
-//!
-//! The general approach used when migrating is:
-//! 1. Clean up old pallet's storage using custom code
-//! 2. Use dedicated dApp staking v3 extrinsic calls for registering dApps & locking funds.
-//!
-//! The main benefits of this approach are that we don't duplicate logic that is already present in dApp staking v3,
-//! and that we ensure proper events are emitted for each action which will make indexers happy. No special handling will
-//! be required to migrate dApps or locked/staked funds over from the old pallet to the new one, from the indexers  perspective.
-//!
-//! ### Final Cleanup
-//!
-//! The pallet doesn't clean after itself, so when it's removed from the runtime,
-//! the old storage should be cleaned up using `RemovePallet` type.
-//!
-
 pub use pallet::*;
 
 use frame_support::{
     dispatch::PostDispatchInfo,
     log,
     pallet_prelude::*,
-    traits::{Get, LockableCurrency, ReservableCurrency},
+    storage_alias,
+    traits::{ConstU32, Get},
+    WeakBoundedVec,
 };
 
-use frame_system::{pallet_prelude::*, RawOrigin};
+use frame_system::pallet_prelude::*;
 use parity_scale_codec::{Decode, Encode};
-use sp_io::{hashing::twox_128, storage::clear_prefix, KillStorageResult};
-use sp_runtime::{
-    traits::{TrailingZeroInput, UniqueSaturatedInto},
-    Saturating,
-};
+use sp_runtime::Saturating;
 
-use pallet_dapps_staking::{Ledger as OldLedger, RegisteredDapps as OldRegisteredDapps};
+use pallet_dapp_staking_v3::{SingularStakingInfo, StakeAmount, StakerInfo};
 
-#[cfg(feature = "try-runtime")]
-use astar_primitives::Balance;
 #[cfg(feature = "try-runtime")]
 use sp_std::vec::Vec;
 
-pub use crate::pallet::DappStakingMigrationHandler;
-
-#[cfg(test)]
-mod mock;
-#[cfg(test)]
-mod tests;
+pub use crate::pallet::SingularStakingInfoTranslationUpgrade;
 
 #[cfg(feature = "runtime-benchmarks")]
 mod benchmarking;
@@ -88,6 +48,35 @@ pub use weights::WeightInfo;
 
 const LOG_TARGET: &str = "dapp-staking-migration";
 
+mod v5 {
+    use super::*;
+
+    #[derive(
+        Encode, Decode, MaxEncodedLen, Copy, Clone, Debug, PartialEq, Eq, TypeInfo, Default,
+    )]
+    #[scale_info(skip_type_params(T))]
+    pub struct SingularStakingInfo {
+        /// Staked amount
+        pub(crate) staked: StakeAmount,
+        /// Indicates whether a staker is a loyal staker or not.
+        pub(crate) loyal_staker: bool,
+    }
+
+    #[storage_alias]
+    pub type StakerInfo<T: Config> = StorageDoubleMap<
+        pallet_dapp_staking_v3::Pallet<T>,
+        Blake2_128Concat,
+        <T as frame_system::Config>::AccountId,
+        Blake2_128Concat,
+        <T as pallet_dapp_staking_v3::Config>::SmartContract,
+        SingularStakingInfo,
+        OptionQuery,
+    >;
+}
+
+const MAX_KEY_SIZE: u32 = 1024;
+type StakingInfoKey = WeakBoundedVec<u8, ConstU32<MAX_KEY_SIZE>>;
+
 #[frame_support::pallet]
 pub mod pallet {
     use super::*;
@@ -98,7 +87,7 @@ pub mod pallet {
     #[pallet::config]
     pub trait Config:
         // Tight coupling, but it's fine since pallet is supposed to be just temporary and will be removed after migration.
-        frame_system::Config + pallet_dapp_staking_v3::Config + pallet_dapps_staking::Config
+        frame_system::Config + pallet_dapp_staking_v3::Config
     {
         /// The overarching event type.
         type RuntimeEvent: From<Event<Self>> + IsType<<Self as frame_system::Config>::RuntimeEvent>;
@@ -114,10 +103,8 @@ pub mod pallet {
     #[pallet::event]
     #[pallet::generate_deposit(pub(crate) fn deposit_event)]
     pub enum Event<T: Config> {
-        /// Number of entries migrated from v2 over to v3
-        EntriesMigrated(u32),
-        /// Number of entries deleted from v2
-        EntriesDeleted(u32),
+        /// Number of staking info entries translated
+        SingularStakingInfoTranslated(u32),
     }
 
     #[pallet::call]
@@ -185,73 +172,38 @@ pub mod pallet {
                 return Err(consumed_weight);
             }
 
-            // Ensure we can call dApp staking v3 extrinsics within this call.
-            consumed_weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2));
-            pallet_dapp_staking_v3::ActiveProtocolState::<T>::mutate(|state| {
-                state.maintenance = false;
-            });
+            consumed_weight.saturating_accrue(T::DbWeight::get().writes(1));
 
             let mut migration_state = init_migration_state;
-            let (mut entries_migrated, mut entries_deleted) = (0_u32, 0_u32);
+            let mut entries_migrated = 0_u32;
 
-            // Execute migration steps only if we have enough weight to do so.
-            //
-            // 1. Migrate registered dApps
-            // 2. Migrate ledgers
-            // 3. Cleanup
             while weight_limit
                 .saturating_sub(consumed_weight)
                 .all_gte(Self::migration_weight_margin())
             {
-                match migration_state {
-                    MigrationState::NotInProgress | MigrationState::RegisteredDApps => {
-                        migration_state = MigrationState::RegisteredDApps;
-
-                        match Self::migrate_dapps() {
-                            Ok(weight) => {
-                                consumed_weight.saturating_accrue(weight);
-                                entries_migrated.saturating_inc();
-                            }
-                            Err(weight) => {
-                                consumed_weight.saturating_accrue(weight);
-                                migration_state = MigrationState::Ledgers;
-                            }
-                        }
-                    }
-                    MigrationState::Ledgers => match Self::migrate_ledger() {
-                        Ok(weight) => {
+                match migration_state.clone() {
+                    MigrationState::NotInProgress => match Self::translate_staking_info(None) {
+                        Ok((last_key, weight)) => {
                             consumed_weight.saturating_accrue(weight);
                             entries_migrated.saturating_inc();
+
+                            migration_state = MigrationState::SingularStakingInfo(last_key);
                         }
                         Err(weight) => {
                             consumed_weight.saturating_accrue(weight);
-                            migration_state = MigrationState::Cleanup;
+                            migration_state = MigrationState::Finished;
                         }
                     },
-                    MigrationState::Cleanup => {
-                        // Ensure we don't attempt to delete too much at once.
-                        const SAFETY_MARGIN: u32 = 2000;
-                        let remaining_weight = weight_limit.saturating_sub(consumed_weight);
-                        let capacity = match remaining_weight.checked_div_per_component(
-                            &<T as Config>::WeightInfo::cleanup_old_storage_success(1),
-                        ) {
-                            Some(entries_to_delete) => {
-                                SAFETY_MARGIN.min(entries_to_delete.unique_saturated_into())
-                            }
-                            None => {
-                                // Not enough weight to delete even a single entry
-                                break;
-                            }
-                        };
-
-                        match Self::cleanup_old_storage(capacity) {
-                            Ok((weight, count)) => {
+                    MigrationState::SingularStakingInfo(last_key) => {
+                        match Self::translate_staking_info(Some(last_key)) {
+                            Ok((last_key, weight)) => {
                                 consumed_weight.saturating_accrue(weight);
-                                entries_deleted.saturating_accrue(count);
+                                entries_migrated.saturating_inc();
+
+                                migration_state = MigrationState::SingularStakingInfo(last_key);
                             }
-                            Err((weight, count)) => {
+                            Err(weight) => {
                                 consumed_weight.saturating_accrue(weight);
-                                entries_deleted.saturating_accrue(count);
                                 migration_state = MigrationState::Finished;
                             }
                         }
@@ -265,199 +217,61 @@ pub mod pallet {
 
             // Deposit events if needed
             if entries_migrated > 0 {
-                Self::deposit_event(Event::<T>::EntriesMigrated(entries_migrated));
-            }
-            if entries_deleted > 0 {
-                Self::deposit_event(Event::<T>::EntriesDeleted(entries_deleted));
+                Self::deposit_event(Event::<T>::SingularStakingInfoTranslated(entries_migrated));
             }
 
-            // Put the pallet back into maintenance mode in case we're still migration the old storage over,
-            // otherwise disable the maintenance mode.
-            pallet_dapp_staking_v3::ActiveProtocolState::<T>::mutate(|state| {
-                state.maintenance = match migration_state {
-                    MigrationState::NotInProgress
-                    | MigrationState::RegisteredDApps
-                    | MigrationState::Ledgers => true,
-                    MigrationState::Cleanup | MigrationState::Finished => false,
-                };
-            });
+            // Update the migration status
+            MigrationStateStorage::<T>::put(migration_state.clone());
 
-            if migration_state != init_migration_state {
-                // Already charged in pessimistic manner at the beginning of the function.
-                MigrationStateStorage::<T>::put(migration_state);
+            // Once migration has been finished, disable the maintenance mode and set correct storage version.
+            if migration_state == MigrationState::Finished {
+                log::trace!(target: LOG_TARGET, "Migration has been finished.");
+
+                pallet_dapp_staking_v3::ActiveProtocolState::<T>::mutate(|state| {
+                    state.maintenance = false;
+                });
+                StorageVersion::new(6).put::<pallet_dapp_staking_v3::Pallet<T>>();
+                consumed_weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2));
             }
 
             Ok(consumed_weight)
         }
 
-        /// Used to migrate `RegisteredDapps` from the _old_ dApps staking v2 pallet over to the new `IntegratedDApps`.
-        ///
-        /// Steps:
-        /// 1. Attempt to `drain` a single DB entry from the old storage. If it's unregistered, move on.
-        /// 2. Unreserve the old `RegisterDeposit` amount from the developer account.
-        /// 2. Re-decode old smart contract type into new one. Operation should be infallible in practice since the same underlying type is used.
-        /// 3. `register` the old-new smart contract into dApp staking v3 pallet.
-        ///
-        /// Returns `Ok(_)` if an entry was migrated, `Err(_)` if there are no more entries to migrate.
-        pub(crate) fn migrate_dapps() -> Result<Weight, Weight> {
-            match OldRegisteredDapps::<T>::drain().next() {
-                Some((smart_contract, old_dapp_info)) => {
-                    // In case dApp was unregistered, nothing more to do here
-                    if old_dapp_info.is_unregistered() {
-                        // Not precise, but happens rarely
-                        return Ok(<T as Config>::WeightInfo::migrate_dapps_success());
-                    }
-
-                    // Release reserved funds from the old dApps staking
-                    <T as pallet_dapps_staking::Config>::Currency::unreserve(
-                        &old_dapp_info.developer,
-                        <T as pallet_dapps_staking::Config>::RegisterDeposit::get(),
-                    );
+        pub(crate) fn translate_staking_info(
+            last_key: Option<StakingInfoKey>,
+        ) -> Result<(StakingInfoKey, Weight), Weight> {
+            // Create an iterator to be used for reading a single entry
+            let mut iter = if let Some(last_key) = last_key {
+                v5::StakerInfo::<T>::iter_from(last_key.into_inner())
+            } else {
+                v5::StakerInfo::<T>::iter()
+            };
 
-                    // Trick to get around different associated types which are essentially the same underlying struct.
-                    let new_smart_contract = match Decode::decode(&mut TrailingZeroInput::new(
-                        smart_contract.encode().as_ref(),
-                    )) {
-                        Ok(new_smart_contract) => new_smart_contract,
-                        Err(_) => {
-                            log::error!(
-                                target: LOG_TARGET,
-                                "Failed to decode smart contract: {:?}.",
-                                smart_contract,
-                            );
-
-                            // This should never happen, but if it does, we want to know about it.
-                            #[cfg(feature = "try-runtime")]
-                            panic!("Failed to decode smart contract: {:?}", smart_contract);
-                            #[cfg(not(feature = "try-runtime"))]
-                            // Not precise, but must never happen in production
-                            return Ok(<T as Config>::WeightInfo::migrate_dapps_success());
-                        }
-                    };
-
-                    match pallet_dapp_staking_v3::Pallet::<T>::register(
-                        RawOrigin::Root.into(),
-                        old_dapp_info.developer.clone(),
-                        new_smart_contract,
-                    ) {
-                        Ok(_) => {}
-                        Err(error) => {
-                            log::error!(
-                                target: LOG_TARGET,
-                                "Failed to register smart contract: {:?} with error: {:?}.",
-                                smart_contract,
-                                error,
-                            );
-                        }
-                    }
+            // Try to read the next entry
+            if let Some((account_id, smart_contract_id, old)) = iter.next() {
+                // Entry exists so it needs to be translated into the new format
+                let new_staking_info = SingularStakingInfo::new_migration(
+                    StakeAmount::default(),
+                    old.staked,
+                    old.loyal_staker,
+                );
+                StakerInfo::<T>::insert(&account_id, &smart_contract_id, new_staking_info);
 
-                    Ok(<T as Config>::WeightInfo::migrate_dapps_success())
-                }
-                None => {
-                    // Nothing more to migrate here
-                    Err(<T as Config>::WeightInfo::migrate_dapps_noop())
-                }
-            }
-        }
+                let hashed_key = StakerInfo::<T>::hashed_key_for(&account_id, &smart_contract_id);
 
-        /// Used to migrate `Ledger` from the _old_ dApps staking v2 pallet over to the new `Ledger`.
-        ///
-        /// Steps:
-        /// 1. Attempt to `drain` a single DB entry from the old storage.
-        /// 2. Release the old lock from the staker account, in full.
-        /// 3. Lock (or freeze) the old _staked_ amount into the new dApp staking v3 pallet.
-        ///
-        /// **NOTE:** the amount that was undergoing the unbonding process is not migrated but is immediately fully released.
-        ///
-        /// Returns `Ok(_)` if an entry was migrated, `Err(_)` if there are no more entries to migrate.
-        pub(crate) fn migrate_ledger() -> Result<Weight, Weight> {
-            match OldLedger::<T>::drain().next() {
-                Some((staker, old_account_ledger)) => {
-                    let old_locked = old_account_ledger.locked;
-
-                    // Old unbonding amount can just be released, to keep things simple.
-                    // Alternative is to re-calculate this into unlocking chunks.
-                    let total_unbonding = old_account_ledger.unbonding_info.sum();
-
-                    <T as pallet_dapps_staking::Config>::Currency::remove_lock(
-                        pallet_dapps_staking::pallet::STAKING_ID,
-                        &staker,
+                if cfg!(feature = "try-runtime") {
+                    assert!(
+                        hashed_key.len() < MAX_KEY_SIZE as usize,
+                        "Key size exceeded max limit!"
                     );
-
-                    let locked = old_locked.saturating_sub(total_unbonding);
-
-                    // No point in attempting to lock the old amount into dApp staking v3 if amount is insufficient.
-                    if locked >= <T as pallet_dapp_staking_v3::Config>::MinimumLockedAmount::get() {
-                        match pallet_dapp_staking_v3::Pallet::<T>::lock(
-                            RawOrigin::Signed(staker.clone()).into(),
-                            locked,
-                        ) {
-                            Ok(_) => {}
-                            Err(error) => {
-                                log::error!(
-                                    target: LOG_TARGET,
-                                    "Failed to lock for staker {:?} with error: {:?}.",
-                                    staker,
-                                    error,
-                                );
-
-                                // This should never happen, but if it does, we want to know about it.
-                                #[cfg(feature = "try-runtime")]
-                                panic!(
-                                    "Failed to lock for staker {:?} with error: {:?}.",
-                                    staker, error,
-                                );
-                            }
-                        }
-                    }
-
-                    // In case no lock action, it will be imprecise but it's fine since this
-                    // isn't expected to happen, and even if it does, it's not a big deal.
-                    Ok(<T as Config>::WeightInfo::migrate_ledger_success())
                 }
-                None => {
-                    // Nothing more to migrate here
-                    Err(<T as Config>::WeightInfo::migrate_ledger_noop())
-                }
-            }
-        }
-
-        /// Used to remove one entry from the old _dapps_staking_v2_ storage.
-        ///
-        /// If there are no more entries to remove, returns `Err(_)` with consumed weight and number of deleted entries.
-        /// Otherwise returns `Ok(_)` with consumed weight and number of consumed entries.
-        pub(crate) fn cleanup_old_storage(limit: u32) -> Result<(Weight, u32), (Weight, u32)> {
-            let hashed_prefix = twox_128(pallet_dapps_staking::Pallet::<T>::name().as_bytes());
-
-            // Repeated calls in the same block don't work, so we set the limit to `Unlimited` in case of `try-runtime` testing.
-            let inner_limit = if cfg!(feature = "try-runtime") {
-                None
-            } else {
-                Some(limit)
-            };
-
-            let (keys_removed, done) = match clear_prefix(&hashed_prefix, inner_limit) {
-                KillStorageResult::AllRemoved(value) => (value, true),
-                KillStorageResult::SomeRemaining(value) => (value, false),
-            };
-
-            log::trace!(
-                target: LOG_TARGET,
-                "Removed {} keys from storage.",
-                keys_removed
-            );
 
-            if !done {
                 Ok((
-                    <T as Config>::WeightInfo::cleanup_old_storage_success(keys_removed),
-                    keys_removed as u32,
+                    WeakBoundedVec::force_from(hashed_key, None),
+                    <T as Config>::WeightInfo::translate_staking_info_success(),
                 ))
             } else {
-                log::trace!(target: LOG_TARGET, "All keys have been removed.",);
-                Err((
-                    <T as Config>::WeightInfo::cleanup_old_storage_noop(),
-                    keys_removed as u32,
-                ))
+                Err(<T as Config>::WeightInfo::translate_staking_info_success_noop())
             }
         }
 
@@ -492,24 +306,18 @@ pub mod pallet {
         /// This is used to ensure we don't go over the limit.
         fn migration_weight_margin() -> Weight {
             // Consider the weight of all steps
-            <T as Config>::WeightInfo::migrate_dapps_success()
-                .max(<T as Config>::WeightInfo::migrate_ledger_success())
-                .max(<T as Config>::WeightInfo::cleanup_old_storage_success(1))
+            <T as Config>::WeightInfo::translate_staking_info_success()
                 // and add the weight of updating migration status
                 .saturating_add(T::DbWeight::get().writes(1))
         }
     }
 
-    #[derive(PartialEq, Eq, Clone, Encode, Decode, Copy, TypeInfo, RuntimeDebug, MaxEncodedLen)]
+    #[derive(PartialEq, Eq, Clone, Encode, Decode, TypeInfo, RuntimeDebug, MaxEncodedLen)]
     pub enum MigrationState {
         /// No migration in progress
         NotInProgress,
-        /// In the middle of `RegisteredDApps` migration.
-        RegisteredDApps,
-        /// In the middle of `Ledgers` migration.
-        Ledgers,
-        /// In the middle of old v2 storage cleanup
-        Cleanup,
+        /// In the middle of `SingularStakingInfo` migration/translation.
+        SingularStakingInfo(StakingInfoKey),
         /// All migrations have been finished
         Finished,
     }
@@ -520,20 +328,18 @@ pub mod pallet {
         }
     }
 
-    pub struct DappStakingMigrationHandler<T: Config>(PhantomData<T>);
-    impl<T: Config> frame_support::traits::OnRuntimeUpgrade for DappStakingMigrationHandler<T> {
+    pub struct SingularStakingInfoTranslationUpgrade<T: Config>(PhantomData<T>);
+    impl<T: Config> frame_support::traits::OnRuntimeUpgrade
+        for SingularStakingInfoTranslationUpgrade<T>
+    {
         fn on_runtime_upgrade() -> Weight {
-            // When upgrade happens, we need to put dApp staking v3 into maintenance mode immediately.
-            // For the old pallet, since the storage cleanup is going to happen, maintenance mode must be ensured
-            // by the runtime config itself.
             let mut consumed_weight = T::DbWeight::get().reads_writes(1, 2);
+
+            // Enable maintenance mode.
             pallet_dapp_staking_v3::ActiveProtocolState::<T>::mutate(|state| {
                 state.maintenance = true;
             });
 
-            // Set the correct init storage version
-            pallet_dapp_staking_v3::STORAGE_VERSION.put::<pallet_dapp_staking_v3::Pallet<T>>();
-
             // In case of try-runtime, we want to execute the whole logic, to ensure it works
             // with on-chain data.
             if cfg!(feature = "try-runtime") {
@@ -565,146 +371,67 @@ pub mod pallet {
 
         #[cfg(feature = "try-runtime")]
         fn pre_upgrade() -> Result<Vec<u8>, sp_runtime::TryRuntimeError> {
-            // Get dev accounts with registered dapps and their total reserved balance
-            let developers: Vec<_> = pallet_dapps_staking::RegisteredDapps::<T>::iter()
-                .filter_map(|(smart_contract, info)| {
-                    if info.state == pallet_dapps_staking::DAppState::Registered {
-                        let reserved =
-                            <T as pallet_dapps_staking::Config>::Currency::reserved_balance(
-                                &info.developer,
-                            );
-                        Some((info.developer, smart_contract, reserved))
-                    } else {
-                        None
-                    }
-                })
-                .collect();
-
-            // Get the stakers and their active locked (staked) amount.
-
-            use sp_runtime::traits::Zero;
-            let mut total_locked = Balance::zero();
-            let min_lock_amount: Balance =
-                <T as pallet_dapp_staking_v3::Config>::MinimumLockedAmount::get();
-            let stakers: Vec<_> = pallet_dapps_staking::Ledger::<T>::iter()
-                .filter_map(|(staker, ledger)| {
-                    total_locked.saturating_accrue(ledger.locked);
-                    total_locked.saturating_reduce(ledger.unbonding_info.sum());
-
-                    let new_lock_amount = ledger.locked.saturating_sub(ledger.unbonding_info.sum());
-                    if new_lock_amount >= min_lock_amount {
-                        Some((staker, new_lock_amount))
-                    } else {
-                        None
-                    }
+            // Get all staker info entries to be used later for verification
+            let staker_info: Vec<_> = v5::StakerInfo::<T>::iter()
+                .map(|(account_id, smart_contract, staking_info)| {
+                    (
+                        account_id,
+                        smart_contract,
+                        staking_info.staked,
+                        staking_info.loyal_staker,
+                    )
                 })
                 .collect();
 
-            log::info!(
-                target: LOG_TARGET,
-                "Total locked amount in the old pallet: {:?}.",
-                total_locked,
-            );
-
-            log::info!(
-                target: LOG_TARGET,
-                "Out of {} stakers, {} have sufficient amount to lock.",
-                pallet_dapps_staking::Ledger::<T>::iter().count(),
-                stakers.len(),
-            );
-
-            let helper = Helper::<T> {
-                developers,
-                stakers,
-            };
+            let helper = Helper::<T> { staker_info };
 
             Ok(helper.encode())
         }
 
         #[cfg(feature = "try-runtime")]
         fn post_upgrade(state: Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
-            use sp_runtime::traits::Zero;
+            use sp_runtime::traits::TrailingZeroInput;
+
+            // 0. Verify that migration state is `Finished`
+            if MigrationStateStorage::<T>::get() != MigrationState::Finished {
+                return Err("Migration state is not `Finished`".into());
+            }
 
             let helper = Helper::<T>::decode(&mut TrailingZeroInput::new(state.as_ref()))
                 .map_err(|_| "Cannot decode data from pre_upgrade")?;
 
-            // 1. Ensure that all entries have been unregistered/removed and all dev accounts have been refunded.
-            //    Also check that dApps have been registered in the new pallet.
-            assert!(pallet_dapps_staking::RegisteredDapps::<T>::iter()
-                .count()
-                .is_zero());
-            assert_eq!(
-                pallet_dapp_staking_v3::IntegratedDApps::<T>::iter().count(),
-                helper.developers.len()
-            );
-
-            let register_deposit = <T as pallet_dapps_staking::Config>::RegisterDeposit::get();
-            for (dev_account, smart_contract, old_reserved) in helper.developers {
-                let new_reserved =
-                    <T as pallet_dapps_staking::Config>::Currency::reserved_balance(&dev_account);
-                assert_eq!(old_reserved, new_reserved + register_deposit);
-
-                let new_smart_contract: <T as pallet_dapp_staking_v3::Config>::SmartContract =
-                    Decode::decode(&mut TrailingZeroInput::new(
-                        smart_contract.encode().as_ref(),
-                    ))
-                    .expect("Must succeed since we're using the same underlying type.");
-
-                let dapp_info =
-                    pallet_dapp_staking_v3::IntegratedDApps::<T>::get(&new_smart_contract)
-                        .expect("Must exist!");
-                assert_eq!(dapp_info.owner, dev_account);
+            // 1. Verify that staker info is essentially same as before
+            for (account_id, smart_contract, staked, loyal_staker) in helper.staker_info {
+                let staking_info = StakerInfo::<T>::get(&account_id, &smart_contract)
+                    .ok_or("Staking info not found but it must exist!")?;
+
+                let expected_staking_info = SingularStakingInfo::new_migration(
+                    StakeAmount::default(),
+                    staked,
+                    loyal_staker,
+                );
+
+                if staking_info != expected_staking_info {
+                    log::error!(target: LOG_TARGET,
+                        "Staking info mismatch for account {:?} and smart contract {:?}. Expected: {:?}, got: {:?}",
+                        account_id, smart_contract, expected_staking_info, staking_info
+                    );
+
+                    return Err("Failed to verify staking info".into());
+                }
+            }
+
+            // 2. Verify pallet is no longer in maintenance mode
+            if pallet_dapp_staking_v3::ActiveProtocolState::<T>::get().maintenance {
+                return Err("Pallet is still in maintenance mode".into());
             }
 
-            // 2. Ensure that all ledger entries have been migrated over to the new pallet.
-            //    Total locked amount in the new pallet must equal the sum of all old locked amounts.
-            assert!(pallet_dapps_staking::Ledger::<T>::iter().count().is_zero());
-            assert_eq!(
-                pallet_dapp_staking_v3::Ledger::<T>::iter().count(),
-                helper.stakers.len()
-            );
-
-            for (staker, old_locked) in &helper.stakers {
-                let new_locked = pallet_dapp_staking_v3::Ledger::<T>::get(&staker).locked;
-                assert!(*old_locked >= new_locked);
+            // 3. Verify on-chain storage version is correct
+            if StorageVersion::get::<pallet_dapp_staking_v3::Pallet<T>>() != 6 {
+                return Err("Storage version is not correct".into());
             }
 
-            let total_locked = helper
-                .stakers
-                .iter()
-                .map(|(_, locked)| locked)
-                .sum::<Balance>();
-            assert_eq!(
-                pallet_dapp_staking_v3::CurrentEraInfo::<T>::get().total_locked,
-                total_locked
-            );
-
-            log::info!(
-                target: LOG_TARGET,
-                "Total locked amount in the new pallet: {:?}.",
-                total_locked,
-            );
-
-            // 3. Check that rest of the storage has been cleaned up.
-            assert!(!pallet_dapps_staking::PalletDisabled::<T>::exists());
-            assert!(!pallet_dapps_staking::CurrentEra::<T>::exists());
-            assert!(!pallet_dapps_staking::BlockRewardAccumulator::<T>::exists());
-            assert!(!pallet_dapps_staking::ForceEra::<T>::exists());
-            assert!(!pallet_dapps_staking::NextEraStartingBlock::<T>::exists());
-            assert!(!pallet_dapps_staking::StorageVersion::<T>::exists());
-
-            assert!(pallet_dapps_staking::RegisteredDevelopers::<T>::iter()
-                .count()
-                .is_zero());
-            assert!(pallet_dapps_staking::GeneralEraInfo::<T>::iter()
-                .count()
-                .is_zero());
-            assert!(pallet_dapps_staking::ContractEraStake::<T>::iter()
-                .count()
-                .is_zero());
-            assert!(pallet_dapps_staking::GeneralStakerInfo::<T>::iter()
-                .count()
-                .is_zero());
+            log::trace!(target: LOG_TARGET, "Post-upgrade checks successful.");
 
             Ok(())
         }
@@ -715,12 +442,5 @@ pub mod pallet {
 /// Used to help with `try-runtime` testing.
 #[derive(Encode, Decode)]
 struct Helper<T: Config> {
-    /// Vec of devs, with their associated smart contract & total reserved balance
-    developers: Vec<(
-        T::AccountId,
-        <T as pallet_dapps_staking::Config>::SmartContract,
-        Balance,
-    )>,
-    /// Stakers with their total active locked amount (not undergoing the unbonding process)
-    stakers: Vec<(T::AccountId, Balance)>,
+    staker_info: Vec<(T::AccountId, T::SmartContract, StakeAmount, bool)>,
 }
diff --git a/pallets/dapp-staking-migration/src/mock.rs b/pallets/dapp-staking-migration/src/mock.rs
deleted file mode 100644
index 64c7adc54c..0000000000
--- a/pallets/dapp-staking-migration/src/mock.rs
+++ /dev/null
@@ -1,293 +0,0 @@
-// This file is part of Astar.
-
-// Copyright (C) 2019-2023 Stake Technologies Pte.Ltd.
-// SPDX-License-Identifier: GPL-3.0-or-later
-
-// Astar is free software: you can redistribute it and/or modify
-// it under the terms of the GNU General Public License as published by
-// the Free Software Foundation, either version 3 of the License, or
-// (at your option) any later version.
-
-// Astar is distributed in the hope that it will be useful,
-// but WITHOUT ANY WARRANTY; without even the implied warranty of
-// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-// GNU General Public License for more details.
-
-// You should have received a copy of the GNU General Public License
-// along with Astar. If not, see <http://www.gnu.org/licenses/>.
-
-use crate::{self as pallet_dapp_staking_migration, *};
-
-use frame_support::{
-    assert_ok, construct_runtime, parameter_types,
-    traits::{fungible::Mutate as FunMutate, ConstBool, ConstU128, ConstU32, Currency},
-    weights::Weight,
-    PalletId,
-};
-use sp_arithmetic::fixed_point::FixedU64;
-use sp_core::H256;
-use sp_io::TestExternalities;
-use sp_runtime::traits::{BlakeTwo256, IdentityLookup};
-
-use astar_primitives::{
-    dapp_staking::{CycleConfiguration, SmartContract, StakingRewardHandler},
-    oracle::PriceProvider,
-    testing::Header,
-    Balance, BlockNumber,
-};
-
-pub(crate) type AccountId = u64;
-
-pub(crate) const UNBONDING_ACCOUNT: AccountId = 10;
-
-pub(crate) const EXISTENTIAL_DEPOSIT: Balance = 2;
-pub(crate) const MINIMUM_LOCK_AMOUNT: Balance = 10;
-
-type UncheckedExtrinsic = frame_system::mocking::MockUncheckedExtrinsic<Test>;
-type Block = frame_system::mocking::MockBlock<Test>;
-
-construct_runtime!(
-    pub struct Test
-    where
-        Block = Block,
-        NodeBlock = Block,
-        UncheckedExtrinsic = UncheckedExtrinsic,
-    {
-        System: frame_system,
-        Balances: pallet_balances,
-        DappStaking: pallet_dapp_staking_v3,
-        DappsStaking: pallet_dapps_staking,
-        DappStakingMigration: pallet_dapp_staking_migration,
-    }
-);
-
-parameter_types! {
-    pub const BlockHashCount: BlockNumber = 250;
-    pub BlockWeights: frame_system::limits::BlockWeights =
-        frame_system::limits::BlockWeights::simple_max(Weight::from_parts(1024, 0));
-}
-
-impl frame_system::Config for Test {
-    type BaseCallFilter = frame_support::traits::Everything;
-    type BlockWeights = ();
-    type BlockLength = ();
-    type RuntimeOrigin = RuntimeOrigin;
-    type Index = u64;
-    type RuntimeCall = RuntimeCall;
-    type BlockNumber = BlockNumber;
-    type Hash = H256;
-    type Hashing = BlakeTwo256;
-    type AccountId = AccountId;
-    type Lookup = IdentityLookup<Self::AccountId>;
-    type Header = Header;
-    type RuntimeEvent = RuntimeEvent;
-    type BlockHashCount = BlockHashCount;
-    type DbWeight = ();
-    type Version = ();
-    type PalletInfo = PalletInfo;
-    type AccountData = pallet_balances::AccountData<Balance>;
-    type OnNewAccount = ();
-    type OnKilledAccount = ();
-    type SystemWeightInfo = ();
-    type SS58Prefix = ();
-    type OnSetCode = ();
-    type MaxConsumers = frame_support::traits::ConstU32<16>;
-}
-
-impl pallet_balances::Config for Test {
-    type MaxLocks = ConstU32<4>;
-    type MaxReserves = ();
-    type ReserveIdentifier = [u8; 8];
-    type Balance = Balance;
-    type RuntimeEvent = RuntimeEvent;
-    type DustRemoval = ();
-    type ExistentialDeposit = ConstU128<EXISTENTIAL_DEPOSIT>;
-    type AccountStore = System;
-    type HoldIdentifier = ();
-    type FreezeIdentifier = RuntimeFreezeReason;
-    type MaxHolds = ConstU32<0>;
-    type MaxFreezes = ConstU32<1>;
-    type WeightInfo = ();
-}
-
-pub struct DummyPriceProvider;
-impl PriceProvider for DummyPriceProvider {
-    fn average_price() -> FixedU64 {
-        FixedU64::from_rational(1, 10)
-    }
-}
-
-pub struct DummyStakingRewardHandler;
-impl StakingRewardHandler<AccountId> for DummyStakingRewardHandler {
-    fn staker_and_dapp_reward_pools(_total_staked_value: Balance) -> (Balance, Balance) {
-        (
-            Balance::from(1_000_000_000_000_u128),
-            Balance::from(1_000_000_000_u128),
-        )
-    }
-
-    fn bonus_reward_pool() -> Balance {
-        Balance::from(3_000_000_u128)
-    }
-
-    fn payout_reward(beneficiary: &AccountId, reward: Balance) -> Result<(), ()> {
-        let _ = Balances::mint_into(beneficiary, reward);
-        Ok(())
-    }
-}
-
-pub(crate) type MockSmartContract = SmartContract<AccountId>;
-
-#[cfg(feature = "runtime-benchmarks")]
-pub struct BenchmarkHelper<SC, ACC>(sp_std::marker::PhantomData<(SC, ACC)>);
-#[cfg(feature = "runtime-benchmarks")]
-impl pallet_dapp_staking_v3::BenchmarkHelper<MockSmartContract, AccountId>
-    for BenchmarkHelper<MockSmartContract, AccountId>
-{
-    fn get_smart_contract(id: u32) -> MockSmartContract {
-        MockSmartContract::Wasm(id as AccountId)
-    }
-
-    fn set_balance(account: &AccountId, amount: Balance) {
-        use frame_support::traits::fungible::Unbalanced as FunUnbalanced;
-        Balances::write_balance(account, amount)
-            .expect("Must succeed in test/benchmark environment.");
-    }
-}
-
-pub struct DummyCycleConfiguration;
-impl CycleConfiguration for DummyCycleConfiguration {
-    fn periods_per_cycle() -> u32 {
-        4
-    }
-
-    fn eras_per_voting_subperiod() -> u32 {
-        8
-    }
-
-    fn eras_per_build_and_earn_subperiod() -> u32 {
-        16
-    }
-
-    fn blocks_per_era() -> u32 {
-        10
-    }
-}
-
-impl pallet_dapp_staking_v3::Config for Test {
-    type RuntimeEvent = RuntimeEvent;
-    type RuntimeFreezeReason = RuntimeFreezeReason;
-    type Currency = Balances;
-    type SmartContract = MockSmartContract;
-    type ManagerOrigin = frame_system::EnsureRoot<AccountId>;
-    type NativePriceProvider = DummyPriceProvider;
-    type StakingRewardHandler = DummyStakingRewardHandler;
-    type CycleConfiguration = DummyCycleConfiguration;
-    type Observers = ();
-    type AccountCheck = ();
-    type EraRewardSpanLength = ConstU32<8>;
-    type RewardRetentionInPeriods = ConstU32<2>;
-    type MaxNumberOfContracts = ConstU32<10>;
-    type MaxUnlockingChunks = ConstU32<5>;
-    type MinimumLockedAmount = ConstU128<MINIMUM_LOCK_AMOUNT>;
-    type UnlockingPeriod = ConstU32<2>;
-    type MaxNumberOfStakedContracts = ConstU32<5>;
-    type MinimumStakeAmount = ConstU128<3>;
-    type NumberOfTiers = ConstU32<4>;
-    type WeightInfo = pallet_dapp_staking_v3::weights::SubstrateWeight<Test>;
-    #[cfg(feature = "runtime-benchmarks")]
-    type BenchmarkHelper = BenchmarkHelper<MockSmartContract, AccountId>;
-}
-
-parameter_types! {
-    pub const DappsStakingPalletId: PalletId = PalletId(*b"mokdpstk");
-}
-
-impl pallet_dapps_staking::Config for Test {
-    type RuntimeEvent = RuntimeEvent;
-    type Currency = Balances;
-    type BlockPerEra = ConstU32<10>;
-    type RegisterDeposit = ConstU128<100>;
-    type SmartContract = MockSmartContract;
-    type WeightInfo = pallet_dapps_staking::weights::SubstrateWeight<Test>;
-    type MaxNumberOfStakersPerContract = ConstU32<10>;
-    type MinimumStakingAmount = ConstU128<MINIMUM_LOCK_AMOUNT>;
-    type PalletId = DappsStakingPalletId;
-    type MinimumRemainingAmount = ConstU128<1>;
-    type MaxUnlockingChunks = ConstU32<5>;
-    type UnbondingPeriod = ConstU32<3>;
-    type MaxEraStakeValues = ConstU32<10>;
-    type UnregisteredDappRewardRetention = ConstU32<10>;
-    type ForcePalletDisabled = ConstBool<false>;
-    type DelegateClaimFee = ConstU128<1>;
-}
-
-impl pallet_dapp_staking_migration::Config for Test {
-    type RuntimeEvent = RuntimeEvent;
-    type WeightInfo = crate::weights::SubstrateWeight<Test>;
-}
-
-pub struct ExtBuilder;
-impl ExtBuilder {
-    pub fn build() -> TestExternalities {
-        // Normal behavior is for reward payout to succeed
-        let mut storage = frame_system::GenesisConfig::default()
-            .build_storage::<Test>()
-            .unwrap();
-
-        let balances = vec![1000; 11]
-            .into_iter()
-            .enumerate()
-            .map(|(idx, amount)| (idx as u64 + 1, amount))
-            .collect();
-
-        pallet_balances::GenesisConfig::<Test> { balances: balances }
-            .assimilate_storage(&mut storage)
-            .ok();
-
-        let mut ext = TestExternalities::from(storage);
-        ext.execute_with(|| {
-            System::set_block_number(1);
-        });
-
-        ext
-    }
-}
-
-/// Initialize old dApps staking storage.
-///
-/// This is kept outside of the test ext creation since the same mock is reused
-/// in the benchmarks code.
-pub fn init() {
-    let dapps_number = 10_u32;
-    let staker = dapps_number.into();
-    Balances::make_free_balance_be(&staker, 1_000_000_000_000_000_000);
-
-    // Add some dummy dApps to the old pallet & stake on them.
-    for idx in 0..dapps_number {
-        let developer = idx.into();
-        Balances::make_free_balance_be(&developer, 1_000_000_000_000);
-        let smart_contract = MockSmartContract::Wasm(idx.into());
-        assert_ok!(pallet_dapps_staking::Pallet::<Test>::register(
-            RawOrigin::Root.into(),
-            developer,
-            smart_contract.clone(),
-        ));
-        assert_ok!(pallet_dapps_staking::Pallet::<Test>::bond_and_stake(
-            RawOrigin::Signed(staker.clone()).into(),
-            smart_contract,
-            1_000,
-        ));
-    }
-
-    assert_ok!(pallet_dapps_staking::Pallet::<Test>::bond_and_stake(
-        RawOrigin::Signed(UNBONDING_ACCOUNT).into(),
-        MockSmartContract::Wasm(0),
-        2000
-    ));
-    assert_ok!(pallet_dapps_staking::Pallet::<Test>::unbond_and_unstake(
-        RawOrigin::Signed(UNBONDING_ACCOUNT).into(),
-        MockSmartContract::Wasm(0),
-        500
-    ));
-}
diff --git a/pallets/dapp-staking-migration/src/tests.rs b/pallets/dapp-staking-migration/src/tests.rs
deleted file mode 100644
index 0245be36db..0000000000
--- a/pallets/dapp-staking-migration/src/tests.rs
+++ /dev/null
@@ -1,198 +0,0 @@
-// This file is part of Astar.
-
-// Copyright (C) 2019-2023 Stake Technologies Pte.Ltd.
-// SPDX-License-Identifier: GPL-3.0-or-later
-
-// Astar is free software: you can redistribute it and/or modify
-// it under the terms of the GNU General Public License as published by
-// the Free Software Foundation, either version 3 of the License, or
-// (at your option) any later version.
-
-// Astar is distributed in the hope that it will be useful,
-// but WITHOUT ANY WARRANTY; without even the implied warranty of
-// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-// GNU General Public License for more details.
-
-// You should have received a copy of the GNU General Public License
-// along with Astar. If not, see <http://www.gnu.org/licenses/>.
-
-use crate::mock::*;
-use crate::*;
-
-use frame_support::{assert_ok, assert_storage_noop};
-use sp_runtime::traits::Zero;
-
-#[test]
-fn sanity_check() {
-    ExtBuilder::build().execute_with(|| {
-        assert!(DappStakingMigration::max_call_weight()
-            .all_gte(DappStakingMigration::min_call_weight()));
-    });
-}
-
-#[test]
-fn migrate_dapps_check() {
-    ExtBuilder::build().execute_with(|| {
-        init();
-
-        // Cleanup single entry, check pre and post states
-        let init_old_count = pallet_dapps_staking::RegisteredDapps::<Test>::iter().count();
-        assert!(init_old_count > 0, "Sanity check.");
-
-        let init_new_count = pallet_dapp_staking_v3::IntegratedDApps::<Test>::iter().count();
-        assert!(init_new_count.is_zero(), "Sanity check.");
-
-        assert_eq!(
-            DappStakingMigration::migrate_dapps(),
-            Ok(<Test as Config>::WeightInfo::migrate_dapps_success())
-        );
-        assert_eq!(
-            init_old_count,
-            pallet_dapps_staking::RegisteredDapps::<Test>::iter().count() + 1,
-            "One entry should have been cleaned up."
-        );
-        assert_eq!(
-            pallet_dapp_staking_v3::IntegratedDApps::<Test>::iter().count(),
-            1,
-            "Single new entry should have been added."
-        );
-
-        // Cleanup the remaining entries.
-        for _ in 1..init_old_count {
-            assert_eq!(
-                DappStakingMigration::migrate_dapps(),
-                Ok(<Test as Config>::WeightInfo::migrate_dapps_success())
-            );
-        }
-
-        // Further calls should result in Err
-        assert_eq!(
-            DappStakingMigration::migrate_dapps(),
-            Err(<Test as Config>::WeightInfo::migrate_dapps_noop())
-        );
-    });
-}
-
-#[test]
-fn migrate_ledgers_check() {
-    ExtBuilder::build().execute_with(|| {
-        init();
-
-        // Check for unbonding accounts
-        let init_unbonding_ledger = pallet_dapps_staking::Ledger::<Test>::get(&UNBONDING_ACCOUNT);
-        let unbonding_account_init_locked_amount = init_unbonding_ledger.locked;
-        let unbonding_account_unbonding_amount = init_unbonding_ledger.unbonding_info.sum();
-
-        // Cleanup all entries, check pre and post states.
-        let init_old_count = pallet_dapps_staking::Ledger::<Test>::iter().count();
-        assert!(init_old_count > 0, "Sanity check.");
-
-        let init_new_count = pallet_dapp_staking_v3::Ledger::<Test>::iter().count();
-        assert!(init_new_count.is_zero(), "Sanity check.");
-
-        assert!(pallet_dapp_staking_v3::CurrentEraInfo::<Test>::get()
-            .total_locked
-            .is_zero());
-
-        for x in 0..init_old_count {
-            assert_eq!(
-                DappStakingMigration::migrate_ledger(),
-                Ok(<Test as Config>::WeightInfo::migrate_ledger_success())
-            );
-
-            assert_eq!(
-                init_old_count - x - 1,
-                pallet_dapps_staking::Ledger::<Test>::iter().count(),
-                "One entry should have been cleaned up."
-            );
-            assert_eq!(
-                x + 1,
-                pallet_dapp_staking_v3::Ledger::<Test>::iter().count(),
-                "Single new entry should have been added."
-            );
-            assert!(pallet_dapp_staking_v3::CurrentEraInfo::<Test>::get().total_locked > 0);
-        }
-
-        // Further calls should result in Err
-        assert_eq!(
-            DappStakingMigration::migrate_ledger(),
-            Err(<Test as Config>::WeightInfo::migrate_ledger_noop())
-        );
-
-        let post_unbonding_ledger = pallet_dapp_staking_v3::Ledger::<Test>::get(&UNBONDING_ACCOUNT);
-        assert_eq!(
-            post_unbonding_ledger.locked,
-            unbonding_account_init_locked_amount - unbonding_account_unbonding_amount
-        );
-    });
-}
-
-// TODO: this doesn't work since clear_prefix doesn't work in tests for some reason.
-#[ignore]
-#[test]
-fn storage_cleanup_check() {
-    let mut ext = ExtBuilder::build();
-    assert_ok!(ext.commit_all());
-
-    ext.execute_with(|| {
-        init();
-
-        let init_count = (pallet_dapps_staking::RegisteredDapps::<Test>::iter().count()
-            + pallet_dapps_staking::Ledger::<Test>::iter().count()) as u32;
-
-        for _ in 0..init_count {
-            assert_ok!(DappStakingMigration::cleanup_old_storage(init_count));
-        }
-    });
-}
-
-#[test]
-fn migrate_call_works() {
-    ExtBuilder::build().execute_with(|| {
-        init();
-        let account = 1;
-
-        // Call enough times to clean everything up.
-        while MigrationStateStorage::<Test>::get() != MigrationState::Finished {
-            assert_ok!(DappStakingMigration::migrate(
-                frame_system::RawOrigin::Signed(account).into(),
-                Some(Weight::from_parts(1, 1))
-            ));
-
-            match MigrationStateStorage::<Test>::get() {
-                MigrationState::RegisteredDApps | MigrationState::Ledgers => {
-                    assert!(
-                        pallet_dapp_staking_v3::ActiveProtocolState::<Test>::get().maintenance,
-                        "Pallet must be in the maintenance mode during old storage migration."
-                    );
-                }
-                _ => {
-                    assert!(
-                        !pallet_dapp_staking_v3::ActiveProtocolState::<Test>::get().maintenance,
-                        "Maintenance mode is disabled during old storage cleanup."
-                    );
-                }
-            }
-        }
-
-        // Check post-state
-        assert!(pallet_dapps_staking::RegisteredDapps::<Test>::iter()
-            .count()
-            .is_zero());
-        assert!(pallet_dapps_staking::Ledger::<Test>::iter()
-            .count()
-            .is_zero());
-        assert!(pallet_dapps_staking::RegisteredDevelopers::<Test>::iter()
-            .count()
-            .is_zero());
-        assert!(pallet_dapps_staking::GeneralEraInfo::<Test>::iter()
-            .count()
-            .is_zero());
-
-        // Migrate call can still be called, but it shouldn't have any effect.
-        assert_storage_noop!(assert_ok!(DappStakingMigration::migrate(
-            frame_system::RawOrigin::Signed(account).into(),
-            None
-        )));
-    });
-}
diff --git a/pallets/dapp-staking-migration/src/weights.rs b/pallets/dapp-staking-migration/src/weights.rs
index 87250d7fa7..8ee919f35d 100644
--- a/pallets/dapp-staking-migration/src/weights.rs
+++ b/pallets/dapp-staking-migration/src/weights.rs
@@ -20,24 +20,24 @@
 //! Autogenerated weights for pallet_dapp_staking_migration
 //!
 //! THIS FILE WAS AUTO-GENERATED USING THE SUBSTRATE BENCHMARK CLI VERSION 4.0.0-dev
-//! DATE: 2024-02-07, STEPS: `50`, REPEAT: `20`, LOW RANGE: `[]`, HIGH RANGE: `[]`
+//! DATE: 2024-03-12, STEPS: `5`, REPEAT: `1`, LOW RANGE: `[]`, HIGH RANGE: `[]`
 //! WORST CASE MAP SIZE: `1000000`
-//! HOSTNAME: `gh-runner-01-ovh`, CPU: `Intel(R) Xeon(R) E-2236 CPU @ 3.40GHz`
-//! EXECUTION: Some(Wasm), WASM-EXECUTION: Compiled, CHAIN: Some("astar-dev"), DB CACHE: 1024
+//! HOSTNAME: `Dinos-MBP.fritz.box`, CPU: `<UNKNOWN>`
+//! EXECUTION: Some(Wasm), WASM-EXECUTION: Compiled, CHAIN: Some("shibuya-dev"), DB CACHE: 1024
 
 // Executed Command:
 // ./target/release/astar-collator
 // benchmark
 // pallet
-// --chain=astar-dev
-// --steps=50
-// --repeat=20
+// --chain=shibuya-dev
+// --steps=5
+// --repeat=1
 // --pallet=pallet_dapp_staking_migration
 // --extrinsic=*
 // --execution=wasm
 // --wasm-execution=compiled
 // --heap-pages=4096
-// --output=./benchmark-results/astar-dev/dapp_staking_migration_weights.rs
+// --output=weights.rs
 // --template=./scripts/templates/weight-template.hbs
 
 #![cfg_attr(rustfmt, rustfmt_skip)]
@@ -49,189 +49,57 @@ use core::marker::PhantomData;
 
 /// Weight functions needed for pallet_dapp_staking_migration.
 pub trait WeightInfo {
-	fn migrate_dapps_success() -> Weight;
-	fn migrate_dapps_noop() -> Weight;
-	fn migrate_ledger_success() -> Weight;
-	fn migrate_ledger_noop() -> Weight;
-	fn cleanup_old_storage_success(x: u32, ) -> Weight;
-	fn cleanup_old_storage_noop() -> Weight;
+	fn translate_staking_info_success() -> Weight;
+	fn translate_staking_info_success_noop() -> Weight;
 }
 
 /// Weights for pallet_dapp_staking_migration using the Substrate node and recommended hardware.
 pub struct SubstrateWeight<T>(PhantomData<T>);
 impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {
-	/// Storage: DappsStaking RegisteredDapps (r:2 w:1)
-	/// Proof: DappsStaking RegisteredDapps (max_values: None, max_size: Some(86), added: 2561, mode: MaxEncodedLen)
-	/// Storage: System Account (r:1 w:1)
-	/// Proof: System Account (max_values: None, max_size: Some(128), added: 2603, mode: MaxEncodedLen)
-	/// Storage: DappStaking IntegratedDApps (r:1 w:1)
-	/// Proof: DappStaking IntegratedDApps (max_values: Some(65535), max_size: Some(116), added: 2096, mode: MaxEncodedLen)
-	/// Storage: DappStaking CounterForIntegratedDApps (r:1 w:1)
-	/// Proof: DappStaking CounterForIntegratedDApps (max_values: Some(1), max_size: Some(4), added: 499, mode: MaxEncodedLen)
-	/// Storage: DappStaking NextDAppId (r:1 w:1)
-	/// Proof: DappStaking NextDAppId (max_values: Some(1), max_size: Some(2), added: 497, mode: MaxEncodedLen)
-	fn migrate_dapps_success() -> Weight {
+	/// Storage: DappStaking StakerInfo (r:2 w:1)
+	/// Proof: DappStaking StakerInfo (max_values: None, max_size: Some(178), added: 2653, mode: MaxEncodedLen)
+	fn translate_staking_info_success() -> Weight {
 		// Proof Size summary in bytes:
-		//  Measured:  `558`
-		//  Estimated: `6112`
-		// Minimum execution time: 47_762_000 picoseconds.
-		Weight::from_parts(48_426_000, 6112)
-			.saturating_add(T::DbWeight::get().reads(6_u64))
-			.saturating_add(T::DbWeight::get().writes(5_u64))
+		//  Measured:  `368`
+		//  Estimated: `6296`
+		// Minimum execution time: 16_000_000 picoseconds.
+		Weight::from_parts(16_000_000, 6296)
+			.saturating_add(T::DbWeight::get().reads(2_u64))
+			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
-	/// Storage: DappsStaking RegisteredDapps (r:1 w:0)
-	/// Proof: DappsStaking RegisteredDapps (max_values: None, max_size: Some(86), added: 2561, mode: MaxEncodedLen)
-	fn migrate_dapps_noop() -> Weight {
+	/// Storage: DappStaking StakerInfo (r:1 w:0)
+	/// Proof: DappStaking StakerInfo (max_values: None, max_size: Some(178), added: 2653, mode: MaxEncodedLen)
+	fn translate_staking_info_success_noop() -> Weight {
 		// Proof Size summary in bytes:
-		//  Measured:  `0`
-		//  Estimated: `3551`
-		// Minimum execution time: 3_150_000 picoseconds.
-		Weight::from_parts(3_368_000, 3551)
+		//  Measured:  `19`
+		//  Estimated: `3643`
+		// Minimum execution time: 4_000_000 picoseconds.
+		Weight::from_parts(4_000_000, 3643)
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 	}
-	/// Storage: DappsStaking Ledger (r:2 w:1)
-	/// Proof: DappsStaking Ledger (max_values: None, max_size: Some(266), added: 2741, mode: MaxEncodedLen)
-	/// Storage: Balances Locks (r:1 w:1)
-	/// Proof: Balances Locks (max_values: None, max_size: Some(1299), added: 3774, mode: MaxEncodedLen)
-	/// Storage: Balances Freezes (r:1 w:1)
-	/// Proof: Balances Freezes (max_values: None, max_size: Some(67), added: 2542, mode: MaxEncodedLen)
-	/// Storage: System Account (r:1 w:1)
-	/// Proof: System Account (max_values: None, max_size: Some(128), added: 2603, mode: MaxEncodedLen)
-	/// Storage: DappStaking Ledger (r:1 w:1)
-	/// Proof: DappStaking Ledger (max_values: None, max_size: Some(310), added: 2785, mode: MaxEncodedLen)
-	/// Storage: CollatorSelection Candidates (r:1 w:0)
-	/// Proof Skipped: CollatorSelection Candidates (max_values: Some(1), max_size: None, mode: Measured)
-	/// Storage: DappStaking CurrentEraInfo (r:1 w:1)
-	/// Proof: DappStaking CurrentEraInfo (max_values: Some(1), max_size: Some(112), added: 607, mode: MaxEncodedLen)
-	fn migrate_ledger_success() -> Weight {
-		// Proof Size summary in bytes:
-		//  Measured:  `1875`
-		//  Estimated: `6472`
-		// Minimum execution time: 70_640_000 picoseconds.
-		Weight::from_parts(72_730_000, 6472)
-			.saturating_add(T::DbWeight::get().reads(8_u64))
-			.saturating_add(T::DbWeight::get().writes(6_u64))
-	}
-	/// Storage: DappsStaking Ledger (r:1 w:0)
-	/// Proof: DappsStaking Ledger (max_values: None, max_size: Some(266), added: 2741, mode: MaxEncodedLen)
-	fn migrate_ledger_noop() -> Weight {
-		// Proof Size summary in bytes:
-		//  Measured:  `0`
-		//  Estimated: `3731`
-		// Minimum execution time: 2_769_000 picoseconds.
-		Weight::from_parts(2_894_000, 3731)
-			.saturating_add(T::DbWeight::get().reads(1_u64))
-	}
-	/// Storage: DappsStaking Ledger (r:6 w:5)
-	/// Proof: DappsStaking Ledger (max_values: None, max_size: Some(266), added: 2741, mode: MaxEncodedLen)
-	/// The range of component `x` is `[1, 5]`.
-	fn cleanup_old_storage_success(x: u32, ) -> Weight {
-		// Proof Size summary in bytes:
-		//  Measured:  `700 + x * (40 ±0)`
-		//  Estimated: `3731 + x * (2741 ±0)`
-		// Minimum execution time: 6_830_000 picoseconds.
-		Weight::from_parts(6_598_680, 3731)
-			// Standard Error: 6_203
-			.saturating_add(Weight::from_parts(686_141, 0).saturating_mul(x.into()))
-			.saturating_add(T::DbWeight::get().reads(1_u64))
-			.saturating_add(T::DbWeight::get().reads((1_u64).saturating_mul(x.into())))
-			.saturating_add(T::DbWeight::get().writes((1_u64).saturating_mul(x.into())))
-			.saturating_add(Weight::from_parts(0, 2741).saturating_mul(x.into()))
-	}
-	fn cleanup_old_storage_noop() -> Weight {
-		// Proof Size summary in bytes:
-		//  Measured:  `0`
-		//  Estimated: `0`
-		// Minimum execution time: 2_017_000 picoseconds.
-		Weight::from_parts(2_116_000, 0)
-	}
 }
 
 // For backwards compatibility and tests
 impl WeightInfo for () {
-	/// Storage: DappsStaking RegisteredDapps (r:2 w:1)
-	/// Proof: DappsStaking RegisteredDapps (max_values: None, max_size: Some(86), added: 2561, mode: MaxEncodedLen)
-	/// Storage: System Account (r:1 w:1)
-	/// Proof: System Account (max_values: None, max_size: Some(128), added: 2603, mode: MaxEncodedLen)
-	/// Storage: DappStaking IntegratedDApps (r:1 w:1)
-	/// Proof: DappStaking IntegratedDApps (max_values: Some(65535), max_size: Some(116), added: 2096, mode: MaxEncodedLen)
-	/// Storage: DappStaking CounterForIntegratedDApps (r:1 w:1)
-	/// Proof: DappStaking CounterForIntegratedDApps (max_values: Some(1), max_size: Some(4), added: 499, mode: MaxEncodedLen)
-	/// Storage: DappStaking NextDAppId (r:1 w:1)
-	/// Proof: DappStaking NextDAppId (max_values: Some(1), max_size: Some(2), added: 497, mode: MaxEncodedLen)
-	fn migrate_dapps_success() -> Weight {
+	/// Storage: DappStaking StakerInfo (r:2 w:1)
+	/// Proof: DappStaking StakerInfo (max_values: None, max_size: Some(178), added: 2653, mode: MaxEncodedLen)
+	fn translate_staking_info_success() -> Weight {
 		// Proof Size summary in bytes:
-		//  Measured:  `558`
-		//  Estimated: `6112`
-		// Minimum execution time: 47_762_000 picoseconds.
-		Weight::from_parts(48_426_000, 6112)
-			.saturating_add(RocksDbWeight::get().reads(6_u64))
-			.saturating_add(RocksDbWeight::get().writes(5_u64))
+		//  Measured:  `368`
+		//  Estimated: `6296`
+		// Minimum execution time: 16_000_000 picoseconds.
+		Weight::from_parts(16_000_000, 6296)
+			.saturating_add(RocksDbWeight::get().reads(2_u64))
+			.saturating_add(RocksDbWeight::get().writes(1_u64))
 	}
-	/// Storage: DappsStaking RegisteredDapps (r:1 w:0)
-	/// Proof: DappsStaking RegisteredDapps (max_values: None, max_size: Some(86), added: 2561, mode: MaxEncodedLen)
-	fn migrate_dapps_noop() -> Weight {
+	/// Storage: DappStaking StakerInfo (r:1 w:0)
+	/// Proof: DappStaking StakerInfo (max_values: None, max_size: Some(178), added: 2653, mode: MaxEncodedLen)
+	fn translate_staking_info_success_noop() -> Weight {
 		// Proof Size summary in bytes:
-		//  Measured:  `0`
-		//  Estimated: `3551`
-		// Minimum execution time: 3_150_000 picoseconds.
-		Weight::from_parts(3_368_000, 3551)
+		//  Measured:  `19`
+		//  Estimated: `3643`
+		// Minimum execution time: 4_000_000 picoseconds.
+		Weight::from_parts(4_000_000, 3643)
 			.saturating_add(RocksDbWeight::get().reads(1_u64))
 	}
-	/// Storage: DappsStaking Ledger (r:2 w:1)
-	/// Proof: DappsStaking Ledger (max_values: None, max_size: Some(266), added: 2741, mode: MaxEncodedLen)
-	/// Storage: Balances Locks (r:1 w:1)
-	/// Proof: Balances Locks (max_values: None, max_size: Some(1299), added: 3774, mode: MaxEncodedLen)
-	/// Storage: Balances Freezes (r:1 w:1)
-	/// Proof: Balances Freezes (max_values: None, max_size: Some(67), added: 2542, mode: MaxEncodedLen)
-	/// Storage: System Account (r:1 w:1)
-	/// Proof: System Account (max_values: None, max_size: Some(128), added: 2603, mode: MaxEncodedLen)
-	/// Storage: DappStaking Ledger (r:1 w:1)
-	/// Proof: DappStaking Ledger (max_values: None, max_size: Some(310), added: 2785, mode: MaxEncodedLen)
-	/// Storage: CollatorSelection Candidates (r:1 w:0)
-	/// Proof Skipped: CollatorSelection Candidates (max_values: Some(1), max_size: None, mode: Measured)
-	/// Storage: DappStaking CurrentEraInfo (r:1 w:1)
-	/// Proof: DappStaking CurrentEraInfo (max_values: Some(1), max_size: Some(112), added: 607, mode: MaxEncodedLen)
-	fn migrate_ledger_success() -> Weight {
-		// Proof Size summary in bytes:
-		//  Measured:  `1875`
-		//  Estimated: `6472`
-		// Minimum execution time: 70_640_000 picoseconds.
-		Weight::from_parts(72_730_000, 6472)
-			.saturating_add(RocksDbWeight::get().reads(8_u64))
-			.saturating_add(RocksDbWeight::get().writes(6_u64))
-	}
-	/// Storage: DappsStaking Ledger (r:1 w:0)
-	/// Proof: DappsStaking Ledger (max_values: None, max_size: Some(266), added: 2741, mode: MaxEncodedLen)
-	fn migrate_ledger_noop() -> Weight {
-		// Proof Size summary in bytes:
-		//  Measured:  `0`
-		//  Estimated: `3731`
-		// Minimum execution time: 2_769_000 picoseconds.
-		Weight::from_parts(2_894_000, 3731)
-			.saturating_add(RocksDbWeight::get().reads(1_u64))
-	}
-	/// Storage: DappsStaking Ledger (r:6 w:5)
-	/// Proof: DappsStaking Ledger (max_values: None, max_size: Some(266), added: 2741, mode: MaxEncodedLen)
-	/// The range of component `x` is `[1, 5]`.
-	fn cleanup_old_storage_success(x: u32, ) -> Weight {
-		// Proof Size summary in bytes:
-		//  Measured:  `700 + x * (40 ±0)`
-		//  Estimated: `3731 + x * (2741 ±0)`
-		// Minimum execution time: 6_830_000 picoseconds.
-		Weight::from_parts(6_598_680, 3731)
-			// Standard Error: 6_203
-			.saturating_add(Weight::from_parts(686_141, 0).saturating_mul(x.into()))
-			.saturating_add(RocksDbWeight::get().reads(1_u64))
-			.saturating_add(RocksDbWeight::get().reads((1_u64).saturating_mul(x.into())))
-			.saturating_add(RocksDbWeight::get().writes((1_u64).saturating_mul(x.into())))
-			.saturating_add(Weight::from_parts(0, 2741).saturating_mul(x.into()))
-	}
-	fn cleanup_old_storage_noop() -> Weight {
-		// Proof Size summary in bytes:
-		//  Measured:  `0`
-		//  Estimated: `0`
-		// Minimum execution time: 2_017_000 picoseconds.
-		Weight::from_parts(2_116_000, 0)
-	}
 }
diff --git a/pallets/dapp-staking-v3/src/lib.rs b/pallets/dapp-staking-v3/src/lib.rs
index bbc2d835ed..83f9fde53c 100644
--- a/pallets/dapp-staking-v3/src/lib.rs
+++ b/pallets/dapp-staking-v3/src/lib.rs
@@ -92,7 +92,7 @@ pub mod pallet {
     use super::*;
 
     /// The current storage version.
-    pub const STORAGE_VERSION: StorageVersion = StorageVersion::new(5);
+    pub const STORAGE_VERSION: StorageVersion = StorageVersion::new(6);
 
     #[pallet::pallet]
     #[pallet::storage_version(STORAGE_VERSION)]
diff --git a/pallets/dapp-staking-v3/src/types.rs b/pallets/dapp-staking-v3/src/types.rs
index dc3f4fbf39..94ca91c866 100644
--- a/pallets/dapp-staking-v3/src/types.rs
+++ b/pallets/dapp-staking-v3/src/types.rs
@@ -997,6 +997,19 @@ impl SingularStakingInfo {
         }
     }
 
+    /// TODO: remove this after migration to v6 has been done.
+    pub fn new_migration(
+        previous_staked: StakeAmount,
+        staked: StakeAmount,
+        loyal_staker: bool,
+    ) -> Self {
+        Self {
+            previous_staked,
+            staked,
+            loyal_staker,
+        }
+    }
+
     /// Stake the specified amount on the contract, for the specified subperiod.
     pub fn stake(&mut self, amount: Balance, current_era: EraNumber, subperiod: Subperiod) {
         // Keep the previous stake amount for future reference
diff --git a/runtime/astar/src/lib.rs b/runtime/astar/src/lib.rs
index b80ce5b058..1c1e6a198a 100644
--- a/runtime/astar/src/lib.rs
+++ b/runtime/astar/src/lib.rs
@@ -28,8 +28,8 @@ use frame_support::{
     dispatch::DispatchClass,
     parameter_types,
     traits::{
-        AsEnsureOriginWithArg, ConstBool, ConstU128, ConstU32, Contains, Currency, FindAuthor, Get,
-        Imbalance, InstanceFilter, Nothing, OnFinalize, OnUnbalanced, Randomness, WithdrawReasons,
+        AsEnsureOriginWithArg, ConstBool, ConstU32, Contains, Currency, FindAuthor, Get, Imbalance,
+        InstanceFilter, Nothing, OnFinalize, OnUnbalanced, Randomness, WithdrawReasons,
     },
     weights::{
         constants::{
@@ -320,39 +320,7 @@ impl pallet_multisig::Config for Runtime {
 }
 
 parameter_types! {
-    pub const BlockPerEra: BlockNumber = DAYS;
-    pub const RegisterDeposit: Balance = 1000 * ASTR;
-    pub const MaxNumberOfStakersPerContract: u32 = 16384;
     pub const MinimumStakingAmount: Balance = 500 * ASTR;
-    pub const MinimumRemainingAmount: Balance = ASTR;
-    pub const MaxEraStakeValues: u32 = 5;
-    pub const MaxUnlockingChunks: u32 = 4;
-    pub const UnbondingPeriod: u32 = 10;
-}
-
-impl pallet_dapps_staking::Config for Runtime {
-    type Currency = Balances;
-    type BlockPerEra = BlockPerEra;
-    type SmartContract = SmartContract<AccountId>;
-    type RegisterDeposit = RegisterDeposit;
-    type RuntimeEvent = RuntimeEvent;
-    type WeightInfo = pallet_dapps_staking::weights::SubstrateWeight<Runtime>;
-    type MaxNumberOfStakersPerContract = MaxNumberOfStakersPerContract;
-    type MinimumStakingAmount = MinimumStakingAmount;
-    type PalletId = DappsStakingPalletId;
-    type MaxUnlockingChunks = MaxUnlockingChunks;
-    type UnbondingPeriod = UnbondingPeriod;
-    type MinimumRemainingAmount = MinimumRemainingAmount;
-    type MaxEraStakeValues = MaxEraStakeValues;
-    // Not allowed on Astar yet
-    type UnregisteredDappRewardRetention = ConstU32<{ u32::MAX }>;
-    // Needed so benchmark can use the pallets extrinsics
-    #[cfg(feature = "runtime-benchmarks")]
-    type ForcePalletDisabled = ConstBool<false>;
-    #[cfg(not(feature = "runtime-benchmarks"))]
-    type ForcePalletDisabled = ConstBool<true>;
-    // Fee required to claim rewards for another account. Calculated & tested manually.
-    type DelegateClaimFee = ConstU128<057_950_348_114_187_155>;
 }
 
 impl pallet_static_price_provider::Config for Runtime {
@@ -1104,8 +1072,6 @@ construct_runtime!(
         StaticPriceProvider: pallet_static_price_provider = 253,
         // To be removed & cleaned up after migration has been finished
         DappStakingMigration: pallet_dapp_staking_migration = 254,
-        // Legacy dApps staking v2, to be removed after migration has been finished
-        DappsStaking: pallet_dapps_staking = 255,
     }
 );
 
@@ -1143,10 +1109,21 @@ pub type Executive = frame_executive::Executive<
     Migrations,
 >;
 
+parameter_types! {
+    pub const DappStakingMigrationName: &'static str = "DappStakingMigration";
+}
 /// All migrations that will run on the next runtime upgrade.
 ///
 /// Once done, migrations should be removed from the tuple.
-pub type Migrations = ();
+pub type Migrations = (
+    // Part of astar-82, need to first cleanup old storage before re-using the pallet
+    frame_support::migrations::RemovePallet<
+        DappStakingMigrationName,
+        <Runtime as frame_system::Config>::DbWeight,
+    >,
+    // Part of astar-82
+    (pallet_dapp_staking_migration::SingularStakingInfoTranslationUpgrade<Runtime>,),
+);
 
 type EventRecord = frame_system::EventRecord<
     <Runtime as frame_system::Config>::RuntimeEvent,
@@ -1223,7 +1200,6 @@ mod benches {
         [pallet_assets, Assets]
         [pallet_balances, Balances]
         [pallet_timestamp, Timestamp]
-        [pallet_dapps_staking, DappsStaking]
         [pallet_dapp_staking_v3, DappStaking]
         [pallet_inflation, Inflation]
         [pallet_dapp_staking_migration, DappStakingMigration]
diff --git a/runtime/shibuya/Cargo.toml b/runtime/shibuya/Cargo.toml
index aa35a772c0..63a6409822 100644
--- a/runtime/shibuya/Cargo.toml
+++ b/runtime/shibuya/Cargo.toml
@@ -101,6 +101,7 @@ astar-xcm-benchmarks = { workspace = true, optional = true }
 pallet-chain-extension-unified-accounts = { workspace = true }
 pallet-chain-extension-xvm = { workspace = true }
 pallet-collator-selection = { workspace = true }
+pallet-dapp-staking-migration = { workspace = true }
 pallet-dapp-staking-v3 = { workspace = true }
 pallet-dynamic-evm-base-fee = { workspace = true }
 pallet-ethereum-checked = { workspace = true }
@@ -153,6 +154,7 @@ std = [
 	"sp-api/std",
 	"sp-core/std",
 	"sp-consensus-aura/std",
+	"pallet-dapp-staking-migration/std",
 	"sp-io/std",
 	"sp-runtime/std",
 	"sp-runtime-interface/std",
@@ -248,6 +250,7 @@ std = [
 ]
 runtime-benchmarks = [
 	"astar-xcm-benchmarks/runtime-benchmarks",
+	"pallet-dapp-staking-migration/runtime-benchmarks",
 	"pallet-xcm-benchmarks/runtime-benchmarks",
 	"frame-benchmarking",
 	"frame-support/runtime-benchmarks",
@@ -295,6 +298,7 @@ try-runtime = [
 	"pallet-assets/try-runtime",
 	"pallet-authorship/try-runtime",
 	"pallet-collator-selection/try-runtime",
+	"pallet-dapp-staking-migration/try-runtime",
 	"pallet-session/try-runtime",
 	"pallet-xcm/try-runtime",
 	"pallet-identity/try-runtime",
diff --git a/runtime/shibuya/src/lib.rs b/runtime/shibuya/src/lib.rs
index c99e1a8965..aad8042fd6 100644
--- a/runtime/shibuya/src/lib.rs
+++ b/runtime/shibuya/src/lib.rs
@@ -1259,6 +1259,11 @@ impl pallet_unified_accounts::Config for Runtime {
     type WeightInfo = pallet_unified_accounts::weights::SubstrateWeight<Self>;
 }
 
+impl pallet_dapp_staking_migration::Config for Runtime {
+    type RuntimeEvent = RuntimeEvent;
+    type WeightInfo = pallet_dapp_staking_migration::weights::SubstrateWeight<Self>;
+}
+
 construct_runtime!(
     pub struct Runtime where
         Block = Block,
@@ -1316,6 +1321,8 @@ construct_runtime!(
 
         Sudo: pallet_sudo = 99,
 
+        // Remove after migrating to v6 storage
+        DappStakingMigration: pallet_dapp_staking_migration = 252,
         // To be removed & cleaned up once proper oracle is implemented
         StaticPriceProvider: pallet_static_price_provider = 253,
     }
@@ -1358,7 +1365,8 @@ pub type Executive = frame_executive::Executive<
 /// All migrations that will run on the next runtime upgrade.
 ///
 /// Once done, migrations should be removed from the tuple.
-pub type Migrations = ();
+pub type Migrations =
+    (pallet_dapp_staking_migration::SingularStakingInfoTranslationUpgrade<Runtime>,);
 
 type EventRecord = frame_system::EventRecord<
     <Runtime as frame_system::Config>::RuntimeEvent,
@@ -1446,6 +1454,7 @@ mod benches {
         [pallet_unified_accounts, UnifiedAccounts]
         [xcm_benchmarks_generic, XcmGeneric]
         [xcm_benchmarks_fungible, XcmFungible]
+        [pallet_dapp_staking_migration, DappStakingMigration]
     );
 }
 
diff --git a/runtime/shiden/src/lib.rs b/runtime/shiden/src/lib.rs
index afa5b49d07..2adb5b98a8 100644
--- a/runtime/shiden/src/lib.rs
+++ b/runtime/shiden/src/lib.rs
@@ -1006,6 +1006,11 @@ impl pallet_proxy::Config for Runtime {
     type AnnouncementDepositFactor = AnnouncementDepositFactor;
 }
 
+impl pallet_dapp_staking_migration::Config for Runtime {
+    type RuntimeEvent = RuntimeEvent;
+    type WeightInfo = pallet_dapp_staking_migration::weights::SubstrateWeight<Self>;
+}
+
 construct_runtime!(
     pub struct Runtime where
         Block = Block,
@@ -1055,6 +1060,8 @@ construct_runtime!(
 
         Sudo: pallet_sudo = 99,
 
+        // Remove after migrating to v6 storage
+        DappStakingMigration: pallet_dapp_staking_migration = 252,
         // To be removed & cleaned up once proper oracle is implemented
         StaticPriceProvider: pallet_static_price_provider = 253,
     }
@@ -1108,6 +1115,8 @@ pub type Migrations = (
     >,
     // Part of shiden-119
     RecalculationEraFix,
+    // Part of shiden-121
+    (pallet_dapp_staking_migration::SingularStakingInfoTranslationUpgrade<Runtime>,),
 );
 
 use frame_support::traits::OnRuntimeUpgrade;

From 986c09bdf3cf4bcb960664368dc8f3e3dba6009c Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Tue, 12 Mar 2024 15:59:01 +0100
Subject: [PATCH 23/28] Merge issue resolution

---
 .../dapp-staking-migration/src/benchmarking.rs |  2 +-
 runtime/shibuya/src/lib.rs                     |  1 +
 runtime/shiden/src/lib.rs                      | 18 +-----------------
 3 files changed, 3 insertions(+), 18 deletions(-)

diff --git a/pallets/dapp-staking-migration/src/benchmarking.rs b/pallets/dapp-staking-migration/src/benchmarking.rs
index 60e7843120..d15796797d 100644
--- a/pallets/dapp-staking-migration/src/benchmarking.rs
+++ b/pallets/dapp-staking-migration/src/benchmarking.rs
@@ -18,8 +18,8 @@
 
 use super::{Pallet as Migration, *};
 
-use frame_benchmarking::{account as benchmark_account, v2::*};
 use astar_primitives::{dapp_staking::SmartContractHandle, Balance};
+use frame_benchmarking::{account as benchmark_account, v2::*};
 use pallet_dapp_staking_v3::StakeAmount;
 
 fn smart_contract<T: Config>(idx: u8) -> T::SmartContract {
diff --git a/runtime/shibuya/src/lib.rs b/runtime/shibuya/src/lib.rs
index 00cd3d1aaa..ae6c4397be 100644
--- a/runtime/shibuya/src/lib.rs
+++ b/runtime/shibuya/src/lib.rs
@@ -1379,6 +1379,7 @@ pub type Executive = frame_executive::Executive<
 ///
 /// Once done, migrations should be removed from the tuple.
 pub type Migrations =
+    // Part of shibuya-124
     (pallet_dapp_staking_migration::SingularStakingInfoTranslationUpgrade<Runtime>,);
 
 type EventRecord = frame_system::EventRecord<
diff --git a/runtime/shiden/src/lib.rs b/runtime/shiden/src/lib.rs
index 85750fc223..d64c48e1a9 100644
--- a/runtime/shiden/src/lib.rs
+++ b/runtime/shiden/src/lib.rs
@@ -1131,28 +1131,12 @@ parameter_types! {
 /// Once done, migrations should be removed from the tuple.
 pub type Migrations = (
     // Part of shiden-121
-    pallet_dapp_staking_migration::SingularStakingInfoTranslationUpgrade<Runtime>,,
+    pallet_dapp_staking_migration::SingularStakingInfoTranslationUpgrade<Runtime>,
     // Part of shiden-121 (added after v5.33.0 release)
     SetNewTierConfig,
 );
 
 use frame_support::traits::OnRuntimeUpgrade;
-pub struct RecalculationEraFix;
-impl OnRuntimeUpgrade for RecalculationEraFix {
-    fn on_runtime_upgrade() -> Weight {
-        let first_dapp_staking_v3_era = 743;
-
-        let expected_recalculation_era =
-            InflationCycleConfig::eras_per_cycle().saturating_add(first_dapp_staking_v3_era);
-
-        pallet_inflation::ActiveInflationConfig::<Runtime>::mutate(|config| {
-            config.recalculation_era = expected_recalculation_era;
-        });
-
-        <Runtime as frame_system::Config>::DbWeight::get().reads_writes(1, 1)
-    }
-}
-
 pub struct SetNewTierConfig;
 impl OnRuntimeUpgrade for SetNewTierConfig {
     fn on_runtime_upgrade() -> Weight {

From 61e9fa2e4c8aadfcce7c3119e1616c6188bcb78d Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Tue, 12 Mar 2024 16:05:23 +0100
Subject: [PATCH 24/28] Vec dep

---
 pallets/dapp-staking-migration/src/benchmarking.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pallets/dapp-staking-migration/src/benchmarking.rs b/pallets/dapp-staking-migration/src/benchmarking.rs
index d15796797d..828f7120ec 100644
--- a/pallets/dapp-staking-migration/src/benchmarking.rs
+++ b/pallets/dapp-staking-migration/src/benchmarking.rs
@@ -18,6 +18,7 @@
 
 use super::{Pallet as Migration, *};
 
+use sp_std::vec;
 use astar_primitives::{dapp_staking::SmartContractHandle, Balance};
 use frame_benchmarking::{account as benchmark_account, v2::*};
 use pallet_dapp_staking_v3::StakeAmount;

From b8c7e702b74356b66120ddc81bc9675d20d94769 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Wed, 13 Mar 2024 11:20:07 +0100
Subject: [PATCH 25/28] Fixes, formatting

---
 pallets/dapp-staking-migration/src/benchmarking.rs | 2 +-
 precompiles/dapp-staking-v3/src/lib.rs             | 4 ++--
 precompiles/dapp-staking-v3/src/test/tests_v2.rs   | 2 +-
 runtime/shibuya/src/lib.rs                         | 1 -
 4 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/pallets/dapp-staking-migration/src/benchmarking.rs b/pallets/dapp-staking-migration/src/benchmarking.rs
index 828f7120ec..39e30f2623 100644
--- a/pallets/dapp-staking-migration/src/benchmarking.rs
+++ b/pallets/dapp-staking-migration/src/benchmarking.rs
@@ -18,10 +18,10 @@
 
 use super::{Pallet as Migration, *};
 
-use sp_std::vec;
 use astar_primitives::{dapp_staking::SmartContractHandle, Balance};
 use frame_benchmarking::{account as benchmark_account, v2::*};
 use pallet_dapp_staking_v3::StakeAmount;
+use sp_std::vec;
 
 fn smart_contract<T: Config>(idx: u8) -> T::SmartContract {
     let address: T::AccountId = benchmark_account("smart_contract", idx.into(), 456);
diff --git a/precompiles/dapp-staking-v3/src/lib.rs b/precompiles/dapp-staking-v3/src/lib.rs
index 231d96d8e8..5c84e25d47 100644
--- a/precompiles/dapp-staking-v3/src/lib.rs
+++ b/precompiles/dapp-staking-v3/src/lib.rs
@@ -553,7 +553,7 @@ where
         // the entire amount from the origin to target contract.
         //
         // In case value comes from the past period, we don't care, since the `unstake` call will fall apart.
-        let stake_amount = if staked_amount > 0
+        let amount = if staked_amount > 0
             && staked_amount.saturating_sub(amount) < minimum_allowed_stake_amount
         {
             staked_amount
@@ -571,7 +571,7 @@ where
         // Then call stake on the target smart contract
         let stake_call = pallet_dapp_staking_v3::Call::<R>::stake {
             smart_contract: target_smart_contract,
-            amount: stake_amount,
+            amount,
         };
         RuntimeHelper::<R>::try_dispatch(handle, Some(origin).into(), stake_call)?;
 
diff --git a/precompiles/dapp-staking-v3/src/test/tests_v2.rs b/precompiles/dapp-staking-v3/src/test/tests_v2.rs
index 2279b97fbf..17fe9db258 100644
--- a/precompiles/dapp-staking-v3/src/test/tests_v2.rs
+++ b/precompiles/dapp-staking-v3/src/test/tests_v2.rs
@@ -768,7 +768,7 @@ fn nomination_transfer_is_ok() {
     ExternalityBuilder::build().execute_with(|| {
         initialize();
 
-        // Register the first dApp, and stke on it.
+        // Register the first dApp, and stake on it.
         let staker_h160 = ALICE;
         let staker_native = AddressMapper::into_account_id(staker_h160);
         let smart_contract_address_1 = H160::repeat_byte(0xFA);
diff --git a/runtime/shibuya/src/lib.rs b/runtime/shibuya/src/lib.rs
index ae6c4397be..00cd3d1aaa 100644
--- a/runtime/shibuya/src/lib.rs
+++ b/runtime/shibuya/src/lib.rs
@@ -1379,7 +1379,6 @@ pub type Executive = frame_executive::Executive<
 ///
 /// Once done, migrations should be removed from the tuple.
 pub type Migrations =
-    // Part of shibuya-124
     (pallet_dapp_staking_migration::SingularStakingInfoTranslationUpgrade<Runtime>,);
 
 type EventRecord = frame_system::EventRecord<

From 27f0bc0fc95bfa8453e21c1e632ccea76b17eae9 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Wed, 13 Mar 2024 14:27:58 +0100
Subject: [PATCH 26/28] Fix

---
 pallets/dapp-staking-v3/src/lib.rs            |  4 +--
 pallets/dapp-staking-v3/src/test/mock.rs      |  1 -
 .../dapp-staking-v3/src/test/testing_utils.rs | 31 ++++++++++---------
 precompiles/dapp-staking-v3/src/lib.rs        |  4 +--
 4 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/lib.rs b/pallets/dapp-staking-v3/src/lib.rs
index cd56a5da13..5e5303b93a 100644
--- a/pallets/dapp-staking-v3/src/lib.rs
+++ b/pallets/dapp-staking-v3/src/lib.rs
@@ -1110,7 +1110,7 @@ pub mod pallet {
 
             // 1.
             // Update `StakerInfo` storage with the reduced stake amount on the specified contract.
-            let (new_staking_info, era_and_amount_pairs) =
+            let (new_staking_info, amount, era_and_amount_pairs) =
                 match StakerInfo::<T>::get(&account, &smart_contract) {
                     Some(mut staking_info) => {
                         ensure!(
@@ -1135,7 +1135,7 @@ pub mod pallet {
                         let era_and_amount_pairs =
                             staking_info.unstake(amount, current_era, protocol_state.subperiod());
 
-                        (staking_info, era_and_amount_pairs)
+                        (staking_info, amount, era_and_amount_pairs)
                     }
                     None => {
                         return Err(Error::<T>::NoStakingInfo.into());
diff --git a/pallets/dapp-staking-v3/src/test/mock.rs b/pallets/dapp-staking-v3/src/test/mock.rs
index a13ee76066..4f8746aadd 100644
--- a/pallets/dapp-staking-v3/src/test/mock.rs
+++ b/pallets/dapp-staking-v3/src/test/mock.rs
@@ -38,7 +38,6 @@ use sp_std::cell::RefCell;
 
 use astar_primitives::{
     dapp_staking::{Observer as DappStakingObserver, SmartContract, StandardTierSlots},
-    testing::Header,
     Balance, BlockNumber,
 };
 
diff --git a/pallets/dapp-staking-v3/src/test/testing_utils.rs b/pallets/dapp-staking-v3/src/test/testing_utils.rs
index df6ee25430..5306656839 100644
--- a/pallets/dapp-staking-v3/src/test/testing_utils.rs
+++ b/pallets/dapp-staking-v3/src/test/testing_utils.rs
@@ -636,7 +636,7 @@ pub(crate) fn assert_unstake(
         pre_staker_info.total_staked_amount().saturating_sub(amount) < minimum_stake_amount;
 
     // Unstake all if we expect to go below the minimum stake amount
-    let amount = if is_full_unstake {
+    let expected_amount = if is_full_unstake {
         pre_staker_info.total_staked_amount()
     } else {
         amount
@@ -651,7 +651,7 @@ pub(crate) fn assert_unstake(
     System::assert_last_event(RuntimeEvent::DappStaking(Event::Unstake {
         account,
         smart_contract: smart_contract.clone(),
-        amount,
+        amount: expected_amount,
     }));
 
     // Verify post-state
@@ -668,12 +668,12 @@ pub(crate) fn assert_unstake(
     // =====================
     assert_eq!(
         post_ledger.staked_amount(unstake_period),
-        pre_ledger.staked_amount(unstake_period) - amount,
+        pre_ledger.staked_amount(unstake_period) - expected_amount,
         "Stake amount must decrease by the 'amount'"
     );
     assert_eq!(
         post_ledger.stakeable_amount(unstake_period),
-        pre_ledger.stakeable_amount(unstake_period) + amount,
+        pre_ledger.stakeable_amount(unstake_period) + expected_amount,
         "Stakeable amount must increase by the 'amount'"
     );
 
@@ -697,14 +697,14 @@ pub(crate) fn assert_unstake(
         assert_eq!(post_staker_info.period_number(), unstake_period);
         assert_eq!(
             post_staker_info.total_staked_amount(),
-            pre_staker_info.total_staked_amount() - amount,
+            pre_staker_info.total_staked_amount() - expected_amount,
             "Total staked amount must decrease by the 'amount'"
         );
         assert_eq!(
             post_staker_info.staked_amount(unstake_subperiod),
             pre_staker_info
                 .staked_amount(unstake_subperiod)
-                .saturating_sub(amount),
+                .saturating_sub(expected_amount),
             "Staked amount must decrease by the 'amount'"
         );
 
@@ -727,14 +727,14 @@ pub(crate) fn assert_unstake(
     let unstaked_amount_era_pairs =
         pre_staker_info
             .clone()
-            .unstake(amount, unstake_period, unstake_subperiod);
+            .unstake(expected_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);
+        assert_eq!(*last_unstake_amount, expected_amount);
     }
 
     // 3. verify contract stake
@@ -742,14 +742,14 @@ pub(crate) fn assert_unstake(
     // =========================
     assert_eq!(
         post_contract_stake.total_staked_amount(unstake_period),
-        pre_contract_stake.total_staked_amount(unstake_period) - amount,
+        pre_contract_stake.total_staked_amount(unstake_period) - expected_amount,
         "Staked amount must decreased by the 'amount'"
     );
     assert_eq!(
         post_contract_stake.staked_amount(unstake_period, unstake_subperiod),
         pre_contract_stake
             .staked_amount(unstake_period, unstake_subperiod)
-            .saturating_sub(amount),
+            .saturating_sub(expected_amount),
         "Staked amount must decreased by the 'amount'"
     );
 
@@ -780,21 +780,22 @@ pub(crate) fn assert_unstake(
     } else {
         assert_eq!(
             post_era_info.total_staked_amount(),
-            pre_era_info.total_staked_amount() - amount,
+            pre_era_info.total_staked_amount() - expected_amount,
             "Total staked amount for the current era must decrease by 'amount'."
         );
     }
     assert_eq!(
         post_era_info.total_staked_amount_next_era(),
-        pre_era_info.total_staked_amount_next_era() - amount,
+        pre_era_info.total_staked_amount_next_era() - expected_amount,
         "Total staked amount for the next era must decrease by 'amount'. No overflow is allowed."
     );
 
     // Check for unstake underflow.
     if unstake_subperiod == Subperiod::BuildAndEarn
-        && pre_era_info.staked_amount_next_era(Subperiod::BuildAndEarn) < amount
+        && pre_era_info.staked_amount_next_era(Subperiod::BuildAndEarn) < expected_amount
     {
-        let overflow = amount - pre_era_info.staked_amount_next_era(Subperiod::BuildAndEarn);
+        let overflow =
+            expected_amount - pre_era_info.staked_amount_next_era(Subperiod::BuildAndEarn);
 
         assert!(post_era_info
             .staked_amount_next_era(Subperiod::BuildAndEarn)
@@ -806,7 +807,7 @@ pub(crate) fn assert_unstake(
     } else {
         assert_eq!(
             post_era_info.staked_amount_next_era(unstake_subperiod),
-            pre_era_info.staked_amount_next_era(unstake_subperiod) - amount
+            pre_era_info.staked_amount_next_era(unstake_subperiod) - expected_amount
         );
     }
 }
diff --git a/precompiles/dapp-staking-v3/src/lib.rs b/precompiles/dapp-staking-v3/src/lib.rs
index 5c84e25d47..231d96d8e8 100644
--- a/precompiles/dapp-staking-v3/src/lib.rs
+++ b/precompiles/dapp-staking-v3/src/lib.rs
@@ -553,7 +553,7 @@ where
         // the entire amount from the origin to target contract.
         //
         // In case value comes from the past period, we don't care, since the `unstake` call will fall apart.
-        let amount = if staked_amount > 0
+        let stake_amount = if staked_amount > 0
             && staked_amount.saturating_sub(amount) < minimum_allowed_stake_amount
         {
             staked_amount
@@ -571,7 +571,7 @@ where
         // Then call stake on the target smart contract
         let stake_call = pallet_dapp_staking_v3::Call::<R>::stake {
             smart_contract: target_smart_contract,
-            amount,
+            amount: stake_amount,
         };
         RuntimeHelper::<R>::try_dispatch(handle, Some(origin).into(), stake_call)?;
 

From c6dda60be0aa91196c652abb532b4155bc8f5d14 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Thu, 14 Mar 2024 13:47:16 +0100
Subject: [PATCH 27/28] Update docs

---
 runtime/astar/src/lib.rs  | 4 ++--
 runtime/shiden/src/lib.rs | 4 +---
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/runtime/astar/src/lib.rs b/runtime/astar/src/lib.rs
index d610c6d889..b71742fb74 100644
--- a/runtime/astar/src/lib.rs
+++ b/runtime/astar/src/lib.rs
@@ -1128,12 +1128,12 @@ parameter_types! {
 ///
 /// Once done, migrations should be removed from the tuple.
 pub type Migrations = (
-    // Part of astar-82, need to first cleanup old storage before re-using the pallet
+    // Part of astar-83, need to first cleanup old storage before re-using the pallet
     frame_support::migrations::RemovePallet<
         DappStakingMigrationName,
         <Runtime as frame_system::Config>::DbWeight,
     >,
-    // Part of astar-82
+    // Part of astar-83
     (pallet_dapp_staking_migration::SingularStakingInfoTranslationUpgrade<Runtime>,),
 );
 
diff --git a/runtime/shiden/src/lib.rs b/runtime/shiden/src/lib.rs
index d64c48e1a9..91786beb45 100644
--- a/runtime/shiden/src/lib.rs
+++ b/runtime/shiden/src/lib.rs
@@ -1130,10 +1130,8 @@ parameter_types! {
 ///
 /// Once done, migrations should be removed from the tuple.
 pub type Migrations = (
-    // Part of shiden-121
+    // Part of shiden-122
     pallet_dapp_staking_migration::SingularStakingInfoTranslationUpgrade<Runtime>,
-    // Part of shiden-121 (added after v5.33.0 release)
-    SetNewTierConfig,
 );
 
 use frame_support::traits::OnRuntimeUpgrade;

From 379f272dc11b7ef852aa45797943773056ff59b0 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Thu, 14 Mar 2024 13:51:06 +0100
Subject: [PATCH 28/28] Comment

---
 runtime/shibuya/src/lib.rs | 1 -
 1 file changed, 1 deletion(-)

diff --git a/runtime/shibuya/src/lib.rs b/runtime/shibuya/src/lib.rs
index 588d3c39b2..12330074f6 100644
--- a/runtime/shibuya/src/lib.rs
+++ b/runtime/shibuya/src/lib.rs
@@ -1374,7 +1374,6 @@ pub type Executive = frame_executive::Executive<
 ///
 /// Once done, migrations should be removed from the tuple.
 pub type Migrations =
-    // Part of shiden-122
     (pallet_dapp_staking_migration::SingularStakingInfoTranslationUpgrade<Runtime>,);
 
 type EventRecord = frame_system::EventRecord<