update ledger to save restaked value for each role

pallets/roles/src/lib.rs

@@ -54,16 +54,29 @@ pub use weights::WeightInfo;
 pub struct RoleStakingLedger<T: Config> {
 	/// The stash account whose balance is actually locked and at stake.
 	pub stash: T::AccountId,
-	/// The total amount of the stash's balance that is re-staked for selected services
+	/// The total amount of the stash's balance that is re-staked for all selected roles.
 	/// This re-staked balance we are currently accounting for new slashing conditions.
 	#[codec(compact)]
 	pub total: BalanceOf<T>,
+	/// The list of roles and their re-staked amounts.
+	pub roles: Vec<RoleStakeInfo<T>>,
+}
+
+/// The information regarding the re-staked amount for a particular role.
+#[derive(PartialEqNoBound, EqNoBound, Encode, Decode, RuntimeDebugNoBound, TypeInfo, Clone)]
+#[scale_info(skip_type_params(T))]
+pub struct RoleStakeInfo<T: Config> {
+	/// Role type
+	pub role: RoleType,
+	/// The total amount of the stash's balance that is re-staked for selected role.
+	#[codec(compact)]
+	pub re_staked: BalanceOf<T>,
 }
 
 impl<T: Config> RoleStakingLedger<T> {
 	/// Initializes the default object using the given `validator`.
 	pub fn default_from(stash: T::AccountId) -> Self {
-		Self { stash, total: Zero::zero() }
+		Self { stash, total: Zero::zero(), roles: vec![] }
 	}
 
 	/// Returns `true` if the stash account has no funds at all.
@@ -196,9 +209,25 @@ pub mod pallet {
 			// Validate re-staking bond, should not exceed active staked bond.
 			ensure!(staking_ledger.active >= re_stake_amount, Error::<T>::InvalidReStakingBond);
 
-			// Update ledger.
-			let item = RoleStakingLedger { stash: stash_account.clone(), total: re_stake_amount };
-			Self::update_ledger(&stash_account, &item);
+			// Check if account is already paired with ledger
+			let maybe_ledger = Ledger::<T>::get(&stash_account);
+			if maybe_ledger.is_none() {
+				// Add stash account to ledger.
+				let role_info = RoleStakeInfo { role, re_staked: re_stake_amount };
+				let item = RoleStakingLedger {
+					stash: stash_account.clone(),
+					total: re_stake_amount,
+					roles: vec![role_info],
+				};
+				Self::update_ledger(&stash_account, &item);
+			} else {
+				// Update ledger and add role info.
+				let mut ledger = maybe_ledger.unwrap();
+				ledger.total += re_stake_amount;
+				let role_info = RoleStakeInfo { role, re_staked: re_stake_amount };
+				ledger.roles.push(role_info);
+				Self::update_ledger(&stash_account, &ledger);
+			}
 
 			// Add role mapping for the stash account.
 			Self::add_role(stash_account.clone(), role)?;

pallets/roles/src/tests.rs

@@ -36,7 +36,14 @@ fn test_assign_role() {
 		// Lets verify role assigned to account.
 		assert_eq!(Roles::has_role(1, RoleType::Tss), true);
 		// Verify ledger mapping
-		assert_eq!(Roles::ledger(1), Some(RoleStakingLedger { stash: 1, total: 5000 }));
+		assert_eq!(
+			Roles::ledger(1),
+			Some(RoleStakingLedger {
+				stash: 1,
+				total: 5000,
+				roles: vec![RoleStakeInfo { role: RoleType::Tss, re_staked: 5000 }]
+			})
+		);
 	});
 }
 
@@ -59,7 +66,14 @@ fn test_assign_role_with_full_staking_option() {
 		// Lets verify role assigned to account.
 		assert_eq!(Roles::has_role(1, RoleType::Tss), true);
 		// Verify ledger mapping
-		assert_eq!(Roles::ledger(1), Some(RoleStakingLedger { stash: 1, total: 10000 }));
+		assert_eq!(
+			Roles::ledger(1),
+			Some(RoleStakingLedger {
+				stash: 1,
+				total: 10000,
+				roles: vec![RoleStakeInfo { role: RoleType::Tss, re_staked: 10000 }]
+			})
+		);
 	});
 }
 
@@ -86,6 +100,18 @@ fn test_assign_multiple_roles() {
 
 		// Lets verify role assigned to account.
 		assert_eq!(Roles::has_role(1, RoleType::ZkSaas), true);
+
+		assert_eq!(
+			Roles::ledger(1),
+			Some(RoleStakingLedger {
+				stash: 1,
+				total: 20000,
+				roles: vec![
+					RoleStakeInfo { role: RoleType::Tss, re_staked: 10000 },
+					RoleStakeInfo { role: RoleType::ZkSaas, re_staked: 10000 },
+				]
+			})
+		);
 	});
 }