Types UTs

frame/dapp-staking-v3/src/lib.rs

@@ -434,7 +434,7 @@ pub mod pallet {
         /// Only the amount that isn't actively used for staking can be unlocked.
         /// If the amount is greater than the available amount for unlocking, everything is unlocked.
         /// If the remaining locked amount would take the account below the minimum locked amount, everything is unlocked.
-        #[pallet::call_index(5)]
+        #[pallet::call_index(6)]
         #[pallet::weight(Weight::zero())]
         pub fn unlock(
             origin: OriginFor<T>,

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

@@ -64,15 +64,18 @@ fn account_ledger_add_lock_amount_works() {
     // First step, sanity checks
     let first_era = 1;
     assert!(acc_ledger.active_locked_amount().is_zero());
+    assert!(acc_ledger.total_locked_amount().is_zero());
     assert!(acc_ledger.add_lock_amount(0, first_era).is_ok());
     assert!(acc_ledger.active_locked_amount().is_zero());
 
     // Adding lock value works as expected
     let init_amount = 20;
     assert!(acc_ledger.add_lock_amount(init_amount, first_era).is_ok());
     assert_eq!(acc_ledger.active_locked_amount(), init_amount);
+    assert_eq!(acc_ledger.total_locked_amount(), init_amount);
     assert_eq!(acc_ledger.lock_era(), first_era);
     assert!(!acc_ledger.is_empty());
+    assert_eq!(acc_ledger.locked.len(), 1);
     assert_eq!(
         acc_ledger.latest_locked_chunk(),
         Some(&LockedChunk::<Balance> {
@@ -85,7 +88,9 @@ fn account_ledger_add_lock_amount_works() {
     let addition = 7;
     assert!(acc_ledger.add_lock_amount(addition, first_era).is_ok());
     assert_eq!(acc_ledger.active_locked_amount(), init_amount + addition);
+    assert_eq!(acc_ledger.total_locked_amount(), init_amount + addition);
     assert_eq!(acc_ledger.lock_era(), first_era);
+    assert_eq!(acc_ledger.locked.len(), 1);
 
     // Add up to storage limit
     for i in 2..=LockedDummy::get() {
@@ -95,11 +100,178 @@ fn account_ledger_add_lock_amount_works() {
             init_amount + addition * i as u128
         );
         assert_eq!(acc_ledger.lock_era(), first_era + i);
+        assert_eq!(acc_ledger.locked.len(), i as usize);
     }
 
     // Any further additions should fail due to exhausting bounded storage capacity
     assert!(acc_ledger
         .add_lock_amount(addition, acc_ledger.lock_era() + 1)
         .is_err());
     assert!(!acc_ledger.is_empty());
+    assert_eq!(acc_ledger.locked.len(), LockedDummy::get() as usize);
+}
+
+#[test]
+fn account_ledger_subtract_lock_amount_basic_usage_works() {
+    get_u32_type!(LockedDummy, 5);
+    get_u32_type!(UnlockingDummy, 5);
+    let mut acc_ledger =
+        AccountLedger::<Balance, BlockNumber, LockedDummy, UnlockingDummy>::default();
+
+    // Sanity check scenario
+    // Cannot reduce if there is nothing locked, should be a noop
+    assert!(acc_ledger.subtract_lock_amount(0, 1).is_ok());
+    assert!(acc_ledger.subtract_lock_amount(10, 1).is_ok());
+    assert!(acc_ledger.locked.len().is_zero());
+    assert!(acc_ledger.is_empty());
+
+    // First basic scenario
+    // Add some lock amount, then reduce it for the same era
+    let first_era = 1;
+    let first_lock_amount = 19;
+    let unlock_amount = 7;
+    assert!(acc_ledger
+        .add_lock_amount(first_lock_amount, first_era)
+        .is_ok());
+    assert!(acc_ledger
+        .subtract_lock_amount(unlock_amount, first_era)
+        .is_ok());
+    assert_eq!(acc_ledger.locked.len(), 1);
+    assert_eq!(
+        acc_ledger.total_locked_amount(),
+        first_lock_amount - unlock_amount
+    );
+    assert_eq!(
+        acc_ledger.active_locked_amount(),
+        first_lock_amount - unlock_amount
+    );
+    assert_eq!(acc_ledger.unlocking_amount(), 0);
+
+    // Second basic scenario
+    // Reduce the lock from the era which isn't latest in the vector
+    let first_lock_amount = first_lock_amount - unlock_amount;
+    let second_lock_amount = 31;
+    let second_era = 2;
+    assert!(acc_ledger
+        .add_lock_amount(second_lock_amount - first_lock_amount, second_era)
+        .is_ok());
+    assert_eq!(acc_ledger.active_locked_amount(), second_lock_amount);
+    assert_eq!(acc_ledger.locked.len(), 2);
+
+    // Subtract from the first era and verify state is as expected
+    assert!(acc_ledger
+        .subtract_lock_amount(unlock_amount, first_era)
+        .is_ok());
+    assert_eq!(acc_ledger.locked.len(), 2);
+    assert_eq!(
+        acc_ledger.active_locked_amount(),
+        second_lock_amount - unlock_amount
+    );
+    assert_eq!(
+        acc_ledger.locked[0].amount,
+        first_lock_amount - unlock_amount
+    );
+    assert_eq!(
+        acc_ledger.locked[1].amount,
+        second_lock_amount - unlock_amount
+    );
+
+    // Third basic scenario
+    // Reduce the the latest era, don't expect the first one to change
+    assert!(acc_ledger
+        .subtract_lock_amount(unlock_amount, second_era)
+        .is_ok());
+    assert_eq!(acc_ledger.locked.len(), 2);
+    assert_eq!(
+        acc_ledger.active_locked_amount(),
+        second_lock_amount - unlock_amount * 2
+    );
+    assert_eq!(
+        acc_ledger.locked[0].amount,
+        first_lock_amount - unlock_amount
+    );
+    assert_eq!(
+        acc_ledger.locked[1].amount,
+        second_lock_amount - unlock_amount * 2
+    );
+}
+
+#[test]
+fn account_ledger_subtract_lock_amount_overflow_fails() {
+    get_u32_type!(LockedDummy, 5);
+    get_u32_type!(UnlockingDummy, 5);
+    let mut acc_ledger =
+        AccountLedger::<Balance, BlockNumber, LockedDummy, UnlockingDummy>::default();
+
+    let first_lock_amount = 17 * 19;
+    let era = 1;
+    let unlock_amount = 5;
+    assert!(acc_ledger.add_lock_amount(first_lock_amount, era).is_ok());
+    for idx in 1..=LockedDummy::get() {
+        assert!(acc_ledger.subtract_lock_amount(unlock_amount, idx).is_ok());
+        assert_eq!(acc_ledger.locked.len(), idx as usize);
+        assert_eq!(
+            acc_ledger.active_locked_amount(),
+            first_lock_amount - unlock_amount * idx as u128
+        );
+    }
+
+    // Updating existing lock should still work
+    for _ in 1..10 {
+        assert!(acc_ledger
+            .subtract_lock_amount(unlock_amount, LockedDummy::get())
+            .is_ok());
+    }
+
+    // Attempt to add additional chunks should fail.
+    assert!(acc_ledger
+        .subtract_lock_amount(unlock_amount, LockedDummy::get() + 1)
+        .is_err());
+}
+
+#[test]
+fn account_ledger_subtract_lock_amount_advanced_example_works() {
+    get_u32_type!(LockedDummy, 5);
+    get_u32_type!(UnlockingDummy, 5);
+    let mut acc_ledger =
+        AccountLedger::<Balance, BlockNumber, LockedDummy, UnlockingDummy>::default();
+
+    // Prepare an example where we have two non-consecutive entries, and we unlock in the era right before the second entry.
+    // This covers a scenario where user has locked in the current era,
+    // creating an entry for the next era, and then decides to immediately unlock.
+    let first_lock_amount = 17;
+    let second_lock_amount = 23;
+    let first_era = 1;
+    let second_era = 5;
+    let unlock_era = second_era - 1;
+    let unlock_amount = 5;
+    assert!(acc_ledger
+        .add_lock_amount(first_lock_amount, first_era)
+        .is_ok());
+    assert!(acc_ledger
+        .add_lock_amount(second_lock_amount, second_era)
+        .is_ok());
+    assert_eq!(acc_ledger.locked.len(), 2);
+
+    assert!(acc_ledger
+        .subtract_lock_amount(unlock_amount, unlock_era)
+        .is_ok());
+    assert_eq!(
+        acc_ledger.active_locked_amount(),
+        first_lock_amount + second_lock_amount - unlock_amount
+    );
+
+    // Check entries in more detail
+    assert_eq!(acc_ledger.locked.len(), 3);
+    assert_eq!(acc_ledger.locked[0].amount, first_lock_amount,);
+    assert_eq!(
+        acc_ledger.locked[2].amount,
+        first_lock_amount + second_lock_amount - unlock_amount
+    );
+    // Verify the new entry
+    assert_eq!(
+        acc_ledger.locked[1].amount,
+        first_lock_amount - unlock_amount
+    );
+    assert_eq!(acc_ledger.locked[1].era, unlock_era);
 }

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

@@ -144,7 +144,6 @@ where
     }
 }
 
-// TODO: would users get better UX if we kept using eras? Using blocks is more precise though.
 /// How much was unlocked in some block.
 #[derive(Encode, Decode, MaxEncodedLen, Clone, Copy, Debug, PartialEq, Eq, TypeInfo)]
 pub struct UnlockingChunk<
@@ -337,18 +336,20 @@ where
 
         // Update existing or insert a new chunk
         let mut inner = self.locked.clone().into_inner();
-        if inner[index].era == era {
+        let relevant_chunk_index = if inner[index].era == era {
             inner[index].amount.saturating_reduce(amount);
+            index
         } else {
             let mut chunk = inner[index];
             chunk.amount.saturating_reduce(amount);
             chunk.era = era;
 
             inner.insert(index + 1, chunk);
-        }
+            index + 1
+        };
 
         // Update all chunks after the relevant one, and remove zero chunks
-        inner[index + 1..]
+        inner[relevant_chunk_index + 1..]
             .iter_mut()
             .for_each(|chunk| chunk.amount.saturating_reduce(amount));
         inner.retain(|chunk| !chunk.amount.is_zero());