1use crate::model::{Clip, Interpolation, Skeleton, Track, TrackValues, Transform};
15use glam::{Mat4, Quat, Vec3};
16
17#[derive(Debug)]
19pub struct PoseGrid {
20 pub times: Vec<f32>,
22 bone_count: usize,
23 local: Vec<Transform>,
25 model: Vec<Mat4>,
26 model_rotation: Vec<Quat>,
27}
28
29impl PoseGrid {
30 fn index(&self, frame: usize, bone: usize) -> usize {
31 assert!(
32 frame < self.frame_count(),
33 "frame index {frame} outside PoseGrid frame count {}",
34 self.frame_count()
35 );
36 assert!(
37 bone < self.bone_count,
38 "bone index {bone} outside PoseGrid bone count {}",
39 self.bone_count
40 );
41 frame * self.bone_count + bone
42 }
43
44 pub fn frame_count(&self) -> usize {
46 self.times.len()
47 }
48
49 pub fn bone_count(&self) -> usize {
51 self.bone_count
52 }
53
54 pub fn local(&self, frame: usize, bone: usize) -> Transform {
60 self.local[self.index(frame, bone)]
61 }
62
63 pub fn model(&self, frame: usize, bone: usize) -> Mat4 {
69 self.model[self.index(frame, bone)]
70 }
71
72 pub fn model_position(&self, frame: usize, bone: usize) -> Vec3 {
78 self.model(frame, bone).w_axis.truncate()
79 }
80
81 pub fn model_rotation(&self, frame: usize, bone: usize) -> Quat {
89 self.model_rotation[self.index(frame, bone)]
90 }
91}
92
93pub fn default_frame_count(clip: &Clip) -> usize {
97 clip.tracks
98 .iter()
99 .map(Track::key_count)
100 .max()
101 .unwrap_or(2)
102 .max(2)
103}
104
105pub fn sample_clip(skeleton: &Skeleton, clip: &Clip, frames: usize) -> PoseGrid {
123 let frames = frames.max(2);
124 let nb = skeleton.bones.len();
125 let duration = clip.duration_s as f32;
126 let times: Vec<f32> = (0..frames)
127 .map(|i| duration * i as f32 / (frames - 1) as f32)
128 .collect();
129
130 let mut local = vec![Transform::IDENTITY; frames * nb];
131 for f in 0..frames {
132 for (b, bone) in skeleton.bones.iter().enumerate() {
133 local[f * nb + b] = bone.rest;
134 }
135 }
136
137 for track in &clip.tracks {
138 if track.times.is_empty() || track.bone >= nb {
139 continue;
140 }
141 for (f, &t) in times.iter().enumerate() {
142 let slot = &mut local[f * nb + track.bone];
143 match &track.values {
144 TrackValues::Vec3s(_) => {
145 let v = sample_vec3(track, t);
146 match track.property {
147 crate::model::Property::Translation => slot.translation = v,
148 crate::model::Property::Scale => slot.scale = v,
149 crate::model::Property::Rotation => {}
150 }
151 }
152 TrackValues::Quats(_) => slot.rotation = sample_quat(track, t),
153 }
154 }
155 }
156
157 let mut model = vec![Mat4::IDENTITY; frames * nb];
158 let mut model_rotation = vec![Quat::IDENTITY; frames * nb];
159 for f in 0..frames {
160 for (b, bone) in skeleton.bones.iter().enumerate() {
161 let m = local[f * nb + b].to_mat4();
162 model[f * nb + b] = match bone.parent {
163 Some(p) => model[f * nb + p] * m,
164 None => m,
165 };
166 let rotation = local[f * nb + b].rotation;
167 model_rotation[f * nb + b] = match bone.parent {
168 Some(p) => model_rotation[f * nb + p] * rotation,
169 None => rotation,
170 };
171 }
172 }
173
174 PoseGrid {
175 times,
176 bone_count: nb,
177 local,
178 model,
179 model_rotation,
180 }
181}
182
183#[derive(Debug, Clone, Copy, PartialEq)]
185pub enum TrackSample {
186 Vec3(Vec3),
188 Quat(Quat),
190}
191
192pub fn sample_track(track: &Track, t: f32) -> TrackSample {
200 match &track.values {
201 TrackValues::Vec3s(_) => TrackSample::Vec3(sample_vec3(track, t)),
202 TrackValues::Quats(_) => TrackSample::Quat(sample_quat(track, t)),
203 }
204}
205
206fn segment(times: &[f32], t: f32) -> (usize, usize, f32) {
214 let n = times.len();
215 if n <= 1 || t <= times[0] {
216 return (0, 0, 0.0);
217 }
218 if t >= times[n - 1] {
219 return (n - 1, n - 1, 0.0);
220 }
221 let k1 = times.partition_point(|&k| k <= t).min(n - 1);
222 if k1 == 0 {
223 return (0, 0, 0.0);
224 }
225 let k0 = k1 - 1;
226 let dt = times[k1] - times[k0];
227 let u = if dt > 0.0 { (t - times[k0]) / dt } else { 0.0 };
228 (k0, k1, u)
229}
230
231fn value_at<T: Copy + Default>(vals: &[T], index: usize) -> T {
235 vals.get(index).copied().unwrap_or_default()
236}
237
238fn hermite(u: f32) -> (f32, f32, f32, f32) {
240 let u2 = u * u;
241 let u3 = u2 * u;
242 (
243 2.0 * u3 - 3.0 * u2 + 1.0, u3 - 2.0 * u2 + u, -2.0 * u3 + 3.0 * u2, u3 - u2, )
248}
249
250fn sample_vec3(track: &Track, t: f32) -> Vec3 {
251 let TrackValues::Vec3s(vals) = &track.values else {
252 return Vec3::ZERO;
253 };
254 let (k0, k1, u) = segment(&track.times, t);
255 let v0 = value_at(vals, track.value_index(k0));
256 if k0 == k1 {
257 return v0;
258 }
259 let v1 = value_at(vals, track.value_index(k1));
260 match track.interpolation {
261 Interpolation::Step => v0,
262 Interpolation::Linear => v0.lerp(v1, u),
263 Interpolation::CubicSpline => {
264 let dt = track.times[k1] - track.times[k0];
265 let m0 = value_at(vals, 3 * k0 + 2) * dt;
267 let m1 = value_at(vals, 3 * k1) * dt;
268 let (h00, h10, h01, h11) = hermite(u);
269 v0 * h00 + m0 * h10 + v1 * h01 + m1 * h11
270 }
271 }
272}
273
274fn sample_quat(track: &Track, t: f32) -> Quat {
275 let TrackValues::Quats(vals) = &track.values else {
276 return Quat::IDENTITY;
277 };
278 let (k0, k1, u) = segment(&track.times, t);
279 let q0: Quat = value_at(vals, track.value_index(k0));
280 if k0 == k1 {
281 return q0.normalize();
282 }
283 let q1 = value_at(vals, track.value_index(k1));
284 match track.interpolation {
285 Interpolation::Step => q0.normalize(),
286 Interpolation::Linear => {
287 let q1 = if q0.dot(q1) < 0.0 { -q1 } else { q1 };
290 q0.normalize().slerp(q1.normalize(), u)
291 }
292 Interpolation::CubicSpline => {
293 let dt = track.times[k1] - track.times[k0];
296 let m0 = value_at(vals, 3 * k0 + 2).to_array();
297 let m1 = value_at(vals, 3 * k1).to_array();
298 let a0 = q0.to_array();
299 let a1 = q1.to_array();
300 let (h00, h10, h01, h11) = hermite(u);
301 let mut out = [0.0f32; 4];
302 for i in 0..4 {
303 out[i] = a0[i] * h00 + m0[i] * dt * h10 + a1[i] * h01 + m1[i] * dt * h11;
304 }
305 Quat::from_array(out).normalize()
306 }
307 }
308}
309
310#[cfg(test)]
311mod tests {
312 use super::*;
313 use crate::model::{Bone, Property};
314
315 fn track(times: Vec<f32>, quats: Vec<Quat>) -> Track {
316 Track {
317 bone: 0,
318 property: Property::Rotation,
319 interpolation: Interpolation::Linear,
320 times,
321 values: TrackValues::Quats(quats),
322 }
323 }
324
325 fn vec3_track(interpolation: Interpolation, times: Vec<f32>, vals: Vec<Vec3>) -> Track {
326 Track {
327 bone: 0,
328 property: Property::Translation,
329 interpolation,
330 times,
331 values: TrackValues::Vec3s(vals),
332 }
333 }
334
335 #[test]
336 #[should_panic(expected = "outside PoseGrid bone count")]
337 fn pose_grid_rejects_bone_index_at_frame_boundary() {
338 let skeleton = Skeleton {
339 bones: vec![Bone {
340 name: "root".into(),
341 parent: None,
342 rest: Transform::IDENTITY,
343 inverse_bind: None,
344 }],
345 };
346 let clip = Clip {
347 name: "idle".into(),
348 duration_s: 1.0,
349 tracks: Vec::new(),
350 };
351 let grid = sample_clip(&skeleton, &clip, 2);
352
353 let _ = grid.local(0, grid.bone_count());
354 }
355
356 #[test]
357 fn out_of_range_clip_track_is_skipped() {
358 let skeleton = Skeleton {
359 bones: vec![Bone {
360 name: "root".into(),
361 parent: None,
362 rest: Transform::IDENTITY,
363 inverse_bind: None,
364 }],
365 };
366 let clip = Clip {
367 name: "hostile".into(),
368 duration_s: 1.0,
369 tracks: vec![Track {
370 bone: 1,
371 property: Property::Translation,
372 interpolation: Interpolation::Linear,
373 times: vec![0.0],
374 values: TrackValues::Vec3s(vec![Vec3::splat(42.0)]),
375 }],
376 };
377
378 let grid = sample_clip(&skeleton, &clip, 2);
379 assert_eq!(grid.local(0, 0), Transform::IDENTITY);
380 assert_eq!(grid.local(1, 0), Transform::IDENTITY);
381 }
382
383 #[test]
386 fn nan_first_key_time_samples_without_panicking() {
387 let t = track(
388 vec![f32::NAN, 0.5, 1.0],
389 vec![Quat::IDENTITY, Quat::IDENTITY, Quat::IDENTITY],
390 );
391 for time in [-1.0, 0.0, 0.25, 0.75, 2.0] {
392 let TrackSample::Quat(q) = sample_track(&t, time) else {
393 panic!("rotation track samples a quat");
394 };
395 assert!(q.is_finite() || q.is_nan()); }
397 }
398
399 #[test]
400 fn all_nan_times_sample_without_panicking() {
401 let t = track(
402 vec![f32::NAN, f32::NAN],
403 vec![Quat::IDENTITY, Quat::IDENTITY],
404 );
405 sample_track(&t, 0.5);
406 }
407
408 #[test]
409 fn empty_track_samples_default_without_panicking() {
410 let t = track(vec![], vec![]);
411 let TrackSample::Quat(q) = sample_track(&t, 0.5) else {
412 panic!("rotation track samples a quat");
413 };
414 assert_eq!(q, Quat::IDENTITY.normalize());
415 }
416
417 #[test]
421 fn short_values_sample_default_without_panicking() {
422 let t = track(vec![0.0, 0.5, 1.0], vec![Quat::IDENTITY, Quat::IDENTITY]);
423 sample_track(&t, 0.75); }
425
426 #[test]
427 fn short_vec3_values_sample_default_without_panicking() {
428 let t = vec3_track(Interpolation::Linear, vec![0.0, 0.5, 1.0], vec![Vec3::ONE]);
429 let TrackSample::Vec3(v) = sample_track(&t, 0.75) else {
430 panic!("translation track samples a vec3");
431 };
432 assert!(v.is_finite());
433 }
434
435 #[test]
438 fn short_cubic_values_sample_default_without_panicking() {
439 let t = vec3_track(
440 Interpolation::CubicSpline,
441 vec![0.0, 1.0],
442 vec![Vec3::ZERO, Vec3::ONE, Vec3::ZERO], );
444 sample_track(&t, 0.5);
445 }
446}