puffin/fuzzer/

mutations.rs

use std::borrow::Cow;

use libafl::prelude::*;
use libafl_bolts::prelude::*;

use super::utils::{
    choose, choose_iter, choose_term, choose_term_filtered_mut, choose_term_path_filtered,
    find_all_sub_term_filtered, find_all_term_filtered, find_term, find_term_mut, reservoir_sample,
    Choosable, TermConstraints, TracePath,
};
use crate::algebra::atoms::Function;
use crate::algebra::signature::Signature;
use crate::algebra::{DYTerm, Subterms, Term, TermType};
use crate::fuzzer::term_zoo::TermZoo;
use crate::protocol::{ProtocolBehavior, ProtocolTypes};
use crate::put_registry::PutRegistry;
use crate::trace::{Action, Spawner, Trace, TraceContext};

#[derive(Clone, Copy, Debug)]
pub struct MutationConfig {
    pub fresh_zoo_after: u64,
    pub max_result_trace_length: usize,
    pub min_trace_length: usize,
    /// Below this term size we no longer mutate. Note that it is possible to reach
    /// smaller terms by having a mutation which removes all symbols in a single mutation.
    /// Above this term size we no longer mutate.
    pub term_constraints: TermConstraints,
    pub with_bit_level: bool,
    pub with_dy: bool,
    /// Focus on one payload at a time for a whole StdMutationalStage
    pub with_focus: bool,
    /// Relative weights for the scope at which a *replacement* mutation is applied, see
    /// [`ScopeWeights`] and [`MutationScope`].
    pub scope_weights: ScopeWeights,
}

impl Default for MutationConfig {
    //  TODO:EVAL: evaluate modif to this config
    fn default() -> Self {
        Self {
            fresh_zoo_after: 100000,
            max_result_trace_length: 20,
            min_trace_length: 2,
            term_constraints: TermConstraints::default(),
            with_bit_level: false,
            with_dy: true,
            with_focus: true,
            scope_weights: ScopeWeights::default(),
        }
    }
}

/// Relative weights for picking a [`MutationScope`] when applying a replacement mutation.
///
/// The default weights the scope inverse to its breadth (`global:1, step:2, individual:3`) so that
/// surgical edits dominate. Setting them allows ablation studies and reproducing historical
/// behaviours, e.g. `(0, 0, 1)` is purely individual replacements (behaviour before global
/// mutations existed).
#[derive(Clone, Copy, Debug)]
pub struct ScopeWeights {
    pub global: usize,
    pub step: usize,
    pub individual: usize,
}

impl Default for ScopeWeights {
    fn default() -> Self {
        Self {
            // Weights INVERSE to the semantic breadth of the scope: a Global replace rewrites the
            // matched sub-term EVERYWHERE in the trace (broadest, most destructive for structured
            // inputs), Step within one step, Individual a single occurrence (most surgical). For
            // grammar/DY fuzzing we want surgical edits to dominate, so global is the least likely.
            // (AFL-style broad stacking suits flat byte inputs, not deep protocol structures.)
            global: 1,
            step: 2,
            individual: 3,
        }
    }
}

impl ScopeWeights {
    #[must_use]
    pub const fn new(global: usize, step: usize, individual: usize) -> Self {
        Self {
            global,
            step,
            individual,
        }
    }

    /// Saturating, so that absurdly large weights keep a meaningful (if degenerate) distribution
    /// instead of overflowing: the dominant weight simply wins.
    const fn total(self) -> usize {
        self.global
            .saturating_add(self.step)
            .saturating_add(self.individual)
    }
}

/// Scope over which a chosen replacement is applied, see [`apply_scoped_mutation`].
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum MutationScope {
    /// Replace *every* occurrence structurally equal to the chosen term, in the whole trace.
    Global,
    /// Replace every such occurrence, but only within the step of the chosen term.
    Step,
    /// Replace only the single chosen occurrence.
    Individual,
}

impl MutationScope {
    /// Draw a scope at random according to `weights`. Falls back to [`Self::Individual`] when all
    /// weights are zero.
    fn choose<R: Rand>(weights: ScopeWeights, rand: &mut R) -> Self {
        let total = weights.total();
        if total == 0 {
            return Self::Individual;
        }
        let draw = rand.between(0, total - 1);
        if draw < weights.global {
            Self::Global
        } else if draw < weights.global.saturating_add(weights.step) {
            Self::Step
        } else {
            Self::Individual
        }
    }
}

/// Apply `new_term` in place of the term at `to_mutate_path`, generalising the replacement to
/// structurally equal terms according to a randomly drawn [`MutationScope`].
///
/// This is the shared engine behind the *replacement* DY mutators (generate / replace-match /
/// replace-reuse): they all have the shape "pick a sub-term, compute a replacement for it, write it
/// back", so broadcasting that replacement to the other occurrences of the *same* term is
/// well-defined. Doing so collapses the cost of goals that require the same edit in several places
/// (e.g. turning every `fn_seq_2` of one action into `fn_seq_5`) from one lucky draw *per
/// occurrence* down to one lucky draw *per distinct sub-term*.
///
/// The generalisation is deliberately performed without any term constraint: skipping a large
/// recipe here would silently turn "replace all occurrences" into a partial replacement, which is
/// precisely what a broader scope is meant to avoid; and the cost stays linear in the size of the
/// trace, i.e. the same order as the selection pass that already ran.
///
/// The chosen occurrence is always part of the targets, so that a broader scope is a strict
/// superset of [`MutationScope::Individual`] whatever the search returns.
// TODO: middle-ground scope: instead of the whole trace or the whole step, pick a position between
// the chosen term and the root of the recipe and replace all occurrences of the chosen term in that
// sub-term only.
pub fn apply_scoped_mutation<PT: ProtocolTypes, R: Rand>(
    trace: &mut Trace<PT>,
    to_mutate_path: &TracePath,
    representative: &Term<PT>,
    new_term: &Term<PT>,
    config: &MutationConfig,
    rand: &mut R,
) -> MutationResult {
    let targets = scoped_targets(trace, to_mutate_path, representative, config, rand);
    apply_to_targets(
        trace,
        &targets,
        new_term,
        representative,
        config.with_bit_level,
        &config.term_constraints,
    )
}

/// Returns `true` iff replacing a term of size `representative_size` by one of size `new_size` at
/// every path in `targets` keeps each affected single step recipe within
/// `constraints.max_result_term_size`.
///
/// **Precondition: the input `trace` is already within the caps.** This is a loop invariant, not
/// something re-checked here: hand-written seeds are authored within the caps, and every
/// replacement mutation is reject-whole, so a trace that reached the corpus is always `<=
/// max_result_term_size`. The size-neutral fast path below relies on it — a shrinking or
/// size-neutral replacement is accepted in O(1) precisely because it cannot push an *already
/// bounded* trace over a cap. The one way to break the invariant is to configure a result cap
/// *below* the size of an already-imported seed: such a seed is out of bounds before any mutation,
/// and a neutral replacement will (correctly, per this contract) keep it that way. Enforcing caps
/// against oversized inputs is a corpus-boundary concern, out of scope for the mutator.
///
/// Efficiency (why we don't recompute every step's size):
/// - a mutation that does not grow the trace (`new_size <= representative_size`) can never exceed a
///   cap, so we accept in O(1) without touching any size;
/// - otherwise we only look at the **affected steps** (those containing a target) — never all steps
///   — recomputing each such step's recipe size once, and derive the new whole-trace size as
///   `old_trace_size + n_targets * delta`. `old_trace_size` is read from the (already bounded, `<=
///   already-bounded) trace, so it is cheap.
fn replacement_within_caps<PT: ProtocolTypes>(
    trace: &Trace<PT>,
    targets: &[TracePath],
    representative_size: usize,
    new_size: usize,
    constraints: &TermConstraints,
) -> bool {
    if new_size <= representative_size {
        return true; // shrink or size-neutral: cannot exceed any cap
    }
    let delta = new_size - representative_size;

    // Per-affected-step check: count targets per step, only recompute those steps' sizes.
    let mut per_step: std::collections::HashMap<usize, usize> = std::collections::HashMap::new();
    for (step_index, _) in targets {
        *per_step.entry(*step_index).or_insert(0) += 1;
    }
    for (step_index, count) in &per_step {
        let cur = match trace.steps.get(*step_index).map(|s| &s.action) {
            Some(Action::Input(input)) => input.recipe.size(),
            _ => 0,
        };
        if cur + count * delta > constraints.max_result_term_size {
            return false;
        }
    }

    // We intentionally cap only the per-step recipe size here. Whole-trace growth is bounded by the
    // number of steps (`max_result_trace_length`, enforced when steps are added) times this cap,
    // which is enough; we therefore avoid the extra whole-trace `Trace::size()` recomputation.
    true
}

/// Returns `true` iff replacing a term carrying `representative_payloads` payloads by one carrying
/// `new_payloads`, at every path in `targets`, keeps each *affected* step recipe within
/// `constraints.threshold_max_payloads_per_term`.
///
/// The budget is enforced **per term** (per step recipe), not as an average over the whole trace:
/// a broad scope that rewrites several occurrences within a single recipe is charged the full
/// growth for that recipe, so one dense recipe cannot hide behind many payload-free ones. Mirrors
/// [`replacement_within_caps`]: only a growing replacement (`new_payloads >
/// representative_payloads`) can exceed the budget, so a payload-neutral or shrinking replacement
/// is accepted in O(1) without inspecting any recipe. Enforced uniformly for every replacement
/// mutator via [`apply_to_targets`] (previously only [`ReplaceReuseMutator`] checked it, and only
/// against a trace-wide average).
fn payloads_within_caps<PT: ProtocolTypes>(
    trace: &Trace<PT>,
    targets: &[TracePath],
    representative_payloads: usize,
    new_payloads: usize,
    constraints: &TermConstraints,
) -> bool {
    if new_payloads <= representative_payloads {
        return true; // no payload growth: cannot exceed the budget
    }
    let delta = new_payloads - representative_payloads;

    // Per-affected-step check: count targets per step, only inspect those steps' recipes. `cur`
    // already includes the payloads of the occurrences about to be replaced, so the projected
    // recipe payload count is `cur + count * delta`.
    let mut per_step: std::collections::HashMap<usize, usize> = std::collections::HashMap::new();
    for (step_index, _) in targets {
        *per_step.entry(*step_index).or_insert(0) += 1;
    }
    for (step_index, count) in &per_step {
        let cur = match trace.steps.get(*step_index).map(|s| &s.action) {
            Some(Action::Input(input)) => input.recipe.count_payloads(),
            _ => 0,
        };
        if cur + count * delta > constraints.threshold_max_payloads_per_term {
            return false;
        }
    }
    true
}

/// Replace the term at each of `targets` (each currently structurally equal to `representative`) by
/// `new_term`, enforcing the result-size caps (`constraints.max_result_{term,trace}_size`) and,
/// when `with_bit` is set, the per-term payload budget
/// (`constraints.threshold_max_payloads_per_term`). If applying to all targets would exceed a cap,
/// the mutation is rejected wholesale (returns [`MutationResult::Skipped`]) so the "replace all
/// equal occurrences" atomicity is preserved. See `replacement_within_caps` and
/// `payloads_within_caps`.
///
/// `with_bit` mirrors whether bit-level mutations are enabled: when they are not, no term carries a
/// payload, so the whole payload accounting is skipped without even walking `new_term`.
pub fn apply_to_targets<PT: ProtocolTypes>(
    trace: &mut Trace<PT>,
    targets: &[TracePath],
    new_term: &Term<PT>,
    representative: &Term<PT>,
    with_bit: bool,
    constraints: &TermConstraints,
) -> MutationResult {
    if !replacement_within_caps(
        trace,
        targets,
        representative.size(),
        new_term.size(),
        constraints,
    ) {
        log::debug!("[Mutation] rejected: would exceed result-size caps");
        return MutationResult::Skipped;
    }
    // Payloads only exist under bit-level mutations; when those are off, skip the accounting
    // entirely (no `new_term` walk). Even under bit-level, only a replacement that *adds* payloads
    // can exceed the budget, so a payload-free `new_term` short-circuits before the per-recipe
    // walks.
    if with_bit {
        let new_payloads = new_term.count_payloads();
        if new_payloads > 0
            && !payloads_within_caps(
                trace,
                targets,
                representative.count_payloads(),
                new_payloads,
                constraints,
            )
        {
            log::debug!("[Mutation] rejected: would exceed per-term payload budget");
            return MutationResult::Skipped;
        }
    }

    let mut mutated = false;
    for target in targets {
        if let Some(term_mut) = find_term_mut(trace, target) {
            term_mut.mutate(new_term.clone());
            mutated = true;
        }
    }

    if mutated {
        MutationResult::Mutated
    } else {
        MutationResult::Skipped
    }
}

/// Draw a [`MutationScope`] and return *all* the occurrences to replace, see
/// [`apply_scoped_mutation`]. The search is unbounded (finds every structurally-equal occurrence);
/// runaway is prevented at *application* time by the result-size caps in [`apply_to_targets`], not
/// by cutting the search short.
///
/// Exposed separately from [`apply_to_targets`] so a caller can inspect *how many* occurrences
/// would be replaced before committing: [`ReplaceReuseMutator`] needs it to bound the payloads it
/// adds.
pub fn scoped_targets<PT: ProtocolTypes, R: Rand>(
    trace: &Trace<PT>,
    to_mutate_path: &TracePath,
    representative: &Term<PT>,
    config: &MutationConfig,
    rand: &mut R,
) -> Vec<TracePath> {
    let scope = MutationScope::choose(config.scope_weights, rand);
    let (chosen_step, _) = to_mutate_path;

    // the type shape is only compared first to speed up the comparison
    let filter = |term: &Term<PT>| {
        term.get_type_shape() == representative.get_type_shape() && term == representative
    };
    // The result-size caps enforced in `apply_to_targets` bound growth, so the search for the
    // other occurrences is itself unconstrained.
    let constraints = TermConstraints::no_constraint();

    let mut targets: Vec<TracePath> = match scope {
        // no search needed, the chosen occurrence is added below
        MutationScope::Individual => vec![],
        // only the recipe of the chosen step is explored, rather than filtering the whole trace
        MutationScope::Step => match trace.steps.get(*chosen_step).map(|step| &step.action) {
            Some(Action::Input(input)) => {
                find_all_sub_term_filtered(&input.recipe, filter, &constraints)
                    .into_iter()
                    .map(|term_path| (*chosen_step, term_path))
                    .collect()
            }
            _ => vec![],
        },
        MutationScope::Global => find_all_term_filtered(trace, filter, &constraints),
    };
    // the chosen occurrence is added last, make sure it is not replaced twice
    targets.retain(|path| path != to_mutate_path);
    targets.push(to_mutate_path.clone());

    log::debug!(
        "[Mutation] scope {scope:?} selected {} occurrence(s)",
        targets.len()
    );
    targets
}

impl MutationConfig {
    pub fn default_with_bit() -> Self {
        MutationConfig {
            with_bit_level: true,
            with_focus: false,
            ..Self::default()
        }
    }
}
pub type DyMutations<'harness, PT, PB, S> = tuple_list_type!(
// DY mutations
    RepeatMutator<S>,
    SkipMutator<S>,
    ReplaceReuseMutator<S>,
    ReplaceMatchMutator<S, PT>,
    RemoveAndLiftMutator<S>,
    GenerateMutator<'harness, S, PB>,
    SwapMutator<S>,
);

#[must_use]
pub fn dy_mutations<'harness, S, PT: ProtocolTypes, PB>(
    mutation_config: MutationConfig,
    signature: &'static Signature<PT>,
    put_registry: &'harness PutRegistry<PB>,
) -> DyMutations<'harness, PT, PB, S>
where
    S: HasCorpus<Trace<PT>> + HasMetadata + HasMaxSize + HasRand,
    PB: ProtocolBehavior<ProtocolTypes = PT>,
{
    let MutationConfig {
        fresh_zoo_after,
        max_result_trace_length,
        min_trace_length,
        term_constraints,
        with_dy,
        ..
    } = mutation_config;

    tuple_list!(
        RepeatMutator::new(max_result_trace_length, with_dy),
        SkipMutator::new(min_trace_length, with_dy),
        ReplaceReuseMutator::new(mutation_config),
        ReplaceMatchMutator::new(mutation_config, signature),
        RemoveAndLiftMutator::new(term_constraints, with_dy),
        GenerateMutator::new(
            0,
            fresh_zoo_after,
            None,
            signature,
            put_registry,
            mutation_config,
        ), /* Refresh zoo after 100000M mutations */
        SwapMutator::new(term_constraints, with_dy),
    )
}

/// SWAP: Swaps a sub-term with a different sub-term which is part of the trace

/// (such that types match).
///
/// Note: this mutation is deliberately *not* scoped (see [`MutationScope`]): it exchanges two
/// sub-terms, so "replace all equal occurrences" has no single well-defined replacement.
// TODO: we might later give it its own scoped variant: swap *all* occurrences of A with *all*
// occurrences of B (with probability 1/3), or only within the current step (with probability 1/3).
pub struct SwapMutator<S>
where
    S: HasRand,
{
    constraints: TermConstraints,
    phantom_s: std::marker::PhantomData<S>,
    with_dy: bool,
}

impl<S> SwapMutator<S>
where
    S: HasRand,
{
    #[must_use]
    pub const fn new(constraints: TermConstraints, with_dy: bool) -> Self {
        Self {
            constraints,
            phantom_s: std::marker::PhantomData,
            with_dy,
        }
    }
}

impl<S, PT: ProtocolTypes> Mutator<Trace<PT>, S> for SwapMutator<S>
where
    S: HasRand,
{
    fn mutate(&mut self, state: &mut S, trace: &mut Trace<PT>) -> Result<MutationResult, Error> {
        log::debug!("[DY] Start mutate with {}", self.name());
        if !self.with_dy {
            return Ok(MutationResult::Skipped);
        }
        let _a = BytesInsertMutator;
        let rand = state.rand_mut();
        if let Some((term_a, trace_path_a)) = choose(trace, &self.constraints, rand) {
            if let Some(trace_path_b) = choose_term_path_filtered(
                trace,
                |term: &Term<PT>| term.get_type_shape() == term_a.get_type_shape(),
                &self.constraints,
                rand,
            ) {
                let step_a_size = match &trace.steps[trace_path_a.0].action {
                    crate::trace::Action::Input(input) => input.recipe.size(),
                    crate::trace::Action::Output(_) => 0,
                };
                let step_b_size = match &trace.steps[trace_path_b.0].action {
                    crate::trace::Action::Input(input) => input.recipe.size(),
                    crate::trace::Action::Output(_) => 0,
                };

                let term_a_cloned = term_a.clone();
                let term_a_size = term_a.size();

                if let Some(term_b_mut) = find_term_mut(trace, &trace_path_b) {
                    let term_b_size = term_b_mut.size();

                    // Post-mutation hard cap of 500 to prevent runaway growth while allowing
                    // natural overshoot (pre-PR 472 behavior).
                    if step_b_size + term_a_size <= 500 + term_b_size
                        && step_a_size + term_b_size <= 500 + term_a_size
                    {
                        log::debug!(
                            "[Mutation] Mutate SwapMutator on terms\n{} and\n {}",
                            term_a_cloned,
                            term_b_mut
                        );
                        let term_b_cloned = term_b_mut.clone();
                        term_b_mut.mutate(term_a_cloned);
                        if let Some(trace_a_mut) = find_term_mut(trace, &trace_path_a) {
                            trace_a_mut.mutate(term_b_cloned);
                        }
                        return Ok(MutationResult::Mutated);
                    } else {
                        log::debug!("[SwapMutator] Skipped as it would exceed max_term_size.");
                    }
                }
            }
        }
        log::debug!("       Skipped {}", self.name());
        Ok(MutationResult::Skipped)
    }

    fn post_exec(&mut self, _state: &mut S, _new_corpus_id: Option<CorpusId>) -> Result<(), Error> {
        Ok(())
    }
}
impl<S> Named for SwapMutator<S>
where
    S: HasRand,
{
    fn name(&self) -> &Cow<'static, str> {
        &Cow::Borrowed("SwapMutator")
    }
}

/// REMOVE AND LIFT: Removes a sub-term from a term and attaches orphaned children to the parent

/// (such that types match). This only works if there is only a single child.
///
/// Note: this mutation is deliberately *not* scoped (see [`MutationScope`]): the replacement is a
/// grand-sub-term of the mutated term, i.e. position-dependent, so there is no single replacement
/// to broadcast.
// TODO: a scoped variant would instead apply the *same transformation* to all terms (which are
// `make_list` terms) that are equal to the impacted term *before* the RemoveAndLift.
// Note: this mutation will eventually be removed in favour of more scoped list-only mutations.
pub struct RemoveAndLiftMutator<S>
where
    S: HasRand,
{
    constraints: TermConstraints,
    phantom_s: std::marker::PhantomData<S>,
    with_dy: bool,
}

impl<S> RemoveAndLiftMutator<S>
where
    S: HasRand,
{
    #[must_use]
    pub const fn new(constraints: TermConstraints, with_dy: bool) -> Self {
        Self {
            constraints,
            phantom_s: std::marker::PhantomData,
            with_dy,
        }
    }
}

impl<S, PT: ProtocolTypes> Mutator<Trace<PT>, S> for RemoveAndLiftMutator<S>
where
    S: HasRand,
{
    fn mutate(&mut self, state: &mut S, trace: &mut Trace<PT>) -> Result<MutationResult, Error> {
        log::debug!("[DY] Start mutate with {}", self.name());
        if !self.with_dy {
            return Ok(MutationResult::Skipped);
        }
        let rand = state.rand_mut();
        let filter = |term: &Term<PT>| {
            term.is_symbolic() && // exclude terms with payloads since we aim to modify its internal structure
            match &term.term {
                DYTerm::Variable(_) => false,
                DYTerm::Application(_, subterms) =>
                    {
                        subterms
                            .find_subterm(|subterm| match &subterm.term {
                                DYTerm::Variable(_) => false,
                                DYTerm::Application(_, grand_subterms) => {
                                    grand_subterms.find_subterm_same_shape(subterm).is_some()
                                }
                            })
                            .is_some()
                    }
            }
        };
        if let Some(to_mutate) = choose_term_filtered_mut(trace, filter, &self.constraints, rand) {
            log::debug!(
                "[Mutation] Mutate RemoveAndLiftMutator on term\n{}",
                to_mutate
            );
            match &mut to_mutate.term {
                // TODO-bitlevel: maybe also SKIP if not(to_mutate.is_symbolic())
                DYTerm::Variable(_) => {
                    log::debug!("       Skipped {}", self.name());
                    Ok(MutationResult::Skipped)
                }
                DYTerm::Application(_, ref mut subterms) => {
                    if let Some(((subterm_index, _), grand_subterm)) = choose_iter(
                        subterms.filter_grand_subterms(|subterm, grand_subterm| {
                            subterm.get_type_shape() == grand_subterm.get_type_shape()
                        }),
                        rand,
                    ) {
                        let grand_subterm_cloned = grand_subterm.clone();
                        subterms.push(grand_subterm_cloned);
                        subterms.swap_remove(subterm_index);
                        return Ok(MutationResult::Mutated);
                    }
                    log::debug!("       Skipped {}", self.name());
                    Ok(MutationResult::Skipped)
                }
            }
        } else {
            log::debug!("       Skipped {}", self.name());
            Ok(MutationResult::Skipped)
        }
    }

    fn post_exec(&mut self, _state: &mut S, _new_corpus_id: Option<CorpusId>) -> Result<(), Error> {
        Ok(())
    }
}

impl<S> Named for RemoveAndLiftMutator<S>
where
    S: HasRand,
{
    fn name(&self) -> &Cow<'static, str> {
        &Cow::Borrowed("RemoveAndLiftMutator")
    }
}

/// REPLACE-MATCH: Replaces a function symbol with a different one (such that types match).
///
/// An example would be to replace a constant with another constant or the binary function
/// `fn_add` with `fn_sub`.
/// It can also replace any variable with a constant.
pub struct ReplaceMatchMutator<S, PT: ProtocolTypes>
where
    S: HasRand,
{
    config: MutationConfig,
    signature: &'static Signature<PT>,
    phantom_s: std::marker::PhantomData<S>,
}

impl<S, PT: ProtocolTypes> ReplaceMatchMutator<S, PT>
where
    S: HasRand,
{
    #[must_use]
    pub const fn new(config: MutationConfig, signature: &'static Signature<PT>) -> Self {
        Self {
            config: MutationConfig {
                term_constraints: TermConstraints {
                    // forbid replacing function symbols in terms with payloads
                    must_be_symbolic: true,
                    ..config.term_constraints
                },
                ..config
            },
            signature,
            phantom_s: std::marker::PhantomData,
        }
    }
}

impl<S, PT: ProtocolTypes> Mutator<Trace<PT>, S> for ReplaceMatchMutator<S, PT>
where
    S: HasRand,
{
    fn mutate(&mut self, state: &mut S, trace: &mut Trace<PT>) -> Result<MutationResult, Error> {
        log::debug!("[DY] Start mutate with {}", self.name());
        if !self.config.with_dy {
            return Ok(MutationResult::Skipped);
        }
        let rand = state.rand_mut();
        let Some(to_mutate_path) =
            choose_term_path_filtered(trace, |_| true, &self.config.term_constraints, rand)
        else {
            log::debug!("       Skipped {}", self.name());
            return Ok(MutationResult::Skipped);
        };
        let Some(to_mutate) = find_term(trace, &to_mutate_path).cloned() else {
            log::debug!("       Skipped {}", self.name());
            return Ok(MutationResult::Skipped);
        };
        log::debug!("[Mutation] ReplaceMatchMutator on term\n{}", to_mutate);

        // Build the fully-formed replacement, of the same type as the chosen term.
        let new_term = match &to_mutate.term {
            DYTerm::Variable(variable) => {
                let Some((shape, dynamic_fn)) = self.signature.functions.choose_filtered(
                    |(shape, _)| variable.typ == shape.return_type && shape.is_constant(),
                    rand,
                ) else {
                    log::debug!("       Skipped {}", self.name());
                    return Ok(MutationResult::Skipped);
                };
                Term::from(DYTerm::Application(
                    Function::new(shape.clone(), dynamic_fn.clone()),
                    Vec::new(),
                ))
            }
            DYTerm::Application(func, _) => {
                let Some((shape, dynamic_fn)) = self.signature.functions.choose_filtered(
                    |(shape, _)| {
                        func.shape() != shape
                            && func.shape().return_type == shape.return_type
                            && func.shape().argument_types == shape.argument_types
                    },
                    rand,
                ) else {
                    log::debug!("       Skipped {}", self.name());
                    return Ok(MutationResult::Skipped);
                };
                // Only the function symbol changes, the sub-terms are kept.
                let mut new_term = to_mutate.clone();
                if let DYTerm::Application(new_func, _) = &mut new_term.term {
                    new_func.change_function(shape.clone(), dynamic_fn.clone());
                }
                new_term
            }
        };

        Ok(apply_scoped_mutation(
            trace,
            &to_mutate_path,
            &to_mutate,
            &new_term,
            &self.config,
            rand,
        ))
    }

    fn post_exec(&mut self, _state: &mut S, _new_corpus_id: Option<CorpusId>) -> Result<(), Error> {
        Ok(())
    }
}

impl<S, PT: ProtocolTypes> Named for ReplaceMatchMutator<S, PT>
where
    S: HasRand,
{
    fn name(&self) -> &Cow<'static, str> {
        &Cow::Borrowed("ReplaceMatchMutator")
    }
}

/// REPLACE-REUSE: Replaces a sub-term with a different sub-term which is part of the trace
/// (such that types match). The new sub-term could come from another step which has a different
/// recipe term.
pub struct ReplaceReuseMutator<S>
where
    S: HasRand,
{
    config: MutationConfig,
    phantom_s: std::marker::PhantomData<S>,
}

impl<S> ReplaceReuseMutator<S>
where
    S: HasRand,
{
    #[must_use]
    pub const fn new(config: MutationConfig) -> Self {
        Self {
            config,
            phantom_s: std::marker::PhantomData,
        }
    }
}

impl<S, PT: ProtocolTypes> Mutator<Trace<PT>, S> for ReplaceReuseMutator<S>
where
    S: HasRand,
{
    fn mutate(&mut self, state: &mut S, trace: &mut Trace<PT>) -> Result<MutationResult, Error> {
        log::debug!("[DY] Start mutate with {}", self.name());
        if !self.config.with_dy {
            return Ok(MutationResult::Skipped);
        }
        let rand = state.rand_mut();
        if let Some(replacement) = choose_term(trace, &self.config.term_constraints, rand).cloned()
        {
            if let Some(to_replace_path) = choose_term_path_filtered(
                trace,
                |term: &Term<PT>| term.get_type_shape() == replacement.get_type_shape(),
                &self.config.term_constraints,
                rand,
            ) {
                let Some(to_replace) = find_term(trace, &to_replace_path).cloned() else {
                    log::debug!("       Skipped {}", self.name());
                    return Ok(MutationResult::Skipped);
                };
                // the scope is drawn first: a broader scope replaces several occurrences, and each
                // of them contributes to the payload budget checked below
                let targets =
                    scoped_targets(trace, &to_replace_path, &to_replace, &self.config, rand);
                // The per-term payload budget (threshold_max_payloads_per_term) is enforced
                // wholesale in `apply_to_targets`, uniformly across all replacement mutators, and
                // skipped entirely when bit-level mutations are off (`self.with_bit`).
                log::debug!(
                    "[Mutation] Mutate ReplaceReuseMutator on terms\n {} and\n{}",
                    to_replace,
                    replacement
                );
                return Ok(apply_to_targets(
                    trace,
                    &targets,
                    &replacement,
                    &to_replace,
                    self.config.with_bit_level,
                    &self.config.term_constraints,
                ));
            }
        }
        log::debug!("       Skipped {}", self.name());
        Ok(MutationResult::Skipped)
    }

    fn post_exec(&mut self, _state: &mut S, _new_corpus_id: Option<CorpusId>) -> Result<(), Error> {
        Ok(())
    }
}

impl<S> Named for ReplaceReuseMutator<S>
where
    S: HasRand,
{
    fn name(&self) -> &Cow<'static, str> {
        &Cow::Borrowed("ReplaceReuseMutator")
    }
}

/// SKIP:  Removes an input step
pub struct SkipMutator<S>
where
    S: HasRand,
{
    min_trace_length: usize,
    phantom_s: std::marker::PhantomData<S>,
    with_dy: bool,
}

impl<S> SkipMutator<S>
where
    S: HasRand,
{
    #[must_use]
    pub const fn new(min_trace_length: usize, with_dy: bool) -> Self {
        Self {
            min_trace_length,
            phantom_s: std::marker::PhantomData,
            with_dy,
        }
    }
}
impl<S, PT: ProtocolTypes> Mutator<Trace<PT>, S> for SkipMutator<S>
where
    S: HasRand,
{
    fn mutate(&mut self, state: &mut S, trace: &mut Trace<PT>) -> Result<MutationResult, Error> {
        log::debug!("[DY] Start mutate with {}", self.name());
        if !self.with_dy {
            return Ok(MutationResult::Skipped);
        }
        let steps = &mut trace.steps;
        let length = steps.len();
        if length <= self.min_trace_length {
            log::debug!("       Skipped {}", self.name());
            return Ok(MutationResult::Skipped);
        }
        if length == 0 {
            log::debug!("       Skipped {}", self.name());
            return Ok(MutationResult::Skipped);
        }
        let remove_index = state.rand_mut().between(0, length - 1);
        log::debug!("[Mutation] Mutate SkipMutator on step {remove_index}");
        steps.remove(remove_index);
        Ok(MutationResult::Mutated)
    }

    fn post_exec(&mut self, _state: &mut S, _new_corpus_id: Option<CorpusId>) -> Result<(), Error> {
        Ok(())
    }
}
impl<S> Named for SkipMutator<S>
where
    S: HasRand,
{
    fn name(&self) -> &Cow<'static, str> {
        &Cow::Borrowed("SkipMutator")
    }
}

/// REPEAT: Repeats an input which is already part of the trace
pub struct RepeatMutator<S>
where
    S: HasRand,
{
    max_result_trace_length: usize,
    phantom_s: std::marker::PhantomData<S>,
    with_dy: bool,
}

impl<S> RepeatMutator<S>
where
    S: HasRand,
{
    #[must_use]
    pub const fn new(max_result_trace_length: usize, with_dy: bool) -> Self {
        Self {
            max_result_trace_length,
            phantom_s: std::marker::PhantomData,
            with_dy,
        }
    }
}
impl<S, PT: ProtocolTypes> Mutator<Trace<PT>, S> for RepeatMutator<S>
where
    S: HasRand,
{
    fn mutate(&mut self, state: &mut S, trace: &mut Trace<PT>) -> Result<MutationResult, Error> {
        log::debug!("[DY] Start mutate with {}", self.name());
        if !self.with_dy {
            return Ok(MutationResult::Skipped);
        }
        let steps = &trace.steps;
        let length = steps.len();
        if length >= self.max_result_trace_length {
            log::debug!("       Skipped {}", self.name());
            return Ok(MutationResult::Skipped);
        }
        if length == 0 {
            log::debug!("       Skipped {}", self.name());
            return Ok(MutationResult::Skipped);
        }
        let insert_index = state.rand_mut().between(0, length);
        let Some(step) = state.rand_mut().choose(steps) else {
            return Ok(MutationResult::Skipped);
        };
        let step = step.clone();
        log::debug!("[Mutation] Mutate RepeatMutator on step {insert_index}");
        trace.steps.insert(insert_index, step);
        Ok(MutationResult::Mutated)
    }

    fn post_exec(&mut self, _state: &mut S, _new_corpus_id: Option<CorpusId>) -> Result<(), Error> {
        Ok(())
    }
}
impl<S> Named for RepeatMutator<S>
where
    S: HasRand,
{
    fn name(&self) -> &Cow<'static, str> {
        &Cow::Borrowed("RepeatMutator")
    }
}

/// GENERATE: Generates a previously-unseen term using a term zoo
pub struct GenerateMutator<'a, S, PB: ProtocolBehavior>
where
    S: HasRand,
{
    mutation_counter: u64,
    refresh_zoo_after: u64,
    config: MutationConfig,
    zoo: Option<TermZoo<PB>>,
    signature: &'static Signature<PB::ProtocolTypes>,
    put_registry: &'a PutRegistry<PB>,
    phantom_s: std::marker::PhantomData<S>,
}
impl<'a, S, PB: ProtocolBehavior> GenerateMutator<'a, S, PB>
where
    S: HasRand,
{
    #[must_use]
    #[allow(clippy::too_many_arguments)]
    pub const fn new(
        mutation_counter: u64,
        refresh_zoo_after: u64,
        zoo: Option<TermZoo<PB>>,
        signature: &'static Signature<PB::ProtocolTypes>,
        put_registry: &'a PutRegistry<PB>,
        config: MutationConfig,
    ) -> Self {
        Self {
            mutation_counter,
            refresh_zoo_after,
            config,
            zoo,
            signature,
            put_registry,
            phantom_s: std::marker::PhantomData,
        }
    }
}
impl<'a, S, PB: ProtocolBehavior> Mutator<Trace<PB::ProtocolTypes>, S>
    for GenerateMutator<'a, S, PB>
where
    S: HasRand,
{
    fn mutate(
        &mut self,
        state: &mut S,
        trace: &mut Trace<PB::ProtocolTypes>,
    ) -> Result<MutationResult, Error> {
        log::debug!("[DY] Start mutate with {}", self.name());
        if !self.config.with_dy {
            return Ok(MutationResult::Skipped);
        }
        let rand = state.rand_mut();
        let (to_mutate_path, to_mutate, new_term) = {
            if let Some((to_mutate, to_mutate_path)) =
                reservoir_sample(trace, |_| true, &self.config.term_constraints, rand)
            {
                log::debug!("[Mutation] Mutate GenerateMutator on term\n{}", to_mutate);
                self.mutation_counter += 1;
                let zoo = if self.mutation_counter % self.refresh_zoo_after == 0 {
                    log::debug!("[Mutation] Mutate GenerateMutator: refresh zoo");
                    let spawner = Spawner::new(self.put_registry.clone());
                    let ctx = TraceContext::new(spawner); // zoo generate symbolic terms
                    self.zoo.insert(TermZoo::generate(
                        &ctx,
                        self.signature,
                        rand,
                        self.config.term_constraints.zoo_gen_how_many,
                        self.config.term_constraints.zoo_max_depth,
                    ))
                } else {
                    self.zoo.get_or_insert_with(|| {
                        let spawner = Spawner::new(self.put_registry.clone());
                        let ctx = TraceContext::new(spawner); // zoo generate symbolic terms
                        TermZoo::generate(
                            &ctx,
                            self.signature,
                            rand,
                            self.config.term_constraints.zoo_gen_how_many,
                            self.config.term_constraints.zoo_max_depth,
                        )
                    })
                };
                if let Some(new_term) = zoo.choose_filtered(
                    |term| {
                        // We seek for a different term with a matching type
                        to_mutate.get_type_shape() == term.get_type_shape() && *to_mutate != **term
                    },
                    rand,
                ) {
                    log::debug!(
                        "Found to_mutate and new_term: {}\n ----------------\n{}",
                        to_mutate,
                        new_term
                    );
                    // clone so that the immutable borrow of `trace` ends here
                    (to_mutate_path, to_mutate.clone(), new_term.clone())
                } else {
                    return Ok(MutationResult::Skipped);
                }
            } else {
                return Ok(MutationResult::Skipped);
            }
        };

        // Apply the replacement at a randomly drawn scope (global / step / individual)
        Ok(apply_scoped_mutation(
            trace,
            &to_mutate_path,
            &to_mutate,
            &new_term,
            &self.config,
            rand,
        ))
    }

    fn post_exec(&mut self, _state: &mut S, _new_corpus_id: Option<CorpusId>) -> Result<(), Error> {
        Ok(())
    }
}

impl<'a, S, PB: ProtocolBehavior> Named for GenerateMutator<'a, S, PB>
where
    S: HasRand,
{
    fn name(&self) -> &Cow<'static, str> {
        &Cow::Borrowed("GenerateMutator")
    }
}

#[cfg(test)]
mod tests {
    use std::collections::{HashMap, HashSet};

    use libafl::corpus::InMemoryCorpus;
    use libafl::mutators::{MutationResult, Mutator};
    use libafl::state::StdState;
    use libafl_bolts::rands::{RomuDuoJrRand, StdRand};

    use super::*;
    use crate::agent::AgentName;
    use crate::algebra::dynamic_function::DescribableFunction;
    use crate::algebra::test_signature::{TestTrace, *};
    use crate::algebra::DYTerm;
    use crate::fuzzer::utils::{choose_term_path, TracePath};
    use crate::term;
    use crate::trace::{Action, InputAction, Step};

    fn create_state(
    ) -> StdState<InMemoryCorpus<TestTrace>, TestTrace, RomuDuoJrRand, InMemoryCorpus<TestTrace>>
    {
        let rand = StdRand::with_seed(1235);
        let corpus: InMemoryCorpus<TestTrace> = InMemoryCorpus::new();
        StdState::new(rand, corpus, InMemoryCorpus::new(), &mut (), &mut ()).unwrap()
    }

    /// Checks whether repeat can repeat the last step
    #[test_log::test]
    fn test_repeat_mutator() {
        let mut state = create_state();

        let mut mutator = RepeatMutator::new(15, true);

        fn check_is_encrypt12(step: &Step<TestProtocolTypes>) -> bool {
            if let Action::Input(input) = &step.action {
                if input.recipe.name() == fn_encrypt12.name() {
                    return true;
                }
            }
            false
        }

        loop {
            let mut trace = setup_simple_trace();
            mutator.mutate(&mut state, &mut trace).unwrap();

            let length = trace.steps.len();
            if let Some(last) = trace.steps.get(length - 1) {
                if check_is_encrypt12(last) {
                    if let Some(step) = trace.steps.get(length - 2) {
                        if check_is_encrypt12(step) {
                            break;
                        }
                    }
                }
            }
        }
    }

    /// The defining behaviour of each scope, on a trace that contains the *same* sub-term several
    /// times inside one step and also in another step:
    /// - `Individual` replaces exactly the chosen occurrence,
    /// - `Step` replaces every occurrence of the chosen step, and only those,
    /// - `Global` replaces every occurrence of the whole trace.
    #[test_log::test]
    fn test_scope_extent() {
        let mut rand = StdRand::with_seed(11);

        // a term that appears 3 times in step 0 and 2 times in step 1
        let duplicated: Term<TestProtocolTypes> = term! { fn_signature_algorithm_extension };
        let recipe_a: Term<TestProtocolTypes> = term! {
            fn_client_extensions_append(
                (fn_client_extensions_append(
                    (fn_client_extensions_append(
                        fn_client_extensions_new,
                        fn_signature_algorithm_extension
                    )),
                    fn_signature_algorithm_extension
                )),
                fn_signature_algorithm_extension
            )
        };
        let recipe_b: Term<TestProtocolTypes> = term! {
            fn_client_extensions_append(
                (fn_client_extensions_append(
                    fn_client_extensions_new,
                    fn_signature_algorithm_extension
                )),
                fn_signature_algorithm_extension
            )
        };
        let build_trace = || Trace {
            steps: vec![
                Step {
                    agent: AgentName::first(),
                    action: Action::Input(InputAction {
                        recipe: recipe_a.clone(),
                        precomputations: vec![],
                    }),
                },
                Step {
                    agent: AgentName::first(),
                    action: Action::Input(InputAction {
                        recipe: recipe_b.clone(),
                        precomputations: vec![],
                    }),
                },
            ],
            ..setup_simple_trace()
        };
        let replacement: Term<TestProtocolTypes> = term! { fn_ec_point_formats_extension };

        // count the occurrences of `duplicated` left in each step
        let remaining = |trace: &Trace<TestProtocolTypes>| {
            (0..2)
                .map(|step| match &trace.steps[step].action {
                    Action::Input(input) => (&input.recipe)
                        .into_iter()
                        .filter(|t| **t == duplicated)
                        .count(),
                    Action::Output(_) => 0,
                })
                .collect::<Vec<_>>()
        };

        let chosen_path: TracePath = (0, vec![1]); // outermost duplicated term of step 0
        let sanity = build_trace();
        assert_eq!(remaining(&sanity), vec![3, 2], "test fixture");

        for (weights, expected, label) in [
            (ScopeWeights::new(0, 0, 1), vec![2, 2], "individual"),
            (ScopeWeights::new(0, 1, 0), vec![0, 2], "step"),
            (ScopeWeights::new(1, 0, 0), vec![0, 0], "global"),
        ] {
            let mut trace = build_trace();
            let targets = scoped_targets(
                &trace,
                &chosen_path,
                &duplicated,
                &MutationConfig {
                    scope_weights: weights,
                    ..MutationConfig::default()
                },
                &mut rand,
            );
            assert_eq!(
                apply_to_targets(
                    &mut trace,
                    &targets,
                    &replacement,
                    &duplicated,
                    true,
                    &TermConstraints::no_constraint(),
                ),
                MutationResult::Mutated,
                "{label}: should have mutated"
            );
            assert_eq!(
                remaining(&trace),
                expected,
                "{label}: unexpected extent of the replacement (per-step occurrences left)"
            );
        }
    }

    /// The payload budget (`threshold_max_payloads_per_term`) is enforced **per term** (per step
    /// recipe), not as a trace-wide average: concentrating the payload growth of a broad scope in a
    /// single recipe is rejected wholesale, even when the same growth spread over the whole trace
    /// would stay under budget on average. Payload-neutral or shrinking replacements are always
    /// accepted. Exercises [`payloads_within_caps`] directly (the guard folded into
    /// [`apply_to_targets`]).
    #[test_log::test]
    fn test_payload_budget_is_per_term_not_average() {
        let trace = setup_simple_trace();
        let step = trace
            .steps
            .iter()
            .position(|s| matches!(s.action, Action::Input(_)))
            .expect("the fixture has at least one input step");
        // baseline payloads already in that recipe (0 for a symbolic seed, but read it to stay
        // robust if the fixture changes)
        let base = match &trace.steps[step].action {
            Action::Input(input) => input.recipe.count_payloads(),
            Action::Output(_) => unreachable!(),
        };

        // one growing replacement (+1 payload) for each of four occurrences, all in ONE recipe
        let concentrated: Vec<TracePath> = (0..4).map(|i| (step, vec![i])).collect();
        let single: Vec<TracePath> = vec![(step, vec![0])];
        // budget leaves room for +3 payloads in a single recipe
        let constraints = TermConstraints {
            threshold_max_payloads_per_term: base + 3,
            ..TermConstraints::no_constraint()
        };

        // per-term: this recipe would reach base + 4 > base + 3 -> reject the whole mutation. A
        // trace-wide average (4 payloads over several steps) would have stayed under budget; that
        // dilution is exactly what this check no longer allows.
        assert!(
            !payloads_within_caps(&trace, &concentrated, base, base + 1, &constraints),
            "four +1-payload replacements in one recipe must exceed the per-term budget"
        );
        // a single occurrence only reaches base + 1 <= base + 3 -> accepted
        assert!(
            payloads_within_caps(&trace, &single, base, base + 1, &constraints),
            "a single growing replacement stays within the per-term budget"
        );
        // payload-neutral and shrinking replacements are accepted regardless of scope width
        assert!(payloads_within_caps(
            &trace,
            &concentrated,
            7,
            7,
            &constraints
        ));
        assert!(payloads_within_caps(
            &trace,
            &concentrated,
            9,
            2,
            &constraints
        ));
    }

    /// `Term::count_payloads` is the alloc-free equivalent of `all_payloads().len()`: it must agree
    /// with it node-for-node, both in the all-symbolic case (0) and once a payload is attached.
    #[test_log::test]
    fn test_count_payloads_matches_all_payloads_len() {
        let trace = setup_simple_trace();
        for step in &trace.steps {
            let Action::Input(input) = &step.action else {
                continue;
            };
            // every sub-term of the recipe: the two traversals must return the same count
            for sub in &input.recipe {
                assert_eq!(sub.count_payloads(), sub.all_payloads().len());
            }
            // and once a payload is attached at the root, both see exactly one
            let mut with_payload = input.recipe.clone();
            with_payload.add_payload(vec![0u8; 4]);
            assert_eq!(
                with_payload.count_payloads(),
                with_payload.all_payloads().len()
            );
            assert_eq!(with_payload.count_payloads(), 1);
        }
    }

    /// The result-size caps are reject-whole: a *growing* replacement that would push either the
    /// affected step recipe over `max_result_term_size` is refused. The size-neutral / shrinking
    /// fast path is accepted in
    /// O(1) even when the input already exceeds the caps — that is the documented bounded-input
    /// precondition of [`replacement_within_caps`] (a cap set below an imported seed is a
    /// corpus-boundary misconfiguration, not the mutator's job to repair).
    #[test_log::test]
    fn test_replacement_within_caps_rejects_growth() {
        let trace = setup_simple_trace();
        let step = trace
            .steps
            .iter()
            .position(|s| matches!(s.action, Action::Input(_)))
            .expect("the fixture has at least one input step");
        let cur = match &trace.steps[step].action {
            Action::Input(input) => input.recipe.size(),
            Action::Output(_) => unreachable!(),
        };
        let targets: Vec<TracePath> = vec![(step, vec![0])]; // one target in this step

        // representative_size/new_size only fix the delta (+5 here); `cur` is read from the recipe.
        let (repr, new) = (1, 6);

        // generous caps: +5 fits in both -> accepted
        let generous = TermConstraints {
            max_result_term_size: cur + 100,
            ..TermConstraints::no_constraint()
        };
        assert!(replacement_within_caps(
            &trace, &targets, repr, new, &generous
        ));

        // per-step cap allows only +2 in this recipe -> +5 rejected
        let tight_step = TermConstraints {
            max_result_term_size: cur + 2,
            ..TermConstraints::no_constraint()
        };
        assert!(!replacement_within_caps(
            &trace,
            &targets,
            repr,
            new,
            &tight_step
        ));

        // NOTE: we no longer enforce a whole-trace node cap here — growth is bounded by the
        // per-step recipe cap times the number of steps (`max_result_trace_length`).

        // documented precondition: with the caps set *below* the already-imported seed, a
        // size-neutral or shrinking replacement is still accepted (the input is out of bounds
        // before any mutation; the mutator does not retroactively enforce the cap).
        let below_seed = TermConstraints {
            max_result_term_size: 0,
            ..TermConstraints::no_constraint()
        };
        assert!(replacement_within_caps(&trace, &targets, 4, 4, &below_seed)); // neutral
        assert!(replacement_within_caps(&trace, &targets, 9, 2, &below_seed)); // shrink
    }

    /// Whatever the scope, the *chosen* occurrence is always replaced. In particular a broader
    /// scope must not silently skip the mutation when the search for the other occurrences finds
    /// nothing.
    #[test_log::test]
    fn test_scoped_mutation_always_mutates_chosen_occurrence() {
        let mut rand = StdRand::with_seed(7);

        for weights in [
            ScopeWeights::new(1, 0, 0), // global only
            ScopeWeights::new(0, 1, 0), // step only
            ScopeWeights::new(0, 0, 1), // individual only
            ScopeWeights::default(),
        ] {
            let mut trace = setup_simple_trace();
            let path =
                choose_term_path_filtered(&trace, |_| true, &TermConstraints::default(), &mut rand)
                    .unwrap();
            let chosen = find_term(&trace, &path).unwrap().clone();

            // any different term of the same type is a valid replacement here
            let Some(replacement) = trace
                .steps
                .iter()
                .filter_map(|step| match &step.action {
                    Action::Input(input) => Some(&input.recipe),
                    Action::Output(_) => None,
                })
                .flat_map(|recipe| recipe.into_iter())
                .find(|term| term.get_type_shape() == chosen.get_type_shape() && **term != chosen)
                .cloned()
            else {
                continue; // no valid replacement in this trace, nothing to assert
            };

            let targets = scoped_targets(
                &trace,
                &path,
                &chosen,
                &MutationConfig {
                    scope_weights: weights,
                    ..MutationConfig::default()
                },
                &mut rand,
            );
            let result = apply_to_targets(
                &mut trace,
                &targets,
                &replacement,
                &chosen,
                true,
                &TermConstraints::no_constraint(),
            );

            assert_eq!(
                result,
                MutationResult::Mutated,
                "weights {weights:?} should have mutated"
            );
            assert_eq!(
                find_term(&trace, &path).unwrap(),
                &replacement,
                "weights {weights:?}: the chosen occurrence must carry the replacement"
            );
        }
    }

    /// The scope drawn by [`MutationScope::choose`] must follow the configured weights: a zeroed
    /// weight is never drawn, and `(0, 0, 0)` degrades gracefully to `Individual`.
    #[test_log::test]
    fn test_scope_weights() {
        let mut rand = StdRand::with_seed(42);

        let draw = |w: ScopeWeights, rand: &mut _| {
            let mut seen = std::collections::HashSet::new();
            for _ in 0..200 {
                seen.insert(MutationScope::choose(w, rand));
            }
            seen
        };

        // uniform: all three scopes must show up
        let all = draw(ScopeWeights::default(), &mut rand);
        assert_eq!(all.len(), 3, "uniform weights should draw all three scopes");

        // (1, 0, 1): the historical behaviour, never per-step
        let no_step = draw(ScopeWeights::new(1, 0, 1), &mut rand);
        assert!(!no_step.contains(&MutationScope::Step));
        assert!(no_step.contains(&MutationScope::Global));
        assert!(no_step.contains(&MutationScope::Individual));

        // (0, 0, 1): purely individual replacements
        let only_individual = draw(ScopeWeights::new(0, 0, 1), &mut rand);
        assert_eq!(only_individual, HashSet::from([MutationScope::Individual]));

        // degenerate weights fall back to Individual instead of panicking
        let zeroed = draw(ScopeWeights::new(0, 0, 0), &mut rand);
        assert_eq!(zeroed, HashSet::from([MutationScope::Individual]));
    }

    #[test_log::test]
    fn test_replace_match_mutator() {
        let _server = AgentName::first();
        let mut state = create_state();
        let mut mutator =
            ReplaceMatchMutator::new(MutationConfig::default_with_bit(), &TEST_SIGNATURE);

        loop {
            let mut trace = setup_simple_trace();
            mutator.mutate(&mut state, &mut trace).unwrap();

            if let Some(last) = trace.steps.iter().last() {
                match &last.action {
                    Action::Input(input) => match &input.recipe.term {
                        DYTerm::Variable(_) => {}
                        DYTerm::Application(_, subterms) => {
                            if let Some(last_subterm) = subterms.iter().last() {
                                if last_subterm.name() == fn_seq_1.name() {
                                    break;
                                }
                            }
                        }
                    },
                    Action::Output(_) => {}
                }
            }
        }
    }

    #[test_log::test]
    fn test_remove_lift_mutator() {
        // Should remove an extension
        let mut state = create_state();
        let _server = AgentName::first();
        let mut mutator = RemoveAndLiftMutator::new(TermConstraints::default(), true);

        // Returns the amount of extensions in the trace
        fn sum_extension_appends(trace: &TestTrace) -> usize {
            trace.count_functions_by_name(fn_client_extensions_append.name())
        }

        loop {
            let mut trace = setup_simple_trace();
            let before_mutation = sum_extension_appends(&trace);
            let result = mutator.mutate(&mut state, &mut trace).unwrap();

            if result == MutationResult::Mutated {
                let after_mutation = sum_extension_appends(&trace);
                if after_mutation < before_mutation {
                    // extension removed
                    break;
                }
            }
        }
    }

    #[test_log::test]
    fn test_replace_reuse_mutator() {
        let mut state = create_state();
        let _server = AgentName::first();
        let mut mutator = ReplaceReuseMutator::new(MutationConfig::default_with_bit());

        fn count_client_hello(trace: &TestTrace) -> usize {
            trace.count_functions_by_name(fn_client_hello.name())
        }

        fn count_finished(trace: &TestTrace) -> usize {
            trace.count_functions_by_name(fn_finished.name())
        }

        loop {
            let mut trace = setup_simple_trace();
            let result = mutator.mutate(&mut state, &mut trace).unwrap();

            if result == MutationResult::Mutated {
                let client_hellos = count_client_hello(&trace);
                let finishes = count_finished(&trace);
                if client_hellos == 2 && finishes == 0 {
                    // finished replaced by client_hello
                    break;
                }
            }
        }
    }

    #[test_log::test]
    fn test_skip_mutator() {
        let mut state = create_state();
        let _server = AgentName::first();
        let mut mutator = SkipMutator::new(2, true);

        loop {
            let mut trace = setup_simple_trace();
            let before_len = trace.steps.len();
            mutator.mutate(&mut state, &mut trace).unwrap();

            if before_len - 1 == trace.steps.len() {
                break;
            }
        }
    }

    #[test_log::test]
    fn test_swap_mutator() {
        let mut state = create_state();
        let mut mutator = SwapMutator::new(TermConstraints::default(), true);

        loop {
            let mut trace = setup_simple_trace();
            mutator.mutate(&mut state, &mut trace).unwrap();

            let is_first_not_ch = if let Some(first) = trace.steps.first() {
                match &first.action {
                    Action::Input(input) => Some(input.recipe.name() != fn_client_hello.name()),
                    Action::Output(_) => None,
                }
            } else {
                None
            };

            let is_next_not_fn_client_key_exchange = if let Some(next) = trace.steps.get(1) {
                match &next.action {
                    Action::Input(input) => {
                        Some(input.recipe.name() != fn_client_key_exchange.name())
                    }
                    Action::Output(_) => None,
                }
            } else {
                None
            };

            if let Some(first) = is_first_not_ch {
                if let Some(second) = is_next_not_fn_client_key_exchange {
                    if first && second {
                        break;
                    }
                }
            }
        }
    }

    #[test_log::test]
    fn test_find_term() {
        let mut rand = StdRand::with_seed(45);
        let mut trace = setup_simple_trace();
        let term_size = trace.count_functions();

        let mut stats: HashSet<TracePath> = HashSet::new();

        for _ in 0..10000 {
            let path = choose_term_path(&trace, &TermConstraints::default(), &mut rand).unwrap();
            find_term_mut(&mut trace, &path).unwrap();
            stats.insert(path);
        }

        assert_eq!(term_size, stats.len());
    }

    #[test_log::test]
    fn test_reservoir_sample_randomness() {
        /// https://rust-lang-nursery.github.io/rust-cookbook/science/mathematics/statistics.html#standard-deviation
        fn std_deviation(data: &[u32]) -> Option<f32> {
            fn mean(data: &[u32]) -> Option<f32> {
                let sum = data.iter().sum::<u32>() as f32;
                let count = data.len();

                match count {
                    positive if positive > 0 => Some(sum / count as f32),
                    _ => None,
                }
            }

            match (mean(data), data.len()) {
                (Some(data_mean), count) if count > 0 => {
                    let variance = data
                        .iter()
                        .map(|value| {
                            let diff = data_mean - (*value as f32);

                            diff * diff
                        })
                        .sum::<f32>()
                        / count as f32;

                    Some(variance.sqrt())
                }
                _ => None,
            }
        }

        let trace = setup_simple_trace();
        let term_size = trace.count_functions();

        let mut rand = StdRand::with_seed(45);
        let mut stats: HashMap<u32, u32> = HashMap::new();

        for _ in 0..10000 {
            let term = choose(&trace, &TermConstraints::default(), &mut rand).unwrap();

            let id = term.0.resistant_id();

            let count: u32 = *stats.get(&id).unwrap_or(&0);
            stats.insert(id, count + 1);
        }

        let std_dev =
            std_deviation(stats.values().cloned().collect::<Vec<u32>>().as_slice()).unwrap();
        /*        println!("{:?}", std_dev);
        println!("{:?}", stats);*/

        assert!(std_dev < 30.0);
        assert_eq!(term_size, stats.len());
    }
}