use crate::dynamics::RapierRigidBodyHandle;
use crate::plugin::context::systemparams::RAPIER_CONTEXT_EXPECT_ERROR;
use crate::plugin::context::{
DefaultRapierContext, RapierContextColliders, RapierContextEntityLink, RapierRigidBodySet,
};
use crate::plugin::{configuration::TimestepMode, RapierConfiguration};
use crate::{dynamics::RigidBody, plugin::context::SimulationToRenderTime};
use crate::{prelude::*, utils};
use bevy::prelude::*;
use rapier::dynamics::{RigidBodyBuilder, RigidBodyHandle, RigidBodyType};
use std::collections::HashMap;
pub type RigidBodyWritebackComponents<'a> = (
&'a RapierRigidBodyHandle,
&'a RapierContextEntityLink,
Option<&'a Parent>,
Option<&'a mut Transform>,
Option<&'a mut TransformInterpolation>,
Option<&'a mut Velocity>,
Option<&'a mut Sleeping>,
);
pub type RigidBodyComponents<'a> = (
(Entity, Option<&'a RapierContextEntityLink>),
&'a RigidBody,
Option<&'a GlobalTransform>,
Option<&'a Velocity>,
Option<&'a AdditionalMassProperties>,
Option<&'a ReadMassProperties>,
Option<&'a LockedAxes>,
Option<&'a ExternalForce>,
Option<&'a GravityScale>,
(Option<&'a Ccd>, Option<&'a SoftCcd>),
Option<&'a Dominance>,
Option<&'a Sleeping>,
Option<&'a Damping>,
Option<&'a RigidBodyDisabled>,
Option<&'a AdditionalSolverIterations>,
);
pub fn apply_rigid_body_user_changes(
mut rigid_body_sets: Query<&mut RapierRigidBodySet>,
config: Query<&RapierConfiguration>,
changed_rb_types: Query<
(&RapierRigidBodyHandle, &RapierContextEntityLink, &RigidBody),
Changed<RigidBody>,
>,
mut changed_transforms: Query<
(
&RapierRigidBodyHandle,
&RapierContextEntityLink,
&GlobalTransform,
Option<&mut TransformInterpolation>,
),
Changed<GlobalTransform>,
>,
changed_velocities: Query<
(&RapierRigidBodyHandle, &RapierContextEntityLink, &Velocity),
Changed<Velocity>,
>,
changed_additional_mass_props: Query<
(
Entity,
&RapierContextEntityLink,
&RapierRigidBodyHandle,
&AdditionalMassProperties,
),
Changed<AdditionalMassProperties>,
>,
changed_locked_axes: Query<
(
&RapierRigidBodyHandle,
&RapierContextEntityLink,
&LockedAxes,
),
Changed<LockedAxes>,
>,
changed_forces: Query<
(
&RapierRigidBodyHandle,
&RapierContextEntityLink,
&ExternalForce,
),
Changed<ExternalForce>,
>,
mut changed_impulses: Query<
(
&RapierRigidBodyHandle,
&RapierContextEntityLink,
&mut ExternalImpulse,
),
Changed<ExternalImpulse>,
>,
changed_gravity_scale: Query<
(
&RapierRigidBodyHandle,
&RapierContextEntityLink,
&GravityScale,
),
Changed<GravityScale>,
>,
(changed_ccd, changed_soft_ccd): (
Query<(&RapierRigidBodyHandle, &RapierContextEntityLink, &Ccd), Changed<Ccd>>,
Query<(&RapierRigidBodyHandle, &RapierContextEntityLink, &SoftCcd), Changed<SoftCcd>>,
),
changed_dominance: Query<
(&RapierRigidBodyHandle, &RapierContextEntityLink, &Dominance),
Changed<Dominance>,
>,
changed_sleeping: Query<
(&RapierRigidBodyHandle, &RapierContextEntityLink, &Sleeping),
Changed<Sleeping>,
>,
changed_damping: Query<
(&RapierRigidBodyHandle, &RapierContextEntityLink, &Damping),
Changed<Damping>,
>,
(changed_disabled, changed_additional_solver_iterations): (
Query<
(
&RapierRigidBodyHandle,
&RapierContextEntityLink,
&RigidBodyDisabled,
),
Changed<RigidBodyDisabled>,
>,
Query<
(
&RapierRigidBodyHandle,
&RapierContextEntityLink,
&AdditionalSolverIterations,
),
Changed<AdditionalSolverIterations>,
>,
),
mut mass_modified: EventWriter<MassModifiedEvent>,
) {
for (handle, link, sleeping) in changed_sleeping.iter() {
let rigidbody_set = rigid_body_sets
.get_mut(link.0)
.expect(RAPIER_CONTEXT_EXPECT_ERROR)
.into_inner();
if let Some(rb) = rigidbody_set.bodies.get_mut(handle.0) {
let activation = rb.activation_mut();
activation.normalized_linear_threshold = sleeping.normalized_linear_threshold;
activation.angular_threshold = sleeping.angular_threshold;
if !sleeping.sleeping && activation.sleeping {
rb.wake_up(true);
} else if sleeping.sleeping && !activation.sleeping {
rb.sleep();
}
}
}
for (handle, link, rb_type) in changed_rb_types.iter() {
let context = rigid_body_sets
.get_mut(link.0)
.expect(RAPIER_CONTEXT_EXPECT_ERROR)
.into_inner();
if let Some(rb) = context.bodies.get_mut(handle.0) {
rb.set_body_type((*rb_type).into(), true);
}
}
let transform_changed_fn =
|handle: &RigidBodyHandle,
config: &RapierConfiguration,
transform: &GlobalTransform,
last_transform_set: &HashMap<RigidBodyHandle, GlobalTransform>| {
if config.force_update_from_transform_changes {
true
} else if let Some(prev) = last_transform_set.get(handle) {
*prev != *transform
} else {
true
}
};
for (handle, link, global_transform, mut interpolation) in changed_transforms.iter_mut() {
let rigidbody_set = rigid_body_sets
.get_mut(link.0)
.expect(RAPIER_CONTEXT_EXPECT_ERROR)
.into_inner();
let config = config
.get(link.0)
.expect("Could not get `RapierConfiguration`");
let mut transform_changed = None;
if let Some(interpolation) = interpolation.as_deref_mut() {
transform_changed = transform_changed.or_else(|| {
Some(transform_changed_fn(
&handle.0,
config,
global_transform,
&rigidbody_set.last_body_transform_set,
))
});
if transform_changed == Some(true) {
interpolation.start = None;
interpolation.end = None;
}
}
if let Some(rb) = rigidbody_set.bodies.get_mut(handle.0) {
transform_changed = transform_changed.or_else(|| {
Some(transform_changed_fn(
&handle.0,
config,
global_transform,
&rigidbody_set.last_body_transform_set,
))
});
match rb.body_type() {
RigidBodyType::KinematicPositionBased => {
if transform_changed == Some(true) {
rb.set_next_kinematic_position(utils::transform_to_iso(
&global_transform.compute_transform(),
));
rigidbody_set
.last_body_transform_set
.insert(handle.0, *global_transform);
}
}
_ => {
if transform_changed == Some(true) {
rb.set_position(
utils::transform_to_iso(&global_transform.compute_transform()),
true,
);
rigidbody_set
.last_body_transform_set
.insert(handle.0, *global_transform);
}
}
}
}
}
for (handle, link, velocity) in changed_velocities.iter() {
let rigidbody_set = rigid_body_sets
.get_mut(link.0)
.expect(RAPIER_CONTEXT_EXPECT_ERROR)
.into_inner();
if let Some(rb) = rigidbody_set.bodies.get_mut(handle.0) {
rb.set_linvel(velocity.linvel.into(), true);
#[allow(clippy::useless_conversion)] rb.set_angvel(velocity.angvel.into(), true);
}
}
for (entity, link, handle, mprops) in changed_additional_mass_props.iter() {
let rigidbody_set = rigid_body_sets
.get_mut(link.0)
.expect(RAPIER_CONTEXT_EXPECT_ERROR)
.into_inner();
if let Some(rb) = rigidbody_set.bodies.get_mut(handle.0) {
match mprops {
AdditionalMassProperties::MassProperties(mprops) => {
rb.set_additional_mass_properties(mprops.into_rapier(), true);
}
AdditionalMassProperties::Mass(mass) => {
rb.set_additional_mass(*mass, true);
}
}
mass_modified.send(entity.into());
}
}
for (handle, link, additional_solver_iters) in changed_additional_solver_iterations.iter() {
let rigidbody_set = rigid_body_sets
.get_mut(link.0)
.expect(RAPIER_CONTEXT_EXPECT_ERROR)
.into_inner();
if let Some(rb) = rigidbody_set.bodies.get_mut(handle.0) {
rb.set_additional_solver_iterations(additional_solver_iters.0);
}
}
for (handle, link, locked_axes) in changed_locked_axes.iter() {
let rigidbody_set = rigid_body_sets
.get_mut(link.0)
.expect(RAPIER_CONTEXT_EXPECT_ERROR)
.into_inner();
if let Some(rb) = rigidbody_set.bodies.get_mut(handle.0) {
rb.set_locked_axes((*locked_axes).into(), true);
}
}
for (handle, link, forces) in changed_forces.iter() {
let rigidbody_set = rigid_body_sets
.get_mut(link.0)
.expect(RAPIER_CONTEXT_EXPECT_ERROR)
.into_inner();
if let Some(rb) = rigidbody_set.bodies.get_mut(handle.0) {
rb.reset_forces(true);
rb.reset_torques(true);
rb.add_force(forces.force.into(), true);
#[allow(clippy::useless_conversion)] rb.add_torque(forces.torque.into(), true);
}
}
for (handle, link, mut impulses) in changed_impulses.iter_mut() {
let rigidbody_set = rigid_body_sets
.get_mut(link.0)
.expect(RAPIER_CONTEXT_EXPECT_ERROR)
.into_inner();
if let Some(rb) = rigidbody_set.bodies.get_mut(handle.0) {
rb.apply_impulse(impulses.impulse.into(), true);
#[allow(clippy::useless_conversion)] rb.apply_torque_impulse(impulses.torque_impulse.into(), true);
impulses.reset();
}
}
for (handle, link, gravity_scale) in changed_gravity_scale.iter() {
let rigidbody_set = rigid_body_sets
.get_mut(link.0)
.expect(RAPIER_CONTEXT_EXPECT_ERROR)
.into_inner();
if let Some(rb) = rigidbody_set.bodies.get_mut(handle.0) {
rb.set_gravity_scale(gravity_scale.0, true);
}
}
for (handle, link, ccd) in changed_ccd.iter() {
let rigidbody_set = rigid_body_sets
.get_mut(link.0)
.expect(RAPIER_CONTEXT_EXPECT_ERROR)
.into_inner();
if let Some(rb) = rigidbody_set.bodies.get_mut(handle.0) {
rb.enable_ccd(ccd.enabled);
}
}
for (handle, link, soft_ccd) in changed_soft_ccd.iter() {
let rigidbody_set = rigid_body_sets
.get_mut(link.0)
.expect(RAPIER_CONTEXT_EXPECT_ERROR)
.into_inner();
if let Some(rb) = rigidbody_set.bodies.get_mut(handle.0) {
rb.set_soft_ccd_prediction(soft_ccd.prediction);
}
}
for (handle, link, dominance) in changed_dominance.iter() {
let rigidbody_set = rigid_body_sets
.get_mut(link.0)
.expect(RAPIER_CONTEXT_EXPECT_ERROR)
.into_inner();
if let Some(rb) = rigidbody_set.bodies.get_mut(handle.0) {
rb.set_dominance_group(dominance.groups);
}
}
for (handle, link, damping) in changed_damping.iter() {
let rigidbody_set = rigid_body_sets
.get_mut(link.0)
.expect(RAPIER_CONTEXT_EXPECT_ERROR)
.into_inner();
if let Some(rb) = rigidbody_set.bodies.get_mut(handle.0) {
rb.set_linear_damping(damping.linear_damping);
rb.set_angular_damping(damping.angular_damping);
}
}
for (handle, link, _) in changed_disabled.iter() {
let rigidbody_set = rigid_body_sets
.get_mut(link.0)
.expect(RAPIER_CONTEXT_EXPECT_ERROR)
.into_inner();
if let Some(co) = rigidbody_set.bodies.get_mut(handle.0) {
co.set_enabled(false);
}
}
}
pub fn writeback_rigid_bodies(
mut rigid_body_sets: Query<&mut RapierRigidBodySet>,
timestep_mode: Res<TimestepMode>,
config: Query<&RapierConfiguration>,
sim_to_render_time: Query<&SimulationToRenderTime>,
global_transforms: Query<&GlobalTransform>,
mut writeback: Query<
RigidBodyWritebackComponents,
(With<RigidBody>, Without<RigidBodyDisabled>),
>,
) {
for (handle, link, parent, transform, mut interpolation, mut velocity, mut sleeping) in
writeback.iter_mut()
{
let config = config
.get(link.0)
.expect("Could not get `RapierConfiguration`");
if !config.physics_pipeline_active {
continue;
}
let handle = handle.0;
let rigid_body_set = rigid_body_sets
.get_mut(link.0)
.expect(RAPIER_CONTEXT_EXPECT_ERROR)
.into_inner();
let sim_to_render_time = sim_to_render_time
.get(link.0)
.expect("Could not get `SimulationToRenderTime`");
if let Some(rb) = rigid_body_set.bodies.get(handle) {
let mut interpolated_pos = utils::iso_to_transform(rb.position());
if let TimestepMode::Interpolated { dt, .. } = *timestep_mode {
if let Some(interpolation) = interpolation.as_deref_mut() {
if interpolation.end.is_none() {
interpolation.end = Some(*rb.position());
}
if let Some(interpolated) =
interpolation.lerp_slerp((dt + sim_to_render_time.diff) / dt)
{
interpolated_pos = utils::iso_to_transform(&interpolated);
}
}
}
if let Some(mut transform) = transform {
interpolated_pos = interpolated_pos.with_scale(transform.scale);
if let Some(parent_global_transform) =
parent.and_then(|p| global_transforms.get(**p).ok())
{
let (_, inverse_parent_rotation, inverse_parent_translation) =
parent_global_transform
.affine()
.inverse()
.to_scale_rotation_translation();
let new_rotation = inverse_parent_rotation * interpolated_pos.rotation;
#[allow(unused_mut)] let mut new_translation = inverse_parent_rotation
* interpolated_pos.translation
+ inverse_parent_translation;
#[cfg(feature = "dim2")]
{
new_translation.z = transform.translation.z;
}
if transform.rotation != new_rotation
|| transform.translation != new_translation
{
transform.rotation = new_rotation;
transform.translation = new_translation;
}
let new_global_transform = parent_global_transform.mul_transform(*transform);
rigid_body_set
.last_body_transform_set
.insert(handle, new_global_transform);
} else {
#[cfg(feature = "dim2")]
{
interpolated_pos.translation.z = transform.translation.z;
}
if transform.rotation != interpolated_pos.rotation
|| transform.translation != interpolated_pos.translation
{
transform.rotation = interpolated_pos.rotation;
transform.translation = interpolated_pos.translation;
}
rigid_body_set
.last_body_transform_set
.insert(handle, GlobalTransform::from(interpolated_pos));
}
}
if let Some(velocity) = &mut velocity {
let new_vel = Velocity {
linvel: (*rb.linvel()).into(),
#[cfg(feature = "dim3")]
angvel: (*rb.angvel()).into(),
#[cfg(feature = "dim2")]
angvel: rb.angvel(),
};
if **velocity != new_vel {
**velocity = new_vel;
}
}
if let Some(sleeping) = &mut sleeping {
if sleeping.sleeping != rb.is_sleeping() {
sleeping.sleeping = rb.is_sleeping();
}
}
}
}
}
pub fn init_rigid_bodies(
mut commands: Commands,
default_context_access: Query<Entity, With<DefaultRapierContext>>,
mut rigidbody_sets: Query<(Entity, &mut RapierRigidBodySet)>,
rigid_bodies: Query<RigidBodyComponents, Without<RapierRigidBodyHandle>>,
) {
for (
(entity, entity_context_link),
rb,
transform,
vel,
additional_mass_props,
_mass_props,
locked_axes,
force,
gravity_scale,
(ccd, soft_ccd),
dominance,
sleep,
damping,
disabled,
additional_solver_iters,
) in rigid_bodies.iter()
{
let mut builder = RigidBodyBuilder::new((*rb).into());
builder = builder.enabled(disabled.is_none());
if let Some(transform) = transform {
builder = builder.position(utils::transform_to_iso(&transform.compute_transform()));
}
#[allow(clippy::useless_conversion)] if let Some(vel) = vel {
builder = builder.linvel(vel.linvel.into()).angvel(vel.angvel.into());
}
if let Some(locked_axes) = locked_axes {
builder = builder.locked_axes((*locked_axes).into())
}
if let Some(gravity_scale) = gravity_scale {
builder = builder.gravity_scale(gravity_scale.0);
}
if let Some(ccd) = ccd {
builder = builder.ccd_enabled(ccd.enabled)
}
if let Some(soft_ccd) = soft_ccd {
builder = builder.soft_ccd_prediction(soft_ccd.prediction)
}
if let Some(dominance) = dominance {
builder = builder.dominance_group(dominance.groups)
}
if let Some(sleep) = sleep {
builder = builder.sleeping(sleep.sleeping);
}
if let Some(damping) = damping {
builder = builder
.linear_damping(damping.linear_damping)
.angular_damping(damping.angular_damping);
}
if let Some(mprops) = additional_mass_props {
builder = match mprops {
AdditionalMassProperties::MassProperties(mprops) => {
builder.additional_mass_properties(mprops.into_rapier())
}
AdditionalMassProperties::Mass(mass) => builder.additional_mass(*mass),
};
}
if let Some(added_iters) = additional_solver_iters {
builder = builder.additional_solver_iterations(added_iters.0);
}
builder = builder.user_data(entity.to_bits() as u128);
let mut rb = builder.build();
#[allow(clippy::useless_conversion)] if let Some(force) = force {
rb.add_force(force.force.into(), false);
rb.add_torque(force.torque.into(), false);
}
if let Some(sleep) = sleep {
let activation = rb.activation_mut();
activation.normalized_linear_threshold = sleep.normalized_linear_threshold;
activation.angular_threshold = sleep.angular_threshold;
}
let context_entity = entity_context_link.map_or_else(
|| {
let context_entity = default_context_access.get_single().ok()?;
commands
.entity(entity)
.insert(RapierContextEntityLink(context_entity));
Some(context_entity)
},
|link| Some(link.0),
);
let Some(context_entity) = context_entity else {
continue;
};
let Ok((_, mut rigidbody_set)) = rigidbody_sets.get_mut(context_entity) else {
log::error!("Could not find entity {context_entity} with rapier context while initializing {entity}");
continue;
};
let handle = rigidbody_set.bodies.insert(rb);
commands
.entity(entity)
.insert(RapierRigidBodyHandle(handle));
rigidbody_set.entity2body.insert(entity, handle);
if let Some(transform) = transform {
rigidbody_set
.last_body_transform_set
.insert(handle, *transform);
}
}
}
pub fn apply_initial_rigid_body_impulses(
mut context: Query<(&mut RapierRigidBodySet, &RapierContextColliders)>,
mut init_impulses: Query<
(Entity, &RapierContextEntityLink, &mut ExternalImpulse),
Without<RapierRigidBodyHandle>,
>,
) {
for (entity, link, mut impulse) in init_impulses.iter_mut() {
let (mut rigidbody_set, context_colliders) =
context.get_mut(link.0).expect(RAPIER_CONTEXT_EXPECT_ERROR);
let rigidbody_set = &mut *rigidbody_set;
let bodies = &mut rigidbody_set.bodies;
if let Some(rb) = rigidbody_set
.entity2body
.get(&entity)
.and_then(|h| bodies.get_mut(*h))
{
rb.recompute_mass_properties_from_colliders(&context_colliders.colliders);
rb.apply_impulse(impulse.impulse.into(), false);
#[allow(clippy::useless_conversion)] rb.apply_torque_impulse(impulse.torque_impulse.into(), false);
impulse.reset();
}
}
}