heron_rapier 5.0.2

Physics behavior for Heron, using rapier
Documentation
use bevy::ecs::prelude::*;
use bevy::math::Affine3A;
use bevy::transform::prelude::*;
use fnv::FnvHashMap;

use heron_core::{Damping, PhysicMaterial, RigidBody, RotationConstraints, Velocity};

use crate::convert::{IntoBevy, IntoRapier};
use crate::rapier::geometry::ColliderSet;
use crate::rapier::{
    dynamics::{IslandManager, RigidBodyBuilder, RigidBodyHandle, RigidBodySet, RigidBodyType},
    prelude::{ImpulseJointSet, MultibodyJointSet},
};

pub(crate) type HandleMap = FnvHashMap<Entity, RigidBodyHandle>;

#[allow(clippy::type_complexity)]
pub(crate) fn create(
    mut commands: Commands<'_, '_>,
    mut bodies: ResMut<'_, RigidBodySet>,
    mut handles: ResMut<'_, HandleMap>,
    query: Query<
        '_,
        '_,
        (
            Entity,
            &GlobalTransform,
            &RigidBody,
            Option<&Velocity>,
            Option<&Damping>,
            Option<&RotationConstraints>,
        ),
        Without<super::RigidBodyHandle>,
    >,
) {
    for (entity, transform, body, velocity, damping, rotation_constraints) in query.iter() {
        let (_, global_rotation, global_translation) = transform.to_scale_rotation_translation();
        let mut builder = RigidBodyBuilder::new(body_status(*body))
            .user_data(entity.to_bits().into())
            .position((global_translation, global_rotation).into_rapier());

        #[allow(unused_variables)]
        if let Some(RotationConstraints {
            allow_x,
            allow_y,
            allow_z,
        }) = rotation_constraints.copied()
        {
            #[cfg(dim2)]
            if !allow_z {
                builder = builder.lock_rotations();
            }
            #[cfg(dim3)]
            {
                builder = builder.restrict_rotations(allow_x, allow_y, allow_z);
            }
        }

        if let Some(v) = velocity {
            builder = builder
                .linvel(v.linear.into_rapier())
                .angvel(v.angular.into_rapier());
        }

        if let Some(d) = damping {
            builder = builder.linear_damping(d.linear).angular_damping(d.angular);
        }

        let rigid_body_handle = bodies.insert(builder.build());

        handles.insert(entity, rigid_body_handle);
        commands
            .entity(entity)
            .insert(super::RigidBodyHandle(rigid_body_handle));
    }
}

#[allow(clippy::too_many_arguments)]
pub(crate) fn remove_invalids_after_components_removed(
    mut commands: Commands<'_, '_>,
    mut handles: ResMut<'_, HandleMap>,
    mut bodies: ResMut<'_, RigidBodySet>,
    mut islands: ResMut<'_, IslandManager>,
    mut colliders: ResMut<'_, ColliderSet>,
    mut impulse_joints: ResMut<'_, ImpulseJointSet>,
    mut multibody_joints: ResMut<'_, MultibodyJointSet>,
    bodies_removed: RemovedComponents<'_, RigidBody>,
    constraints_removed: RemovedComponents<'_, RotationConstraints>,
    materials_removed: RemovedComponents<'_, PhysicMaterial>,
    rb_entities: Query<'_, '_, Entity, With<super::RigidBodyHandle>>,
    collider_entities: Query<'_, '_, Entity, With<super::ColliderHandle>>,
) {
    bodies_removed
        .iter()
        .chain(constraints_removed.iter())
        .chain(materials_removed.iter())
        .for_each(|entity| {
            if let Some(handle) = handles.remove(&entity) {
                remove_collider_handles(
                    &mut commands,
                    &collider_entities,
                    &bodies,
                    &colliders,
                    handle,
                );
                bodies.remove(
                    handle,
                    &mut islands,
                    &mut colliders,
                    &mut impulse_joints,
                    &mut multibody_joints,
                    true,
                );
                if rb_entities.get(entity).is_ok() {
                    commands.entity(entity).remove::<super::RigidBodyHandle>();
                }
            }
        });
}

#[allow(clippy::type_complexity)]
pub(crate) fn remove_invalids_after_component_changed(
    mut commands: Commands<'_, '_>,
    mut handles: ResMut<'_, HandleMap>,
    mut bodies: ResMut<'_, RigidBodySet>,
    mut islands: ResMut<'_, IslandManager>,
    mut colliders: ResMut<'_, ColliderSet>,
    mut impulse_joints: ResMut<'_, ImpulseJointSet>,
    mut multibody_joints: ResMut<'_, MultibodyJointSet>,
    collider_entities: Query<'_, '_, Entity, With<super::ColliderHandle>>,
    rigidbody_entities: Query<'_, '_, Entity, With<super::RigidBodyHandle>>,
    changed: Query<
        '_,
        '_,
        (Entity, &super::RigidBodyHandle),
        Or<(
            Changed<RigidBody>,
            Changed<RotationConstraints>,
            Changed<PhysicMaterial>,
        )>,
    >,
) {
    for (entity, handle) in changed.iter() {
        remove_collider_handles(
            &mut commands,
            &collider_entities,
            &bodies,
            &colliders,
            handle.0,
        );
        bodies.remove(
            handle.0,
            &mut islands,
            &mut colliders,
            &mut impulse_joints,
            &mut multibody_joints,
            true,
        );
        if rigidbody_entities.get(entity).is_ok() {
            commands.entity(entity).remove::<super::RigidBodyHandle>();
        }
        handles.remove(&entity);
    }
}

#[allow(clippy::manual_filter_map)]
fn remove_collider_handles(
    commands: &mut Commands<'_, '_>,
    entities: &Query<'_, '_, Entity, With<super::ColliderHandle>>,
    bodies: &RigidBodySet,
    colliders: &ColliderSet,
    handle: RigidBodyHandle,
) {
    bodies
        .get(handle)
        .iter()
        .flat_map(|it| it.colliders().iter())
        .filter_map(|it| colliders.get(*it))
        .map(|it| {
            #[allow(clippy::cast_possible_truncation)]
            Entity::from_bits(it.user_data as u64)
        })
        .filter(|e| entities.get(*e).is_ok())
        .for_each(|collider_entity| {
            commands
                .entity(collider_entity)
                .remove::<super::ColliderHandle>();
        });
}

pub(crate) fn update_rapier_position(
    mut bodies: ResMut<'_, RigidBodySet>,
    query: Query<'_, '_, (&GlobalTransform, &super::RigidBodyHandle), Changed<GlobalTransform>>,
) {
    for (transform, handle) in query.iter() {
        if let Some(body) = bodies.get_mut(handle.0) {
            let (_, global_rotation, global_translation) =
                transform.to_scale_rotation_translation();
            let isometry = (global_translation, global_rotation).into_rapier();

            if body.is_kinematic() {
                body.set_next_kinematic_position(isometry);
            } else {
                body.set_position(isometry, true);
            }
        }
    }
}

pub(crate) fn update_bevy_transform(
    bodies: Res<'_, RigidBodySet>,
    mut query: Query<
        '_,
        '_,
        (
            Option<&mut Transform>,
            &mut GlobalTransform,
            &super::RigidBodyHandle,
            Option<&RigidBody>,
        ),
    >,
) {
    for (mut local, mut global, handle, body_type) in query.iter_mut() {
        if !body_type.copied().unwrap_or_default().can_have_velocity() {
            continue;
        }

        let body = match bodies.get(handle.0) {
            None => continue,
            Some(body) => body,
        };

        #[cfg(dim3)]
        let (translation, rotation) = body.position().into_bevy();
        #[cfg(dim2)]
        let (mut translation, rotation) = body.position().into_bevy();

        let (global_scale, global_rotation, global_translation) =
            global.to_scale_rotation_translation();

        #[cfg(dim2)]
        {
            // In 2D, preserve the transform `z` component that may have been set by the user
            translation.z = global_translation.z;
        }

        if translation == global_translation && rotation == global_rotation {
            continue;
        }

        if let Some(local) = &mut local {
            if local.translation == global_translation {
                local.translation = translation;
            } else {
                local.translation = translation - (global_translation - local.translation);
            }

            if local.rotation == global_rotation {
                local.rotation = rotation;
            } else {
                local.rotation =
                    rotation * (global_rotation * local.rotation.conjugate()).conjugate();
            }
        }

        *global = GlobalTransform::from(Affine3A::from_scale_rotation_translation(
            global_scale,
            rotation,
            translation,
        ));
    }
}

fn body_status(body_type: RigidBody) -> RigidBodyType {
    match body_type {
        RigidBody::Dynamic => RigidBodyType::Dynamic,
        RigidBody::Static | RigidBody::Sensor => RigidBodyType::Fixed,
        RigidBody::KinematicPositionBased => RigidBodyType::KinematicPositionBased,
        RigidBody::KinematicVelocityBased => RigidBodyType::KinematicVelocityBased,
    }
}