feat: add support for --include-partial-seats

Makefile

@@ -3,10 +3,11 @@ run:
 	cargo run -p sim-validator-assignment -- \
 		run \
 		--num-blocks 1000 \
-		--num-shards 1 \
+		--num-shards 4 \
 		--seats-per-shard 250 \
 		--stake-per-seat 1 \
-		--max-malicious-stake-per-shard 1/2
+		--max-malicious-stake-per-shard 1/2 \
+		--include-partial-seats
 
 .PHONY: download
 download:
@@ -28,5 +29,6 @@ run-with:
 		--seats-per-shard 250 \
 		--stake-per-seat 50000000000000000000000000000 \
 		--max-malicious-stake-per-shard 2/3 \
+		--include-partial-seats \
 		--validator-data ./validator_data.json
 

sim-validator-assignment/src/config.rs

@@ -3,7 +3,7 @@ use num_rational::Ratio;
 use serde::Serialize;
 use std::path::PathBuf;
 
-use crate::seat::Seat;
+use crate::{partial_seat::PartialSeat, seat::Seat};
 
 #[derive(Args, Serialize, Debug)]
 pub struct Config {
@@ -30,18 +30,24 @@ pub struct Config {
     /// data will be used in the simulation.
     #[arg(long)]
     pub validator_data: Option<PathBuf>,
+    /// A validator's stake might not entirely cover seats given a particular `stake_per_seat`. This
+    /// option controls whether remaining stake (not covering a full seat) should be assigned to a
+    /// partial seat or ignored.
+    #[arg(long, default_value_t = false)]
+    pub include_partial_seats: bool,
 }
 
 impl Config {
     #[cfg(test)]
-    pub fn new_mock() -> Self {
+    pub fn new_mock(include_partial_seats: bool) -> Self {
         Self {
             num_blocks: 1_000,
             num_shards: 4,
             seats_per_shard: 2,
             stake_per_seat: 100,
             max_malicious_stake_per_shard: Ratio::new(1, 3),
             validator_data: None,
+            include_partial_seats,
         }
     }
 
@@ -93,31 +99,70 @@ impl Config {
 
         Ok(shard_seats)
     }
+
+    /// Collect partials seats for `shard_idx` by picking seats from positions with `position %
+    /// num_shards == shard_idx`.
+    ///
+    /// # Motivation for assignment via modulo
+    ///
+    /// Every validator may hold at most 1 partial seat. If a validator's stake covers full seats
+    /// without remainder, it holds 0 partial seats. Hence `0 <= num_partial_seats <=
+    /// num_validators`. Assigning seats to shards with `%` allows distributing the existing number
+    /// of partial seats to shards as evenly as possible.
+    ///
+    /// # No minimum number of required _partial_ seats per shard
+    ///
+    /// It is not needed since partial seats are included only to distribute leftover stake (not
+    /// covering full seats) to shards.
+    pub fn collect_partial_seats_for_shard<'seats>(
+        &self,
+        shard_idx: usize,
+        partial_seats: &'seats [PartialSeat],
+    ) -> anyhow::Result<Vec<&PartialSeat<'seats>>> {
+        if shard_idx >= usize::from(self.num_shards) {
+            anyhow::bail!(
+                "shard_idx {} is an invalid index for {} shards",
+                shard_idx,
+                self.num_shards
+            )
+        }
+
+        let mut shard_partial_seats = vec![];
+        for idx in (shard_idx..partial_seats.len()).step_by(self.num_shards.into()) {
+            shard_partial_seats.push(&partial_seats[idx]);
+        }
+
+        Ok(shard_partial_seats)
+    }
 }
 
 #[cfg(test)]
 mod tests {
+    use std::collections::BTreeMap;
+
     use super::Config;
     use crate::validator::tests::new_test_raw_validator_data;
-    use crate::validator::{new_ordered_seats, parse_raw_validator_data};
+    use crate::validator::{
+        new_ordered_partial_seats, new_ordered_seats, parse_raw_validator_data,
+    };
 
     #[test]
-    pub fn test_seats_per_stake() {
-        let config = Config::new_mock();
+    fn test_seats_per_stake() {
+        let config = Config::new_mock(false);
         assert_eq!(config.seats_per_stake(20), 0);
         assert_eq!(config.seats_per_stake(100), 1);
         assert_eq!(config.seats_per_stake(530), 5);
     }
 
     #[test]
-    pub fn test_total_seats() {
-        let config = Config::new_mock();
+    fn test_total_seats() {
+        let config = Config::new_mock(false);
         assert_eq!(config.total_seats(), 8);
     }
 
     #[test]
-    pub fn test_collect_seats_for_shard() {
-        let config = Config::new_mock();
+    fn test_collect_seats_for_shard() {
+        let config = Config::new_mock(false);
         let (_, validators) = parse_raw_validator_data(&config, &new_test_raw_validator_data());
         // Using ordered seats as input to have a deterministic result of `collect_seats_for_shard`.
         let seats = new_ordered_seats(&validators);
@@ -137,12 +182,52 @@ mod tests {
     }
 
     #[test]
-    pub fn test_collect_seats_for_shard_errors() {
-        let config = Config::new_mock();
+    fn test_collect_seats_for_shard_errors() {
+        let config = Config::new_mock(false);
         let (_, validators) = parse_raw_validator_data(&config, &new_test_raw_validator_data());
         let seats = new_ordered_seats(&validators);
 
         insta::assert_debug_snapshot!(config.collect_seats_for_shard(4, &seats));
         insta::assert_debug_snapshot!(config.collect_seats_for_shard(0, &[]));
     }
+
+    #[test]
+    fn test_collect_partial_seats_for_shard() {
+        let mut config = Config::new_mock(true);
+        config.stake_per_seat = 90;
+        let (_, validators) = parse_raw_validator_data(&config, &new_test_raw_validator_data());
+        // Using ordered partial seats as input to have a deterministic result of
+        // `collect_partial_seats_for_shard`.
+        let partial_seats = new_ordered_partial_seats(&validators, config.stake_per_seat);
+
+        // Using `BTreeMap` for deterministic ordering of keys.
+        let mut assignments = BTreeMap::new();
+        for shard_idx in 0..config.num_shards {
+            let assignment = config
+                .collect_partial_seats_for_shard(shard_idx.into(), &partial_seats)
+                .unwrap();
+            assignments.insert(format!("shard_{shard_idx}"), assignment);
+        }
+
+        insta::with_settings!({
+            info => &(
+                &config,
+                "partial_seats:",
+                &partial_seats
+            )
+        }, {
+            insta::assert_yaml_snapshot!(assignments);
+        })
+    }
+
+    #[test]
+    fn test_collect_partial_seats_for_shard_errors() {
+        let config = Config::new_mock(true);
+        let (_, validators) = parse_raw_validator_data(&config, &new_test_raw_validator_data());
+        let partial_seats = new_ordered_partial_seats(&validators, config.stake_per_seat);
+
+        insta::assert_debug_snapshot!(
+            config.collect_partial_seats_for_shard(config.num_shards.into(), &partial_seats)
+        );
+    }
 }

sim-validator-assignment/src/main.rs

@@ -5,6 +5,7 @@ use config::Config;
 mod download;
 use download::{download, DownloadConfig};
 mod mocks;
+mod partial_seat;
 mod run;
 mod seat;
 mod shard;

sim-validator-assignment/src/partial_seat.rs

@@ -0,0 +1,49 @@
+use serde::Serialize;
+
+use crate::validator::Validator;
+
+/// Represents a partial seat filled by a particular validator.
+///
+/// A partial seat may not outlive the validator it is referrencing.
+#[derive(Serialize, PartialEq, Debug, Clone)]
+pub struct PartialSeat<'validator> {
+    /// Reference to the validator holding this particular partial seat.
+    validator: &'validator Validator,
+    /// The stake attributed to the partial seat. Note that `0 < weight < stake_per_seat`.
+    ///
+    /// For example, let validator `V` have a stake of 12 and let `stake_per_seat = 5`. Then `V`
+    /// holds 2 full seats and a partial seat with `weight = 2`.
+    weight: u128,
+}
+
+impl<'validator> PartialSeat<'validator> {
+    /// Constructs the partial seat filled by `validator`.
+    pub fn new(validator: &'validator Validator, weight: u128) -> Self {
+        Self { validator, weight }
+    }
+
+    pub fn get_is_malicious(&self) -> bool {
+        self.validator.get_is_malicious()
+    }
+
+    pub fn get_weight(&self) -> u128 {
+        self.weight
+    }
+}
+
+pub struct ShuffledPartialSeats<'seats> {
+    partial_seats: &'seats [PartialSeat<'seats>],
+}
+
+impl<'seats> ShuffledPartialSeats<'seats> {
+    /// Shuffled the input `partial_seats`.
+    // TODO(rand) do all shuffling operations in one generic fn
+    pub fn new(partial_seats: &'seats mut [PartialSeat<'seats>]) -> Self {
+        fastrand::shuffle(partial_seats);
+        Self { partial_seats }
+    }
+
+    pub fn get_partial_seats(&self) -> &[PartialSeat] {
+        self.partial_seats
+    }
+}

sim-validator-assignment/src/run.rs

@@ -1,7 +1,10 @@
 use crate::config::Config;
+use crate::partial_seat::{PartialSeat, ShuffledPartialSeats};
 use crate::seat::ShuffledSeats;
 use crate::shard::Shard;
-use crate::validator::{new_ordered_seats, parse_raw_validator_data, RawValidatorData};
+use crate::validator::{
+    new_ordered_partial_seats, new_ordered_seats, parse_raw_validator_data, RawValidatorData,
+};
 use num_rational::Ratio;
 use num_traits::ToPrimitive;
 use std::fs::read_to_string;
@@ -13,7 +16,7 @@ pub fn run(config: &Config) -> anyhow::Result<()> {
         None => mock_validator_data(),
     };
 
-    let (population_stats, validators) = parse_raw_validator_data(&config, &raw_validator_data);
+    let (population_stats, validators) = parse_raw_validator_data(config, &raw_validator_data);
 
     println!("population_stats: {:?}", population_stats);
     println!(
@@ -40,15 +43,24 @@ pub fn run(config: &Config) -> anyhow::Result<()> {
     let mut num_corrupted_shards = 0;
 
     for block_height in 0..config.num_blocks {
-        let mut ordered_seats = new_ordered_seats(&validators);
-        let shuffled_seats = ShuffledSeats::new(&mut ordered_seats);
+        let mut seats = new_ordered_seats(&validators);
+        let shuffled_seats = ShuffledSeats::new(&mut seats);
+
+        let mut partial_seats = config
+            .include_partial_seats
+            .then(|| new_ordered_partial_seats(&validators, config.stake_per_seat));
+        let shuffled_partial_seats = partial_seats
+            .as_mut()
+            .map(|ps| ShuffledPartialSeats::new(ps));
 
         for shard_idx in 0..config.num_shards {
             let shard_idx = usize::from(shard_idx);
             let shard_seats =
                 config.collect_seats_for_shard(shard_idx, shuffled_seats.get_seats())?;
-            let shard = Shard::new(&config, shard_seats)?;
-            if shard.is_corrupted(&config) {
+            let shard_partial_seats =
+                get_partial_seats_for_shard(config, shuffled_partial_seats.as_ref(), shard_idx)?;
+            let shard = Shard::new(config, shard_seats, shard_partial_seats)?;
+            if shard.is_corrupted(config) {
                 num_corrupted_shards += 1;
             }
         }
@@ -87,3 +99,18 @@ fn mock_validator_data() -> Vec<RawValidatorData> {
 fn log_heartbeat(block_height: u64, num_simulated_shards: u64, num_corrupted_shards: u64) {
     println!("heartbeat(block_height: {block_height}): {num_corrupted_shards} / {num_simulated_shards} shards corrupted");
 }
+
+fn get_partial_seats_for_shard<'a>(
+    config: &'a Config,
+    partial_seats: Option<&'a ShuffledPartialSeats>,
+    shard_idx: usize,
+) -> anyhow::Result<Option<Vec<&'a PartialSeat<'a>>>> {
+    match partial_seats {
+        Some(partial_seats) => {
+            let shard_seats = config
+                .collect_partial_seats_for_shard(shard_idx, partial_seats.get_partial_seats())?;
+            Ok(Some(shard_seats))
+        }
+        None => Ok(None),
+    }
+}

sim-validator-assignment/src/seat.rs

@@ -28,12 +28,13 @@ pub struct ShuffledSeats<'seats> {
 
 impl<'seats> ShuffledSeats<'seats> {
     /// Shuffles the input `seats`.
+    // TODO(rand) do all shuffling operations in one generic fn
     pub fn new(seats: &'seats mut [Seat<'seats>]) -> Self {
         fastrand::shuffle(seats);
         Self { seats }
     }
 
     pub fn get_seats(&self) -> &[Seat] {
-        &self.seats
+        self.seats
     }
 }

sim-validator-assignment/src/shard.rs

@@ -1,17 +1,28 @@
 use crate::config::Config;
+use crate::partial_seat::PartialSeat;
 use crate::seat::Seat;
 use num_rational::Ratio;
 
 #[derive(Debug, Default)]
 pub struct Shard<'seats> {
     seats: Vec<&'seats Seat<'seats>>,
+    /// Contains partial seats in case they are enabled by configuration.
+    partial_seats: Option<Vec<&'seats PartialSeat<'seats>>>,
     stake: u128,
     malicious_stake: u128,
 }
 
 impl<'seats> Shard<'seats> {
-    pub fn new(config: &Config, seats: Vec<&'seats Seat>) -> anyhow::Result<Self> {
+    pub fn new(
+        config: &Config,
+        seats: Vec<&'seats Seat>,
+        partial_seats: Option<Vec<&'seats PartialSeat>>,
+    ) -> anyhow::Result<Self> {
         if seats.len() != usize::try_from(config.seats_per_shard).unwrap() {
+            // Count only _full_ seats for the minimum number of required seats, since it is not
+            // clear how a _partial_ seat should be weighted for that concern.
+            // Validator assignment frameworks might try to minimize the number of partial seats or
+            // try to ignore them entirely.
             anyhow::bail!(
                 "Shard requires {} seats, received {} seats",
                 config.seats_per_shard,
@@ -26,7 +37,19 @@ impl<'seats> Shard<'seats> {
                 shard.malicious_stake += config.stake_per_seat;
             }
         }
+
+        if let Some(partial_seats) = partial_seats.as_ref() {
+            for &ps in partial_seats.iter() {
+                let weight = ps.get_weight();
+                shard.stake += weight;
+                if ps.get_is_malicious() {
+                    shard.malicious_stake += weight;
+                }
+            }
+        }
+
         shard.seats = seats;
+        shard.partial_seats = partial_seats;
         Ok(shard)
     }