wgrapier2d/dynamics/
body.rs

1//! Rigid-body definition and set.
2
3use encase::ShaderType;
4use wgcore::tensor::GpuVector;
5use wgcore::Shader;
6use wgebra::{WgQuat, WgSim2, WgSim3};
7use wgparry::cuboid::GpuCuboid;
8use wgparry::math::{AngVector, AngularInertia, GpuSim, Vector};
9use wgparry::{dim_shader_defs, substitute_aliases};
10use wgpu::{BufferUsages, Device};
11
12#[derive(ShaderType, Copy, Clone, PartialEq)]
13#[repr(C)]
14/// Linear and angular forces with a layout compatible with the corresponding WGSL struct.
15pub struct GpuForce {
16    /// The linear part of the force.
17    pub linear: Vector<f32>,
18    /// The angular part of the force (aka. the torque).
19    pub angular: AngVector<f32>,
20}
21
22#[derive(ShaderType, Copy, Clone, PartialEq, Default)]
23#[repr(C)]
24/// Linear and angular velocities with a layout compatible with the corresponding WGSL struct.
25pub struct GpuVelocity {
26    /// The linear (translational) velocity.
27    pub linear: Vector<f32>,
28    /// The angular (rotational) velocity.
29    pub angular: AngVector<f32>,
30}
31
32#[derive(ShaderType, Copy, Clone, PartialEq)]
33#[repr(C)]
34/// Rigid-body mass-properties, with a layout compatible with the corresponding WGSL struct.
35pub struct GpuMassProperties {
36    /// The inverse angular inertia tensor.
37    pub inv_inertia: AngularInertia<f32>,
38    /// The inverse mass.
39    pub inv_mass: Vector<f32>,
40    /// The center-of-mass.
41    pub com: Vector<f32>, // ShaderType isn’t implemented for Point
42}
43
44impl Default for GpuMassProperties {
45    fn default() -> Self {
46        GpuMassProperties {
47            #[rustfmt::skip]
48            #[cfg(feature = "dim2")]
49            inv_inertia: 1.0,
50            #[cfg(feature = "dim3")]
51            inv_inertia: AngularInertia::identity(),
52            inv_mass: Vector::repeat(1.0),
53            com: Vector::zeros(),
54        }
55    }
56}
57
58/// A set of rigid-bodies stored on the gpu.
59pub struct GpuBodySet {
60    len: u32,
61    pub(crate) mprops: GpuVector<GpuMassProperties>,
62    pub(crate) local_mprops: GpuVector<GpuMassProperties>,
63    pub(crate) vels: GpuVector<GpuVelocity>,
64    pub(crate) poses: GpuVector<GpuSim>,
65    // TODO: support other shape types.
66    // TODO: support a shape with a shift relative to the body.
67    pub(crate) shapes: GpuVector<GpuCuboid>,
68}
69
70#[derive(Copy, Clone)]
71/// Helper struct for defining a rigid-body to be added to a [`GpuBodySet`].
72pub struct BodyDesc {
73    /// The rigid-body’s mass-properties.
74    pub mprops: GpuMassProperties,
75    /// The rigid-body’s linear and angular velocities.
76    pub vel: GpuVelocity,
77    /// The rigid-body’s world-space pose.
78    pub pose: GpuSim,
79    /// The rigid-body’s shape.
80    pub shape: GpuCuboid,
81}
82
83impl Default for BodyDesc {
84    fn default() -> Self {
85        Self {
86            mprops: Default::default(),
87            vel: Default::default(),
88            pose: Default::default(),
89            shape: GpuCuboid::new(Vector::repeat(0.5)),
90        }
91    }
92}
93
94impl GpuBodySet {
95    /// Is this set empty?
96    pub fn is_empty(&self) -> bool {
97        self.len == 0
98    }
99
100    /// Number of rigid-bodies in this set.
101    pub fn len(&self) -> u32 {
102        self.len
103    }
104
105    /// Create a set of `bodies` on the gpu.
106    pub fn new(device: &Device, bodies: &[BodyDesc]) -> Self {
107        let (mprops, (vels, (poses, shapes))): (Vec<_>, (Vec<_>, (Vec<_>, Vec<_>))) = bodies
108            .iter()
109            .copied()
110            // NOTE: Looks silly, but we can’t just collect into (Vec, Vec, Vec).
111            .map(|b| (b.mprops, (b.vel, (b.pose, b.shape))))
112            .collect();
113        Self {
114            len: bodies.len() as u32,
115            mprops: GpuVector::encase(device, &mprops, BufferUsages::STORAGE),
116            local_mprops: GpuVector::encase(device, &mprops, BufferUsages::STORAGE),
117            vels: GpuVector::encase(device, &vels, BufferUsages::STORAGE),
118            poses: GpuVector::init(device, &poses, BufferUsages::STORAGE),
119            shapes: GpuVector::encase(device, &shapes, BufferUsages::STORAGE),
120        }
121    }
122
123    /// GPU storage buffer containing the poses of every rigid-body.
124    pub fn poses(&self) -> &GpuVector<GpuSim> {
125        &self.poses
126    }
127
128    /// GPU storage buffer containing the velocities of every rigid-body.
129    pub fn vels(&self) -> &GpuVector<GpuVelocity> {
130        &self.vels
131    }
132
133    /// GPU storage buffer containing the world-space mass-properties of every rigid-body.
134    pub fn mprops(&self) -> &GpuVector<GpuMassProperties> {
135        &self.mprops
136    }
137
138    /// GPU storage buffer containing the local-space mass-properties of every rigid-body.
139    pub fn local_mprops(&self) -> &GpuVector<GpuMassProperties> {
140        &self.local_mprops
141    }
142
143    /// GPU storage buffer containing the shape of every rigid-body.
144    pub fn shapes(&self) -> &GpuVector<GpuCuboid> {
145        &self.shapes
146    }
147}
148
149#[derive(Shader)]
150#[shader(
151    derive(WgQuat, WgSim3, WgSim2),
152    src = "body.wgsl",
153    src_fn = "substitute_aliases",
154    shader_defs = "dim_shader_defs"
155)]
156/// Shader defining structs related to rigid-bodies, as well as functions to compute point velocities
157/// and update world-space mass-properties.
158pub struct WgBody;
159
160// TODO: this test won’t pass due to the lack of `substitute_aliases`
161//       and `dim_shader_defs` in the macro. Figure out a way to make this work.
162// wgcore::test_shader_compilation!(WgBody);