From dd85da39b3042b41213d2dbbf38b28d513425643 Mon Sep 17 00:00:00 2001
From: 1xstj <106580853+1xstj@users.noreply.github.com>
Date: Fri, 24 Nov 2023 12:08:44 +0000
Subject: [PATCH] first pass : simple distribution

---
 Cargo.lock                    |  1 +
 pallets/roles/src/impls.rs    | 79 ++++++++++++++++++++++++++++++++++-
 pallets/roles/src/lib.rs      | 14 ++++++-
 primitives/Cargo.toml         |  1 +
 primitives/src/types/roles.rs | 26 ++++++++++++
 5 files changed, 119 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 706b22215..ecb58ac71 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -15423,6 +15423,7 @@ dependencies = [
  "sp-arithmetic 16.0.0 (git+https://github.com/paritytech/substrate?branch=polkadot-v1.0.0)",
  "sp-core 21.0.0 (git+https://github.com/paritytech/substrate?branch=polkadot-v1.0.0)",
  "sp-runtime 24.0.0 (git+https://github.com/paritytech/substrate?branch=polkadot-v1.0.0)",
+ "sp-std 8.0.0 (git+https://github.com/paritytech/substrate?branch=polkadot-v1.0.0)",
 ]
 
 [[package]]
diff --git a/pallets/roles/src/impls.rs b/pallets/roles/src/impls.rs
index c64a1bf35..5ac9f175e 100644
--- a/pallets/roles/src/impls.rs
+++ b/pallets/roles/src/impls.rs
@@ -15,8 +15,13 @@
 // along with Tangle.  If not, see <http://www.gnu.org/licenses/>.
 
 use super::*;
-use frame_support::pallet_prelude::DispatchResult;
+use frame_support::{
+	log,
+	pallet_prelude::DispatchResult,
+	traits::{Currency, ExistenceRequirement::KeepAlive, OneSessionHandler},
+};
 use sp_runtime::{Percent, Saturating};
+use sp_std::ops::Div;
 use tangle_primitives::{jobs::JobKey, traits::roles::RolesHandler};
 
 /// Implements RolesHandler for the pallet.
@@ -180,4 +185,76 @@ impl<T: Config> Pallet<T> {
 	pub(crate) fn kill_stash(stash: &T::AccountId) {
 		<Ledger<T>>::remove(&stash);
 	}
+
+	pub fn distribute_rewards() -> DispatchResult {
+		let total_rewards = T::InflationRewardPerEra::get();
+
+		let mut tss_validators: Vec<T::AccountId> = Default::default();
+		let mut zksaas_validators: Vec<T::AccountId> = Default::default();
+
+		for (acc, role_types) in AccountRolesMapping::<T>::iter() {
+			if role_types.contains(&RoleType::Tss) {
+				tss_validators.push(acc.clone())
+			}
+
+			if role_types.contains(&RoleType::ZkSaas) {
+				zksaas_validators.push(acc)
+			}
+		}
+
+		log::debug!("Found {:?} tss validators", tss_validators.len());
+		log::debug!("Found {:?} zksaas validators", zksaas_validators.len());
+
+		let reward_distribution = T::ValidatorRewardDistribution::get();
+
+		let dist = reward_distribution.get_reward_distribution();
+		let tss_validator_reward =
+			dist.0.mul_floor(total_rewards).div((tss_validators.len() as u32).into());
+		let zksaas_validators_reward =
+			dist.1.mul_floor(total_rewards).div((zksaas_validators.len() as u32).into());
+
+		log::debug!("Reward for tss validator : {:?}", tss_validator_reward);
+		log::debug!("Reward for zksaas validator : {:?}", zksaas_validators_reward);
+
+		for validator in tss_validators {
+			T::Currency::deposit_creating(&validator, tss_validator_reward);
+		}
+
+		for validator in zksaas_validators {
+			T::Currency::deposit_creating(&validator, zksaas_validators_reward);
+		}
+
+		Ok(())
+	}
+}
+
+impl<T: Config> sp_runtime::BoundToRuntimeAppPublic for Pallet<T> {
+	type Public = T::AuthorityId;
+}
+
+impl<T: Config> OneSessionHandler<T::AccountId> for Pallet<T> {
+	type Key = T::AuthorityId;
+
+	fn on_genesis_session<'a, I: 'a>(_validators: I)
+	where
+		I: Iterator<Item = (&'a T::AccountId, T::AuthorityId)>,
+	{
+		// nothing to be done
+	}
+
+	fn on_new_session<'a, I: 'a>(_changed: bool, _validators: I, _queued_validators: I)
+	where
+		I: Iterator<Item = (&'a T::AccountId, T::AuthorityId)>,
+	{
+		// nothing to be done
+	}
+
+	fn on_disabled(_i: u32) {
+		// ignore
+	}
+
+	// Distribute the inflation rewards
+	fn on_before_session_ending() {
+		let _ = Self::distribute_rewards();
+	}
 }
diff --git a/pallets/roles/src/lib.rs b/pallets/roles/src/lib.rs
index 826b4f51d..4c48749a5 100644
--- a/pallets/roles/src/lib.rs
+++ b/pallets/roles/src/lib.rs
@@ -24,6 +24,7 @@ use frame_support::{
 	traits::{Currency, Get},
 	CloneNoBound, EqNoBound, PartialEqNoBound, RuntimeDebugNoBound,
 };
+use tangle_primitives::roles::ValidatorRewardDistribution;
 
 pub use pallet::*;
 use parity_scale_codec::{Decode, Encode};
@@ -42,6 +43,8 @@ mod tests;
 mod weights;
 pub use weights::WeightInfo;
 
+use sp_runtime::RuntimeAppPublic;
+
 /// The ledger of a (bonded) stash.
 #[derive(
 	PartialEqNoBound,
@@ -122,6 +125,15 @@ pub mod pallet {
 		/// The config that verifies MPC related functions
 		type MPCHandler: MPCHandler<Self::AccountId, BlockNumberFor<Self>, BalanceOf<Self>>;
 
+		/// The inflation reward to distribute per era
+		type InflationRewardPerEra: Get<BalanceOf<Self>>;
+
+		/// The inflation distribution based on validator type
+		type ValidatorRewardDistribution: Get<ValidatorRewardDistribution>;
+
+		/// The type used to identify an authority
+		type AuthorityId: RuntimeAppPublic + Decode;
+
 		type WeightInfo: WeightInfo;
 	}
 
@@ -171,7 +183,7 @@ pub mod pallet {
 	/// Mapping of resource to bridge index
 	pub type AccountRolesMapping<T: Config> = StorageMap<
 		_,
-		Blake2_256,
+		Blake2_128Concat,
 		T::AccountId,
 		BoundedVec<RoleType, T::MaxRolesPerAccount>,
 		ValueQuery,
diff --git a/primitives/Cargo.toml b/primitives/Cargo.toml
index b26338a84..f5d9626a2 100644
--- a/primitives/Cargo.toml
+++ b/primitives/Cargo.toml
@@ -16,6 +16,7 @@ serde = { workspace = true }
 smallvec = { workspace = true }
 sp-arithmetic = { workspace = true }
 sp-core = { workspace = true }
+sp-std = { workspace = true }
 sp-runtime = { workspace = true }
 
 [features]
diff --git a/primitives/src/types/roles.rs b/primitives/src/types/roles.rs
index f1d935757..a21438e7a 100644
--- a/primitives/src/types/roles.rs
+++ b/primitives/src/types/roles.rs
@@ -14,6 +14,8 @@
 // You should have received a copy of the GNU General Public License
 // along with Tangle.  If not, see <http://www.gnu.org/licenses/>.
 use frame_support::{dispatch::Vec, pallet_prelude::*};
+use sp_arithmetic::Percent;
+use sp_std::ops::Add;
 
 /// Role type to be used in the system.
 #[derive(Encode, Decode, Clone, Debug, PartialEq, Eq, TypeInfo)]
@@ -81,3 +83,27 @@ pub enum ReStakingOption {
 	// Re-stake only the given amount of funds for selected role.
 	Custom(u64),
 }
+
+/// Represents the reward distribution percentages for validators in a key generation process.
+pub struct ValidatorRewardDistribution {
+	/// The percentage share of the reward allocated for TSS
+	tss_share : Percent,
+	/// The percentage share of the reward allocated for the ZK-SaaS
+	zksaas_share : Percent
+}
+
+impl ValidatorRewardDistribution {
+	pub fn try_new(tss_share : Percent, zksaas_share: Percent) -> Result<Self, String> {
+		if !tss_share.add(zksaas_share).is_one() {
+			return Err("Shares must add to One".to_string())
+		}
+
+		Ok(Self { tss_share , zksaas_share})
+	}
+
+	pub fn get_reward_distribution(self) -> (Percent, Percent) {
+		(self.tss_share, self.zksaas_share)
+	}
+}
+
+