From ea00f4890469848b08934796820bfd3e209e10c5 Mon Sep 17 00:00:00 2001 From: Erick Cestari Date: Fri, 11 Sep 2026 15:03:12 -0300 Subject: [PATCH] ir: weight generator selection Add Generator::weight (default 10) and AnyGenerator::choose, which picks by weight instead of uniformly. The AFL mutator uses it for fresh programs and generator insertion. This lets future generators such as error, warning or pong be picked less often, since they stall the peer rather than advance the protocol. All existing generators keep the default, so selection is unchanged. rand needs the alloc feature for choose_weighted. --- Cargo.toml | 2 +- smite-ir-mutator/src/lib.rs | 19 ++++++------------- smite-ir/src/generators.rs | 34 ++++++++++++++++++++++++++++++++++ smite-ir/src/tests.rs | 32 ++++++++++++++++++++++++++++++++ 4 files changed, 73 insertions(+), 14 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index bf448bc9..05567431 100644 --- a/Cargo.toml +++ b/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" diff --git a/smite-ir-mutator/src/lib.rs b/smite-ir-mutator/src/lib.rs index 30acdd3b..67775a2c 100644 --- a/smite-ir-mutator/src/lib.rs +++ b/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" } diff --git a/smite-ir/src/generators.rs b/smite-ir/src/generators.rs index 58d1b2df..1abf2c2e 100644 --- a/smite-ir/src/generators.rs +++ b/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(), + } + } } diff --git a/smite-ir/src/tests.rs b/smite-ir/src/tests.rs index b627d63c..5c21b401 100644 --- a/smite-ir/src/tests.rs +++ b/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