henad-compute 0.3.0

//! The engine that runs an [`AgentModel`] as a [`SimState`], and the parameters it prepends to the model's own
//! parameters.

use henad_core::action::action_seed;
use henad_core::authoring::model::agent_model::{
    AgentLanes as _, AgentModel, ChunkTally as _, NeighborIndex as _, StepCtx,
};
use henad_core::authoring::model::field::{Extent, FieldLayer};
use henad_core::helpers::{extract_f32, extract_u32, f32_param, u32_param};
use henad_core::model::SimState;
use henad_core::params::{ParamDescriptor, ParamStore, ParamValue};
use henad_core::view::{GridView, PointView, StatEntry, stat_entries};

/// State an agent model's random number generator (RNG) starts from when a run has no seed.
pub const AGENT_INIT_SEED: u64 = 0xA175_F01A_6ED5_0001;

/// Returns the state an agent model's RNG starts from: `seed` mixed, or [`AGENT_INIT_SEED`] when it is `None`.
///
/// A GPU port that reproduces its CPU model's tick 0 starts from the same state.
pub fn agent_init_rng(seed: Option<u64>) -> u64 {
    seed.map_or(AGENT_INIT_SEED, henad_core::authoring::primitives::rng::mix_seed)
}

/// Engine wrapper that implements `SimState` for any `AgentModel`.
pub struct AgentModelState<A: AgentModel> {
    lanes: A::Lanes,
    field: A::Field,
    index: A::Index,
    deposits: <A::Field as FieldLayer>::DepositLanes,
    params: ParamStore,
    /// Number of the model's own params, read once from [`AgentModel::param_descriptors`].
    own_params: usize,
    extent: Extent,
    tally: A::Tally,
    seed: u64,
    /// RNG stream for actions, separate from the stream the ticks draw from.
    action_seed: u64,
    tick: u64,
}

impl<A: AgentModel> std::fmt::Debug for AgentModelState<A> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AgentModelState")
            .field("model", &A::ID)
            .field("tick", &self.tick)
            .field("extent", &self.extent)
            .field("params", &self.params)
            .finish_non_exhaustive()
    }
}

impl<A: AgentModel> AgentModelState<A> {
    /// Builds a state from the full parameter list, with its RNG starting from [`AGENT_INIT_SEED`].
    pub fn from_params(params: &[ParamValue]) -> Self {
        Self::from_params_seeded(params, None)
    }

    /// Builds a state whose RNG starts from `seed`, or [`AGENT_INIT_SEED`] when it is `None`.
    pub fn from_params_seeded(params: &[ParamValue], seed: Option<u64>) -> Self {
        let n = extract_u32(params, NUM_AGENTS, 10_000) as usize;
        let extent = Extent {
            w: extract_f32(params, WORLD_WIDTH, 1_000.0),
            h: extract_f32(params, WORLD_HEIGHT, 1_000.0),
        };

        let (own, field_params) = split_params::<A>(params);
        let mut lanes = A::Lanes::alloc(n);
        let actions = action_seed(seed);
        let mut seed = agent_init_rng(seed);
        A::init(&mut lanes, extent, own, &mut seed);

        let field = A::Field::new(extent, field_params);
        let deposits = field.alloc_deposits(n);
        let hot = A::from_params(own, extent);
        let (pos_x, pos_y) = lanes.positions();
        let mut index = A::Index::new(extent, A::index_cell_size(&hot));
        index.rebuild(pos_x, pos_y, A::index_cell_size(&hot));

        Self {
            lanes,
            field,
            index,
            deposits,
            params: ParamStore::new(&agent_model_param_descriptors::<A>(), params),
            own_params: A::param_descriptors().len(),
            extent,
            tally: A::Tally::default(),
            seed,
            action_seed: actions,
            tick: 0,
        }
    }

    /// Builds a state whose agents come from `seed_lanes`, for reproducing a particular run.
    pub fn from_agents(params: &[ParamValue], seed_lanes: impl FnOnce(&mut A::Lanes, Extent)) -> Self {
        Self::from_agents_and_field(params, None, seed_lanes, |_field, _extent| {})
    }

    /// Builds a state as [`Self::from_agents`] does, with the field layer seeded too and the RNG stream chosen.
    ///
    /// A model whose agents read a field needs both halves fixed before its kernel means anything.
    /// `seed_field` receives the concrete field type, so it reaches whatever that layer offers. With
    /// `seed` a test can show a scenario reaches the same answer from any stream.
    pub fn from_agents_and_field(
        params: &[ParamValue],
        seed: Option<u64>,
        seed_lanes: impl FnOnce(&mut A::Lanes, Extent),
        seed_field: impl FnOnce(&mut A::Field, Extent),
    ) -> Self {
        let n = extract_u32(params, NUM_AGENTS, 10_000) as usize;
        let extent = Extent {
            w: extract_f32(params, WORLD_WIDTH, 1_000.0),
            h: extract_f32(params, WORLD_HEIGHT, 1_000.0),
        };

        let mut lanes = A::Lanes::alloc(n);
        let actions = action_seed(seed);
        let mut seed = agent_init_rng(seed);
        let (own, field_params) = split_params::<A>(params);

        // Init still runs, to advance the seed to the value the first step starts from.
        A::init(&mut lanes, extent, own, &mut seed);

        seed_lanes(&mut lanes, extent);

        let mut field = A::Field::new(extent, field_params);
        seed_field(&mut field, extent);
        let deposits = field.alloc_deposits(n);
        let hot = A::from_params(own, extent);
        let (pos_x, pos_y) = lanes.positions();
        let mut index = A::Index::new(extent, A::index_cell_size(&hot));
        index.rebuild(pos_x, pos_y, A::index_cell_size(&hot));

        Self {
            lanes,
            field,
            index,
            deposits,
            params: ParamStore::new(&agent_model_param_descriptors::<A>(), params),
            own_params: A::param_descriptors().len(),
            extent,
            tally: A::Tally::default(),
            seed,
            action_seed: actions,
            tick: 0,
        }
    }

    /// Agent lanes.
    pub fn lanes(&self) -> &A::Lanes {
        &self.lanes
    }

    /// Field layer under the agents.
    pub fn field(&self) -> &A::Field {
        &self.field
    }

    /// Tally merged over every tick so far.
    pub fn tally(&self) -> &A::Tally {
        &self.tally
    }
}

// Indices of the params the engine prepends before a model's own params. A GPU port reads them too,
// since it composes the same list.
/// Index of the population.
pub const NUM_AGENTS: usize = 0;
/// Index of the world's width.
pub const WORLD_WIDTH: usize = 1;
/// Index of the world's height.
pub const WORLD_HEIGHT: usize = 2;

/// Number of params the engine prepends, and the index of a model's first own param.
pub const AGENT_PARAM_BASE: usize = 3;

/// Splits a composed parameter list into the model's own parameters and its field layer's parameters.
///
/// The split comes from the descriptor lists. Adding a parameter to the model or to the field layer cannot shift the
/// indices of the other list.
pub fn split_params<A: AgentModel>(params: &[ParamValue]) -> (&[ParamValue], &[ParamValue]) {
    split_at_own(params, A::param_descriptors().len())
}

/// Splits a composed list as [`split_params`] does, given the model's own param count.
fn split_at_own(params: &[ParamValue], own: usize) -> (&[ParamValue], &[ParamValue]) {
    let start = AGENT_PARAM_BASE.min(params.len());
    let mid = (start + own).min(params.len());
    (&params[start..mid], &params[mid..])
}

/// Returns the full descriptor list: `num_agents`, `world_width` and `world_height`, then the model's own params,
/// then its field layer's params.
pub fn agent_model_param_descriptors<A: AgentModel>() -> Vec<ParamDescriptor> {
    let extent = A::DEFAULT_EXTENT;
    let mut descs = vec![
        u32_param("num_agents", "Number of Agents", A::DEFAULT_AGENTS, 1, A::MAX_AGENTS).on_reload(),
        f32_param("world_width", "World Width", extent.w, 1.0, 10_000.0, Some(50.0)).on_reload(),
        f32_param("world_height", "World Height", extent.h, 1.0, 10_000.0, Some(50.0)).on_reload(),
    ];
    descs.extend(A::param_descriptors());
    descs.extend(<A::Field as FieldLayer>::param_descriptors());
    descs
}

impl<A: AgentModel> SimState for AgentModelState<A> {
    fn step(&mut self) {
        let (own, field_slice) = split_at_own(self.params.values(), self.own_params);
        let hot = A::from_params(own, self.extent);
        let field_params = <A::Field as FieldLayer>::from_params(field_slice);

        let (pos_x, pos_y) = self.lanes.positions();
        self.index.rebuild(pos_x, pos_y, A::index_cell_size(&hot));

        {
            let ctx = StepCtx::<A> {
                field: self.field.read(),
                index: &self.index,
                params: &hot,
                extent: self.extent,
            };
            A::run_deposit_pass(&self.lanes, &mut self.deposits, &ctx);
        }

        let tallied = {
            let ctx = StepCtx::<A> {
                field: self.field.read(),
                index: &self.index,
                params: &hot,
                extent: self.extent,
            };
            A::run_step_pass(&mut self.lanes, &ctx, self.seed, self.tick)
        };
        self.tally = std::mem::take(&mut self.tally).merge(tallied);

        self.field.update(&self.deposits, &field_params, self.tick);
        self.lanes.swap();
        self.seed = crate::cpu::primitives::chunked::advance_tick_seed(self.seed, self.tick);
        self.tick += 1;
    }

    fn tick(&self) -> u64 {
        self.tick
    }

    fn grid_view(&self) -> Option<GridView<'_>> {
        self.field.grid_view()
    }

    fn point_view(&self) -> Option<PointView<'_>> {
        let (pos_x, pos_y) = self.lanes.positions();
        Some(PointView {
            pos_x,
            pos_y,
            world_w: self.extent.w,
            world_h: self.extent.h,
            color: self.lanes.colors(),
            palette: A::PALETTE,
        })
    }

    fn prepare_view(&mut self) {
        self.field.prepare_view();
    }

    fn stats(&self) -> Vec<StatEntry> {
        stat_entries(A::STATS, A::stats(&self.lanes, &self.field, &self.tally))
    }

    fn set_param(&mut self, index: usize, value: &ParamValue) -> bool {
        self.params.set(index, value)
    }

    fn act(&mut self, index: usize) -> bool {
        if index >= A::ACTIONS.len() {
            return false;
        }
        // The state is destructured, so the params borrow and the lane borrows cover different fields.
        let Self {
            lanes,
            field,
            params,
            own_params,
            extent,
            action_seed,
            ..
        } = self;
        let (own, _) = split_at_own(params.values(), *own_params);
        A::act(index, lanes, field, *extent, own, action_seed);
        true
    }

    fn population(&self) -> u64 {
        self.lanes.len() as u64
    }

    fn heap_bytes(&self) -> usize {
        self.lanes.heap_bytes() + self.field.heap_bytes() + self.index.heap_bytes()
    }

    /// Number of chunks of the step pass.
    fn parallel_jobs(&self) -> Option<usize> {
        Some(self.lanes.len().div_ceil(A::CHUNK.max(1)))
    }
}