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)]
{
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,
}
}