bevy_tnua_avian3d/

lib.rs

//! # avian3d Integration for bevy-tnua
//!
//! In addition to the instruction in bevy-tnua's documentation:
//!
//! * Add [`TnuaAvian3dPlugin`] to the Bevy app.
//! * Optionally: Add [`TnuaAvian3dSensorShape`] to either entity of the character controller by
//!   Tnua or to the to the sensor entities. If exists, the shape on the sensor entity overrides
//!   the one on the character entity for that specific sensor.
mod spatial_ext;

use avian3d::{prelude::*, schedule::PhysicsStepSystems};
use bevy::ecs::schedule::{InternedScheduleLabel, ScheduleLabel};
use bevy::prelude::*;
use ordered_float::OrderedFloat;
pub use spatial_ext::TnuaSpatialExtAvian3d;

use bevy_tnua_physics_integration_layer::{
    TnuaPipelineSystems, TnuaSystems,
    data_for_backends::{
        TnuaGhostPlatform, TnuaGhostSensor, TnuaGravity, TnuaMotor, TnuaNotPlatform,
        TnuaProximitySensor, TnuaProximitySensorOutput, TnuaRigidBodyTracker, TnuaSensorOf,
        TnuaToggle,
    },
    math::{AdjustPrecision, AsF32, Float, Quaternion, Vector3},
    obstacle_radar::TnuaObstacleRadar,
};

pub mod prelude {
    pub use crate::{TnuaAvian3dPlugin, TnuaAvian3dSensorShape, TnuaSpatialExtAvian3d};
}

/// Add this plugin to use avian3d as a physics backend.
///
/// This plugin should be used in addition to `TnuaControllerPlugin`.
/// Note that you should make sure both of these plugins use the same schedule.
/// This should usually be `PhysicsSchedule`, which by default is `FixedUpdate`.
///
/// # Example
///
/// ```ignore
/// App::new()
///     .add_plugins((
///         DefaultPlugins,
///         PhysicsPlugins::default(),
///         TnuaControllerPlugin::new(PhysicsSchedule),
///         TnuaAvian3dPlugin::new(PhysicsSchedule),
///     ));
/// ```
pub struct TnuaAvian3dPlugin {
    schedule: InternedScheduleLabel,
}

impl TnuaAvian3dPlugin {
    pub fn new(schedule: impl ScheduleLabel) -> Self {
        Self {
            schedule: schedule.intern(),
        }
    }
}

impl Plugin for TnuaAvian3dPlugin {
    fn build(&self, app: &mut App) {
        app.configure_sets(
            self.schedule,
            TnuaSystems
                // Need to run _before_ `First`, not after it. The documentation is misleading. See
                // https://github.com/Jondolf/avian/issues/675
                .before(PhysicsStepSystems::First)
                .run_if(|physics_time: Res<Time<Physics>>| !physics_time.is_paused()),
        );
        app.add_systems(
            self.schedule,
            (
                update_rigid_body_trackers_system,
                update_proximity_sensors_system,
                update_obstacle_radars_system
                    // Both use SpatialQuery, which uses a ResMut
                    .ambiguous_with(update_proximity_sensors_system),
            )
                .in_set(TnuaPipelineSystems::Sensors),
        );
        app.add_systems(
            self.schedule,
            apply_motors_system.in_set(TnuaPipelineSystems::Motors),
        );
        app.register_required_components_with::<TnuaGravity, GravityScale>(|| GravityScale(0.0));
    }
}

/// Add this component to make [`TnuaProximitySensor`] cast a shape instead of a ray.
#[derive(Component)]
pub struct TnuaAvian3dSensorShape(pub Collider);

impl FromTemplate for TnuaAvian3dSensorShape {
    type Template = TnuaAvian3dSensorShapeTemplate;
}

#[derive(Default)]
pub struct TnuaAvian3dSensorShapeTemplate(pub Option<Collider>);

impl Template for TnuaAvian3dSensorShapeTemplate {
    type Output = TnuaAvian3dSensorShape;

    fn build_template(&self, _: &mut bevy::ecs::template::TemplateContext) -> Result<Self::Output> {
        Ok(TnuaAvian3dSensorShape(self.0.clone().ok_or_else(|| {
            BevyError::error(
                "TnuaAvian3dSensorShape does not have a default value. Must provide shape",
            )
        })?))
    }

    fn clone_template(&self) -> Self {
        Self(self.0.clone())
    }
}

#[allow(clippy::type_complexity)]
fn update_rigid_body_trackers_system(
    gravity: Res<Gravity>,
    mut query: Query<(
        &Position,
        &Rotation,
        &LinearVelocity,
        &AngularVelocity,
        &mut TnuaRigidBodyTracker,
        Option<&TnuaToggle>,
        Option<&TnuaGravity>,
    )>,
) {
    for (
        position,
        rotation,
        linaer_velocity,
        angular_velocity,
        mut tracker,
        tnua_toggle,
        tnua_gravity,
    ) in query.iter_mut()
    {
        match tnua_toggle.copied().unwrap_or_default() {
            TnuaToggle::Disabled => continue,
            TnuaToggle::SenseOnly => {}
            TnuaToggle::Enabled => {}
        }
        *tracker = TnuaRigidBodyTracker {
            translation: position.adjust_precision(),
            rotation: rotation.adjust_precision(),
            velocity: linaer_velocity.0.adjust_precision(),
            angvel: angular_velocity.0.adjust_precision(),
            gravity: tnua_gravity
                .map(|g| g.0)
                .unwrap_or(gravity.0.adjust_precision()),
        };
    }
}

#[allow(clippy::type_complexity)]
fn update_proximity_sensors_system(
    spatial_query: SpatialQuery,
    mut sensor_query: Query<(
        &mut TnuaProximitySensor,
        &TnuaSensorOf,
        Option<&TnuaAvian3dSensorShape>,
        Option<&mut TnuaGhostSensor>,
    )>,
    owner_query: Query<(
        &Position,
        &Rotation,
        Option<&Collider>,
        Option<&TnuaAvian3dSensorShape>,
        Option<&TnuaToggle>,
    )>,
    collision_layers_query: Query<&CollisionLayers>,
    other_object_query: Query<(
        Option<(
            &Position,
            &LinearVelocity,
            &AngularVelocity,
            Option<&RigidBody>,
        )>,
        Option<&CollisionLayers>,
        Option<&ColliderOf>,
        Has<TnuaGhostPlatform>,
        Has<Sensor>,
        Has<TnuaNotPlatform>,
    )>,
) {
    sensor_query.par_iter_mut().for_each(
        |(mut sensor, &TnuaSensorOf(owner_entity), shape, mut ghost_sensor)| {
            let Ok((position, rotation, collider, owner_shape, tnua_toggle)) =
                owner_query.get(owner_entity)
            else {
                return;
            };
            let shape = shape.or(owner_shape);
            match tnua_toggle.copied().unwrap_or_default() {
                TnuaToggle::Disabled => return,
                TnuaToggle::SenseOnly => {}
                TnuaToggle::Enabled => {}
            }
            let transform = Transform {
                translation: position.0.f32(),
                rotation: rotation.0.f32(),
                scale: collider
                    .map(|collider| collider.scale().f32())
                    .unwrap_or(Vec3::ONE),
            };

            // TODO: is there any point in doing these transformations as f64 when that feature
            // flag is active?
            let cast_origin = transform
                .transform_point(sensor.cast_origin.f32())
                .adjust_precision();
            let cast_direction = sensor.cast_direction;

            struct CastResult {
                entity: Entity,
                proximity: Float,
                intersection_point: Vector3,
                normal: Dir3,
            }

            let collision_layers = collision_layers_query.get(owner_entity).ok();

            let mut final_sensor_output: Option<TnuaProximitySensorOutput> = None;
            if let Some(ghost_sensor) = ghost_sensor.as_mut() {
                ghost_sensor.0.clear();
            }
            let mut apply_cast = |cast_result: CastResult| {
                let CastResult {
                    entity,
                    proximity,
                    intersection_point,
                    normal,
                } = cast_result;

                let Ok((
                    mut entity_kinematic_data,
                    mut entity_collision_layers,
                    entity_collider_of,
                    mut entity_is_ghost,
                    mut entity_is_sensor,
                    mut entity_is_not_platform,
                )) = other_object_query.get(entity)
                else {
                    return false;
                };

                if let Some(collider_of) = entity_collider_of {
                    let parent_entity = collider_of.body;

                    // Collider is child of our rigid body. ignore.
                    if parent_entity == owner_entity {
                        return true;
                    }

                    if let Ok((
                        parent_kinematic_data,
                        parent_collision_layers,
                        _,
                        parent_is_ghost,
                        parent_is_sensor,
                        parent_is_not_platform,
                    )) = other_object_query.get(parent_entity)
                    {
                        if entity_kinematic_data.is_none() {
                            entity_kinematic_data = parent_kinematic_data;
                        }
                        if entity_collision_layers.is_none() {
                            entity_collision_layers = parent_collision_layers;
                        }
                        entity_is_ghost = entity_is_ghost || parent_is_ghost;
                        entity_is_sensor = entity_is_sensor || parent_is_sensor;
                        entity_is_not_platform = entity_is_not_platform || parent_is_not_platform;
                    }
                }

                if entity_is_not_platform {
                    return true;
                }

                let entity_linvel;
                let entity_angvel;
                if let Some((
                    entity_position,
                    entity_linear_velocity,
                    entity_angular_velocity,
                    rigid_body,
                )) = entity_kinematic_data
                {
                    if rigid_body == Some(&RigidBody::Static) {
                        entity_angvel = Vector3::ZERO;
                        entity_linvel = Vector3::ZERO;
                    } else {
                        entity_angvel = entity_angular_velocity.0.adjust_precision();
                        entity_linvel = entity_linear_velocity.0.adjust_precision()
                            + if 0.0 < entity_angvel.length_squared() {
                                let relative_point =
                                    intersection_point - entity_position.adjust_precision();
                                // NOTE: no need to project relative_point on the
                                // rotation plane, it will not affect the cross
                                // product.
                                entity_angvel.cross(relative_point)
                            } else {
                                Vector3::ZERO
                            };
                    }
                } else {
                    entity_angvel = Vector3::ZERO;
                    entity_linvel = Vector3::ZERO;
                }
                let sensor_output = TnuaProximitySensorOutput {
                    entity,
                    proximity,
                    normal,
                    entity_linvel,
                    entity_angvel,
                };

                let excluded_by_collision_layers = || {
                    let collision_layers = collision_layers.copied().unwrap_or_default();
                    let entity_collision_layers =
                        entity_collision_layers.copied().unwrap_or_default();
                    !collision_layers.interacts_with(entity_collision_layers)
                };

                if entity_is_ghost {
                    if let Some(ghost_sensor) = ghost_sensor.as_mut() {
                        ghost_sensor.0.push(sensor_output);
                    }
                    true
                } else if entity_is_sensor || excluded_by_collision_layers() {
                    true
                } else {
                    if final_sensor_output.as_ref().is_none_or(|current_output| {
                        sensor_output.proximity < current_output.proximity
                    }) {
                        // Hits are not guaranteed to be ordered, so we need to make them ordered.
                        // See https://github.com/idanarye/bevy-tnua/issues/123
                        final_sensor_output = Some(sensor_output);
                    }
                    false
                }
            };

            let query_filter = SpatialQueryFilter::from_excluded_entities([owner_entity]);
            if let Some(TnuaAvian3dSensorShape(shape)) = shape {
                // TODO: can I bake `owner_rotation` into
                // `sensor.cast_shape_rotation`?
                let owner_rotation = Quaternion::from_axis_angle(
                    cast_direction.adjust_precision(),
                    rotation
                        .to_scaled_axis()
                        .adjust_precision()
                        .dot(cast_direction.adjust_precision()),
                );
                spatial_query.shape_hits_callback(
                    shape,
                    cast_origin,
                    owner_rotation
                        .mul_quat(sensor.cast_shape_rotation.adjust_precision())
                        .f32(),
                    cast_direction,
                    &ShapeCastConfig {
                        max_distance: sensor.cast_range.f32(),
                        ignore_origin_penetration: true,
                        ..default()
                    },
                    &query_filter,
                    |shape_hit_data| {
                        apply_cast(CastResult {
                            entity: shape_hit_data.entity,
                            proximity: shape_hit_data.distance.adjust_precision(),
                            intersection_point: shape_hit_data.point1,
                            normal: Dir3::new(shape_hit_data.normal1.f32())
                                .unwrap_or_else(|_| -cast_direction),
                        })
                    },
                );
            } else {
                spatial_query.ray_hits_callback(
                    cast_origin,
                    cast_direction,
                    sensor.cast_range.f32(),
                    true,
                    &query_filter,
                    |ray_hit_data| {
                        apply_cast(CastResult {
                            entity: ray_hit_data.entity,
                            proximity: ray_hit_data.distance.adjust_precision(),
                            intersection_point: cast_origin
                                + ray_hit_data.distance.adjust_precision()
                                    * cast_direction.adjust_precision(),
                            normal: Dir3::new(ray_hit_data.normal.f32())
                                .unwrap_or_else(|_| -cast_direction),
                        })
                    },
                );
            }
            if let Some(ghost_sensor) = ghost_sensor.as_mut() {
                // Hits are not guaranteed to be ordered, so we need to make them ordered.
                // See https://github.com/idanarye/bevy-tnua/issues/123
                if let Some(final_sensor_output) = final_sensor_output.as_ref() {
                    ghost_sensor
                        .0
                        .retain(|ghost_hit| ghost_hit.proximity < final_sensor_output.proximity);
                }
                ghost_sensor
                    .0
                    .sort_by_key(|ghost_hit| OrderedFloat(ghost_hit.proximity));
            }
            sensor.output = final_sensor_output;
        },
    );
}

fn update_obstacle_radars_system(
    spatial_query: SpatialQuery,
    gravity: Res<Gravity>,
    mut radars_query: Query<(Entity, &mut TnuaObstacleRadar, &Position)>,
) {
    if radars_query.is_empty() {
        return;
    }
    for (radar_owner_entity, mut radar, radar_position) in radars_query.iter_mut() {
        radar.pre_marking_update(
            radar_owner_entity,
            radar_position.0,
            Dir3::new(gravity.0.f32()).unwrap_or(Dir3::Y),
        );
        spatial_query.shape_intersections_callback(
            &Collider::cylinder(radar.radius.f32(), radar.height.f32()),
            radar_position.0,
            Rotation::default(),
            &SpatialQueryFilter::DEFAULT,
            |obstacle_entity| {
                if radar_owner_entity == obstacle_entity {
                    return true;
                }
                radar.mark_seen(obstacle_entity);
                true
            },
        );
    }
}

#[allow(clippy::type_complexity)]
fn apply_motors_system(
    mut query: Query<(
        &TnuaMotor,
        Forces,
        Option<&TnuaToggle>,
        Option<&TnuaGravity>,
    )>,
) {
    for (motor, mut forces, tnua_toggle, tnua_gravity) in query.iter_mut() {
        match tnua_toggle.copied().unwrap_or_default() {
            TnuaToggle::Disabled | TnuaToggle::SenseOnly => {
                return;
            }
            TnuaToggle::Enabled => {}
        }

        if motor.lin.boost.is_finite() {
            *forces.linear_velocity_mut() += motor.lin.boost.f32();
        }
        if motor.lin.acceleration.is_finite() {
            forces.apply_linear_acceleration(motor.lin.acceleration.f32());
        }
        if motor.ang.boost.is_finite() {
            *forces.angular_velocity_mut() += motor.ang.boost.f32();
        }
        if motor.ang.acceleration.is_finite() {
            forces.apply_torque(motor.ang.acceleration.f32());
        }
        if let Some(gravity) = tnua_gravity {
            forces.apply_linear_acceleration(gravity.0.f32());
        }
    }
}