From 9aeece5bf6074948202e2ab148ce84b52e78588c Mon Sep 17 00:00:00 2001
From: Luke Parker <lukeparker5132@gmail.com>
Date: Fri, 13 Oct 2023 02:29:22 -0400
Subject: [PATCH] Give one weight per key share to validators in Tributary

---
 coordinator/src/substrate/mod.rs             | 21 ++++++++++++++++++--
 coordinator/src/tributary/mod.rs             |  7 +++----
 substrate/client/src/serai/validator_sets.rs | 19 +++++++++++++++++-
 3 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/coordinator/src/substrate/mod.rs b/coordinator/src/substrate/mod.rs
index 529273f8..f54ff9eb 100644
--- a/coordinator/src/substrate/mod.rs
+++ b/coordinator/src/substrate/mod.rs
@@ -9,7 +9,7 @@ use serai_client::{
   SeraiError, Block, Serai,
   primitives::{BlockHash, NetworkId},
   validator_sets::{
-    primitives::{ValidatorSet, KeyPair},
+    primitives::{ValidatorSet, KeyPair, amortize_excess_key_shares},
     ValidatorSetsEvent,
   },
   in_instructions::InInstructionsEvent,
@@ -61,6 +61,23 @@ async fn handle_new_set<D: Db, CNT: Clone + Fn(&mut D, TributarySpec)>(
       .await?
       .expect("NewSet for set which doesn't exist");
 
+    let allocation_per_key_share = serai
+      .get_allocation_per_key_share(set.network, block.hash())
+      .await?
+      .expect("NewSet for set which didn't have an allocation per key share")
+      .0;
+
+    let mut set_data = vec![];
+    for participant in set_participants {
+      let allocation = serai
+        .get_allocation(set.network, participant, block.hash())
+        .await?
+        .expect("validator selected for set yet didn't have an allocation")
+        .0;
+      set_data.push((participant, allocation / allocation_per_key_share));
+    }
+    amortize_excess_key_shares(&mut set_data);
+
     let time = if let Ok(time) = block.time() {
       time
     } else {
@@ -82,7 +99,7 @@ async fn handle_new_set<D: Db, CNT: Clone + Fn(&mut D, TributarySpec)>(
     const SUBSTRATE_TO_TRIBUTARY_TIME_DELAY: u64 = 120;
     let time = time + SUBSTRATE_TO_TRIBUTARY_TIME_DELAY;
 
-    let spec = TributarySpec::new(block.hash(), time, set, set_participants);
+    let spec = TributarySpec::new(block.hash(), time, set, set_data);
     create_new_tributary(db, spec.clone());
   } else {
     log::info!("not present in set {:?}", set);
diff --git a/coordinator/src/tributary/mod.rs b/coordinator/src/tributary/mod.rs
index ab7e3d42..e0533792 100644
--- a/coordinator/src/tributary/mod.rs
+++ b/coordinator/src/tributary/mod.rs
@@ -51,16 +51,15 @@ impl TributarySpec {
     serai_block: [u8; 32],
     start_time: u64,
     set: ValidatorSet,
-    set_participants: Vec<PublicKey>,
+    set_participants: Vec<(PublicKey, u64)>,
   ) -> TributarySpec {
     let mut validators = vec![];
-    for participant in set_participants {
+    for (participant, shares) in set_participants {
       // TODO: Ban invalid keys from being validators on the Serai side
       // (make coordinator key a session key?)
       let participant = <Ristretto as Ciphersuite>::read_G::<&[u8]>(&mut participant.0.as_ref())
         .expect("invalid key registered as participant");
-      // TODO: Give one weight on Tributary per bond instance
-      validators.push((participant, 1));
+      validators.push((participant, shares));
     }
 
     Self { serai_block, start_time, set, validators }
diff --git a/substrate/client/src/serai/validator_sets.rs b/substrate/client/src/serai/validator_sets.rs
index 1c34ff68..0b2362e3 100644
--- a/substrate/client/src/serai/validator_sets.rs
+++ b/substrate/client/src/serai/validator_sets.rs
@@ -1,6 +1,6 @@
 use sp_core::sr25519::{Public, Signature};
 
-use serai_runtime::{validator_sets, ValidatorSets, Runtime};
+use serai_runtime::{primitives::Amount, validator_sets, ValidatorSets, Runtime};
 pub use validator_sets::primitives;
 use primitives::{Session, ValidatorSet, KeyPair};
 
@@ -47,6 +47,23 @@ impl Serai {
     self.storage(PALLET, "Participants", Some(vec![scale_value(network)]), at_hash).await
   }
 
+  pub async fn get_allocation_per_key_share(
+    &self,
+    network: NetworkId,
+    at_hash: [u8; 32],
+  ) -> Result<Option<Amount>, SeraiError> {
+    self.storage(PALLET, "AllocationPerKeyShare", Some(vec![scale_value(network)]), at_hash).await
+  }
+
+  pub async fn get_allocation(
+    &self,
+    network: NetworkId,
+    key: Public,
+    at_hash: [u8; 32],
+  ) -> Result<Option<Amount>, SeraiError> {
+    self.storage(PALLET, "Allocations", Some(vec![scale_value(network), scale_value(key)]), at_hash).await
+  }
+
   pub async fn get_validator_set_musig_key(
     &self,
     set: ValidatorSet,