rustsim 0.0.1

//! Persistent device-side SoA storage for GPU-accelerated stepping.
//!
//! Inspired by FlameGPU2's `CUDAAgent` / `CUDAFatAgentStateList`, which keeps
//! agent variables on the GPU across steps to avoid per-step host-device
//! round-trips.
//!
//! [`DeviceSoaStore`] maintains a mirror of the agent data in SoA layout.
//! After initial upload, subsequent steps operate entirely on-device. Data
//! is only downloaded when explicitly requested (e.g. for data collection
//! or checkpointing).
//!
//! This module is **cross-platform**: on systems without CUDA the store
//! operates in CPU-only mode using the same SoA layout for cache-friendly
//! batch processing.
//!
//! # CUDA safety and failure surfaces
//!
//! The only `unsafe` operations in this module are CUDA kernel launches via
//! `cudarc`. Those launches rely on the following invariants:
//! - `block_size > 0`
//! - the PTX kernel signature matches the launched argument tuple
//! - uploaded device buffers correspond exactly to the stored SoA columns
//! - the kernel performs bounds checks for `idx < n`
//! - the kernel does not read or write out of bounds
//!
//! `step_cuda` returns explicit `Err(String)` values for:
//! - invalid `block_size`
//! - CUDA device initialization
//! - PTX load or kernel lookup
//! - host/device transfer failures
//! - unsupported SoA arity outside `1..=8`
//! - kernel launch / synchronization failures
//!
//! CPU-side `step_cpu` remains the always-available fallback execution path.
//!
//! # Architecture (mirrors FlameGPU2)
//!
//! ```text
//! ???????????????          ????????????????????
//! ? AgentStore   ?  upload  ?  DeviceSoaStore   ?
//! ? (AoS, host)  ? ???????> ?  (SoA, device/cpu)?
//! ?              ?          ?  columns: Vec<f32> ?
//! ?              ? <??????? ?  ids: Vec<AgentId> ?
//! ?              ? download ?                    ?
//! ???????????????          ????????????????????
//! ```

use rustsim_core::soa::{self, SoaExtractable, SoaExtractableF64};
use rustsim_core::store::AgentStore;
use rustsim_core::types::AgentId;
use thiserror::Error;

/// Host-visible checkpoint of a [`DeviceSoaStore`].
///
/// The checkpoint contains only deterministic simulation state: row-ordered
/// agent IDs, SoA columns, and column names. CUDA contexts, streams, modules,
/// pinned staging buffers, and dirty flags are intentionally not serialized;
/// restored stores restart as host-authoritative and can lazily re-enter CUDA
/// residency on the next runtime phase.
#[derive(Debug, Clone, PartialEq)]
pub struct DeviceSoaCheckpoint {
    /// Agent IDs in row order.
    pub ids: Vec<AgentId>,
    /// Column data in SoA layout. `columns[c][i]` = column `c`, agent `i`.
    pub columns: Vec<Vec<f32>>,
    /// Column names for diagnostics and schema compatibility checks.
    pub column_names: Vec<&'static str>,
    /// Typed column schema for validation and kernel ABI planning.
    pub schema: Vec<SoaColumnSchema>,
}

impl DeviceSoaCheckpoint {
    /// Number of checkpointed agents.
    pub fn agent_count(&self) -> usize {
        self.ids.len()
    }

    /// Number of checkpointed columns.
    pub fn num_columns(&self) -> usize {
        self.columns.len()
    }

    /// Approximate bytes held by checkpointed IDs and columns.
    pub fn resident_bytes(&self) -> usize {
        let ids_bytes = self.ids.len() * std::mem::size_of::<AgentId>();
        let cols_bytes: usize = self
            .columns
            .iter()
            .map(|c| c.len() * std::mem::size_of::<f32>())
            .sum();
        ids_bytes + cols_bytes
    }
}

/// Numeric precision of a persistent SoA column.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum SoaColumnType {
    /// 32-bit floating-point column.
    F32,
    /// 64-bit floating-point column.
    F64,
}

/// Stable schema entry for one persistent SoA column.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SoaColumnSchema {
    /// Human-readable column name.
    pub name: &'static str,
    /// Numeric precision stored by the runtime.
    pub column_type: SoaColumnType,
}

impl SoaColumnSchema {
    /// Build an `f32` schema entry.
    pub const fn f32(name: &'static str) -> Self {
        Self {
            name,
            column_type: SoaColumnType::F32,
        }
    }

    /// Build an `f64` schema entry.
    pub const fn f64(name: &'static str) -> Self {
        Self {
            name,
            column_type: SoaColumnType::F64,
        }
    }
}

/// Host-visible checkpoint of a double-precision authoritative SoA store.
#[derive(Debug, Clone, PartialEq)]
pub struct DeviceSoaCheckpointF64 {
    /// Agent IDs in row order.
    pub ids: Vec<AgentId>,
    /// Column data in SoA layout. `columns[c][i]` = column `c`, agent `i`.
    pub columns: Vec<Vec<f64>>,
    /// Column names for diagnostics and schema compatibility checks.
    pub column_names: Vec<&'static str>,
    /// Typed column schema for validation and kernel ABI planning.
    pub schema: Vec<SoaColumnSchema>,
}

impl DeviceSoaCheckpointF64 {
    /// Number of checkpointed agents.
    pub fn agent_count(&self) -> usize {
        self.ids.len()
    }

    /// Number of checkpointed columns.
    pub fn num_columns(&self) -> usize {
        self.columns.len()
    }

    /// Approximate bytes held by checkpointed IDs and columns.
    pub fn resident_bytes(&self) -> usize {
        let ids_bytes = self.ids.len() * std::mem::size_of::<AgentId>();
        let cols_bytes: usize = self
            .columns
            .iter()
            .map(|c| c.len() * std::mem::size_of::<f64>())
            .sum();
        ids_bytes + cols_bytes
    }
}

/// Errors returned when restoring a [`DeviceSoaStore`] from a checkpoint.
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum DeviceSoaRestoreError {
    /// The number of column names does not match the number of columns.
    #[error("checkpoint has {columns} columns but {names} column names")]
    ColumnNameCountMismatch { columns: usize, names: usize },
    /// A column has a different row count from the ID array.
    #[error("checkpoint column {column} has {len} rows but ids contain {agent_count} agents")]
    ColumnLengthMismatch {
        column: usize,
        len: usize,
        agent_count: usize,
    },
    /// Agent IDs must be unique inside a checkpoint.
    #[error("checkpoint contains duplicate agent id {0}")]
    DuplicateAgentId(AgentId),
    /// The typed schema does not match the checkpoint columns.
    #[error("checkpoint schema has {schema} columns but data contains {columns} columns")]
    SchemaColumnCountMismatch { columns: usize, schema: usize },
    /// The typed schema precision is incompatible with the store being restored.
    #[error("checkpoint column {column} has precision {actual:?}, expected {expected:?}")]
    SchemaTypeMismatch {
        column: usize,
        expected: SoaColumnType,
        actual: SoaColumnType,
    },
}

/// Errors returned by fallible SoA lifecycle mutations.
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum DeviceSoaMutationError {
    /// Caller supplied the wrong number of column slices.
    #[error("insert supplied {actual} columns, expected {expected}")]
    ColumnCountMismatch { expected: usize, actual: usize },
    /// A column slice length does not match the inserted ID count.
    #[error("insert column {column} has {actual} rows, expected {expected}")]
    ColumnLengthMismatch {
        column: usize,
        expected: usize,
        actual: usize,
    },
    /// Inserted IDs must be unique within the insert batch.
    #[error("insert batch contains duplicate agent id {0}")]
    DuplicateInsertedId(AgentId),
    /// Inserted IDs must not already exist in the authoritative store.
    #[error("agent id {0} already exists in authoritative SoA store")]
    ExistingAgentId(AgentId),
    /// Removed IDs must be unique within one lifecycle command.
    #[error("remove command contains duplicate agent id {0}")]
    DuplicateRemovedId(AgentId),
    /// Removed IDs must already exist in the authoritative store.
    #[error("agent id {0} does not exist in authoritative SoA store")]
    MissingAgentId(AgentId),
}

/// Persistent SoA storage that eliminates per-step extraction overhead.
///
/// After the initial [`upload`](Self::upload), the SoA columns live in the
/// store and can be stepped repeatedly via [`step_cpu`](Self::step_cpu) or
/// (when CUDA is available) GPU kernels.
///
/// # Zero-copy stepping
///
/// Unlike `cpu_batch_step` (which extracts SoA ? processes ? writes back
/// every step), `DeviceSoaStore` keeps the SoA columns persistent. This
/// saves ~40 ms/step for 1M agents by avoiding the AoS?SoA?AoS conversion.
///
/// Mirrors FlameGPU2's model where agent data lives on the GPU and is only
/// copied to host when needed for logging or I/O.
#[derive(Debug)]
pub struct DeviceSoaStore {
    /// Agent IDs in row order.
    ids: Vec<AgentId>,
    /// Column data in SoA layout. `columns[c][i]` = column `c`, agent `i`.
    columns: Vec<Vec<f32>>,
    /// Column names for debugging.
    column_names: Vec<&'static str>,
    /// Typed column schema for validation and kernel ABI planning.
    schema: Vec<SoaColumnSchema>,
    /// Number of active agents.
    agent_count: usize,
    /// Whether the device data is dirty (modified since last download).
    dirty: bool,
    /// Lazily-initialized CUDA resident state for the persistent-device
    /// stepping API (`step_cuda_resident` / `sync_to_host`).
    ///
    /// When `Some`, the CUDA column buffers, module, kernel function, and
    /// streams are kept alive across calls so that successive steps launch
    /// without any host/device transfer.
    #[cfg(feature = "cuda")]
    cuda_resident: Option<CudaResident>,
}

/// CUDA-resident companion state for [`DeviceSoaStore`].
///
/// Holds the persistent device buffers plus the CUDA context, copy/compute
/// streams, module, and kernel function. Living inside the store means the
/// host-side SoA arrays and the device-side column buffers share a single
/// lifecycle and `device_dirty` / `host_dirty` flags coordinate when a
/// transfer is actually required.
#[cfg(feature = "cuda")]
struct CudaResident {
    #[allow(dead_code)]
    ctx: std::sync::Arc<cudarc::driver::CudaContext>,
    copy_stream: std::sync::Arc<cudarc::driver::CudaStream>,
    compute_stream: std::sync::Arc<cudarc::driver::CudaStream>,
    #[allow(dead_code)]
    module: std::sync::Arc<cudarc::driver::CudaModule>,
    func: cudarc::driver::CudaFunction,
    kernel_name: String,
    d_cols: Vec<cudarc::driver::CudaSlice<f32>>,
    pinned: Vec<cudarc::driver::PinnedHostSlice<f32>>,
    /// Device buffers contain newer data than `self.columns`.
    device_dirty: bool,
    /// Host `self.columns` contain newer data than the device buffers.
    host_dirty: bool,
}

#[cfg(feature = "cuda")]
impl std::fmt::Debug for CudaResident {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CudaResident")
            .field("kernel_name", &self.kernel_name)
            .field("num_columns", &self.d_cols.len())
            .field("device_dirty", &self.device_dirty)
            .field("host_dirty", &self.host_dirty)
            .finish()
    }
}

impl DeviceSoaStore {
    /// Upload agent data from an `AgentStore` into persistent SoA layout.
    ///
    /// This is the initial data transfer. After this call, stepping
    /// operates on the SoA columns directly.
    pub fn upload<A, S>(store: &S) -> Self
    where
        A: SoaExtractable,
        S: AgentStore<A>,
    {
        let (ids, columns) = soa::extract_soa::<A, S>(store);
        let agent_count = ids.len();
        let column_names = A::column_names();
        let schema = schema_from_names(&column_names, SoaColumnType::F32);
        Self {
            ids,
            columns,
            column_names,
            agent_count,
            dirty: false,
            schema,
            #[cfg(feature = "cuda")]
            cuda_resident: None,
        }
    }

    /// Rebuild a persistent SoA store from a host-visible checkpoint.
    ///
    /// Restored stores start in CPU/host-authoritative mode. Any CUDA-resident
    /// state is re-created lazily by the runtime when a CUDA phase executes.
    pub fn from_checkpoint(checkpoint: DeviceSoaCheckpoint) -> Result<Self, DeviceSoaRestoreError> {
        let agent_count = checkpoint.ids.len();
        if checkpoint.column_names.len() != checkpoint.columns.len() {
            return Err(DeviceSoaRestoreError::ColumnNameCountMismatch {
                columns: checkpoint.columns.len(),
                names: checkpoint.column_names.len(),
            });
        }
        validate_schema(
            &checkpoint.schema,
            checkpoint.columns.len(),
            SoaColumnType::F32,
        )?;

        validate_column_lengths(&checkpoint.columns, agent_count)?;
        validate_unique_ids(&checkpoint.ids)?;

        Ok(Self {
            ids: checkpoint.ids,
            columns: checkpoint.columns,
            column_names: checkpoint.column_names,
            agent_count,
            dirty: false,
            schema: checkpoint.schema,
            #[cfg(feature = "cuda")]
            cuda_resident: None,
        })
    }

    /// Capture a deterministic host-visible checkpoint of the current SoA state.
    ///
    /// If CUDA residency is active, callers must synchronize first via
    /// [`sync_to_host`](Self::sync_to_host) or a higher-level runtime wrapper.
    pub fn checkpoint(&self) -> DeviceSoaCheckpoint {
        DeviceSoaCheckpoint {
            ids: self.ids.clone(),
            columns: self.columns.clone(),
            column_names: self.column_names.clone(),
            schema: self.schema.clone(),
        }
    }

    /// Number of active agents.
    pub fn agent_count(&self) -> usize {
        self.agent_count
    }

    /// Number of SoA columns.
    pub fn num_columns(&self) -> usize {
        self.columns.len()
    }

    /// Approximate resident bytes held by the persistent SoA buffers.
    ///
    /// This includes the active ID array and the active lengths of all columns,
    /// but not allocator over-capacity, metadata, or background CUDA allocations.
    pub fn resident_bytes(&self) -> usize {
        let ids_bytes = self.ids.len() * std::mem::size_of::<AgentId>();
        let cols_bytes: usize = self
            .columns
            .iter()
            .map(|c| c.len() * std::mem::size_of::<f32>())
            .sum();
        ids_bytes + cols_bytes
    }

    /// Column names.
    pub fn column_names(&self) -> &[&'static str] {
        &self.column_names
    }

    /// Typed column schema.
    pub fn schema(&self) -> &[SoaColumnSchema] {
        &self.schema
    }

    /// Read-only access to a column by index.
    pub fn column(&self, index: usize) -> &[f32] {
        &self.columns[index]
    }

    /// Mutable access to all columns (for kernel closures).
    pub fn columns_mut(&mut self) -> &mut [Vec<f32>] {
        self.dirty = true;
        &mut self.columns
    }

    /// Read-only access to agent IDs in row order.
    pub fn ids(&self) -> &[AgentId] {
        &self.ids
    }

    /// Return the row currently occupied by `id`, if present.
    pub fn row_of(&self, id: AgentId) -> Option<usize> {
        self.ids.iter().position(|existing| *existing == id)
    }

    /// Whether the authoritative store currently contains `id`.
    pub fn contains_id(&self, id: AgentId) -> bool {
        self.row_of(id).is_some()
    }

    /// Smallest non-negative [`AgentId`] not currently present in the store.
    pub fn next_available_id(&self) -> AgentId {
        next_available_id(&self.ids)
    }

    /// Whether the SoA data has been modified since the last download.
    pub fn is_dirty(&self) -> bool {
        self.dirty
    }

    /// Execute a CPU kernel on the SoA columns, in-place.
    ///
    /// The kernel receives `(columns, agent_count)`. This is equivalent to
    /// `cpu_batch_step` but without the extract/write-back overhead.
    pub fn step_cpu<F>(&mut self, mut kernel: F) -> u128
    where
        F: FnMut(&mut [Vec<f32>], usize),
    {
        let t0 = std::time::Instant::now();
        kernel(&mut self.columns, self.agent_count);
        self.dirty = true;
        t0.elapsed().as_micros()
    }

    /// Execute a named sequence of CPU kernels over the persistent SoA columns.
    ///
    /// This is the CPU mirror of FlameGPU2's agent-function / layer model:
    /// each kernel in the sequence sees the SoA columns after all previous
    /// kernels in the same call, without any intermediate host?device copies
    /// or extract/write-back.
    ///
    /// Every kernel is invoked with `(columns, agent_count)`. Timing of each
    /// individual kernel is captured and returned as
    /// `[(name, kernel_us)]` in the same order as `kernels`.
    ///
    /// The `dirty` flag is set once at the end, not per-kernel, so a
    /// subsequent `download` flushes all kernel effects in one pass.
    pub fn run_sequence_cpu<'a, I, K>(&mut self, kernels: I) -> Vec<(&'static str, u128)>
    where
        I: IntoIterator<Item = (&'static str, K)>,
        K: FnOnce(&mut [Vec<f32>], usize) + 'a,
    {
        let mut timings = Vec::new();
        for (name, kernel) in kernels {
            let t0 = std::time::Instant::now();
            kernel(&mut self.columns, self.agent_count);
            timings.push((name, t0.elapsed().as_micros()));
        }
        if !timings.is_empty() {
            self.dirty = true;
        }
        timings
    }

    /// Download the SoA data back into an `AgentStore`.
    ///
    /// This is the inverse of `upload` -- writes SoA column values back
    /// into the AoS agent structs. Only needed for data collection,
    /// checkpointing, or when switching from GPU to CPU stepping.
    pub fn download<A, S>(&mut self, store: &S)
    where
        A: SoaExtractable,
        S: AgentStore<A>,
    {
        soa::write_back_soa::<A, S>(store, &self.ids, &self.columns);
        self.dirty = false;
    }

    /// Execute a CUDA kernel on the SoA columns.
    ///
    /// Uploads columns to GPU, launches the kernel, downloads results.
    /// Columns remain in `DeviceSoaStore` for subsequent steps.
    ///
    /// This is more efficient than `cuda_batch_step` for multi-step runs
    /// because the SoA extraction is not repeated.
    ///
    /// Failure surfaces returned as `Err(String)` include:
    /// - invalid `block_size`
    /// - CUDA context initialization
    /// - PTX compile/load or kernel lookup
    /// - host/device transfer failures
    /// - kernel launch or stream synchronization failures
    ///
    /// # Safety requirements for the PTX kernel
    /// - the kernel signature must match the launched argument list
    /// - the kernel must bounds-check against `n`
    /// - the kernel must not read or write outside the provided column buffers
    #[cfg(feature = "cuda")]
    pub fn step_cuda(
        &mut self,
        ptx_source: &str,
        _module_name: &str,
        kernel_name: &str,
        block_size: u32,
    ) -> Result<u128, String> {
        use cudarc::driver::{LaunchConfig, PushKernelArg};

        if block_size == 0 {
            return Err("block_size must be positive".to_string());
        }

        let n = self.agent_count;
        if n == 0 {
            return Ok(0);
        }

        let ctx = crate::cuda_context::new_context(0)?;
        let stream = ctx.default_stream();

        let ptx = cudarc::nvrtc::Ptx::from_src(ptx_source);
        let module = ctx.load_module(ptx).map_err(|e| format!("PTX load: {e}"))?;
        let func = module
            .load_function(kernel_name)
            .map_err(|e| format!("kernel '{kernel_name}' not found: {e}"))?;

        // Upload columns on the compute stream.
        let mut d_cols = Vec::with_capacity(self.columns.len());
        for col in &self.columns {
            let d = stream
                .clone_htod(col.as_slice())
                .map_err(|e| format!("htod: {e}"))?;
            d_cols.push(d);
        }

        let n_u32 = n as u32;
        let grid_size = n.div_ceil(block_size as usize) as u32;
        let cfg = LaunchConfig {
            grid_dim: (grid_size, 1, 1),
            block_dim: (block_size, 1, 1),
            shared_mem_bytes: 0,
        };

        let t0 = std::time::Instant::now();

        // SAFETY:
        // - each `d_cols[i]` is a valid device buffer allocated on `stream`
        // - `n_u32` is the logical row count passed to the kernel for bounds checking
        // - callers are responsible for providing a PTX kernel whose signature matches
        //   the argument list built below and whose implementation does not access out
        //   of bounds
        // - `block_size > 0` has been validated above
        // - work is scheduled on `stream` only, and `stream.synchronize()` is called
        //   below before any host-visible read
        unsafe {
            let mut builder = stream.launch_builder(&func);
            for d in d_cols.iter_mut() {
                builder.arg(d);
            }
            builder.arg(&n_u32);
            builder.launch(cfg).map_err(|e| format!("launch: {e}"))?;
        }

        stream.synchronize().map_err(|e| format!("sync: {e}"))?;
        let kernel_us = t0.elapsed().as_micros();

        // Download results.
        for (i, d_col) in d_cols.iter().enumerate() {
            stream
                .memcpy_dtoh(d_col, &mut self.columns[i])
                .map_err(|e| format!("dtoh: {e}"))?;
        }

        self.dirty = true;
        Ok(kernel_us)
    }

    /// Execute a CUDA kernel over the SoA columns using **pinned host memory
    /// and dedicated non-default CUDA streams** for host/device transfer
    /// overlap.
    ///
    /// This is structurally the same as [`Self::step_cuda`] but differs in how
    /// the transfers are scheduled:
    ///
    /// - SoA columns are first copied into page-locked (pinned) host buffers
    ///   allocated via [`cudarc::driver::CudaContext::alloc_pinned`]. Pinned
    ///   memory lets the CUDA driver issue truly asynchronous `memcpy_htod` /
    ///   `memcpy_dtoh` without going through an internal bounce buffer.
    /// - Host-to-device uploads and device-to-host downloads run on a
    ///   dedicated *copy* stream ([`cudarc::driver::CudaContext::new_stream`]).
    /// - The kernel launch runs on a dedicated *compute* stream.
    /// - The two streams are serialized via
    ///   [`cudarc::driver::CudaStream::join`], which records an event on the
    ///   other stream and makes `self` wait on it — the standard CUDA pattern
    ///   for overlapping transfer and compute.
    ///
    /// With only one step in flight at a time this still runs in-order on the
    /// host timeline; the benefit materializes when multiple `step_cuda_pinned`
    /// calls or a multi-kernel sequence are issued back-to-back and the
    /// driver is free to overlap the download of step *N* with the upload of
    /// step *N+1*.
    ///
    /// Failure surfaces returned as `Err(String)` include:
    /// - invalid `block_size`
    /// - CUDA context / stream / pinned allocation failures
    /// - PTX compile/load or kernel lookup
    /// - host/device transfer failures
    /// - kernel launch or stream synchronization failures
    ///
    /// # Safety requirements for the PTX kernel
    /// - the kernel signature must match the launched argument list
    /// - the kernel must bounds-check against `n`
    /// - the kernel must not read or write outside the provided column buffers
    #[cfg(feature = "cuda")]
    pub fn step_cuda_pinned(
        &mut self,
        ptx_source: &str,
        _module_name: &str,
        kernel_name: &str,
        block_size: u32,
    ) -> Result<u128, String> {
        use cudarc::driver::{LaunchConfig, PushKernelArg};

        if block_size == 0 {
            return Err("block_size must be positive".to_string());
        }

        let n = self.agent_count;
        if n == 0 {
            return Ok(0);
        }

        let ctx = crate::cuda_context::new_context(0)?;

        // Dedicated non-default streams: one for host<->device transfers, one
        // for compute. Creating a new stream implicitly synchronizes the
        // context (see `CudaContext::new_stream` docs) which is fine here
        // because `step_cuda_pinned` is itself a synchronous boundary.
        let copy_stream = ctx
            .new_stream()
            .map_err(|e| format!("copy stream init: {e}"))?;
        let compute_stream = ctx
            .new_stream()
            .map_err(|e| format!("compute stream init: {e}"))?;

        let ptx = cudarc::nvrtc::Ptx::from_src(ptx_source);
        let module = ctx.load_module(ptx).map_err(|e| format!("PTX load: {e}"))?;
        let func = module
            .load_function(kernel_name)
            .map_err(|e| format!("kernel '{kernel_name}' not found: {e}"))?;

        // Allocate pinned (page-locked) host staging buffers and fill them
        // from the SoA columns. The pinned allocation is unsafe because the
        // driver does not initialize it; we overwrite every element
        // immediately via `copy_from_slice` so no uninitialized read occurs.
        let mut pinned: Vec<cudarc::driver::PinnedHostSlice<f32>> =
            Vec::with_capacity(self.columns.len());
        for col in &self.columns {
            // SAFETY: `alloc_pinned` returns uninitialized pinned host memory.
            // We immediately fill the entire slice via `copy_from_slice` below
            // before any read, so no uninitialized byte is ever observed.
            let mut p = unsafe { ctx.alloc_pinned::<f32>(col.len()) }
                .map_err(|e| format!("pinned alloc: {e}"))?;
            p.as_mut_slice()
                .map_err(|e| format!("pinned access: {e}"))?
                .copy_from_slice(col);
            pinned.push(p);
        }

        // Host -> device on the copy stream.
        let mut d_cols: Vec<cudarc::driver::CudaSlice<f32>> = Vec::with_capacity(pinned.len());
        for p in &pinned {
            let d = copy_stream
                .clone_htod(p)
                .map_err(|e| format!("htod: {e}"))?;
            d_cols.push(d);
        }

        // The compute stream must wait until the uploads on `copy_stream`
        // have completed before it can access `d_cols`.
        compute_stream
            .join(&copy_stream)
            .map_err(|e| format!("compute.join(copy): {e}"))?;

        let n_u32 = n as u32;
        let grid_size = n.div_ceil(block_size as usize) as u32;
        let cfg = LaunchConfig {
            grid_dim: (grid_size, 1, 1),
            block_dim: (block_size, 1, 1),
            shared_mem_bytes: 0,
        };

        let t0 = std::time::Instant::now();

        // SAFETY:
        // - each `d_cols[i]` is a valid device buffer allocated on `copy_stream`
        //   and made visible to `compute_stream` via `compute_stream.join(&copy_stream)`
        // - `n_u32` is the logical row count passed to the kernel for bounds checking
        // - callers are responsible for providing a PTX kernel whose signature matches
        //   the argument list built below and whose implementation does not access out
        //   of bounds
        // - `block_size > 0` has been validated above
        // - work is scheduled on `compute_stream`; the copy stream waits on it via
        //   `copy_stream.join(&compute_stream)` below before issuing dtoh, and the
        //   copy stream is then synchronized before any host-visible read
        unsafe {
            let mut builder = compute_stream.launch_builder(&func);
            for d in d_cols.iter_mut() {
                builder.arg(d);
            }
            builder.arg(&n_u32);
            builder.launch(cfg).map_err(|e| format!("launch: {e}"))?;
        }

        // The copy stream must wait for the kernel to finish before reading
        // the buffers back.
        copy_stream
            .join(&compute_stream)
            .map_err(|e| format!("copy.join(compute): {e}"))?;

        // Device -> pinned host on the copy stream.
        for (i, d_col) in d_cols.iter().enumerate() {
            copy_stream
                .memcpy_dtoh(d_col, &mut pinned[i])
                .map_err(|e| format!("dtoh: {e}"))?;
        }

        copy_stream
            .synchronize()
            .map_err(|e| format!("sync: {e}"))?;
        let kernel_us = t0.elapsed().as_micros();

        // Copy the pinned results back into the owning SoA columns.
        for (i, p) in pinned.iter().enumerate() {
            self.columns[i]
                .copy_from_slice(p.as_slice().map_err(|e| format!("pinned readback: {e}"))?);
        }

        self.dirty = true;
        Ok(kernel_us)
    }

    // ------------------------------------------------------------------
    // Truly-persistent CUDA lifecycle
    //
    // `step_cuda` / `step_cuda_pinned` upload and download on every call,
    // which is useful as a stateless entry point but costs ~2n bytes of
    // PCIe traffic per step. The API below keeps CUDA context, streams,
    // module, kernel function, device buffers, and pinned staging all
    // alive across calls — the FlameGPU2 "agent data lives on the GPU"
    // pattern — so that a back-to-back sequence of kernel launches
    // issues **zero** host/device transfers beyond the initial upload
    // and an explicit `sync_to_host` on demand.
    // ------------------------------------------------------------------

    /// Initialize the CUDA-resident pipeline.
    ///
    /// Allocates persistent device buffers, compiles / loads the PTX
    /// module, looks up the kernel function, creates dedicated copy and
    /// compute streams, allocates pinned host staging, and performs the
    /// initial host→device upload. After this call, subsequent
    /// [`step_cuda_resident`](Self::step_cuda_resident) invocations launch
    /// the kernel **without any host/device transfer**, and the host
    /// columns are only refreshed when [`sync_to_host`](Self::sync_to_host)
    /// is called explicitly.
    ///
    /// Re-initializing with a different kernel replaces the cached state.
    ///
    /// Failure surfaces returned as `Err(String)` include:
    /// - CUDA context / stream / module load / kernel lookup failures
    /// - pinned allocation failures
    /// - host→device upload failures
    #[cfg(feature = "cuda")]
    pub fn init_cuda(&mut self, ptx_source: &str, kernel_name: &str) -> Result<(), String> {
        let ctx = crate::cuda_context::new_context(0)?;
        let copy_stream = ctx
            .new_stream()
            .map_err(|e| format!("copy stream init: {e}"))?;
        let compute_stream = ctx
            .new_stream()
            .map_err(|e| format!("compute stream init: {e}"))?;

        let ptx = cudarc::nvrtc::Ptx::from_src(ptx_source);
        let module = ctx.load_module(ptx).map_err(|e| format!("PTX load: {e}"))?;
        let func = module
            .load_function(kernel_name)
            .map_err(|e| format!("kernel '{kernel_name}' not found: {e}"))?;

        // Allocate pinned staging and fill it from the host columns, then
        // upload into device buffers on the copy stream.
        let mut pinned: Vec<cudarc::driver::PinnedHostSlice<f32>> =
            Vec::with_capacity(self.columns.len());
        let mut d_cols: Vec<cudarc::driver::CudaSlice<f32>> =
            Vec::with_capacity(self.columns.len());
        for col in &self.columns {
            // SAFETY: `alloc_pinned` returns uninitialized pinned host memory.
            // We immediately fill the entire slice via `copy_from_slice` below
            // before any read, so no uninitialized byte is ever observed.
            let mut p = unsafe { ctx.alloc_pinned::<f32>(col.len()) }
                .map_err(|e| format!("pinned alloc: {e}"))?;
            p.as_mut_slice()
                .map_err(|e| format!("pinned access: {e}"))?
                .copy_from_slice(col);
            let d = copy_stream
                .clone_htod(&p)
                .map_err(|e| format!("htod: {e}"))?;
            pinned.push(p);
            d_cols.push(d);
        }

        // Ensure the initial upload is visible before any subsequent launch.
        copy_stream
            .synchronize()
            .map_err(|e| format!("sync: {e}"))?;

        self.cuda_resident = Some(CudaResident {
            ctx,
            copy_stream,
            compute_stream,
            module,
            func,
            kernel_name: kernel_name.to_string(),
            d_cols,
            pinned,
            device_dirty: false,
            host_dirty: false,
        });
        Ok(())
    }

    /// Launch the initialized CUDA kernel over the persistent device
    /// buffers.
    ///
    /// Requires a prior call to [`init_cuda`](Self::init_cuda). Unlike
    /// [`step_cuda`](Self::step_cuda) and
    /// [`step_cuda_pinned`](Self::step_cuda_pinned), this method:
    /// - does **not** upload host columns (unless they are host-dirty, in
    ///   which case it uploads once and clears the flag);
    /// - does **not** download device results;
    /// - runs entirely on the dedicated compute stream.
    ///
    /// The device buffers are marked dirty so that a later
    /// [`sync_to_host`](Self::sync_to_host) or [`download`](Self::download)
    /// performs the device→host copy exactly once.
    ///
    /// Failure surfaces returned as `Err(String)` include:
    /// - `init_cuda` was not called
    /// - invalid `block_size`
    /// - kernel launch or stream synchronization failures
    /// - on-demand re-upload failures
    #[cfg(feature = "cuda")]
    pub fn step_cuda_resident(&mut self, block_size: u32) -> Result<u128, String> {
        use cudarc::driver::{LaunchConfig, PushKernelArg};

        if block_size == 0 {
            return Err("block_size must be positive".to_string());
        }

        let n = self.agent_count;
        if n == 0 {
            return Ok(0);
        }

        let resident = self
            .cuda_resident
            .as_mut()
            .ok_or_else(|| "init_cuda must be called before step_cuda_resident".to_string())?;

        // If the host copied fresh data into `self.columns`, flush it to
        // the device before launching.
        if resident.host_dirty {
            for (i, col) in self.columns.iter().enumerate() {
                resident.pinned[i]
                    .as_mut_slice()
                    .map_err(|e| format!("pinned access: {e}"))?
                    .copy_from_slice(col);
                resident
                    .copy_stream
                    .memcpy_htod(&resident.pinned[i], &mut resident.d_cols[i])
                    .map_err(|e| format!("htod: {e}"))?;
            }
            resident
                .compute_stream
                .join(&resident.copy_stream)
                .map_err(|e| format!("compute.join(copy): {e}"))?;
            resident.host_dirty = false;
        }

        let n_u32 = n as u32;
        let grid_size = n.div_ceil(block_size as usize) as u32;
        let cfg = LaunchConfig {
            grid_dim: (grid_size, 1, 1),
            block_dim: (block_size, 1, 1),
            shared_mem_bytes: 0,
        };

        let t0 = std::time::Instant::now();

        // SAFETY:
        // - each `d_cols[i]` is a valid device buffer owned by `resident`
        //   and was last touched on `compute_stream` (or `copy_stream`
        //   followed by `compute_stream.join(copy_stream)`), so its
        //   contents are visible to the compute stream
        // - `n_u32` is the logical row count passed to the kernel for
        //   bounds checking
        // - callers are responsible for providing a PTX kernel whose
        //   signature matches the argument list built below and whose
        //   implementation does not access out of bounds
        // - `block_size > 0` has been validated above
        // - `sync_to_host` / `download` synchronize the compute stream
        //   before any host-visible read
        unsafe {
            let mut builder = resident.compute_stream.launch_builder(&resident.func);
            for d in resident.d_cols.iter_mut() {
                builder.arg(d);
            }
            builder.arg(&n_u32);
            builder.launch(cfg).map_err(|e| format!("launch: {e}"))?;
        }

        // We don't synchronize here — the whole point of the resident API
        // is to let successive launches queue up without a host barrier.
        // We capture the elapsed time after the launch was submitted; it
        // reflects submission latency, not kernel completion. Callers who
        // need completion timing can call `sync_to_host()` which does
        // synchronize.
        let kernel_us = t0.elapsed().as_micros();

        resident.device_dirty = true;
        self.dirty = true;
        Ok(kernel_us)
    }

    /// Launch a sequence of persistent CUDA kernels back-to-back without
    /// any intermediate host transfer.
    ///
    /// Each entry in `kernels` is a `(name, kernel_name)` pair; the kernel
    /// is looked up in the module loaded by [`init_cuda`](Self::init_cuda).
    /// The first entry *must* match the kernel passed to `init_cuda` (the
    /// persistent `CudaFunction`); other names are resolved from the same
    /// module on each call. Every kernel shares the same persistent device
    /// buffers, so this is a direct mirror of FlameGPU2's layered agent
    /// functions.
    ///
    /// Returns `[(name, submit_us)]` in the same order as `kernels`.
    ///
    /// Failure surfaces returned as `Err(String)` include everything
    /// [`step_cuda_resident`](Self::step_cuda_resident) can return plus
    /// kernel lookup failures for subsequent kernels in the sequence.
    #[cfg(feature = "cuda")]
    pub fn run_sequence_cuda_resident(
        &mut self,
        kernels: &[&str],
        block_size: u32,
    ) -> Result<Vec<(String, u128)>, String> {
        use cudarc::driver::{LaunchConfig, PushKernelArg};

        if block_size == 0 {
            return Err("block_size must be positive".to_string());
        }
        let n = self.agent_count;
        if n == 0 {
            return Ok(Vec::new());
        }

        let resident = self.cuda_resident.as_mut().ok_or_else(|| {
            "init_cuda must be called before run_sequence_cuda_resident".to_string()
        })?;

        if resident.host_dirty {
            for (i, col) in self.columns.iter().enumerate() {
                resident.pinned[i]
                    .as_mut_slice()
                    .map_err(|e| format!("pinned access: {e}"))?
                    .copy_from_slice(col);
                resident
                    .copy_stream
                    .memcpy_htod(&resident.pinned[i], &mut resident.d_cols[i])
                    .map_err(|e| format!("htod: {e}"))?;
            }
            resident
                .compute_stream
                .join(&resident.copy_stream)
                .map_err(|e| format!("compute.join(copy): {e}"))?;
            resident.host_dirty = false;
        }

        let n_u32 = n as u32;
        let grid_size = n.div_ceil(block_size as usize) as u32;
        let cfg = LaunchConfig {
            grid_dim: (grid_size, 1, 1),
            block_dim: (block_size, 1, 1),
            shared_mem_bytes: 0,
        };

        let mut timings = Vec::with_capacity(kernels.len());
        for name in kernels {
            let func = if *name == resident.kernel_name {
                resident.func.clone()
            } else {
                resident
                    .module
                    .load_function(name)
                    .map_err(|e| format!("kernel '{name}' not found: {e}"))?
            };

            let t0 = std::time::Instant::now();

            // SAFETY: same invariants as `step_cuda_resident` — the kernel
            // signature must match the launched argument list, operate on
            // the persistent device buffers, and bounds-check against `n`.
            // The previous launch is implicitly serialized by the shared
            // compute stream.
            unsafe {
                let mut builder = resident.compute_stream.launch_builder(&func);
                for d in resident.d_cols.iter_mut() {
                    builder.arg(d);
                }
                builder.arg(&n_u32);
                builder.launch(cfg).map_err(|e| format!("launch: {e}"))?;
            }

            timings.push((name.to_string(), t0.elapsed().as_micros()));
        }

        resident.device_dirty = true;
        self.dirty = true;
        Ok(timings)
    }

    /// Synchronize the compute stream and copy the persistent device
    /// buffers back into the host SoA columns.
    ///
    /// No-op when there is no CUDA-resident state or when the device is
    /// not dirty. This is the one place device→host traffic happens for
    /// the resident API, so a typical run issues `init_cuda` once,
    /// `step_cuda_resident` repeatedly, then `sync_to_host` only when
    /// data collection / checkpointing is required.
    #[cfg(feature = "cuda")]
    pub fn sync_to_host(&mut self) -> Result<(), String> {
        let resident = match self.cuda_resident.as_mut() {
            Some(r) => r,
            None => return Ok(()),
        };
        if !resident.device_dirty {
            return Ok(());
        }

        // Ensure the compute stream is visible to the copy stream.
        resident
            .copy_stream
            .join(&resident.compute_stream)
            .map_err(|e| format!("copy.join(compute): {e}"))?;
        for (i, d_col) in resident.d_cols.iter().enumerate() {
            resident
                .copy_stream
                .memcpy_dtoh(d_col, &mut resident.pinned[i])
                .map_err(|e| format!("dtoh: {e}"))?;
        }
        resident
            .copy_stream
            .synchronize()
            .map_err(|e| format!("sync: {e}"))?;
        for (i, p) in resident.pinned.iter().enumerate() {
            self.columns[i]
                .copy_from_slice(p.as_slice().map_err(|e| format!("pinned readback: {e}"))?);
        }
        resident.device_dirty = false;
        Ok(())
    }

    /// Mark the host columns as dirty so that the next resident launch
    /// re-uploads them before launching.
    ///
    /// Call this after directly mutating the host columns (e.g. via
    /// [`columns_mut`](Self::columns_mut)) when a CUDA-resident pipeline
    /// is active.
    #[cfg(feature = "cuda")]
    pub fn mark_host_dirty(&mut self) {
        if let Some(r) = self.cuda_resident.as_mut() {
            r.host_dirty = true;
        }
    }

    /// Drop the CUDA-resident state (releasing device buffers, streams,
    /// the module, and the kernel function).
    ///
    /// A subsequent call to [`init_cuda`](Self::init_cuda) re-initializes
    /// the pipeline from scratch.
    #[cfg(feature = "cuda")]
    pub fn release_cuda(&mut self) {
        self.cuda_resident = None;
    }

    /// Whether the store currently owns a CUDA-resident pipeline.
    #[cfg(feature = "cuda")]
    pub fn has_cuda_resident(&self) -> bool {
        self.cuda_resident.is_some()
    }

    /// Remove agents by marking them dead and compacting the SoA arrays.
    ///
    /// This is the CPU equivalent of FlameGPU2's `CUDAScatter` -- a prefix-sum
    /// based compaction that removes dead agents from the SoA buffers without
    /// full re-extraction.
    ///
    /// `dead_ids` must be sorted.
    pub fn scatter_remove(&mut self, dead_ids: &[AgentId]) {
        if dead_ids.is_empty() {
            return;
        }

        // Build a set for O(1) lookup
        let dead_set: std::collections::HashSet<AgentId> = dead_ids.iter().copied().collect();

        // Compact: keep only alive agents
        let mut write = 0;
        for read in 0..self.agent_count {
            if !dead_set.contains(&self.ids[read]) {
                if write != read {
                    self.ids[write] = self.ids[read];
                    for col in &mut self.columns {
                        col[write] = col[read];
                    }
                }
                write += 1;
            }
        }

        self.agent_count = write;
        self.ids.truncate(write);
        for col in &mut self.columns {
            col.truncate(write);
        }
        self.dirty = true;
    }

    /// Remove agents with typed validation.
    pub fn try_scatter_remove(
        &mut self,
        dead_ids: &[AgentId],
    ) -> Result<(), DeviceSoaMutationError> {
        validate_remove_ids(dead_ids, &self.ids)?;
        self.scatter_remove(dead_ids);
        Ok(())
    }

    /// Append new agents to the SoA buffers.
    ///
    /// This is the CPU equivalent of FlameGPU2's scatter-based agent creation.
    /// New agent data is provided as `(ids, column_values)` where
    /// `column_values[c]` is a slice of values for column `c`.
    pub fn scatter_insert(&mut self, new_ids: &[AgentId], new_columns: &[&[f32]]) {
        self.try_scatter_insert(new_ids, new_columns)
            .expect("valid SoA scatter insert");
    }

    /// Append new agents to the SoA buffers with typed validation.
    pub fn try_scatter_insert(
        &mut self,
        new_ids: &[AgentId],
        new_columns: &[&[f32]],
    ) -> Result<(), DeviceSoaMutationError> {
        validate_insert_shape(new_ids, new_columns, self.columns.len())?;
        validate_insert_ids(new_ids, &self.ids)?;

        let new_count = new_ids.len();
        if new_count == 0 {
            return Ok(());
        }

        self.ids.extend_from_slice(new_ids);
        for (i, col) in self.columns.iter_mut().enumerate() {
            col.extend_from_slice(new_columns[i]);
        }
        self.agent_count += new_count;
        self.dirty = true;
        Ok(())
    }

    /// Apply remove + insert as one validated lifecycle mutation.
    pub fn try_scatter_replace(
        &mut self,
        dead_ids: &[AgentId],
        new_ids: &[AgentId],
        new_columns: &[&[f32]],
    ) -> Result<(), DeviceSoaMutationError> {
        validate_remove_ids(dead_ids, &self.ids)?;
        validate_insert_shape(new_ids, new_columns, self.columns.len())?;
        validate_insert_ids_after_removal(new_ids, &self.ids, dead_ids)?;

        self.scatter_remove(dead_ids);
        append_rows(&mut self.ids, &mut self.columns, new_ids, new_columns);
        self.agent_count = self.ids.len();
        if !(dead_ids.is_empty() && new_ids.is_empty()) {
            self.dirty = true;
        }
        Ok(())
    }
}

/// Persistent double-precision SoA storage.
#[derive(Debug)]
pub struct DeviceSoaStoreF64 {
    ids: Vec<AgentId>,
    columns: Vec<Vec<f64>>,
    column_names: Vec<&'static str>,
    schema: Vec<SoaColumnSchema>,
    agent_count: usize,
    dirty: bool,
}

impl DeviceSoaStoreF64 {
    /// Upload agent data from an `AgentStore` into persistent f64 SoA layout.
    pub fn upload<A, S>(store: &S) -> Self
    where
        A: SoaExtractableF64,
        S: AgentStore<A>,
    {
        let (ids, columns) = soa::extract_soa_f64::<A, S>(store);
        let agent_count = ids.len();
        let column_names = <A as SoaExtractableF64>::column_names();
        let schema = schema_from_names(&column_names, SoaColumnType::F64);
        Self {
            ids,
            columns,
            column_names,
            schema,
            agent_count,
            dirty: false,
        }
    }

    /// Rebuild a persistent f64 SoA store from a host-visible checkpoint.
    pub fn from_checkpoint(
        checkpoint: DeviceSoaCheckpointF64,
    ) -> Result<Self, DeviceSoaRestoreError> {
        let agent_count = checkpoint.ids.len();
        if checkpoint.column_names.len() != checkpoint.columns.len() {
            return Err(DeviceSoaRestoreError::ColumnNameCountMismatch {
                columns: checkpoint.columns.len(),
                names: checkpoint.column_names.len(),
            });
        }
        validate_schema(
            &checkpoint.schema,
            checkpoint.columns.len(),
            SoaColumnType::F64,
        )?;
        validate_column_lengths(&checkpoint.columns, agent_count)?;
        validate_unique_ids(&checkpoint.ids)?;

        Ok(Self {
            ids: checkpoint.ids,
            columns: checkpoint.columns,
            column_names: checkpoint.column_names,
            schema: checkpoint.schema,
            agent_count,
            dirty: false,
        })
    }

    /// Capture a deterministic host-visible checkpoint.
    pub fn checkpoint(&self) -> DeviceSoaCheckpointF64 {
        DeviceSoaCheckpointF64 {
            ids: self.ids.clone(),
            columns: self.columns.clone(),
            column_names: self.column_names.clone(),
            schema: self.schema.clone(),
        }
    }

    /// Number of active agents.
    pub fn agent_count(&self) -> usize {
        self.agent_count
    }

    /// Number of SoA columns.
    pub fn num_columns(&self) -> usize {
        self.columns.len()
    }

    /// Approximate resident bytes held by IDs and f64 columns.
    pub fn resident_bytes(&self) -> usize {
        let ids_bytes = self.ids.len() * std::mem::size_of::<AgentId>();
        let cols_bytes: usize = self
            .columns
            .iter()
            .map(|c| c.len() * std::mem::size_of::<f64>())
            .sum();
        ids_bytes + cols_bytes
    }

    /// Column names.
    pub fn column_names(&self) -> &[&'static str] {
        &self.column_names
    }

    /// Typed column schema.
    pub fn schema(&self) -> &[SoaColumnSchema] {
        &self.schema
    }

    /// Read-only access to a column by index.
    pub fn column(&self, index: usize) -> &[f64] {
        &self.columns[index]
    }

    /// Mutable access to all columns.
    pub fn columns_mut(&mut self) -> &mut [Vec<f64>] {
        self.dirty = true;
        &mut self.columns
    }

    /// Read-only access to agent IDs in row order.
    pub fn ids(&self) -> &[AgentId] {
        &self.ids
    }

    /// Return the row currently occupied by `id`, if present.
    pub fn row_of(&self, id: AgentId) -> Option<usize> {
        self.ids.iter().position(|existing| *existing == id)
    }

    /// Whether the authoritative store currently contains `id`.
    pub fn contains_id(&self, id: AgentId) -> bool {
        self.row_of(id).is_some()
    }

    /// Smallest non-negative [`AgentId`] not currently present in the store.
    pub fn next_available_id(&self) -> AgentId {
        next_available_id(&self.ids)
    }

    /// Whether the SoA data has been modified since the last download.
    pub fn is_dirty(&self) -> bool {
        self.dirty
    }

    /// Execute a CPU kernel on the f64 SoA columns, in-place.
    pub fn step_cpu<F>(&mut self, mut kernel: F) -> u128
    where
        F: FnMut(&mut [Vec<f64>], usize),
    {
        let t0 = std::time::Instant::now();
        kernel(&mut self.columns, self.agent_count);
        self.dirty = true;
        t0.elapsed().as_micros()
    }

    /// Execute a named sequence of CPU kernels over persistent f64 columns.
    pub fn run_sequence_cpu<'a, I, K>(&mut self, kernels: I) -> Vec<(&'static str, u128)>
    where
        I: IntoIterator<Item = (&'static str, K)>,
        K: FnOnce(&mut [Vec<f64>], usize) + 'a,
    {
        let mut timings = Vec::new();
        for (name, kernel) in kernels {
            let t0 = std::time::Instant::now();
            kernel(&mut self.columns, self.agent_count);
            timings.push((name, t0.elapsed().as_micros()));
        }
        if !timings.is_empty() {
            self.dirty = true;
        }
        timings
    }

    /// Download the f64 SoA data back into an `AgentStore`.
    pub fn download<A, S>(&mut self, store: &S)
    where
        A: SoaExtractableF64,
        S: AgentStore<A>,
    {
        soa::write_back_soa_f64::<A, S>(store, &self.ids, &self.columns);
        self.dirty = false;
    }

    /// Remove agents by compacting the f64 SoA arrays.
    pub fn scatter_remove(&mut self, dead_ids: &[AgentId]) {
        if dead_ids.is_empty() {
            return;
        }
        let dead_set: std::collections::HashSet<AgentId> = dead_ids.iter().copied().collect();
        let mut write = 0;
        for read in 0..self.agent_count {
            if !dead_set.contains(&self.ids[read]) {
                if write != read {
                    self.ids[write] = self.ids[read];
                    for col in &mut self.columns {
                        col[write] = col[read];
                    }
                }
                write += 1;
            }
        }
        self.agent_count = write;
        self.ids.truncate(write);
        for col in &mut self.columns {
            col.truncate(write);
        }
        self.dirty = true;
    }

    /// Remove agents with typed validation.
    pub fn try_scatter_remove(
        &mut self,
        dead_ids: &[AgentId],
    ) -> Result<(), DeviceSoaMutationError> {
        validate_remove_ids(dead_ids, &self.ids)?;
        self.scatter_remove(dead_ids);
        Ok(())
    }

    /// Append new agents to the f64 SoA buffers with typed validation.
    pub fn try_scatter_insert(
        &mut self,
        new_ids: &[AgentId],
        new_columns: &[&[f64]],
    ) -> Result<(), DeviceSoaMutationError> {
        validate_insert_shape(new_ids, new_columns, self.columns.len())?;
        validate_insert_ids(new_ids, &self.ids)?;

        let new_count = new_ids.len();
        if new_count == 0 {
            return Ok(());
        }
        self.ids.extend_from_slice(new_ids);
        for (i, col) in self.columns.iter_mut().enumerate() {
            col.extend_from_slice(new_columns[i]);
        }
        self.agent_count += new_count;
        self.dirty = true;
        Ok(())
    }

    /// Apply remove + insert as one validated lifecycle mutation.
    pub fn try_scatter_replace(
        &mut self,
        dead_ids: &[AgentId],
        new_ids: &[AgentId],
        new_columns: &[&[f64]],
    ) -> Result<(), DeviceSoaMutationError> {
        validate_remove_ids(dead_ids, &self.ids)?;
        validate_insert_shape(new_ids, new_columns, self.columns.len())?;
        validate_insert_ids_after_removal(new_ids, &self.ids, dead_ids)?;

        self.scatter_remove(dead_ids);
        append_rows(&mut self.ids, &mut self.columns, new_ids, new_columns);
        self.agent_count = self.ids.len();
        if !(dead_ids.is_empty() && new_ids.is_empty()) {
            self.dirty = true;
        }
        Ok(())
    }
}

fn schema_from_names(names: &[&'static str], column_type: SoaColumnType) -> Vec<SoaColumnSchema> {
    names
        .iter()
        .map(|name| SoaColumnSchema { name, column_type })
        .collect()
}

fn validate_schema(
    schema: &[SoaColumnSchema],
    columns: usize,
    expected: SoaColumnType,
) -> Result<(), DeviceSoaRestoreError> {
    if schema.len() != columns {
        return Err(DeviceSoaRestoreError::SchemaColumnCountMismatch {
            columns,
            schema: schema.len(),
        });
    }
    for (column, entry) in schema.iter().enumerate() {
        if entry.column_type != expected {
            return Err(DeviceSoaRestoreError::SchemaTypeMismatch {
                column,
                expected,
                actual: entry.column_type,
            });
        }
    }
    Ok(())
}

fn validate_column_lengths<T>(
    columns: &[Vec<T>],
    agent_count: usize,
) -> Result<(), DeviceSoaRestoreError> {
    for (column, values) in columns.iter().enumerate() {
        if values.len() != agent_count {
            return Err(DeviceSoaRestoreError::ColumnLengthMismatch {
                column,
                len: values.len(),
                agent_count,
            });
        }
    }
    Ok(())
}

fn validate_unique_ids(ids: &[AgentId]) -> Result<(), DeviceSoaRestoreError> {
    let mut seen = std::collections::HashSet::with_capacity(ids.len());
    for id in ids {
        if !seen.insert(*id) {
            return Err(DeviceSoaRestoreError::DuplicateAgentId(*id));
        }
    }
    Ok(())
}

fn validate_insert_shape<T>(
    new_ids: &[AgentId],
    new_columns: &[&[T]],
    expected_columns: usize,
) -> Result<(), DeviceSoaMutationError> {
    if new_columns.len() != expected_columns {
        return Err(DeviceSoaMutationError::ColumnCountMismatch {
            expected: expected_columns,
            actual: new_columns.len(),
        });
    }
    for (column, values) in new_columns.iter().enumerate() {
        if values.len() != new_ids.len() {
            return Err(DeviceSoaMutationError::ColumnLengthMismatch {
                column,
                expected: new_ids.len(),
                actual: values.len(),
            });
        }
    }
    Ok(())
}

fn validate_insert_ids(
    new_ids: &[AgentId],
    existing_ids: &[AgentId],
) -> Result<(), DeviceSoaMutationError> {
    validate_unique_insert_ids(new_ids)?;
    let existing: std::collections::HashSet<AgentId> = existing_ids.iter().copied().collect();
    for id in new_ids {
        if existing.contains(id) {
            return Err(DeviceSoaMutationError::ExistingAgentId(*id));
        }
    }
    Ok(())
}

fn validate_unique_insert_ids(new_ids: &[AgentId]) -> Result<(), DeviceSoaMutationError> {
    let mut inserted = std::collections::HashSet::with_capacity(new_ids.len());
    for id in new_ids {
        if !inserted.insert(*id) {
            return Err(DeviceSoaMutationError::DuplicateInsertedId(*id));
        }
    }
    Ok(())
}

fn validate_remove_ids(
    dead_ids: &[AgentId],
    existing_ids: &[AgentId],
) -> Result<(), DeviceSoaMutationError> {
    let mut removed = std::collections::HashSet::with_capacity(dead_ids.len());
    for id in dead_ids {
        if !removed.insert(*id) {
            return Err(DeviceSoaMutationError::DuplicateRemovedId(*id));
        }
    }
    let existing: std::collections::HashSet<AgentId> = existing_ids.iter().copied().collect();
    for id in dead_ids {
        if !existing.contains(id) {
            return Err(DeviceSoaMutationError::MissingAgentId(*id));
        }
    }
    Ok(())
}

fn validate_insert_ids_after_removal(
    new_ids: &[AgentId],
    existing_ids: &[AgentId],
    dead_ids: &[AgentId],
) -> Result<(), DeviceSoaMutationError> {
    validate_unique_insert_ids(new_ids)?;
    let removed: std::collections::HashSet<AgentId> = dead_ids.iter().copied().collect();
    let existing: std::collections::HashSet<AgentId> = existing_ids.iter().copied().collect();
    for id in new_ids {
        if existing.contains(id) && !removed.contains(id) {
            return Err(DeviceSoaMutationError::ExistingAgentId(*id));
        }
    }
    Ok(())
}

fn append_rows<T: Copy>(
    ids: &mut Vec<AgentId>,
    columns: &mut [Vec<T>],
    new_ids: &[AgentId],
    new_columns: &[&[T]],
) {
    ids.extend_from_slice(new_ids);
    for (i, col) in columns.iter_mut().enumerate() {
        col.extend_from_slice(new_columns[i]);
    }
}

fn next_available_id(ids: &[AgentId]) -> AgentId {
    let existing: std::collections::HashSet<AgentId> = ids.iter().copied().collect();
    let mut candidate = 0;
    while existing.contains(&candidate) {
        candidate = candidate.saturating_add(1);
    }
    candidate
}

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

    #[derive(Debug, Clone)]
    struct TestAgent {
        id: AgentId,
        x: f32,
        vx: f32,
    }

    impl Agent for TestAgent {
        fn id(&self) -> AgentId {
            self.id
        }
    }

    impl SoaExtractable for TestAgent {
        fn num_columns() -> usize {
            2
        }
        fn column_names() -> Vec<&'static str> {
            vec!["x", "vx"]
        }
        fn extract_row(&self, columns: &mut [Vec<f32>]) {
            columns[0].push(self.x);
            columns[1].push(self.vx);
        }
        fn write_back_row(&mut self, columns: &[&[f32]], row: usize) {
            self.x = columns[0][row];
            self.vx = columns[1][row];
        }
    }

    #[test]
    fn persistent_soa_step() {
        let mut store = HashMapStore::new();
        for i in 1..=100 {
            store.insert(TestAgent {
                id: i,
                x: 0.0,
                vx: 1.0,
            });
        }

        let mut device = DeviceSoaStore::upload::<TestAgent, _>(&store);
        assert_eq!(device.agent_count(), 100);
        assert_eq!(device.num_columns(), 2);

        // Step 10 times without re-extraction
        for _ in 0..10 {
            device.step_cpu(|columns, n| {
                let (x_col, rest) = columns.split_at_mut(1);
                let x = &mut x_col[0];
                let vx = &rest[0];
                for i in 0..n {
                    x[i] += vx[i];
                }
            });
        }

        // Verify
        assert!(device.is_dirty());
        // x should be 10.0 for all agents
        for i in 0..device.agent_count() {
            assert!((device.column(0)[i] - 10.0).abs() < 1e-5);
        }

        // Download back
        device.download::<TestAgent, _>(&store);
        assert!(!device.is_dirty());
    }

    #[test]
    fn run_sequence_cpu_applies_kernels_in_order() {
        let mut store = HashMapStore::new();
        for i in 1..=4 {
            store.insert(TestAgent {
                id: i,
                x: 0.0,
                vx: 2.0,
            });
        }

        let mut device = DeviceSoaStore::upload::<TestAgent, _>(&store);

        let timings = device.run_sequence_cpu(vec![
            (
                "advance",
                Box::new(|cols: &mut [Vec<f32>], n: usize| {
                    let (x_col, rest) = cols.split_at_mut(1);
                    let x = &mut x_col[0];
                    let vx = &rest[0];
                    for i in 0..n {
                        x[i] += vx[i];
                    }
                }) as Box<dyn FnOnce(&mut [Vec<f32>], usize)>,
            ),
            (
                "scale_vx",
                Box::new(|cols: &mut [Vec<f32>], n: usize| {
                    for v in cols[1].iter_mut().take(n) {
                        *v *= 3.0;
                    }
                }),
            ),
            (
                "advance_again",
                Box::new(|cols: &mut [Vec<f32>], n: usize| {
                    let (x_col, rest) = cols.split_at_mut(1);
                    let x = &mut x_col[0];
                    let vx = &rest[0];
                    for i in 0..n {
                        x[i] += vx[i];
                    }
                }),
            ),
        ]);

        assert_eq!(timings.len(), 3);
        assert_eq!(timings[0].0, "advance");
        assert_eq!(timings[1].0, "scale_vx");
        assert_eq!(timings[2].0, "advance_again");

        // Start: x=0, vx=2 -> after advance: x=2, vx=2
        // After scale_vx: vx=6 -> after advance_again: x=8, vx=6
        assert!(device.is_dirty());
        for i in 0..device.agent_count() {
            assert!((device.column(0)[i] - 8.0).abs() < 1e-5);
            assert!((device.column(1)[i] - 6.0).abs() < 1e-5);
        }
    }

    #[test]
    fn scatter_remove_and_insert() {
        let mut store = HashMapStore::new();
        for i in 1..=10 {
            store.insert(TestAgent {
                id: i,
                x: i as f32,
                vx: 0.0,
            });
        }

        let mut device = DeviceSoaStore::upload::<TestAgent, _>(&store);
        assert_eq!(device.agent_count(), 10);

        // Remove agents 3 and 7
        device.scatter_remove(&[3, 7]);
        assert_eq!(device.agent_count(), 8);
        assert!(!device.ids().contains(&3));
        assert!(!device.ids().contains(&7));

        // Insert 2 new agents
        device.scatter_insert(&[11, 12], &[&[11.0, 12.0], &[0.0, 0.0]]);
        assert_eq!(device.agent_count(), 10);
        assert!(device.ids().contains(&11));
        assert!(device.ids().contains(&12));
    }
}