rustsim_core/

scheduler.rs

//! Scheduler trait and built-in implementations.
//!
//! Schedulers control the order in which agents are activated each step.
//! The built-in schedulers mirror Julia Agents.jl `Schedulers.*`:
//!
//! | Rust type | Julia equivalent | Behavior |
//! |-----------|-----------------|----------|
//! | [`Fastest`] | `Schedulers.fastest` | Iteration order (no sorting) |
//! | [`ById`] | `Schedulers.ByID` | Sorted ascending by agent ID |
//! | [`Randomly`] | `Schedulers.Randomly` | Shuffled randomly each step |
//! | [`ByProperty`] | `Schedulers.ByProperty` | Sorted by a user-defined key (greatest first) |
//! | [`PartiallyRandom`] | `Schedulers.Partially` | Random subset with a given probability |
//!
//! # Determinism notes
//!
//! Deterministic replay depends on both the scheduler and the order of agent IDs
//! collected from the store.
//!
//! - [`ById`] imposes an explicit deterministic order.
//! - [`Fastest`] is deterministic only if the underlying store iteration order is deterministic.
//! - [`Randomly`] and [`PartiallyRandom`] are reproducible for a fixed RNG seed only when the
//!   pre-randomization ID collection order is itself deterministic.
//! - [`ByProperty`] is reproducible when the selected key is deterministic and either ties do not
//!   occur or the source ID order is deterministic.

use crate::types::AgentId;
use rand::seq::SliceRandom;
use rand::Rng;

use crate::{model::Model, standard::HasAgentIds};

/// Trait for agent activation schedulers.
///
/// Each step, the model calls [`schedule_into`] to populate a buffer with
/// the agent IDs to activate, in the desired order. The buffer is cleared
/// by the caller before being passed in.
///
/// Schedulers are stateless - the buffer is owned by the model and reused
/// across steps to avoid per-step allocations.
///
/// [`schedule_into`]: Scheduler::schedule_into
pub trait Scheduler<M> {
    /// Fill `buf` with the agent IDs to activate this step, in the desired order.
    ///
    /// The buffer is cleared by the caller before this method is invoked.
    fn schedule_into(&mut self, model: &M, buf: &mut Vec<AgentId>);
}

/// Fastest scheduler - returns agents in store iteration order.
///
/// No sorting or shuffling is performed. This is the cheapest scheduler
/// and is appropriate when activation order does not matter.
///
/// Reproducibility inherits whatever ordering guarantees the underlying store
/// provides.
#[derive(Debug, Default)]
pub struct Fastest;

impl Fastest {
    pub const fn new() -> Self {
        Self
    }
}

impl<M> Scheduler<M> for Fastest
where
    M: HasAgentIds,
{
    fn schedule_into(&mut self, model: &M, buf: &mut Vec<AgentId>) {
        model.agent_ids_into(buf);
    }
}

/// ID-ordered scheduler - returns agents sorted ascending by [`AgentId`].
///
/// Produces a deterministic, reproducible activation order regardless of
/// store internals.
#[derive(Debug, Default)]
pub struct ById;

impl ById {
    pub fn new() -> Self {
        Self
    }
}

impl<M> Scheduler<M> for ById
where
    M: HasAgentIds,
{
    fn schedule_into(&mut self, model: &M, buf: &mut Vec<AgentId>) {
        model.agent_ids_into(buf);
        buf.sort_unstable();
    }
}

/// Random scheduler - shuffles agent IDs each step using the model's RNG.
///
/// Produces a different activation order every step while remaining
/// reproducible for a given RNG seed **when the collected source ID order is
/// deterministic**.
#[derive(Debug, Default)]
pub struct Randomly;

impl Randomly {
    pub fn new() -> Self {
        Self
    }
}

impl<M> Scheduler<M> for Randomly
where
    M: HasAgentIds + Model,
{
    fn schedule_into(&mut self, model: &M, buf: &mut Vec<AgentId>) {
        model.agent_ids_into(buf);
        buf.shuffle(&mut *model.rng_mut());
    }
}

/// Partial random scheduler - activates a random subset of agents each step.
///
/// Each agent is independently included with probability `p`. Useful for
/// modelling intermittent activity or sampling.
///
/// Reproducibility for a fixed RNG seed still depends on deterministic source
/// ID enumeration before filtering.
#[derive(Debug)]
pub struct PartiallyRandom {
    probability: f64,
}

impl PartiallyRandom {
    /// Create a new `PartiallyRandom` scheduler with the given inclusion probability.
    ///
    /// # Panics
    ///
    /// Does not validate `probability` - values outside `[0, 1]` will produce
    /// unexpected results from [`rand::Rng::gen_bool`].
    pub fn new(probability: f64) -> Self {
        Self { probability }
    }
}

impl<M> Scheduler<M> for PartiallyRandom
where
    M: HasAgentIds + Model,
{
    fn schedule_into(&mut self, model: &M, buf: &mut Vec<AgentId>) {
        model.agent_ids_into(buf);
        let mut rng = model.rng_mut();
        buf.retain(|_| rng.gen_bool(self.probability));
    }
}

/// Property-sorted scheduler - activates agents sorted by a user-defined key.
///
/// The `selector` closure extracts an [`Ord`]-comparable key from each agent.
/// Agents with the **greatest** key value are activated first (descending order).
///
/// When multiple agents produce the same key, reproducibility falls back to the
/// relative order already present in the collected ID buffer.
///
/// # Example
///
/// ```ignore
/// // Activate highest-energy agents first
/// let sched = ByProperty::new(|a: &MyAgent| a.energy);
/// ```
#[derive(Debug)]
pub struct ByProperty<F> {
    selector: F,
}

impl<F> ByProperty<F> {
    /// Create a new `ByProperty` scheduler with the given key selector.
    pub fn new(selector: F) -> Self {
        Self { selector }
    }
}

impl<M, F, K> Scheduler<M> for ByProperty<F>
where
    M: HasAgentIds + Model,
    F: Fn(&M::Agent) -> K,
    K: Ord,
{
    fn schedule_into(&mut self, model: &M, buf: &mut Vec<AgentId>) {
        model.agent_ids_into(buf);

        // Drop any IDs that vanished since collection.
        buf.retain(|id| model.agent(*id).is_some());

        let selector = &self.selector;
        buf.sort_by(|a, b| {
            let agent_a = model.agent(*a);
            let agent_b = model.agent(*b);
            match (agent_a, agent_b) {
                (Some(agent_a), Some(agent_b)) => {
                    let ka = selector(&agent_a);
                    let kb = selector(&agent_b);
                    kb.cmp(&ka)
                }
                (Some(_), None) => std::cmp::Ordering::Less,
                (None, Some(_)) => std::cmp::Ordering::Greater,
                (None, None) => std::cmp::Ordering::Equal,
            }
        });
    }
}