khora_data/ecs/

soa.rs

// Copyright 2025 eraflo
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Field-split (Structure-of-Arrays) component storage — the physical layout
//! the AGDF compute kernels want.
//!
//! A normal CRPECS column is `Vec<T>` (Array-of-Structures *within* the
//! component): one struct after another. For a compute-heavy hot loop that
//! sweeps a few fields over many entities, that wastes the CPU's SIMD slots and
//! cache. A [`FieldSoaColumn`] instead stores **one contiguous `f32` array per
//! field**, so each field streams without gather and tiles cleanly into
//! `f32x8` (the resident layout the bench showed reaches ~4×).
//!
//! This is the data-layer realisation of "adapt the HOW, not the WHAT": the
//! component's *semantics* are unchanged, only its column's physical mapping.
//! It is **opt-in per component** (`#[component(layout = "soa")]`) and lives
//! behind the same `Box<dyn AnyVec>` as any column, so the rest of CRPECS —
//! pages, domains, queries on other components, serialization — is untouched.
//! A field-SoA component can't yield `&T` (its bytes aren't a contiguous `T`),
//! so it is read by value via the `Soa<T>` query / `World::clone_component`, or
//! in bulk via the per-field arrays for SIMD.
//!
//! The generic column needs only a tiny amount of per-type knowledge — how to
//! scatter a value into the field arrays and gather it back — supplied by
//! [`SoaLayout`], which `#[derive(Component)]` generates for `layout = "soa"`
//! structs whose fields are all `f32`.
//!
//! (Note: "field array" here is the per-component-field `Vec<f32>`; it is
//! unrelated to the engine's [`Lane`](crate::ecs) negotiation concept.)

use std::any::Any;
use std::marker::PhantomData;

use crate::ecs::component::Component;
use crate::ecs::page::{AnyVec, SetFromBytesError, MAX_COLUMN_PAYLOAD_BYTES};

/// Per-type knowledge the generic [`FieldSoaColumn`] needs to scatter a
/// component into its `f32` field arrays and gather it back.
///
/// Generated by `#[derive(Component)]` + `#[component(layout = "soa")]` for
/// structs whose fields are all `f32`. The array order matches field
/// declaration order.
pub trait SoaLayout: Component {
    /// Number of `f32` fields this component decomposes into.
    const FIELD_COUNT: usize;
    /// Field names in declaration order (for the glass-box / debugging).
    const FIELD_NAMES: &'static [&'static str];

    /// Appends each of `self`'s fields onto the matching field array (one new
    /// row).
    ///
    /// `fields.len() == FIELD_COUNT` is guaranteed by the caller.
    fn scatter_push(&self, fields: &mut [Vec<f32>]);

    /// Overwrites row `row` of every field array with `self`'s fields.
    ///
    /// `fields.len() == FIELD_COUNT` and `row` in bounds are guaranteed.
    fn scatter_set(&self, fields: &mut [Vec<f32>], row: usize);

    /// Reconstructs a value from row `row` across the field arrays.
    fn gather(fields: &[Vec<f32>], row: usize) -> Self;
}

/// A field-split component column: one contiguous `Vec<f32>` per component
/// field.
///
/// Stored behind `Box<dyn AnyVec>` exactly like an AoS `Vec<T>` column, so it
/// participates in the normal page lifecycle (spawn / despawn / migration /
/// serialization). Access is by value ([`get`](Self::get)) or in bulk through
/// the per-field [`field`](Self::field) slices for SIMD kernels.
#[derive(Debug)]
pub struct FieldSoaColumn<T: SoaLayout> {
    /// `FIELD_COUNT` field arrays, each holding one `f32` per row. All arrays
    /// have the same length (the row count).
    fields: Vec<Vec<f32>>,
    _marker: PhantomData<fn() -> T>,
}

impl<T: SoaLayout> Default for FieldSoaColumn<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T: SoaLayout> FieldSoaColumn<T> {
    /// Creates an empty column with one field array per component field.
    pub fn new() -> Self {
        Self {
            fields: (0..T::FIELD_COUNT).map(|_| Vec::new()).collect(),
            _marker: PhantomData,
        }
    }

    /// The number of rows (entities) stored.
    pub fn len(&self) -> usize {
        self.fields.first().map_or(0, Vec::len)
    }

    /// Whether the column holds no rows.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Appends a value as a new row.
    pub fn push(&mut self, value: T) {
        value.scatter_push(&mut self.fields);
    }

    /// Reconstructs the value stored at `row`.
    ///
    /// # Panics
    /// Panics if `row >= len()`.
    pub fn get(&self, row: usize) -> T {
        T::gather(&self.fields, row)
    }

    /// Overwrites the value at `row`.
    ///
    /// # Panics
    /// Panics if `row >= len()`.
    pub fn set(&mut self, row: usize, value: T) {
        value.scatter_set(&mut self.fields, row);
    }

    /// Read-only access to one field's array (`field < FIELD_COUNT`) — the
    /// contiguous `f32` stream a SIMD kernel tiles over.
    pub fn field(&self, field: usize) -> &[f32] {
        &self.fields[field]
    }

    /// Mutable access to one field's array.
    pub fn field_mut(&mut self, field: usize) -> &mut [f32] {
        &mut self.fields[field]
    }

    /// All field arrays (read-only).
    pub fn fields(&self) -> &[Vec<f32>] {
        &self.fields
    }

    /// All field arrays (mutable) — for a kernel that writes several fields.
    pub fn fields_mut(&mut self) -> &mut [Vec<f32>] {
        &mut self.fields
    }
}

impl<T: SoaLayout> AnyVec for FieldSoaColumn<T> {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn as_any_mut(&mut self) -> &mut dyn Any {
        self
    }

    fn swap_remove_any(&mut self, index: usize) {
        for field in &mut self.fields {
            field.swap_remove(index);
        }
    }

    fn to_bytes(&self) -> Vec<u8> {
        // Field-major little-endian: field 0's f32s, then field 1's, … This is
        // the column's own self-consistent format; `set_from_bytes` reverses it.
        let mut out = Vec::with_capacity(self.len() * T::FIELD_COUNT * 4);
        for field in &self.fields {
            for &v in field {
                out.extend_from_slice(&v.to_le_bytes());
            }
        }
        out
    }

    unsafe fn set_from_bytes(&mut self, bytes: &[u8]) -> Result<(), SetFromBytesError> {
        let field_count = T::FIELD_COUNT;
        // One row occupies `field_count` little-endian f32s. A zero-field layout
        // has no payload, mirroring the ZST handling in the `Vec<T>` column.
        let row_size = field_count.saturating_mul(4);
        if row_size == 0 {
            return Ok(());
        }

        // Validate the untrusted length before reserving any rows: it must be an
        // exact multiple of a full row and stay under the payload ceiling.
        if !bytes.len().is_multiple_of(row_size) {
            return Err(SetFromBytesError::MisalignedLength {
                len: bytes.len(),
                elem_size: row_size,
            });
        }
        if bytes.len() > MAX_COLUMN_PAYLOAD_BYTES {
            return Err(SetFromBytesError::PayloadTooLarge {
                len: bytes.len(),
                max: MAX_COLUMN_PAYLOAD_BYTES,
            });
        }

        let rows = bytes.len() / row_size;
        let mut idx = 0usize;
        for field in &mut self.fields {
            field.clear();
            field.reserve(rows);
            for _ in 0..rows {
                let b = [
                    bytes[idx * 4],
                    bytes[idx * 4 + 1],
                    bytes[idx * 4 + 2],
                    bytes[idx * 4 + 3],
                ];
                field.push(f32::from_le_bytes(b));
                idx += 1;
            }
        }
        Ok(())
    }
}

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

    // A hand-written SoaLayout to test the generic column without the macro.
    #[derive(Debug, Clone, Copy, PartialEq)]
    struct P {
        x: f32,
        y: f32,
        z: f32,
    }
    impl Component for P {}
    impl SoaLayout for P {
        const FIELD_COUNT: usize = 3;
        const FIELD_NAMES: &'static [&'static str] = &["x", "y", "z"];
        fn scatter_push(&self, fields: &mut [Vec<f32>]) {
            fields[0].push(self.x);
            fields[1].push(self.y);
            fields[2].push(self.z);
        }
        fn scatter_set(&self, fields: &mut [Vec<f32>], row: usize) {
            fields[0][row] = self.x;
            fields[1][row] = self.y;
            fields[2][row] = self.z;
        }
        fn gather(fields: &[Vec<f32>], row: usize) -> Self {
            P {
                x: fields[0][row],
                y: fields[1][row],
                z: fields[2][row],
            }
        }
    }

    #[test]
    fn push_get_set_roundtrip() {
        let mut col = FieldSoaColumn::<P>::new();
        col.push(P {
            x: 1.0,
            y: 2.0,
            z: 3.0,
        });
        col.push(P {
            x: 4.0,
            y: 5.0,
            z: 6.0,
        });
        assert_eq!(col.len(), 2);
        assert_eq!(
            col.get(0),
            P {
                x: 1.0,
                y: 2.0,
                z: 3.0
            }
        );
        assert_eq!(
            col.get(1),
            P {
                x: 4.0,
                y: 5.0,
                z: 6.0
            }
        );
        col.set(
            0,
            P {
                x: 7.0,
                y: 8.0,
                z: 9.0,
            },
        );
        assert_eq!(
            col.get(0),
            P {
                x: 7.0,
                y: 8.0,
                z: 9.0
            }
        );
        // Each field array is contiguous (the SIMD-friendly property).
        assert_eq!(col.field(0), &[7.0, 4.0]);
        assert_eq!(col.field(2), &[9.0, 6.0]);
    }

    #[test]
    fn swap_remove_keeps_fields_aligned() {
        let mut col = FieldSoaColumn::<P>::new();
        for i in 0..3 {
            let f = i as f32;
            col.push(P {
                x: f,
                y: f + 10.0,
                z: f + 20.0,
            });
        }
        col.swap_remove_any(0); // row 0 replaced by last row (index 2)
        assert_eq!(col.len(), 2);
        assert_eq!(
            col.get(0),
            P {
                x: 2.0,
                y: 12.0,
                z: 22.0
            }
        );
        assert_eq!(
            col.get(1),
            P {
                x: 1.0,
                y: 11.0,
                z: 21.0
            }
        );
    }

    #[test]
    fn bytes_roundtrip() {
        let mut col = FieldSoaColumn::<P>::new();
        col.push(P {
            x: 1.5,
            y: -2.5,
            z: 3.25,
        });
        col.push(P {
            x: 4.0,
            y: 5.0,
            z: 6.0,
        });
        let bytes = col.to_bytes();
        let mut restored = FieldSoaColumn::<P>::new();
        // SAFETY: `bytes` came from `to_bytes` on the same `FieldSoaColumn<P>`
        // type, so the field-major f32 layout matches exactly.
        unsafe { restored.set_from_bytes(&bytes) }.expect("valid bytes must round-trip");
        assert_eq!(restored.len(), 2);
        assert_eq!(
            restored.get(0),
            P {
                x: 1.5,
                y: -2.5,
                z: 3.25
            }
        );
        assert_eq!(
            restored.get(1),
            P {
                x: 4.0,
                y: 5.0,
                z: 6.0
            }
        );
    }
}

/// End-to-end test of a real `#[component(layout = "soa")]` component flowing
/// through the full CRPECS lifecycle: spawn, by-value query, bulk SIMD-style
/// mutation, by-value get/set, and despawn — proving the field-SoA column is a
/// first-class CRPECS citizen, not a parallel store.
#[cfg(test)]
mod world_integration {
    use crate::ecs::{SemanticDomain, Soa, World};
    use khora_macros::Component;

    #[derive(Debug, Clone, Copy, PartialEq, Default, Component)]
    #[component(no_serializable, layout = "soa")]
    struct Velocity {
        x: f32,
        y: f32,
        z: f32,
    }

    fn sorted_by_x(world: &World) -> Vec<Velocity> {
        let mut v: Vec<Velocity> = world.query::<Soa<Velocity>>().collect();
        v.sort_by(|a, b| a.x.partial_cmp(&b.x).unwrap());
        v
    }

    #[test]
    fn soa_component_full_lifecycle() {
        let mut world = World::default();
        world.register_component::<Velocity>(SemanticDomain::Physics);

        let e0 = world.spawn(Velocity {
            x: 1.0,
            y: 2.0,
            z: 3.0,
        });
        let e1 = world.spawn(Velocity {
            x: 4.0,
            y: 5.0,
            z: 6.0,
        });

        // By-value query reconstructs each value from the field arrays.
        assert_eq!(
            sorted_by_x(&world),
            vec![
                Velocity {
                    x: 1.0,
                    y: 2.0,
                    z: 3.0
                },
                Velocity {
                    x: 4.0,
                    y: 5.0,
                    z: 6.0
                },
            ]
        );

        // Bulk SIMD-style access: the `x` field array is contiguous — double it.
        world.for_each_soa_column_mut::<Velocity>(|col| {
            for x in col.field_mut(0) {
                *x *= 2.0;
            }
        });
        assert_eq!(world.clone_component::<Velocity>(e0).unwrap().x, 2.0);
        assert_eq!(world.clone_component::<Velocity>(e1).unwrap().x, 8.0);

        // By-value write.
        assert!(world.set_component(
            e0,
            Velocity {
                x: 9.0,
                y: 9.0,
                z: 9.0
            }
        ));
        assert_eq!(
            world.clone_component::<Velocity>(e0),
            Some(Velocity {
                x: 9.0,
                y: 9.0,
                z: 9.0
            })
        );

        // Despawn one; the other's field arrays stay aligned and correct.
        world.despawn(e0);
        assert_eq!(
            world.clone_component::<Velocity>(e1),
            Some(Velocity {
                x: 8.0,
                y: 5.0,
                z: 6.0
            })
        );
        assert_eq!(world.query::<Soa<Velocity>>().count(), 1);
    }
}