avian3d/dynamics/rigid_body/mass_properties/
system_param.rs1use crate::prelude::*;
2use bevy::{
3    ecs::system::{
4        SystemParam,
5        lifetimeless::{Read, Write},
6    },
7    prelude::*,
8};
9
10#[derive(SystemParam)]
14pub struct MassPropertyHelper<'w, 's> {
15    query: Query<
16        'w,
17        's,
18        (
19            Option<Read<Mass>>,
20            Option<Read<AngularInertia>>,
21            Option<Read<CenterOfMass>>,
22            Option<Read<ColliderMassProperties>>,
23            Option<Read<ColliderTransform>>,
24            Has<Sensor>,
25        ),
26        Or<(
27            With<Mass>,
28            With<AngularInertia>,
29            With<CenterOfMass>,
30            (
31                With<ColliderMassProperties>,
32                With<ColliderTransform>,
33                Without<Sensor>,
34            ),
35        )>,
36    >,
37    computed_mass_properties_query: Query<
38        'w,
39        's,
40        (
41            Write<ComputedMass>,
42            Write<ComputedAngularInertia>,
43            Write<ComputedCenterOfMass>,
44            Option<Read<Mass>>,
45            Option<Read<AngularInertia>>,
46            Option<Read<CenterOfMass>>,
47            Has<NoAutoMass>,
48            Has<NoAutoAngularInertia>,
49            Has<NoAutoCenterOfMass>,
50        ),
51    >,
52    children: Query<'w, 's, Read<Children>>,
53}
54
55impl MassPropertyHelper<'_, '_> {
56    pub fn update_mass_properties(&mut self, entity: Entity) {
61        let mut mass_props = self.total_mass_properties(entity);
63
64        let Ok((
65            mut computed_mass,
66            mut computed_inertia,
67            mut computed_com,
68            mass,
69            angular_inertia,
70            center_of_mass,
71            no_auto_mass,
72            no_auto_inertia,
73            no_auto_com,
74        )) = self.computed_mass_properties_query.get_mut(entity)
75        else {
76            return;
77        };
78
79        if no_auto_mass {
83            if let Some(mass) = mass {
84                mass_props.set_mass(mass.0, !no_auto_inertia);
85                computed_mass.set(mass_props.mass as Scalar);
86            } else if !no_auto_inertia {
87                #[allow(clippy::unnecessary_cast)]
89                mass_props.set_mass(computed_mass.value() as f32, true);
90            }
91        } else {
92            computed_mass.set(mass_props.mass as Scalar);
93        }
94
95        if no_auto_inertia {
96            if let Some(angular_inertia) = angular_inertia {
97                #[cfg(feature = "2d")]
98                {
99                    mass_props.angular_inertia = angular_inertia.0;
100                    computed_inertia.set(mass_props.angular_inertia as Scalar);
101                }
102                #[cfg(feature = "3d")]
103                {
104                    mass_props.principal_angular_inertia = angular_inertia.principal;
105                    mass_props.local_inertial_frame = angular_inertia.local_frame;
106                    *computed_inertia = ComputedAngularInertia::new_with_local_frame(
107                        mass_props.principal_angular_inertia.adjust_precision(),
108                        mass_props.local_inertial_frame.adjust_precision(),
109                    );
110                }
111            }
112        } else {
113            #[cfg(feature = "2d")]
114            {
115                computed_inertia.set(mass_props.angular_inertia as Scalar);
116            }
117            #[cfg(feature = "3d")]
118            {
119                *computed_inertia = ComputedAngularInertia::new_with_local_frame(
120                    mass_props.principal_angular_inertia.adjust_precision(),
121                    mass_props.local_inertial_frame.adjust_precision(),
122                );
123            }
124        }
125
126        if no_auto_com {
127            if let Some(center_of_mass) = center_of_mass {
128                mass_props.center_of_mass = center_of_mass.0;
129                computed_com.0 = mass_props.center_of_mass.adjust_precision();
130            }
131        } else {
132            computed_com.0 = mass_props.center_of_mass.adjust_precision();
133        }
134    }
135
136    pub fn total_mass_properties(&self, entity: Entity) -> MassProperties {
141        core::iter::once(self.local_mass_properties(entity))
142            .chain(
143                self.children
144                    .iter_descendants(entity)
145                    .map(|child| self.local_mass_properties(child)),
146            )
147            .flatten()
148            .sum()
149    }
150
151    pub fn descendants_mass_properties(&self, entity: Entity) -> MassProperties {
155        self.children
156            .iter_descendants(entity)
157            .filter_map(|child| self.local_mass_properties(child))
158            .sum()
159    }
160
161    pub fn local_mass_properties(&self, entity: Entity) -> Option<MassProperties> {
168        let (mass, angular_inertia, center_of_mass, collider_mass, collider_transform, is_sensor) =
169            self.query.get(entity).ok()?;
170
171        let mut mass_props = collider_mass
173            .filter(|_| !is_sensor)
174            .map_or(MassProperties::ZERO, |m| **m);
175
176        if let Some(mass) = mass {
178            let update_angular_inertia = angular_inertia.is_none();
181            mass_props.set_mass(mass.0, update_angular_inertia);
182        }
183
184        if let Some(angular_inertia) = angular_inertia {
186            #[cfg(feature = "2d")]
187            {
188                mass_props.angular_inertia = angular_inertia.0;
189            }
190            #[cfg(feature = "3d")]
191            {
192                mass_props.principal_angular_inertia = angular_inertia.principal;
193                mass_props.local_inertial_frame = angular_inertia.local_frame;
194            }
195        }
196
197        if let Some(center_of_mass) = center_of_mass {
199            mass_props.center_of_mass = center_of_mass.0;
200        }
201
202        if let Some(collider_transform) = collider_mass.and(collider_transform) {
203            #[cfg(feature = "2d")]
204            {
205                mass_props.transform_by(Isometry2d::new(
206                    collider_transform.translation.f32(),
207                    Rot2::from(collider_transform.rotation),
208                ));
209            }
210            #[cfg(feature = "3d")]
211            {
212                mass_props.transform_by(Isometry3d::new(
213                    collider_transform.translation.f32(),
214                    collider_transform.rotation.f32(),
215                ));
216            }
217        }
218
219        Some(mass_props)
220    }
221}