ir: weight generator selection

Cargo.toml

@@ -25,7 +25,7 @@ warnings = "deny"
 [workspace.dependencies]
 libc = "0.2"
 log = "0.4"
-rand = { version = "0.10", default-features = false }
+rand = { version = "0.10", default-features = false, features = ["alloc"] }
 simple_logger = { version = "5", default-features = false }
 postcard = { version = "1.1", default-features = false, features = ["alloc"] }
 bitcoin = "0.32"

smite-ir-mutator/src/lib.rs

@@ -37,7 +37,7 @@ use std::os::raw::{c_char, c_uint, c_void};
 use std::slice;
 
 use rand::rngs::SmallRng;
-use rand::{RngExt, SeedableRng, seq::IteratorRandom};
+use rand::{RngExt, SeedableRng};
 
 use smite_ir::generators::AnyGenerator;
 use smite_ir::minimizers::{CommonSubexpressionEliminator, DeadCodeEliminator, Minimizer};
@@ -72,15 +72,11 @@ impl MutatorState {
         }
     }
 
-    /// Generates a fresh program from scratch by randomly delegating to one of
-    /// the registered generators.
+    /// Generates a fresh program from scratch by delegating to one of the
+    /// registered generators, picked by weight.
     fn generate_fresh(&mut self) -> Program {
         let mut builder = ProgramBuilder::new();
-        AnyGenerator::ALL
-            .iter()
-            .choose(&mut self.rng)
-            .expect("AnyGenerator::ALL is non-empty")
-            .generate(&mut builder, &mut self.rng);
+        AnyGenerator::choose(&mut self.rng).generate(&mut builder, &mut self.rng);
         self.last_sequence.clear();
         self.last_sequence.push("fresh");
         builder.build()
@@ -114,11 +110,8 @@ impl MutatorState {
                     "instr-reorder"
                 }
                 4 => {
-                    let generator = *AnyGenerator::ALL
-                        .iter()
-                        .choose(&mut self.rng)
-                        .expect("AnyGenerator::ALL is non-empty");
-                    let mutator = GeneratorInsertionMutator::new(generator);
+                    let mutator =
+                        GeneratorInsertionMutator::new(AnyGenerator::choose(&mut self.rng));
                     mutator.mutate(program, &mut self.rng);
                     "gen-insert"
                 }

smite-ir/src/generators.rs

@@ -22,13 +22,23 @@ pub use node_announcement::NodeAnnouncementGenerator;
 pub use open_channel::OpenChannelGenerator;
 
 use rand::Rng;
+use rand::seq::IndexedRandom;
 
 use super::builder::ProgramBuilder;
 
+/// Weight of a generator that does not override [`Generator::weight`].
+pub const DEFAULT_WEIGHT: u32 = 10;
+
 /// A generator that emits instructions into a `ProgramBuilder`.
 pub trait Generator {
     /// Emits instructions for this generator's protocol interaction.
     fn generate(&self, builder: &mut ProgramBuilder, rng: &mut impl Rng);
+
+    /// Relative pick weight for [`AnyGenerator::choose`]; 0 disables.
+    /// Lower it for generators that should be picked less often.
+    fn weight(&self) -> u32 {
+        DEFAULT_WEIGHT
+    }
 }
 
 /// A list of all the available generators. Any generators included
@@ -55,6 +65,18 @@ impl AnyGenerator {
         Self::ChannelReady(ChannelReadyGenerator),
         Self::FundingFlow(FundingFlowGenerator),
     ];
+
+    /// Picks a generator from `ALL` with probability proportional to its
+    /// [`Generator::weight`].
+    ///
+    /// # Panics
+    ///
+    /// Panics if every generator reports a weight of zero.
+    pub fn choose(rng: &mut impl Rng) -> Self {
+        *Self::ALL
+            .choose_weighted(rng, Generator::weight)
+            .expect("at least one generator must have non-zero weight")
+    }
 }
 
 impl Generator for AnyGenerator {
@@ -69,4 +91,16 @@ impl Generator for AnyGenerator {
             Self::FundingFlow(generator) => generator.generate(builder, rng),
         }
     }
+
+    fn weight(&self) -> u32 {
+        match self {
+            Self::ChannelAnnouncement(generator) => generator.weight(),
+            Self::ChannelUpdate(generator) => generator.weight(),
+            Self::NodeAnnouncement(generator) => generator.weight(),
+            Self::OpenChannel(generator) => generator.weight(),
+            Self::FundingCreated(generator) => generator.weight(),
+            Self::ChannelReady(generator) => generator.weight(),
+            Self::FundingFlow(generator) => generator.weight(),
+        }
+    }
 }

smite-ir/src/tests.rs

@@ -921,6 +921,38 @@ fn any_generator_all_is_complete() {
     assert_eq!(AnyGenerator::ALL.len(), variant_count(AnyGenerator::ALL[0]));
 }
 
+// Weighted selection must cover every generator with non-zero weight and
+// skew towards heavier ones in proportion to their weight.
+#[test]
+fn any_generator_choose_respects_weights() {
+    let mut rng = SmallRng::seed_from_u64(7);
+    let mut counts = vec![0u32; AnyGenerator::ALL.len()];
+    let rounds = 20_000;
+    for _ in 0..rounds {
+        let picked = AnyGenerator::choose(&mut rng);
+        let idx = AnyGenerator::ALL
+            .iter()
+            .position(|g| std::mem::discriminant(g) == std::mem::discriminant(&picked))
+            .expect("chosen generator is in ALL");
+        counts[idx] += 1;
+    }
+
+    let total: u32 = AnyGenerator::ALL.iter().map(Generator::weight).sum();
+    for (generator, &count) in AnyGenerator::ALL.iter().zip(&counts) {
+        let weight = generator.weight();
+        if weight == 0 {
+            assert_eq!(count, 0, "zero-weight generator must never be picked");
+            continue;
+        }
+        let expected = f64::from(rounds) * f64::from(weight) / f64::from(total);
+        let ratio = f64::from(count) / expected;
+        assert!(
+            (0.85..=1.15).contains(&ratio),
+            "generator picked {count} times, expected ~{expected:.0} (ratio {ratio:.2})"
+        );
+    }
+}
+
 // -- ShutdownScriptVariant tests --
 
 // Ensure ShutdownScriptVariant and ShutdownScriptVariant::VARIANT_COUNT stay in