rustsim 0.0.1

//! Layer-based execution model inspired by FlameGPU2.
//!
//! FlameGPU2 organizes agent functions into **layers**. rustsim preserves the
//! idea of explicit execution stages, but the current `LayerExecutor`
//! implementation is intentionally **sequential**:
//!
//! - layers execute in insertion order
//! - functions stored in a layer execute in stored order
//! - no actual in-layer parallelism is performed today
//!
//! This makes layer boundaries a semantic ordering tool first. Future backends
//! may map the same structure to parallel execution, but concurrency is **not**
//! part of the current guarantee.
//!
//! # Example
//!
//! ```ignore
//! let mut executor = LayerExecutor::new();
//!
//! // Layer 0: output stage
//! executor.add_agent_layer("output_locations", output_locations_fn);
//!
//! // Layer 1: update stage
//! executor.add_agent_layer("update_positions", update_positions_fn);
//!
//! // Layer 2: model-level housekeeping
//! executor.add_model_layer("cleanup", cleanup_fn);
//!
//! // Execute all layers in order
//! executor.execute(columns, agent_count, agent_ids, space, properties);
//! ```

use rustsim_core::types::AgentId;

/// A function within a layer.
///
/// This is descriptive metadata only. The current public `LayerExecutor`
/// helpers add one function per layer, and execution is sequential.
pub struct LayerFunction<F> {
    /// Human-readable name for this function.
    pub name: &'static str,
    /// The function to execute.
    pub function: F,
}

/// Timing information for a single layer execution.
#[derive(Debug, Clone)]
pub struct LayerTiming {
    /// Layer index.
    pub layer_index: usize,
    /// Names of functions in this layer.
    pub function_names: Vec<&'static str>,
    /// Wall-clock time for this layer in microseconds.
    pub elapsed_us: u128,
}

/// Timing information for a complete step.
#[derive(Debug, Clone)]
pub struct StepTiming {
    /// Per-layer timing.
    pub layers: Vec<LayerTiming>,
    /// Total wall-clock time for the step in microseconds.
    pub total_us: u128,
}

/// Layer-based execution engine.
///
/// Organizes computation into ordered layers. In the current implementation,
/// the executor provides **deterministic sequential ordering**, not parallel
/// execution. If a backend later introduces concurrency, that will be an
/// additional capability rather than an implicit guarantee of this type.
pub struct LayerExecutor<S, A, Props> {
    layers: Vec<Vec<LayerEntry<S, A, Props>>>,
}

/// Internal layer entry.
#[allow(clippy::type_complexity)]
enum LayerEntry<S, A, Props> {
    /// A function that operates on SoA columns (batch step).
    BatchStep {
        name: &'static str,
        function: Box<dyn FnMut(&mut [Vec<f32>], usize)>,
    },
    /// A function that operates on individual agents by ID.
    AgentStep {
        name: &'static str,
        function: Box<dyn FnMut(AgentId, &mut S, &mut Props)>,
    },
    /// A function that operates on the model state.
    ModelStep {
        name: &'static str,
        function: Box<dyn FnMut(&mut S, &mut Props)>,
    },
    /// Phantom to satisfy unused generics.
    _Phantom(std::marker::PhantomData<A>),
}

impl<S, A, Props> LayerExecutor<S, A, Props> {
    /// Create a new, empty layer executor.
    pub fn new() -> Self {
        Self { layers: Vec::new() }
    }

    /// Add a new layer containing a batch-step function.
    ///
    /// Batch-step functions operate on SoA columns, mirroring FlameGPU2's
    /// agent functions that run as CUDA kernels.
    pub fn add_batch_layer(
        &mut self,
        name: &'static str,
        function: impl FnMut(&mut [Vec<f32>], usize) + 'static,
    ) {
        self.layers.push(vec![LayerEntry::BatchStep {
            name,
            function: Box::new(function),
        }]);
    }

    /// Add a new layer containing an agent-step function.
    pub fn add_agent_layer(
        &mut self,
        name: &'static str,
        function: impl FnMut(AgentId, &mut S, &mut Props) + 'static,
    ) {
        self.layers.push(vec![LayerEntry::AgentStep {
            name,
            function: Box::new(function),
        }]);
    }

    /// Add a new layer containing a model-step function.
    pub fn add_model_layer(
        &mut self,
        name: &'static str,
        function: impl FnMut(&mut S, &mut Props) + 'static,
    ) {
        self.layers.push(vec![LayerEntry::ModelStep {
            name,
            function: Box::new(function),
        }]);
    }

    /// Number of layers.
    pub fn num_layers(&self) -> usize {
        self.layers.len()
    }

    /// Execute all layers in order on the given SoA columns and model state.
    ///
    /// Current execution guarantees:
    /// - layers run sequentially in insertion order
    /// - entries in a layer run sequentially in stored order
    /// - agent-step layers iterate `agent_ids` in the order supplied by the caller
    ///
    /// The returned timing values are observational measurements, not semantic
    /// guarantees about determinism or backend scheduling.
    pub fn execute(
        &mut self,
        columns: &mut [Vec<f32>],
        agent_count: usize,
        agent_ids: &[AgentId],
        space: &mut S,
        properties: &mut Props,
    ) -> StepTiming {
        let t_total = std::time::Instant::now();
        let mut layer_timings = Vec::with_capacity(self.layers.len());

        for (layer_idx, layer) in self.layers.iter_mut().enumerate() {
            let t_layer = std::time::Instant::now();
            let mut names = Vec::new();

            for entry in layer.iter_mut() {
                match entry {
                    LayerEntry::BatchStep { name, function } => {
                        names.push(*name);
                        function(columns, agent_count);
                    }
                    LayerEntry::AgentStep { name, function } => {
                        names.push(*name);
                        for &id in agent_ids {
                            function(id, space, properties);
                        }
                    }
                    LayerEntry::ModelStep { name, function } => {
                        names.push(*name);
                        function(space, properties);
                    }
                    LayerEntry::_Phantom(_) => unreachable!(),
                }
            }

            layer_timings.push(LayerTiming {
                layer_index: layer_idx,
                function_names: names,
                elapsed_us: t_layer.elapsed().as_micros(),
            });
        }

        StepTiming {
            layers: layer_timings,
            total_us: t_total.elapsed().as_micros(),
        }
    }
}

impl<S, A, Props> Default for LayerExecutor<S, A, Props> {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn basic_layer_execution() {
        let mut executor: LayerExecutor<(), (), ()> = LayerExecutor::new();

        // Layer 0: batch step x[i] += 1.0
        executor.add_batch_layer("increment_x", |columns, n| {
            for v in columns[0].iter_mut().take(n) {
                *v += 1.0;
            }
        });

        // Layer 1: batch step x[i] *= 2.0
        executor.add_batch_layer("double_x", |columns, n| {
            for v in columns[0].iter_mut().take(n) {
                *v *= 2.0;
            }
        });

        let mut columns = vec![vec![0.0f32; 10]];
        let ids: Vec<AgentId> = (0..10).collect();
        let mut space = ();
        let mut props = ();

        let timing = executor.execute(&mut columns, 10, &ids, &mut space, &mut props);

        assert_eq!(timing.layers.len(), 2);
        // After layer 0: x = 1.0, after layer 1: x = 2.0
        for &v in &columns[0] {
            assert!((v - 2.0).abs() < 1e-5);
        }
    }
}