finite_light_bevy 0.1.2

#![doc = include_str!("../README.md")]
#![expect(
    clippy::too_many_arguments,
    clippy::type_complexity,
    reason = "Bevy systems commonly have many parameters and complex query types."
)]

mod render;
mod skybox;

use std::ops::Deref;

use bevy::{
    camera::visibility::NoFrustumCulling, prelude::*, render::sync_world::SyncToRenderWorld,
};
pub use finite_light_math as math;
use finite_light_math::{PoincareTransform, SpacetimeEvent, Vec3, Vec4};
pub use skybox::RelativisticSkybox;

#[derive(Component)]
pub struct HideWorldLine;

/// Accumulated proper time along the camera's worldline.
#[derive(Resource, Default)]
pub struct ProperTime {
    elapsed: f64,
    delta: f64,
}

impl ProperTime {
    pub fn elapsed_secs(&self) -> f32 {
        self.elapsed as f32
    }

    pub fn delta_secs(&self) -> f32 {
        self.delta as f32
    }
}

/// Bevy resource wrapping [`math::Metric`].
#[derive(Resource)]
pub struct RelativisticMetric(pub finite_light_math::Metric);

impl Deref for RelativisticMetric {
    type Target = finite_light_math::Metric;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

/// Tracks an entity's spacetime state and history.
///
/// [`Mesh3d`] entities are enrolled automatically -- the plugin adds
/// [`Relativistic`] when it first sees them. For non-mesh entities
/// (cameras, player controllers), add `Relativistic::default()` with a
/// [`Transform`]; the Poincaré state is seeded from the transform on the
/// first frame.
///
/// Access the [`PoincareTransform`] via [`transform`](Self::transform)/
/// [`transform_mut`](Self::transform_mut). Add [`NonRelativistic`] to
/// opt out.
#[derive(Component, Default)]
#[require(Transform)]
pub struct Relativistic {
    /// Current spacetime pose (position, rotation, boost).
    poincare: PoincareTransform,
    /// Spacetime history of this entity's [`PoincareTransform`].
    world_line: finite_light_math::WorldLine,

    /// Conservative bounding radius for world line GC.
    ///
    /// Automatically derived from the mesh at init time. [`None`] for entities
    /// without a mesh.
    bounding_radius: Option<f32>,

    /// Whether the Poincaré has been seeded from the entity's [`Transform`].
    seeded: bool,

    /// Last (translation, rotation) written by [`sync_transforms`]. Used to
    /// detect external [`Transform`] modification.
    last_synced: Option<(bevy::math::Vec3, Quat)>,
}

/// Opt-out: prevents the plugin from auto-adding [`Relativistic`] to this
/// entity or any [`Mesh3d`] descendant.
#[derive(Component)]
pub struct NonRelativistic;

impl Relativistic {
    /// Immutable access to the current [`PoincareTransform`].
    pub fn transform(&self) -> &PoincareTransform {
        &self.poincare
    }

    /// Mutable access to the current [`PoincareTransform`].
    pub fn transform_mut(&mut self) -> &mut PoincareTransform {
        &mut self.poincare
    }

    /// Construct with an initial velocity.
    pub fn with_velocity(mut self, metric: finite_light_math::Metric, v: Vec3) -> Self {
        self.set_velocity(metric, v);
        self
    }

    /// Set the entity's velocity, updating the Lorentz boost.
    pub fn set_velocity(&mut self, metric: finite_light_math::Metric, v: Vec3) {
        self.poincare.lorentz.boost = finite_light_math::Boost::from_velocity(metric, v);
    }

    /// Advance the entity's spatial position by its current velocity times
    /// `dt`.
    pub fn update_position(&mut self, metric: finite_light_math::Metric, dt: f32) {
        let v = self.poincare.lorentz.boost.velocity(metric);
        self.poincare
            .translation
            .set_spatial(self.poincare.translation.spatial() + v * dt);
    }

    /// Read-only access to this entity's spacetime history.
    pub fn world_line(&self) -> &finite_light_math::WorldLine {
        &self.world_line
    }
}

/// Scaled local-space vertex data awaiting extraction to the render world.
/// Created by [`init_mesh_data`], removed after one frame (giving the
/// extract system time to read it).
#[derive(Component)]
pub(crate) struct PendingMeshData {
    pub(crate) vertices: Vec<Vec4>,
    pub(crate) normals: Vec<Vec4>,
}

/// Link a [`Relativistic`] entity to a source. Each frame the entity's
/// [`PoincareTransform`] is composed from the source's transform plus
/// the stored local offset. Added automatically by `init_relativistic`
/// for [`Mesh3d`] descendants of a [`Relativistic`] entity.
#[derive(Component)]
pub struct RelativisticChild {
    /// Entity whose [`PoincareTransform`] to follow.
    pub source: Entity,
    /// Local position offset in the source entity's frame.
    pub offset: Vec3,
    /// Local rotation relative to the source entity.
    pub rotation: Quat,
}

/// Marker for entities whose mesh vertices are transformed by the
/// render-world compute shader. Added by [`init_mesh_data`] when vertex
/// data is first prepared. Their [`Transform`] is reset to identity each
/// frame by [`sync_transforms`].
#[derive(Component)]
pub(crate) struct GpuTransformed;

pub struct RelativisticPlugin {
    pub metric: finite_light_math::Metric,
    pub debug: bool,
}

impl RelativisticPlugin {
    /// Construct with a custom speed of light.
    pub fn with_speed_of_light(speed_of_light: f32) -> Self {
        Self {
            metric: finite_light_math::Metric { speed_of_light },
            debug: false,
        }
    }

    /// Enable or disable debug visualization of world lines.
    pub fn with_debug(mut self, debug: bool) -> Self {
        self.debug = debug;
        self
    }
}

impl Plugin for RelativisticPlugin {
    fn build(&self, app: &mut App) {
        app.insert_resource(RelativisticMetric(self.metric))
            .init_resource::<ProperTime>()
            .add_plugins(render::RelativisticRenderPlugin)
            .add_systems(Update, skybox::assemble_skybox)
            .add_systems(
                PostUpdate,
                (
                    init_relativistic,
                    ApplyDeferred,
                    update_poincare,
                    advance_proper_time,
                    sync_children,
                    record_and_gc,
                    sync_transforms,
                    init_mesh_data,
                )
                    .chain()
                    .after(TransformSystems::Propagate),
            );
        if self.debug {
            app.add_systems(Update, draw_world_lines).add_systems(
                Startup,
                |mut config_store: ResMut<GizmoConfigStore>| {
                    config_store
                        .config_mut::<DefaultGizmoConfigGroup>()
                        .0
                        .depth_bias = -1.;
                },
            );
        }
    }
}

// ---------------------------------------------------------------------------
// Systems
// ---------------------------------------------------------------------------

/// Auto-add [`Relativistic`] to [`Mesh3d`] entities.
///
/// For each [`Mesh3d`] without [`Relativistic`] or [`NonRelativistic`]:
/// * If a [`Relativistic`] ancestor exists, add [`Relativistic`] +
///   [`RelativisticChild`] so the mesh follows the ancestor.
/// * If a [`NonRelativistic`] ancestor exists, skip.
/// * Otherwise, walk to the hierarchy root and add [`Relativistic`] there.
///   [`RelativisticChild`] wiring happens on the next frame. Standalone meshes
///   (no parent) get [`Relativistic`] directly.
fn init_relativistic(
    mut commands: Commands,
    parent_query: Query<&ChildOf>,
    relativistic_query: Query<&GlobalTransform, With<Relativistic>>,
    non_rel_query: Query<(), With<NonRelativistic>>,
    mesh_query: Query<
        (Entity, &GlobalTransform),
        (
            With<Mesh3d>,
            Without<Relativistic>,
            Without<NonRelativistic>,
        ),
    >,
) {
    let mut roots_handled = std::collections::HashSet::new();

    for (mesh_entity, mesh_global) in &mesh_query {
        match walk_ancestors(
            mesh_entity,
            &parent_query,
            &relativistic_query,
            &non_rel_query,
        ) {
            AncestorResult::Relativistic(source, source_global) => {
                // Wire up as a child of the Relativistic ancestor.
                let (_, mesh_rot, mesh_trans) = mesh_global.to_scale_rotation_translation();
                let mut relativistic = Relativistic::default();
                relativistic.poincare.translation = SpacetimeEvent::new(mesh_trans, 0.);
                relativistic.poincare.lorentz.rotation = mesh_rot;
                relativistic.seeded = true;

                let (_, source_rot, source_trans) = source_global.to_scale_rotation_translation();
                let inv_source_rot = source_rot.inverse();

                commands.entity(mesh_entity).insert((
                    relativistic,
                    RelativisticChild {
                        source,
                        offset: inv_source_rot * (mesh_trans - source_trans),
                        rotation: inv_source_rot * mesh_rot,
                    },
                ));
            }
            AncestorResult::NonRelativistic => {}
            AncestorResult::None => {
                // Find the hierarchy root.
                let root = find_root(mesh_entity, &parent_query);
                if root == mesh_entity {
                    // Standalone mesh. Seed Poincaré from GlobalTransform.
                    let (_, rot, trans) = mesh_global.to_scale_rotation_translation();
                    let mut relativistic = Relativistic::default();
                    relativistic.poincare.translation = SpacetimeEvent::new(trans, 0.);
                    relativistic.poincare.lorentz.rotation = rot;
                    relativistic.seeded = true;
                    commands.entity(mesh_entity).insert(relativistic);
                } else if roots_handled.insert(root) {
                    // Add `Relativistic` to the root. `RelativisticChild`
                    // wiring happens on the next frame once the root's
                    // `Relativistic` is visible.
                    commands.entity(root).insert(Relativistic::default());
                }
            }
        }
    }
}

enum AncestorResult {
    /// Found a [`Relativistic`] ancestor: (entity, its [`GlobalTransform`]).
    Relativistic(Entity, GlobalTransform),
    /// Found a [`NonRelativistic`] ancestor: skip this mesh.
    NonRelativistic,
    /// Reached the root without finding either.
    None,
}

/// Walk up from `start` (exclusive) looking for [`Relativistic`] or
/// [`NonRelativistic`] ancestors.
fn walk_ancestors(
    start: Entity,
    parent_query: &Query<&ChildOf>,
    relativistic_query: &Query<&GlobalTransform, With<Relativistic>>,
    non_rel_query: &Query<(), With<NonRelativistic>>,
) -> AncestorResult {
    let mut current = start;
    while let Ok(child_of) = parent_query.get(current) {
        let parent = child_of.0;
        if let Ok(global) = relativistic_query.get(parent) {
            return AncestorResult::Relativistic(parent, *global);
        }
        if non_rel_query.get(parent).is_ok() {
            return AncestorResult::NonRelativistic;
        }
        current = parent;
    }
    AncestorResult::None
}

/// Walk up from `start` to the topmost ancestor (no [`ChildOf`]).
fn find_root(start: Entity, parent_query: &Query<&ChildOf>) -> Entity {
    let mut current = start;
    while let Ok(child_of) = parent_query.get(current) {
        current = child_of.0;
    }
    current
}

/// Seed newly-added [`Relativistic`] entities from their [`Transform`]
/// and stamp the current game time onto every entity's Poincaré.
fn update_poincare(
    time: Res<Time>,
    mut non_child: Query<(&mut Relativistic, &Transform), Without<RelativisticChild>>,
    mut children: Query<&mut Relativistic, With<RelativisticChild>>,
) {
    let t = time.elapsed_secs();

    for (mut relativistic, transform) in &mut non_child {
        // Seed from `Transform` on the first frame.
        if !relativistic.seeded {
            relativistic
                .poincare
                .translation
                .set_spatial(transform.translation);
            relativistic.poincare.lorentz.rotation = transform.rotation;
            relativistic.seeded = true;
        }
        relativistic.poincare.translation.set_t(t);
    }

    for mut relativistic in &mut children {
        relativistic.poincare.translation.set_t(t);
    }
}

/// Advance [`ProperTime`] by `dt_game/gamma`.
fn advance_proper_time(
    time: Res<Time>,
    camera: Query<&Relativistic, With<Camera3d>>,
    mut proper_time: ResMut<ProperTime>,
) {
    let gamma = camera
        .single()
        .map_or(1., |rel| rel.transform().lorentz.boost.gamma() as f64);
    let delta = time.delta_secs_f64() / gamma;
    proper_time.delta = delta;
    proper_time.elapsed += delta;
}

/// Compose the source entity's [`PoincareTransform`] with the stored
/// local offset for each [`RelativisticChild`].
fn sync_children(
    source_query: Query<&Relativistic, Without<RelativisticChild>>,
    mut follower_query: Query<(&mut Relativistic, &RelativisticChild)>,
) {
    for (mut follower, follow) in &mut follower_query {
        let Ok(source) = source_query.get(follow.source) else {
            continue;
        };
        let parent = source.transform();
        let world_offset = parent.lorentz.rotation * follow.offset;
        follower
            .poincare
            .translation
            .set_spatial(parent.translation.spatial() + world_offset);
        follower.poincare.lorentz.rotation = parent.lorentz.rotation * follow.rotation;
        follower.poincare.lorentz.boost = parent.lorentz.boost;
    }
}

/// Record each source entity's current [`PoincareTransform`] as a keyframe,
/// and garbage-collect keyframes that can no longer affect any vertex.
///
/// [`RelativisticChild`] entities are excluded -- they use the source's world
/// line in the shader, so only source entities need keyframe history. The
/// source's GC bounding radius accounts for all children.
fn record_and_gc(
    metric: Res<RelativisticMetric>,
    camera_query: Query<&Relativistic, With<Camera3d>>,
    mut source_query: Query<
        (Entity, &mut Relativistic),
        (Without<Camera3d>, Without<RelativisticChild>),
    >,
    child_query: Query<(&Relativistic, &RelativisticChild), Without<Camera3d>>,
    mut child_extents: Local<std::collections::HashMap<Entity, f32>>,
) {
    // Extract camera position for GC. The camera's world line is not
    // recorded because it is never sent to the render world.
    let camera_position = camera_query
        .single()
        .ok()
        .map(|camera| camera.poincare.translation);

    // Precompute the max effective bounding radius each child contributes
    // to its source entity. Reuse the map across frames to avoid per-frame
    // allocation on WASM where linear memory never shrinks.
    child_extents.clear();
    for (child, follow) in &child_query {
        let extent = follow.offset.length() + child.bounding_radius.unwrap_or(0.);
        let entry = child_extents.entry(follow.source).or_insert(0.);
        *entry = entry.max(extent);
    }

    for (entity, mut relativistic) in &mut source_query {
        let poincare = relativistic.poincare;
        relativistic.world_line.push(poincare);

        if let Some(cam_pos) = camera_position {
            let own_radius = relativistic.bounding_radius.unwrap_or(0.);
            let child_radius = child_extents.get(&entity).copied().unwrap_or(0.);
            relativistic
                .world_line
                .gc(metric.0, cam_pos, own_radius.max(child_radius));
        }
    }
}

fn draw_world_lines(
    query: Query<&Relativistic, (Without<HideWorldLine>, Without<RelativisticChild>)>,
    mut gizmos: Gizmos,
) {
    for relativistic in &query {
        let keyframes = relativistic.world_line.keyframes();
        for window in keyframes.iter().collect::<Vec<_>>().windows(2) {
            let a = window[0].translation.spatial();
            let b = window[1].translation.spatial();
            gizmos.line(a, b, Color::srgb(0., 1., 0.));
        }
    }
}

/// Sync the [`PoincareTransform`] to Bevy's [`Transform`] for rendering.
///
/// * Entities with [`GpuTransformed`] get identity transforms because their
///   mesh vertices are already in world space.
/// * Non-GPU entities without [`RelativisticChild`] have their [`Transform`]
///   driven by the Poincaré state.
/// * [`RelativisticChild`] entities keep their hierarchy-based [`Transform`]
///   until [`GpuTransformed`] is added.
fn sync_transforms(
    mut gpu: Query<(&mut Transform, &mut GlobalTransform), With<GpuTransformed>>,
    mut non_gpu: Query<
        (&mut Relativistic, &mut Transform, &mut GlobalTransform),
        (Without<GpuTransformed>, Without<RelativisticChild>),
    >,
) {
    // The compute shader writes world-space vertices, so the model matrix
    // must be identity. Override both `Transform` and `GlobalTransform` --
    // the latter may include the parent chain from hierarchy propagation.
    // Only write when the value differs to avoid triggering Bevy's change
    // detection every frame (which causes unnecessary re-extraction and
    // bind group re-creation on WebGPU).
    for (mut transform, mut global_transform) in &mut gpu {
        if *transform != Transform::IDENTITY {
            *transform = Transform::IDENTITY;
        }
        if *global_transform != GlobalTransform::IDENTITY {
            *global_transform = GlobalTransform::default();
        }
    }

    for (mut relativistic, mut transform, mut global_transform) in &mut non_gpu {
        // Detect external `Transform` modification. The plugin owns
        // `Transform` for `Relativistic` entities -- mutate the
        // `PoincareTransform` via `Relativistic::transform_mut` instead.
        if let Some((last_t, last_r)) = relativistic.last_synced {
            assert!(
                transform.translation == last_t && transform.rotation == last_r,
                "Transform on a `Relativistic` entity was modified externally. Use \
                 `Relativistic::transform_mut()` to change the spacetime pose.",
            );
        }

        transform.translation = relativistic.poincare.translation.spatial();
        transform.rotation = relativistic.poincare.lorentz.rotation;
        relativistic.last_synced = Some((transform.translation, transform.rotation));

        // Also update `GlobalTransform` because this system runs after
        // `TransformSystems::Propagate`. Without this, cameras (and other
        // non-GPU entities) would have a one-frame-stale view matrix.
        *global_transform = GlobalTransform::from(*transform);
    }
}

/// Read mesh vertices and normals, scale by the entity's
/// [`GlobalTransform`], and create [`PendingMeshData`] for extraction to
/// the render world. Also compute the
/// [`bounding_radius`](Relativistic::bounding_radius). Cleans up
/// [`PendingMeshData`] from the previous frame first.
///
/// Each [`Relativistic`] entity must have a unique [`Mesh3d`] handle.
/// The compute shader writes transformed vertices directly into the
/// vertex buffer keyed by mesh asset ID, so shared handles would cause
/// entities to overwrite each other.
fn init_mesh_data(
    mut commands: Commands,
    meshes: Res<Assets<Mesh>>,
    pending_query: Query<Entity, With<PendingMeshData>>,
    existing_query: Query<&Mesh3d, With<GpuTransformed>>,
    mut query: Query<
        (Entity, &mut Relativistic, &Mesh3d, &GlobalTransform),
        Without<GpuTransformed>,
    >,
) {
    // Clean up `PendingMeshData` from last frame. The render-world
    // extraction system will have already copied the data.
    for entity in &pending_query {
        commands.entity(entity).remove::<PendingMeshData>();
    }

    // Build set of mesh asset IDs already claimed by existing entities.
    let mut claimed_meshes = std::collections::HashSet::new();
    for mesh_handle in &existing_query {
        claimed_meshes.insert(mesh_handle.id());
    }

    for (entity, mut relativistic, mesh_handle, global_transform) in &mut query {
        assert!(
            claimed_meshes.insert(mesh_handle.id()),
            "Multiple Relativistic entities share mesh asset {:?}. Each entity needs a unique \
             mesh handle because the compute shader writes directly into the vertex buffer keyed \
             by asset ID.",
            mesh_handle.id(),
        );

        let Some(mesh) = meshes.get(mesh_handle) else {
            continue;
        };
        let Some(positions) = mesh.attribute(Mesh::ATTRIBUTE_POSITION) else {
            continue;
        };
        // Bake scale from the hierarchy into vertex positions.
        let (scale, _, _) = global_transform.to_scale_rotation_translation();
        let vertices: Vec<Vec4> = positions
            .as_float3()
            .unwrap()
            .iter()
            .map(|p| Vec4::new(p[0] * scale.x, p[1] * scale.y, p[2] * scale.z, 1.))
            .collect();
        // Compute bounding radius as the max vertex distance from the origin.
        let bounding_radius = vertices
            .iter()
            .map(|p| p.truncate().length())
            .fold(0f32, f32::max);
        relativistic.bounding_radius = Some(bounding_radius);

        // Apply inverse-transpose of scale to normals and renormalize.
        let inv_scale = Vec3::new(1. / scale.x, 1. / scale.y, 1. / scale.z);
        let normals: Vec<Vec4> = mesh
            .attribute(Mesh::ATTRIBUTE_NORMAL)
            .and_then(|n| n.as_float3())
            .map(|slice| {
                slice
                    .iter()
                    .map(|n| {
                        let normal =
                            Vec3::new(n[0] * inv_scale.x, n[1] * inv_scale.y, n[2] * inv_scale.z)
                                .normalize();
                        Vec4::new(normal.x, normal.y, normal.z, 0.)
                    })
                    .collect()
            })
            .unwrap_or_else(|| vec![Vec4::new(0., 1., 0., 0.); vertices.len()]);

        commands.entity(entity).insert((
            PendingMeshData { vertices, normals },
            GpuTransformed,
            SyncToRenderWorld,
            NoFrustumCulling,
        ));
    }
}