avian2d/dynamics/rigid_body/mass_properties/
system_param.rs

1use crate::prelude::*;
2use bevy::{
3    ecs::system::{
4        SystemParam,
5        lifetimeless::{Read, Write},
6    },
7    prelude::*,
8};
9
10/// A [`SystemParam`] that provides helper methods for computing and updating [mass properties].
11///
12/// [mass properties]: crate::dynamics::rigid_body::mass_properties
13#[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    /// Updates the [`ComputedMass`], [`ComputedAngularInertia`], and [`ComputedCenterOfMass`] of the given entity.
57    ///
58    /// This takes into account the mass properties of descendants, unless the given entity has the [`NoAutoMass`],
59    /// [`NoAutoAngularInertia`], or [`NoAutoCenterOfMass`] marker components.
60    pub fn update_mass_properties(&mut self, entity: Entity) {
61        // Compute the total mass properties of the entity and its descendants.
62        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 automatic computation of mass properties is disabled, set them to the local `Mass`, `AngularInertia`, and `CenterOfMass`.
80        // Otherwise, use the computed total mass properties.
81
82        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                // Make sure the angular inertia is scaled to match the existing computed mass.
88                #[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    /// Computes the total mass properties of the given entity,
137    /// taking into account the mass properties of descendants.
138    ///
139    /// This ignores the [`NoAutoMass`], [`NoAutoAngularInertia`], and [`NoAutoCenterOfMass`] marker components.
140    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    /// Computes the total mass properties of the descendants of the given entity.
152    ///
153    /// This ignores the [`NoAutoMass`], [`NoAutoAngularInertia`], and [`NoAutoCenterOfMass`] marker components.
154    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    /// Computes the local mass properties of the given entity.
162    ///
163    /// This only considers the entity's own [`Mass`], [`AngularInertia`], [`CenterOfMass`],
164    /// and/or [`ColliderMassProperties`] if present, not those of its children.
165    ///
166    /// If the entity has no mass properties or the entity does not exist, `None` is returned.
167    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        // Initialize the mass properties with the collider's mass properties or zero.
172        let mut mass_props = collider_mass
173            .filter(|_| !is_sensor)
174            .map_or(MassProperties::ZERO, |m| **m);
175
176        // Set the mass if the `Mass` component is present.
177        if let Some(mass) = mass {
178            // TODO: This needs to consider `NoAutoMass`.
179            // Only adjust the angular inertia if it is not exlicitly set with `AngularInertia`.
180            let update_angular_inertia = angular_inertia.is_none();
181            mass_props.set_mass(mass.0, update_angular_inertia);
182        }
183
184        // Set the angular inertia if the `AngularInertia` component is present.
185        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        // Set the center of mass if the `CenterOfMass` component is present.
198        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}