1use bevy::ecs::prelude::*;
2use bevy::math::Affine3A;
3use bevy::transform::prelude::*;
4use fnv::FnvHashMap;
5
6use heron_core::{Damping, PhysicMaterial, RigidBody, RotationConstraints, Velocity};
7
8use crate::convert::{IntoBevy, IntoRapier};
9use crate::rapier::geometry::ColliderSet;
10use crate::rapier::{
11 dynamics::{IslandManager, RigidBodyBuilder, RigidBodyHandle, RigidBodySet, RigidBodyType},
12 prelude::{ImpulseJointSet, MultibodyJointSet},
13};
14
15pub(crate) type HandleMap = FnvHashMap<Entity, RigidBodyHandle>;
16
17#[allow(clippy::type_complexity)]
18pub(crate) fn create(
19 mut commands: Commands<'_, '_>,
20 mut bodies: ResMut<'_, RigidBodySet>,
21 mut handles: ResMut<'_, HandleMap>,
22 query: Query<
23 '_,
24 '_,
25 (
26 Entity,
27 &GlobalTransform,
28 &RigidBody,
29 Option<&Velocity>,
30 Option<&Damping>,
31 Option<&RotationConstraints>,
32 ),
33 Without<super::RigidBodyHandle>,
34 >,
35) {
36 for (entity, transform, body, velocity, damping, rotation_constraints) in query.iter() {
37 let (_, global_rotation, global_translation) = transform.to_scale_rotation_translation();
38 let mut builder = RigidBodyBuilder::new(body_status(*body))
39 .user_data(entity.to_bits().into())
40 .position((global_translation, global_rotation).into_rapier());
41
42 #[allow(unused_variables)]
43 if let Some(RotationConstraints {
44 allow_x,
45 allow_y,
46 allow_z,
47 }) = rotation_constraints.copied()
48 {
49 #[cfg(dim2)]
50 if !allow_z {
51 builder = builder.lock_rotations();
52 }
53 #[cfg(dim3)]
54 {
55 builder = builder.restrict_rotations(allow_x, allow_y, allow_z);
56 }
57 }
58
59 if let Some(v) = velocity {
60 builder = builder
61 .linvel(v.linear.into_rapier())
62 .angvel(v.angular.into_rapier());
63 }
64
65 if let Some(d) = damping {
66 builder = builder.linear_damping(d.linear).angular_damping(d.angular);
67 }
68
69 let rigid_body_handle = bodies.insert(builder.build());
70
71 handles.insert(entity, rigid_body_handle);
72 commands
73 .entity(entity)
74 .insert(super::RigidBodyHandle(rigid_body_handle));
75 }
76}
77
78#[allow(clippy::too_many_arguments)]
79pub(crate) fn remove_invalids_after_components_removed(
80 mut commands: Commands<'_, '_>,
81 mut handles: ResMut<'_, HandleMap>,
82 mut bodies: ResMut<'_, RigidBodySet>,
83 mut islands: ResMut<'_, IslandManager>,
84 mut colliders: ResMut<'_, ColliderSet>,
85 mut impulse_joints: ResMut<'_, ImpulseJointSet>,
86 mut multibody_joints: ResMut<'_, MultibodyJointSet>,
87 bodies_removed: RemovedComponents<'_, RigidBody>,
88 constraints_removed: RemovedComponents<'_, RotationConstraints>,
89 materials_removed: RemovedComponents<'_, PhysicMaterial>,
90 rb_entities: Query<'_, '_, Entity, With<super::RigidBodyHandle>>,
91 collider_entities: Query<'_, '_, Entity, With<super::ColliderHandle>>,
92) {
93 bodies_removed
94 .iter()
95 .chain(constraints_removed.iter())
96 .chain(materials_removed.iter())
97 .for_each(|entity| {
98 if let Some(handle) = handles.remove(&entity) {
99 remove_collider_handles(
100 &mut commands,
101 &collider_entities,
102 &bodies,
103 &colliders,
104 handle,
105 );
106 bodies.remove(
107 handle,
108 &mut islands,
109 &mut colliders,
110 &mut impulse_joints,
111 &mut multibody_joints,
112 true,
113 );
114 if rb_entities.get(entity).is_ok() {
115 commands.entity(entity).remove::<super::RigidBodyHandle>();
116 }
117 }
118 });
119}
120
121#[allow(clippy::type_complexity)]
122pub(crate) fn remove_invalids_after_component_changed(
123 mut commands: Commands<'_, '_>,
124 mut handles: ResMut<'_, HandleMap>,
125 mut bodies: ResMut<'_, RigidBodySet>,
126 mut islands: ResMut<'_, IslandManager>,
127 mut colliders: ResMut<'_, ColliderSet>,
128 mut impulse_joints: ResMut<'_, ImpulseJointSet>,
129 mut multibody_joints: ResMut<'_, MultibodyJointSet>,
130 collider_entities: Query<'_, '_, Entity, With<super::ColliderHandle>>,
131 rigidbody_entities: Query<'_, '_, Entity, With<super::RigidBodyHandle>>,
132 changed: Query<
133 '_,
134 '_,
135 (Entity, &super::RigidBodyHandle),
136 Or<(
137 Changed<RigidBody>,
138 Changed<RotationConstraints>,
139 Changed<PhysicMaterial>,
140 )>,
141 >,
142) {
143 for (entity, handle) in changed.iter() {
144 remove_collider_handles(
145 &mut commands,
146 &collider_entities,
147 &bodies,
148 &colliders,
149 handle.0,
150 );
151 bodies.remove(
152 handle.0,
153 &mut islands,
154 &mut colliders,
155 &mut impulse_joints,
156 &mut multibody_joints,
157 true,
158 );
159 if rigidbody_entities.get(entity).is_ok() {
160 commands.entity(entity).remove::<super::RigidBodyHandle>();
161 }
162 handles.remove(&entity);
163 }
164}
165
166#[allow(clippy::manual_filter_map)]
167fn remove_collider_handles(
168 commands: &mut Commands<'_, '_>,
169 entities: &Query<'_, '_, Entity, With<super::ColliderHandle>>,
170 bodies: &RigidBodySet,
171 colliders: &ColliderSet,
172 handle: RigidBodyHandle,
173) {
174 bodies
175 .get(handle)
176 .iter()
177 .flat_map(|it| it.colliders().iter())
178 .filter_map(|it| colliders.get(*it))
179 .map(|it| {
180 #[allow(clippy::cast_possible_truncation)]
181 Entity::from_bits(it.user_data as u64)
182 })
183 .filter(|e| entities.get(*e).is_ok())
184 .for_each(|collider_entity| {
185 commands
186 .entity(collider_entity)
187 .remove::<super::ColliderHandle>();
188 });
189}
190
191pub(crate) fn update_rapier_position(
192 mut bodies: ResMut<'_, RigidBodySet>,
193 query: Query<'_, '_, (&GlobalTransform, &super::RigidBodyHandle), Changed<GlobalTransform>>,
194) {
195 for (transform, handle) in query.iter() {
196 if let Some(body) = bodies.get_mut(handle.0) {
197 let (_, global_rotation, global_translation) =
198 transform.to_scale_rotation_translation();
199 let isometry = (global_translation, global_rotation).into_rapier();
200
201 if body.is_kinematic() {
202 body.set_next_kinematic_position(isometry);
203 } else {
204 body.set_position(isometry, true);
205 }
206 }
207 }
208}
209
210pub(crate) fn update_bevy_transform(
211 bodies: Res<'_, RigidBodySet>,
212 mut query: Query<
213 '_,
214 '_,
215 (
216 Option<&mut Transform>,
217 &mut GlobalTransform,
218 &super::RigidBodyHandle,
219 Option<&RigidBody>,
220 ),
221 >,
222) {
223 for (mut local, mut global, handle, body_type) in query.iter_mut() {
224 if !body_type.copied().unwrap_or_default().can_have_velocity() {
225 continue;
226 }
227
228 let body = match bodies.get(handle.0) {
229 None => continue,
230 Some(body) => body,
231 };
232
233 #[cfg(dim3)]
234 let (translation, rotation) = body.position().into_bevy();
235 #[cfg(dim2)]
236 let (mut translation, rotation) = body.position().into_bevy();
237
238 let (global_scale, global_rotation, global_translation) =
239 global.to_scale_rotation_translation();
240
241 #[cfg(dim2)]
242 {
243 translation.z = global_translation.z;
245 }
246
247 if translation == global_translation && rotation == global_rotation {
248 continue;
249 }
250
251 if let Some(local) = &mut local {
252 if local.translation == global_translation {
253 local.translation = translation;
254 } else {
255 local.translation = translation - (global_translation - local.translation);
256 }
257
258 if local.rotation == global_rotation {
259 local.rotation = rotation;
260 } else {
261 local.rotation =
262 rotation * (global_rotation * local.rotation.conjugate()).conjugate();
263 }
264 }
265
266 *global = GlobalTransform::from(Affine3A::from_scale_rotation_translation(
267 global_scale,
268 rotation,
269 translation,
270 ));
271 }
272}
273
274fn body_status(body_type: RigidBody) -> RigidBodyType {
275 match body_type {
276 RigidBody::Dynamic => RigidBodyType::Dynamic,
277 RigidBody::Static | RigidBody::Sensor => RigidBodyType::Fixed,
278 RigidBody::KinematicPositionBased => RigidBodyType::KinematicPositionBased,
279 RigidBody::KinematicVelocityBased => RigidBodyType::KinematicVelocityBased,
280 }
281}