sim-rs: support mocking lottery runs

Author: SupernaviX

sim-rs/Cargo.lock

@@ -7,6 +7,7 @@ rust-version = "1.88"
 [dependencies]
 anyhow = "1"
 async-stream = "0.3"
+dashmap = "6"
 futures = "0.3"
 netsim-async = { git = "https://github.com/input-output-hk/ce-netsim.git", rev = "9d1e26c" }
 num-traits = "0.2"

sim-rs/sim-core/Cargo.toml

@@ -28,7 +28,7 @@ use crate::{
     sim::{
         MiniProtocol, NodeImpl, SimCpuTask, SimMessage,
         linear_leios::attackers::{EBWithholdingEvent, EBWithholdingSender},
-        lottery,
+        lottery::{LotteryConfig, LotteryKind, MockLotteryResults, vrf_probabilities},
     },
 };
 
@@ -306,7 +306,7 @@ pub struct LinearLeiosNode {
     tracker: EventTracker,
     rng: ChaChaRng,
     clock: Clock,
-    stake: u64,
+    lottery: LotteryConfig,
     consumers: Vec<NodeId>,
     txs: HashMap<TransactionId, TransactionView>,
     ledger_states: BTreeMap<BlockId, Arc<LedgerState>>,
@@ -333,14 +333,19 @@ impl NodeImpl for LinearLeiosNode {
         rng: ChaChaRng,
         clock: Clock,
     ) -> Self {
+        let lottery = LotteryConfig::Random {
+            stake: config.stake,
+            total_stake: sim_config.total_stake,
+        };
+
         Self {
             id: config.id,
             sim_config,
             queued: EventResult::default(),
             tracker,
             rng,
             clock,
-            stake: config.stake,
+            lottery,
             consumers: config.consumers.clone(),
             txs: HashMap::new(),
             ledger_states: BTreeMap::new(),
@@ -512,7 +517,10 @@ impl LinearLeiosNode {
 // Ranking block propagation
 impl LinearLeiosNode {
     fn try_generate_rb(&mut self, slot: u64) {
-        let Some(vrf) = self.run_vrf(self.sim_config.block_generation_probability) else {
+        let Some(vrf) = self.run_vrf(
+            LotteryKind::GenerateRB,
+            self.sim_config.block_generation_probability,
+        ) else {
             return;
         };
 
@@ -1204,8 +1212,8 @@ impl LinearLeiosNode {
     }
 
     fn try_vote_for_endorser_block(&mut self, eb: &Arc<EndorserBlock>, seen: Timestamp) -> bool {
-        let vrf_wins = lottery::vrf_probabilities(self.sim_config.vote_probability)
-            .filter_map(|f| self.run_vrf(f))
+        let vrf_wins = vrf_probabilities(self.sim_config.vote_probability)
+            .filter_map(|f| self.run_vrf(LotteryKind::GenerateVote, f))
             .count();
         if vrf_wins == 0 {
             return false;
@@ -1487,18 +1495,12 @@ impl LinearLeiosNode {
 
 // Common utilities
 impl LinearLeiosNode {
+    #[allow(unused)]
+    pub fn mock_lottery(&mut self, results: Arc<MockLotteryResults>) {
+        self.lottery = LotteryConfig::Mock { results };
+    }
     // Simulates the output of a VRF using this node's stake (if any).
-    fn run_vrf(&mut self, success_rate: f64) -> Option<u64> {
-        let target_vrf_stake = lottery::compute_target_vrf_stake(
-            self.stake,
-            self.sim_config.total_stake,
-            success_rate,
-        );
-        let result = self.rng.random_range(0..self.sim_config.total_stake);
-        if result < target_vrf_stake {
-            Some(result)
-        } else {
-            None
-        }
+    fn run_vrf(&mut self, kind: LotteryKind, success_rate: f64) -> Option<u64> {
+        self.lottery.run(kind, success_rate, &mut self.rng)
     }
 }

sim-rs/sim-core/src/sim/linear_leios.rs

@@ -1,3 +1,9 @@
+use std::{collections::VecDeque, sync::Arc};
+
+use dashmap::DashMap;
+use rand::Rng;
+use rand_chacha::ChaChaRng;
+
 pub fn compute_target_vrf_stake(stake: u64, total_stake: u64, success_rate: f64) -> u64 {
     let ratio = stake as f64 / total_stake as f64;
     (total_stake as f64 * ratio * success_rate) as u64
@@ -7,3 +13,51 @@ pub fn vrf_probabilities(probability: f64) -> impl Iterator<Item = f64> {
     let final_success_rate = Some(probability.fract()).filter(|f| *f > 0.0);
     std::iter::repeat_n(1.0, probability.trunc() as usize).chain(final_success_rate)
 }
+
+#[derive(Clone, Copy, PartialEq, Eq, Hash)]
+pub enum LotteryKind {
+    GenerateRB,
+    GenerateVote,
+}
+
+#[derive(Default)]
+pub struct MockLotteryResults {
+    results: DashMap<LotteryKind, VecDeque<u64>>,
+}
+impl MockLotteryResults {
+    pub fn run(&self, kind: LotteryKind) -> Option<u64> {
+        self.results.entry(kind).or_default().pop_front()
+    }
+    #[allow(unused)]
+    pub fn configure_win(&self, kind: LotteryKind, result: u64) {
+        self.results.entry(kind).or_default().push_back(result);
+    }
+}
+
+pub enum LotteryConfig {
+    Random {
+        stake: u64,
+        total_stake: u64,
+    },
+    #[allow(unused)]
+    Mock {
+        results: Arc<MockLotteryResults>,
+    },
+}
+
+impl LotteryConfig {
+    pub fn run(&self, kind: LotteryKind, success_rate: f64, rng: &mut ChaChaRng) -> Option<u64> {
+        match self {
+            Self::Random { stake, total_stake } => {
+                let target_vrf_stake = compute_target_vrf_stake(*stake, *total_stake, success_rate);
+                let result = rng.random_range(0..*total_stake);
+                if result < target_vrf_stake {
+                    Some(result)
+                } else {
+                    None
+                }
+            }
+            Self::Mock { results } => results.run(kind),
+        }
+    }
+}