rustsim_core/

messaging.rs

//! Inter-agent messaging system inspired by FlameGPU2.
//!
//! FlameGPU2 achieves massive GPU throughput by decoupling agent communication
//! from direct neighbor queries. Agents **output messages** in one phase, then
//! **read messages** in the next. This two-phase approach eliminates data races
//! and enables efficient GPU-side spatial indexing.
//!
//! rustsim mirrors this with a CPU-first messaging layer that can be backed
//! by GPU buffers when CUDA is available.
//!
//! # Message Types
//!
//! | Type | Description | FlameGPU2 equivalent |
//! |------|-------------|---------------------|
//! | [`BruteForceMessages`] | All-to-all; every agent reads every message | `MessageBruteForce` |
//! | [`SpatialMessages2D`] | Spatial-hashing 2D; agents read messages within a radius | `MessageSpatial2D` |
//! | [`SpatialMessages3D`] | Spatial-hashing 3D; agents read messages within a radius | `MessageSpatial3D` |
//!
//! # Two-Phase Pattern
//!
//! ```ignore
//! // Phase 1: agents output messages
//! for id in agent_ids {
//!     let agent = store.get(id);
//!     messages.output(MyMessage { x: agent.x, y: agent.y, vx: agent.vx });
//! }
//! messages.finalize(); // build spatial index
//!
//! // Phase 2: agents read messages and update
//! for id in agent_ids {
//!     let mut agent = store.get_mut(id);
//!     for msg in messages.read_nearby(agent.x, agent.y, radius) {
//!         // accumulate forces, etc.
//!     }
//! }
//! messages.clear();
//! ```

use std::collections::HashMap;
use thiserror::Error;

/// Errors returned by messaging configuration validation.
#[derive(Debug, Clone, Copy, PartialEq, Error)]
pub enum MessageConfigError {
    #[error("radius must be positive")]
    InvalidRadius,
}

/// Errors returned when a message buffer is used outside its valid phase.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
pub enum MessagePhaseError {
    /// Messages can only be output before finalization.
    #[error("cannot output messages after finalize")]
    AlreadyFinalized,
    /// Messages can only be read after finalization.
    #[error("must finalize messages before reading")]
    NotFinalized,
}

/// A message that can be exchanged between agents.
///
/// Messages are value types -- they are copied into the message buffer
/// during the output phase. Implement this trait for your message struct.
pub trait Message: Clone + Send + Sync + 'static {}

// Blanket impl: anything Clone + Send + Sync + 'static is a Message.
impl<T: Clone + Send + Sync + 'static> Message for T {}

/// Brute-force message list -- every agent can read every message.
///
/// This is the simplest messaging pattern. Each step, agents push messages
/// in the output phase, then iterate all messages in the input phase.
///
/// Performance: O(N) per agent for N messages. Suitable for small populations
/// or global broadcasts.
///
/// Mirrors FlameGPU2 `MessageBruteForce`.
#[derive(Debug, Clone)]
pub struct BruteForceMessages<M> {
    buffer: Vec<M>,
    finalized: bool,
}

impl<M: Clone> BruteForceMessages<M> {
    /// Create an empty brute-force message list.
    pub fn new() -> Self {
        Self {
            buffer: Vec::new(),
            finalized: false,
        }
    }

    /// Create with pre-allocated capacity.
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            buffer: Vec::with_capacity(capacity),
            finalized: false,
        }
    }

    /// Output a message (phase 1).
    pub fn output(&mut self, message: M) {
        self.try_output(message)
            .expect("cannot output after finalize");
    }

    /// Try to output a message (phase 1), returning a typed phase error on misuse.
    pub fn try_output(&mut self, message: M) -> Result<(), MessagePhaseError> {
        if self.finalized {
            return Err(MessagePhaseError::AlreadyFinalized);
        }
        self.buffer.push(message);
        Ok(())
    }

    /// Finalize the message list for reading (phase boundary).
    ///
    /// For brute-force, this is a no-op but enforces the two-phase protocol.
    pub fn finalize(&mut self) {
        self.finalized = true;
    }

    /// Iterate all messages (phase 2).
    ///
    /// # Panics
    ///
    /// Panics if called before [`finalize`](Self::finalize). Use
    /// [`try_read_all`](Self::try_read_all) to handle phase errors explicitly.
    pub fn read_all(&self) -> &[M] {
        self.try_read_all()
            .expect("must finalize messages before reading")
    }

    /// Try to iterate all messages (phase 2), returning a typed phase error on misuse.
    pub fn try_read_all(&self) -> Result<&[M], MessagePhaseError> {
        if !self.finalized {
            return Err(MessagePhaseError::NotFinalized);
        }
        Ok(&self.buffer)
    }

    /// Whether the message buffer has been finalized for reading.
    pub fn is_finalized(&self) -> bool {
        self.finalized
    }

    /// Number of messages in the buffer.
    pub fn len(&self) -> usize {
        self.buffer.len()
    }

    /// Returns `true` if no messages have been output.
    pub fn is_empty(&self) -> bool {
        self.buffer.is_empty()
    }

    /// Clear the message buffer for the next step.
    pub fn clear(&mut self) {
        self.buffer.clear();
        self.finalized = false;
    }
}

impl<M: Clone> Default for BruteForceMessages<M> {
    fn default() -> Self {
        Self::new()
    }
}

/// 2D spatial message list with spatial-hashing neighbor lookup.
///
/// Messages carry an `(x, y)` position. During finalization, messages are
/// sorted into a uniform grid of bins (side length = `radius`). During
/// the read phase, agents query nearby messages by position and only
/// iterate bins within the search radius.
///
/// This mirrors FlameGPU2's `MessageSpatial2D` and its **Partition Boundary
/// Matrix (PBM)** -- a sorted index of messages by spatial bin.
///
/// Performance: O(k) per agent where k = messages in neighboring bins.
#[derive(Debug, Clone)]
pub struct SpatialMessages2D<M> {
    messages: Vec<M>,
    positions: Vec<(f32, f32)>,
    radius: f32,
    // PBM-like: bin -> (start_index, count) into the sorted message arrays
    bin_map: HashMap<(i32, i32), (usize, usize)>,
    // Sorted indices for spatial iteration
    sorted_indices: Vec<usize>,
    finalized: bool,
}

impl<M: Clone> SpatialMessages2D<M> {
    /// Create a new 2D spatial message list.
    ///
    /// `radius` is both the search radius and the bin side length.
    pub fn new(radius: f32) -> Result<Self, MessageConfigError> {
        if radius <= 0.0 {
            return Err(MessageConfigError::InvalidRadius);
        }
        Ok(Self {
            messages: Vec::new(),
            positions: Vec::new(),
            radius,
            bin_map: HashMap::new(),
            sorted_indices: Vec::new(),
            finalized: false,
        })
    }

    /// Output a message at a 2D position (phase 1).
    pub fn output(&mut self, message: M, x: f32, y: f32) {
        self.try_output(message, x, y)
            .expect("cannot output after finalize");
    }

    /// Try to output a message at a 2D position, returning a typed phase error on misuse.
    pub fn try_output(&mut self, message: M, x: f32, y: f32) -> Result<(), MessagePhaseError> {
        if self.finalized {
            return Err(MessagePhaseError::AlreadyFinalized);
        }
        self.messages.push(message);
        self.positions.push((x, y));
        Ok(())
    }

    /// Build the spatial index (phase boundary).
    ///
    /// This constructs a PBM-like structure: messages are logically sorted
    /// by bin, and a map records the (start, count) of each bin.
    pub fn finalize(&mut self) {
        let n = self.messages.len();
        let inv_radius = 1.0 / self.radius;

        // Assign each message to a bin
        let mut bin_assignments: Vec<(i32, i32)> = Vec::with_capacity(n);
        for &(x, y) in &self.positions {
            let bx = (x * inv_radius).floor() as i32;
            let by = (y * inv_radius).floor() as i32;
            bin_assignments.push((bx, by));
        }

        // Create sorted index by bin (mimics PBM sort)
        self.sorted_indices.clear();
        self.sorted_indices.extend(0..n);
        self.sorted_indices
            .sort_unstable_by(|&a, &b| bin_assignments[a].cmp(&bin_assignments[b]));

        // Build bin_map: (bin) -> (start_in_sorted, count)
        self.bin_map.clear();
        if !self.sorted_indices.is_empty() {
            let mut current_bin = bin_assignments[self.sorted_indices[0]];
            let mut start = 0;

            for i in 1..n {
                let bin = bin_assignments[self.sorted_indices[i]];
                if bin != current_bin {
                    self.bin_map.insert(current_bin, (start, i - start));
                    current_bin = bin;
                    start = i;
                }
            }
            self.bin_map.insert(current_bin, (start, n - start));
        }

        self.finalized = true;
    }

    /// Iterate messages within `radius` of `(x, y)` (phase 2).
    ///
    /// Returns an iterator yielding `(message_ref, distance_squared)`.
    /// The caller should filter by `distance_squared <= radius * radius`
    /// for exact circular queries (the method returns a superset from
    /// the bin neighborhood).
    ///
    /// # Panics
    ///
    /// Panics if called before [`finalize`](Self::finalize). Use
    /// [`try_read_nearby`](Self::try_read_nearby) to handle phase errors explicitly.
    pub fn read_nearby(&self, x: f32, y: f32, radius: f32) -> SpatialIter2D<'_, M> {
        self.try_read_nearby(x, y, radius)
            .expect("must finalize messages before reading")
    }

    /// Try to iterate messages within `radius` of `(x, y)` (phase 2).
    pub fn try_read_nearby(
        &self,
        x: f32,
        y: f32,
        radius: f32,
    ) -> Result<SpatialIter2D<'_, M>, MessagePhaseError> {
        if !self.finalized {
            return Err(MessagePhaseError::NotFinalized);
        }
        let inv = 1.0 / self.radius;
        let center_bx = (x * inv).floor() as i32;
        let center_by = (y * inv).floor() as i32;
        let grid_r = ((radius / self.radius).ceil() as i32).max(1);

        // Seed the first bin
        let first_dx = -grid_r;
        let first_dy = -grid_r;
        let bx = center_bx + first_dx;
        let by = center_by + first_dy;
        let (bin_start, bin_count) = self.bin_map.get(&(bx, by)).copied().unwrap_or((0, 0));

        Ok(SpatialIter2D {
            messages: &self.messages,
            positions: &self.positions,
            sorted_indices: &self.sorted_indices,
            bin_map: &self.bin_map,
            query_x: x,
            query_y: y,
            radius_sq: radius * radius,
            center_bx,
            center_by,
            grid_r,
            cur_dx: first_dx,
            cur_dy: first_dy,
            bin_start,
            bin_offset: 0,
            bin_count,
        })
    }

    /// Whether the message buffer has been finalized for reading.
    pub fn is_finalized(&self) -> bool {
        self.finalized
    }

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

    /// Returns `true` if no messages have been output.
    pub fn is_empty(&self) -> bool {
        self.messages.is_empty()
    }

    /// Clear all messages and the spatial index for the next step.
    pub fn clear(&mut self) {
        self.messages.clear();
        self.positions.clear();
        self.bin_map.clear();
        self.sorted_indices.clear();
        self.finalized = false;
    }
}

/// Iterator over spatially-nearby messages.
pub struct SpatialIter2D<'a, M> {
    messages: &'a [M],
    positions: &'a [(f32, f32)],
    sorted_indices: &'a [usize],
    bin_map: &'a HashMap<(i32, i32), (usize, usize)>,
    query_x: f32,
    query_y: f32,
    radius_sq: f32,
    center_bx: i32,
    center_by: i32,
    grid_r: i32,
    cur_dx: i32,
    cur_dy: i32,
    // Current bin's start offset in sorted_indices and remaining count
    bin_start: usize,
    bin_offset: usize,
    bin_count: usize,
}

impl<'a, M> Iterator for SpatialIter2D<'a, M> {
    type Item = (&'a M, f32);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            // Drain entries from the current bin
            while self.bin_offset < self.bin_count {
                let idx = self.sorted_indices[self.bin_start + self.bin_offset];
                self.bin_offset += 1;

                let (px, py) = self.positions[idx];
                let dx = self.query_x - px;
                let dy = self.query_y - py;
                let dist_sq = dx * dx + dy * dy;
                if dist_sq <= self.radius_sq {
                    return Some((&self.messages[idx], dist_sq));
                }
            }

            // Advance to next bin in the grid neighborhood
            loop {
                self.cur_dx += 1;
                if self.cur_dx > self.grid_r {
                    self.cur_dx = -self.grid_r;
                    self.cur_dy += 1;
                    if self.cur_dy > self.grid_r {
                        return None; // exhausted all bins
                    }
                }

                let bx = self.center_bx + self.cur_dx;
                let by = self.center_by + self.cur_dy;
                if let Some(&(start, count)) = self.bin_map.get(&(bx, by)) {
                    self.bin_start = start;
                    self.bin_offset = 0;
                    self.bin_count = count;
                    break; // found a non-empty bin, go drain it
                }
            }
        }
    }
}

/// 3D spatial message list with spatial-hashing neighbor lookup.
///
/// Same concept as [`SpatialMessages2D`] but for 3D positions.
///
/// Mirrors FlameGPU2 `MessageSpatial3D` and its 3D Partition Boundary Matrix.
#[derive(Debug, Clone)]
pub struct SpatialMessages3D<M> {
    messages: Vec<M>,
    positions: Vec<(f32, f32, f32)>,
    radius: f32,
    bin_map: HashMap<(i32, i32, i32), (usize, usize)>,
    sorted_indices: Vec<usize>,
    finalized: bool,
}

impl<M: Clone> SpatialMessages3D<M> {
    /// Create a new 3D spatial message list.
    pub fn new(radius: f32) -> Result<Self, MessageConfigError> {
        if radius <= 0.0 {
            return Err(MessageConfigError::InvalidRadius);
        }
        Ok(Self {
            messages: Vec::new(),
            positions: Vec::new(),
            radius,
            bin_map: HashMap::new(),
            sorted_indices: Vec::new(),
            finalized: false,
        })
    }

    /// Output a message at a 3D position (phase 1).
    pub fn output(&mut self, message: M, x: f32, y: f32, z: f32) {
        self.try_output(message, x, y, z)
            .expect("cannot output after finalize");
    }

    /// Try to output a message at a 3D position, returning a typed phase error on misuse.
    pub fn try_output(
        &mut self,
        message: M,
        x: f32,
        y: f32,
        z: f32,
    ) -> Result<(), MessagePhaseError> {
        if self.finalized {
            return Err(MessagePhaseError::AlreadyFinalized);
        }
        self.messages.push(message);
        self.positions.push((x, y, z));
        Ok(())
    }

    /// Build the spatial index (phase boundary).
    pub fn finalize(&mut self) {
        let n = self.messages.len();
        let inv_radius = 1.0 / self.radius;

        let mut bin_assignments: Vec<(i32, i32, i32)> = Vec::with_capacity(n);
        for &(x, y, z) in &self.positions {
            let bx = (x * inv_radius).floor() as i32;
            let by = (y * inv_radius).floor() as i32;
            let bz = (z * inv_radius).floor() as i32;
            bin_assignments.push((bx, by, bz));
        }

        self.sorted_indices.clear();
        self.sorted_indices.extend(0..n);
        self.sorted_indices
            .sort_unstable_by(|&a, &b| bin_assignments[a].cmp(&bin_assignments[b]));

        self.bin_map.clear();
        if !self.sorted_indices.is_empty() {
            let mut current_bin = bin_assignments[self.sorted_indices[0]];
            let mut start = 0;

            for i in 1..n {
                let bin = bin_assignments[self.sorted_indices[i]];
                if bin != current_bin {
                    self.bin_map.insert(current_bin, (start, i - start));
                    current_bin = bin;
                    start = i;
                }
            }
            self.bin_map.insert(current_bin, (start, n - start));
        }

        self.finalized = true;
    }

    /// Iterate messages within `radius` of `(x, y, z)` (phase 2).
    ///
    /// Returns an iterator yielding `(message_ref, distance_squared)`.
    ///
    /// # Panics
    ///
    /// Panics if called before [`finalize`](Self::finalize). Use
    /// [`try_read_nearby`](Self::try_read_nearby) to handle phase errors explicitly.
    pub fn read_nearby(&self, x: f32, y: f32, z: f32, radius: f32) -> SpatialIter3D<'_, M> {
        self.try_read_nearby(x, y, z, radius)
            .expect("must finalize messages before reading")
    }

    /// Try to iterate messages within `radius` of `(x, y, z)` (phase 2).
    pub fn try_read_nearby(
        &self,
        x: f32,
        y: f32,
        z: f32,
        radius: f32,
    ) -> Result<SpatialIter3D<'_, M>, MessagePhaseError> {
        if !self.finalized {
            return Err(MessagePhaseError::NotFinalized);
        }
        let inv = 1.0 / self.radius;
        let center_bx = (x * inv).floor() as i32;
        let center_by = (y * inv).floor() as i32;
        let center_bz = (z * inv).floor() as i32;
        let grid_r = ((radius / self.radius).ceil() as i32).max(1);

        // Seed the first bin
        let first_dx = -grid_r;
        let first_dy = -grid_r;
        let first_dz = -grid_r;
        let bx = center_bx + first_dx;
        let by = center_by + first_dy;
        let bz = center_bz + first_dz;
        let (bin_start, bin_count) = self.bin_map.get(&(bx, by, bz)).copied().unwrap_or((0, 0));

        Ok(SpatialIter3D {
            messages: &self.messages,
            positions: &self.positions,
            sorted_indices: &self.sorted_indices,
            bin_map: &self.bin_map,
            query_x: x,
            query_y: y,
            query_z: z,
            radius_sq: radius * radius,
            center_bx,
            center_by,
            center_bz,
            grid_r,
            cur_dx: first_dx,
            cur_dy: first_dy,
            cur_dz: first_dz,
            bin_start,
            bin_offset: 0,
            bin_count,
        })
    }

    /// Whether the message buffer has been finalized for reading.
    pub fn is_finalized(&self) -> bool {
        self.finalized
    }

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

    /// Returns `true` if no messages have been output.
    pub fn is_empty(&self) -> bool {
        self.messages.is_empty()
    }

    /// Clear all messages and the spatial index for the next step.
    pub fn clear(&mut self) {
        self.messages.clear();
        self.positions.clear();
        self.bin_map.clear();
        self.sorted_indices.clear();
        self.finalized = false;
    }
}

/// Iterator over spatially-nearby 3D messages.
pub struct SpatialIter3D<'a, M> {
    messages: &'a [M],
    positions: &'a [(f32, f32, f32)],
    sorted_indices: &'a [usize],
    bin_map: &'a HashMap<(i32, i32, i32), (usize, usize)>,
    query_x: f32,
    query_y: f32,
    query_z: f32,
    radius_sq: f32,
    center_bx: i32,
    center_by: i32,
    center_bz: i32,
    grid_r: i32,
    cur_dx: i32,
    cur_dy: i32,
    cur_dz: i32,
    bin_start: usize,
    bin_offset: usize,
    bin_count: usize,
}

impl<'a, M> Iterator for SpatialIter3D<'a, M> {
    type Item = (&'a M, f32);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            // Drain entries from the current bin
            while self.bin_offset < self.bin_count {
                let idx = self.sorted_indices[self.bin_start + self.bin_offset];
                self.bin_offset += 1;

                let (px, py, pz) = self.positions[idx];
                let dx = self.query_x - px;
                let dy = self.query_y - py;
                let dz = self.query_z - pz;
                let dist_sq = dx * dx + dy * dy + dz * dz;
                if dist_sq <= self.radius_sq {
                    return Some((&self.messages[idx], dist_sq));
                }
            }

            // Advance to next bin in the 3D grid neighborhood
            loop {
                self.cur_dx += 1;
                if self.cur_dx > self.grid_r {
                    self.cur_dx = -self.grid_r;
                    self.cur_dy += 1;
                    if self.cur_dy > self.grid_r {
                        self.cur_dy = -self.grid_r;
                        self.cur_dz += 1;
                        if self.cur_dz > self.grid_r {
                            return None; // exhausted all bins
                        }
                    }
                }

                let bx = self.center_bx + self.cur_dx;
                let by = self.center_by + self.cur_dy;
                let bz = self.center_bz + self.cur_dz;
                if let Some(&(start, count)) = self.bin_map.get(&(bx, by, bz)) {
                    self.bin_start = start;
                    self.bin_offset = 0;
                    self.bin_count = count;
                    break; // found a non-empty bin, go drain it
                }
            }
        }
    }
}

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

    #[test]
    fn brute_force_basic() {
        let mut msgs = BruteForceMessages::new();
        msgs.output(42i32);
        msgs.output(99);
        msgs.finalize();
        assert_eq!(msgs.read_all(), &[42, 99]);
        msgs.clear();
        assert!(msgs.is_empty());
    }

    #[test]
    fn brute_force_phase_errors_are_typed() {
        let mut msgs = BruteForceMessages::new();
        assert_eq!(msgs.try_read_all(), Err(MessagePhaseError::NotFinalized));
        msgs.try_output(1).unwrap();
        msgs.finalize();
        assert_eq!(msgs.try_output(2), Err(MessagePhaseError::AlreadyFinalized));
        assert_eq!(msgs.try_read_all().unwrap(), &[1]);
    }

    #[test]
    fn spatial_2d_basic() {
        let mut msgs = SpatialMessages2D::new(1.0).unwrap();
        msgs.output("a", 0.0, 0.0);
        msgs.output("b", 0.5, 0.5);
        msgs.output("c", 10.0, 10.0);
        msgs.finalize();

        // Query near origin -- should find "a" and "b" but not "c"
        let nearby: Vec<_> = msgs.read_nearby(0.0, 0.0, 1.0).collect();
        assert_eq!(nearby.len(), 2);
        let labels: Vec<&str> = nearby.iter().map(|(&m, _)| m).collect();
        assert!(labels.contains(&"a"));
        assert!(labels.contains(&"b"));
    }

    #[test]
    fn spatial_2d_phase_errors_are_typed() {
        let mut msgs = SpatialMessages2D::new(1.0).unwrap();
        assert_eq!(
            msgs.try_read_nearby(0.0, 0.0, 1.0).err(),
            Some(MessagePhaseError::NotFinalized)
        );
        msgs.try_output("a", 0.0, 0.0).unwrap();
        msgs.finalize();
        assert_eq!(
            msgs.try_output("b", 1.0, 1.0),
            Err(MessagePhaseError::AlreadyFinalized)
        );
    }

    #[test]
    fn spatial_3d_basic() {
        let mut msgs = SpatialMessages3D::new(1.0).unwrap();
        msgs.output("a", 0.0, 0.0, 0.0);
        msgs.output("b", 0.5, 0.5, 0.5);
        msgs.output("c", 10.0, 10.0, 10.0);
        msgs.finalize();

        let nearby: Vec<_> = msgs.read_nearby(0.0, 0.0, 0.0, 1.0).collect();
        assert_eq!(nearby.len(), 2);
        let labels: Vec<&str> = nearby.iter().map(|(&m, _)| m).collect();
        assert!(labels.contains(&"a"));
        assert!(labels.contains(&"b"));
    }

    #[test]
    fn spatial_3d_phase_errors_are_typed() {
        let mut msgs = SpatialMessages3D::new(1.0).unwrap();
        assert_eq!(
            msgs.try_read_nearby(0.0, 0.0, 0.0, 1.0).err(),
            Some(MessagePhaseError::NotFinalized)
        );
        msgs.try_output("a", 0.0, 0.0, 0.0).unwrap();
        msgs.finalize();
        assert_eq!(
            msgs.try_output("b", 1.0, 1.0, 1.0),
            Err(MessagePhaseError::AlreadyFinalized)
        );
    }
}