From 985c75384ac48ddec64205a1fcd33fc1e91a9b3b Mon Sep 17 00:00:00 2001
From: Simon Gellis <simongellis@gmail.com>
Date: Thu, 7 Nov 2024 15:13:26 -0500
Subject: [PATCH] sim-rs: make uniform IB generation an option

---
 sim-rs/sim-cli/src/bin/gen-test-data.rs | 1 +
 sim-rs/sim-core/src/config.rs           | 3 +++
 sim-rs/sim-core/src/sim/node.rs         | 8 +++++++-
 sim-rs/test_data/realistic.toml         | 1 +
 sim-rs/test_data/simple.toml            | 1 +
 sim-rs/test_data/small.toml             | 1 +
 6 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/sim-rs/sim-cli/src/bin/gen-test-data.rs b/sim-rs/sim-cli/src/bin/gen-test-data.rs
index 8c2da07..ecdc525 100644
--- a/sim-rs/sim-cli/src/bin/gen-test-data.rs
+++ b/sim-rs/sim-cli/src/bin/gen-test-data.rs
@@ -123,6 +123,7 @@ fn main() -> Result<()> {
         ib_shards: 8,
         max_block_size: 90112,
         stage_length: 2,
+        uniform_ib_generation: true,
         max_ib_requests_per_peer: 1,
         max_ib_size: 327680,
         max_tx_size: 16384,
diff --git a/sim-rs/sim-core/src/config.rs b/sim-rs/sim-core/src/config.rs
index ec00f00..8be6365 100644
--- a/sim-rs/sim-core/src/config.rs
+++ b/sim-rs/sim-core/src/config.rs
@@ -61,6 +61,7 @@ pub struct RawConfig {
     pub max_block_size: u64,
     pub max_tx_size: u64,
     pub stage_length: u64,
+    pub uniform_ib_generation: bool,
     pub max_ib_size: u64,
     pub max_ib_requests_per_peer: usize,
     pub ib_shards: u64,
@@ -116,6 +117,7 @@ impl From<RawConfig> for SimConfiguration {
             max_block_size: value.max_block_size,
             max_tx_size: value.max_tx_size,
             stage_length: value.stage_length,
+            uniform_ib_generation: value.uniform_ib_generation,
             max_ib_size: value.max_ib_size,
             max_ib_requests_per_peer: value.max_ib_requests_per_peer,
             ib_shards: value.ib_shards,
@@ -138,6 +140,7 @@ pub struct SimConfiguration {
     pub max_block_size: u64,
     pub max_tx_size: u64,
     pub stage_length: u64,
+    pub uniform_ib_generation: bool,
     pub max_ib_size: u64,
     pub max_ib_requests_per_peer: usize,
     pub ib_shards: u64,
diff --git a/sim-rs/sim-core/src/sim/node.rs b/sim-rs/sim-core/src/sim/node.rs
index 07cc671..c67e000 100644
--- a/sim-rs/sim-core/src/sim/node.rs
+++ b/sim-rs/sim-core/src/sim/node.rs
@@ -210,7 +210,13 @@ impl Node {
         while probability > 0.0 {
             let next_p = f64::min(probability, 1.0);
             if let Some(vrf) = self.run_vrf(next_p) {
-                let vrf_slot = slot + self.rng.gen_range(0..self.sim_config.stage_length);
+                let vrf_slot = if self.sim_config.uniform_ib_generation {
+                    // IBs are generated at the start of any slot within this stage
+                    slot + self.rng.gen_range(0..self.sim_config.stage_length)
+                } else {
+                    // IBs are generated at the start of the first slot of this stage
+                    slot
+                };
                 slot_vrfs.entry(vrf_slot).or_default().push(vrf);
             }
             probability -= 1.0;
diff --git a/sim-rs/test_data/realistic.toml b/sim-rs/test_data/realistic.toml
index 9f12f03..78c836c 100644
--- a/sim-rs/test_data/realistic.toml
+++ b/sim-rs/test_data/realistic.toml
@@ -6,6 +6,7 @@ max_ib_requests_per_peer = 1
 max_ib_size = 327680
 max_tx_size = 16384
 stage_length = 3
+uniform_ib_generation = true
 
 [transaction_frequency_ms]
 distribution = "exp"
diff --git a/sim-rs/test_data/simple.toml b/sim-rs/test_data/simple.toml
index cabf661..42f8a1c 100644
--- a/sim-rs/test_data/simple.toml
+++ b/sim-rs/test_data/simple.toml
@@ -51,6 +51,7 @@ max_block_size = 90112
 
 # IB parameters
 stage_length = 2
+uniform_ib_generation = true
 ib_generation_probability = 0.5 # corresponds to 𝑓I in the model
 ib_shards = 2
 max_ib_size = 327680
diff --git a/sim-rs/test_data/small.toml b/sim-rs/test_data/small.toml
index 246e88a..91acae5 100644
--- a/sim-rs/test_data/small.toml
+++ b/sim-rs/test_data/small.toml
@@ -6,6 +6,7 @@ max_ib_size = 327680
 max_ib_requests_per_peer = 1
 ib_shards = 8
 stage_length = 2
+uniform_ib_generation = true
 
 [[nodes]]
 location = [