use crate::model::{Clip, Interpolation, Skeleton, Track, TrackValues, Transform};
use glam::{Mat4, Quat, Vec3};
#[derive(Debug)]
pub struct PoseGrid {
pub times: Vec<f32>,
bone_count: usize,
local: Vec<Transform>,
model: Vec<Mat4>,
}
impl PoseGrid {
fn index(&self, frame: usize, bone: usize) -> usize {
assert!(
frame < self.frame_count(),
"frame index {frame} outside PoseGrid frame count {}",
self.frame_count()
);
assert!(
bone < self.bone_count,
"bone index {bone} outside PoseGrid bone count {}",
self.bone_count
);
frame * self.bone_count + bone
}
pub fn frame_count(&self) -> usize {
self.times.len()
}
pub fn bone_count(&self) -> usize {
self.bone_count
}
pub fn local(&self, frame: usize, bone: usize) -> Transform {
self.local[self.index(frame, bone)]
}
pub fn model(&self, frame: usize, bone: usize) -> Mat4 {
self.model[self.index(frame, bone)]
}
pub fn model_position(&self, frame: usize, bone: usize) -> Vec3 {
self.model(frame, bone).w_axis.truncate()
}
}
pub fn default_frame_count(clip: &Clip) -> usize {
clip.tracks
.iter()
.map(Track::key_count)
.max()
.unwrap_or(2)
.max(2)
}
pub fn sample_clip(skeleton: &Skeleton, clip: &Clip, frames: usize) -> PoseGrid {
let frames = frames.max(2);
let nb = skeleton.bones.len();
let duration = clip.duration_s as f32;
let times: Vec<f32> = (0..frames)
.map(|i| duration * i as f32 / (frames - 1) as f32)
.collect();
let mut local = vec![Transform::IDENTITY; frames * nb];
for f in 0..frames {
for (b, bone) in skeleton.bones.iter().enumerate() {
local[f * nb + b] = bone.rest;
}
}
for track in &clip.tracks {
if track.times.is_empty() || track.bone >= nb {
continue;
}
for (f, &t) in times.iter().enumerate() {
let slot = &mut local[f * nb + track.bone];
match &track.values {
TrackValues::Vec3s(_) => {
let v = sample_vec3(track, t);
match track.property {
crate::model::Property::Translation => slot.translation = v,
crate::model::Property::Scale => slot.scale = v,
crate::model::Property::Rotation => {}
}
}
TrackValues::Quats(_) => slot.rotation = sample_quat(track, t),
}
}
}
let mut model = vec![Mat4::IDENTITY; frames * nb];
for f in 0..frames {
for (b, bone) in skeleton.bones.iter().enumerate() {
let m = local[f * nb + b].to_mat4();
model[f * nb + b] = match bone.parent {
Some(p) => model[f * nb + p] * m,
None => m,
};
}
}
PoseGrid {
times,
bone_count: nb,
local,
model,
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum TrackSample {
Vec3(Vec3),
Quat(Quat),
}
pub fn sample_track(track: &Track, t: f32) -> TrackSample {
match &track.values {
TrackValues::Vec3s(_) => TrackSample::Vec3(sample_vec3(track, t)),
TrackValues::Quats(_) => TrackSample::Quat(sample_quat(track, t)),
}
}
fn segment(times: &[f32], t: f32) -> (usize, usize, f32) {
let n = times.len();
if n <= 1 || t <= times[0] {
return (0, 0, 0.0);
}
if t >= times[n - 1] {
return (n - 1, n - 1, 0.0);
}
let k1 = times.partition_point(|&k| k <= t).min(n - 1);
if k1 == 0 {
return (0, 0, 0.0);
}
let k0 = k1 - 1;
let dt = times[k1] - times[k0];
let u = if dt > 0.0 { (t - times[k0]) / dt } else { 0.0 };
(k0, k1, u)
}
fn value_at<T: Copy + Default>(vals: &[T], index: usize) -> T {
vals.get(index).copied().unwrap_or_default()
}
fn hermite(u: f32) -> (f32, f32, f32, f32) {
let u2 = u * u;
let u3 = u2 * u;
(
2.0 * u3 - 3.0 * u2 + 1.0, u3 - 2.0 * u2 + u, -2.0 * u3 + 3.0 * u2, u3 - u2, )
}
fn sample_vec3(track: &Track, t: f32) -> Vec3 {
let TrackValues::Vec3s(vals) = &track.values else {
return Vec3::ZERO;
};
let (k0, k1, u) = segment(&track.times, t);
let v0 = value_at(vals, track.value_index(k0));
if k0 == k1 {
return v0;
}
let v1 = value_at(vals, track.value_index(k1));
match track.interpolation {
Interpolation::Step => v0,
Interpolation::Linear => v0.lerp(v1, u),
Interpolation::CubicSpline => {
let dt = track.times[k1] - track.times[k0];
let m0 = value_at(vals, 3 * k0 + 2) * dt;
let m1 = value_at(vals, 3 * k1) * dt;
let (h00, h10, h01, h11) = hermite(u);
v0 * h00 + m0 * h10 + v1 * h01 + m1 * h11
}
}
}
fn sample_quat(track: &Track, t: f32) -> Quat {
let TrackValues::Quats(vals) = &track.values else {
return Quat::IDENTITY;
};
let (k0, k1, u) = segment(&track.times, t);
let q0: Quat = value_at(vals, track.value_index(k0));
if k0 == k1 {
return q0.normalize();
}
let q1 = value_at(vals, track.value_index(k1));
match track.interpolation {
Interpolation::Step => q0.normalize(),
Interpolation::Linear => {
let q1 = if q0.dot(q1) < 0.0 { -q1 } else { q1 };
q0.normalize().slerp(q1.normalize(), u)
}
Interpolation::CubicSpline => {
let dt = track.times[k1] - track.times[k0];
let m0 = value_at(vals, 3 * k0 + 2).to_array();
let m1 = value_at(vals, 3 * k1).to_array();
let a0 = q0.to_array();
let a1 = q1.to_array();
let (h00, h10, h01, h11) = hermite(u);
let mut out = [0.0f32; 4];
for i in 0..4 {
out[i] = a0[i] * h00 + m0[i] * dt * h10 + a1[i] * h01 + m1[i] * dt * h11;
}
Quat::from_array(out).normalize()
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::model::{Bone, Property};
fn track(times: Vec<f32>, quats: Vec<Quat>) -> Track {
Track {
bone: 0,
property: Property::Rotation,
interpolation: Interpolation::Linear,
times,
values: TrackValues::Quats(quats),
}
}
fn vec3_track(interpolation: Interpolation, times: Vec<f32>, vals: Vec<Vec3>) -> Track {
Track {
bone: 0,
property: Property::Translation,
interpolation,
times,
values: TrackValues::Vec3s(vals),
}
}
#[test]
#[should_panic(expected = "outside PoseGrid bone count")]
fn pose_grid_rejects_bone_index_at_frame_boundary() {
let skeleton = Skeleton {
bones: vec![Bone {
name: "root".into(),
parent: None,
rest: Transform::IDENTITY,
inverse_bind: None,
}],
};
let clip = Clip {
name: "idle".into(),
duration_s: 1.0,
tracks: Vec::new(),
};
let grid = sample_clip(&skeleton, &clip, 2);
let _ = grid.local(0, grid.bone_count());
}
#[test]
fn nan_first_key_time_samples_without_panicking() {
let t = track(
vec![f32::NAN, 0.5, 1.0],
vec![Quat::IDENTITY, Quat::IDENTITY, Quat::IDENTITY],
);
for time in [-1.0, 0.0, 0.25, 0.75, 2.0] {
let TrackSample::Quat(q) = sample_track(&t, time) else {
panic!("rotation track samples a quat");
};
assert!(q.is_finite() || q.is_nan()); }
}
#[test]
fn all_nan_times_sample_without_panicking() {
let t = track(
vec![f32::NAN, f32::NAN],
vec![Quat::IDENTITY, Quat::IDENTITY],
);
sample_track(&t, 0.5);
}
#[test]
fn empty_track_samples_default_without_panicking() {
let t = track(vec![], vec![]);
let TrackSample::Quat(q) = sample_track(&t, 0.5) else {
panic!("rotation track samples a quat");
};
assert_eq!(q, Quat::IDENTITY.normalize());
}
#[test]
fn short_values_sample_default_without_panicking() {
let t = track(vec![0.0, 0.5, 1.0], vec![Quat::IDENTITY, Quat::IDENTITY]);
sample_track(&t, 0.75); }
#[test]
fn short_vec3_values_sample_default_without_panicking() {
let t = vec3_track(Interpolation::Linear, vec![0.0, 0.5, 1.0], vec![Vec3::ONE]);
let TrackSample::Vec3(v) = sample_track(&t, 0.75) else {
panic!("translation track samples a vec3");
};
assert!(v.is_finite());
}
#[test]
fn short_cubic_values_sample_default_without_panicking() {
let t = vec3_track(
Interpolation::CubicSpline,
vec![0.0, 1.0],
vec![Vec3::ZERO, Vec3::ONE, Vec3::ZERO], );
sample_track(&t, 0.5);
}
}