1use std::sync::Arc;
2
3use glam::{Quat, Vec3A};
4
5use crate::ik_primitive::ChainLinkState;
6use crate::{AnimationClip, ModelArena, PoseArena};
7
8mod ik;
9mod morph;
10mod physics;
11mod world;
12
13#[cfg(test)]
14use crate::ik_primitive::{
15 LimitedAxesLinkStepInput, PlaneLinkStepInput, axis_vec, decompose_euler_xyz, euler_xyz_to_quat,
16 limit_axis_bounds, quat_to_rotation_mat3, signed_projected_angle, solve_limited_axes_link_step,
17 solve_plane_link_step,
18};
19
20#[derive(Debug)]
21struct IkScratch {
22 links: Vec<crate::IkLink>,
23 base_rotations: Vec<Quat>,
24 base_ik_rotations: Vec<Quat>,
25 ik_rotations: Vec<Quat>,
26 best_ik_rotations: Vec<Quat>,
27 chain_states: Vec<ChainLinkState>,
28}
29
30impl IkScratch {
31 fn new(model: &ModelArena) -> Self {
32 let max_links = model
33 .ik_solvers()
34 .iter()
35 .map(|s| s.links.len())
36 .max()
37 .unwrap_or(0);
38 IkScratch {
39 links: Vec::with_capacity(max_links),
40 base_rotations: Vec::with_capacity(max_links),
41 base_ik_rotations: Vec::with_capacity(max_links),
42 ik_rotations: Vec::with_capacity(max_links),
43 best_ik_rotations: Vec::with_capacity(max_links),
44 chain_states: Vec::with_capacity(max_links),
45 }
46 }
47}
48
49#[derive(Debug)]
50struct MorphScratch {
51 expanded_weights: Vec<f32>,
52}
53
54impl MorphScratch {
55 fn new(morph_count: usize) -> Self {
56 Self {
57 expanded_weights: vec![0.0; morph_count],
58 }
59 }
60}
61
62#[derive(Clone, Copy, Debug, Default, PartialEq)]
63pub struct IkSolverRuntimeStats {
64 pub solver_evaluations: u64,
65 pub configured_iterations: u64,
66 pub executed_iterations: u64,
67 pub tolerance_precheck_breaks: u64,
68 pub tolerance_post_iteration_breaks: u64,
69 pub rollback_breaks: u64,
70 pub max_iteration_exhaustions: u64,
71 pub link_visits: u64,
72 pub link_steps: u64,
73 pub final_distance_sum: f64,
74 pub final_distance_max: f32,
75 pub exhausted_final_distance_sum: f64,
76 pub exhausted_final_distance_max: f32,
77}
78
79impl IkSolverRuntimeStats {
80 fn reset(&mut self) {
81 *self = Self::default();
82 }
83}
84
85#[derive(Clone, Copy, Debug, PartialEq)]
86pub struct IkSolveOptions {
87 pub tolerance: f32,
88 pub max_iterations_cap: Option<u32>,
89}
90
91pub use physics::{PhysicsMode, PhysicsStepStats, PhysicsTickConfig};
92
93impl Default for IkSolveOptions {
94 fn default() -> Self {
95 Self {
96 tolerance: 0.0,
97 max_iterations_cap: None,
98 }
99 }
100}
101
102#[cfg(test)]
103#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
104pub(super) enum WorldMatrixBoneUpdateCategory {
105 LeadingBookend,
106 PhaseLoop,
107 TrailingBookend,
108 IkLinkChange,
109 #[default]
110 Other,
111}
112
113#[derive(Clone, Copy, Debug)]
114pub struct HostPoseView<'a> {
115 pub local_position_offsets: &'a [Vec3A],
116 pub local_rotations: &'a [Quat],
117 pub local_scales: &'a [Vec3A],
118 pub morph_weights: &'a [f32],
119 pub ik_enabled: &'a [u8],
120}
121
122#[derive(Clone, Debug, PartialEq, Eq)]
123pub enum HostPoseError {
124 BoneCountMismatch { expected: usize, got: usize },
125 MorphCountMismatch { expected: usize, got: usize },
126 IkCountMismatch { expected: usize, got: usize },
127 NonFiniteValue { field: &'static str, index: usize },
128 NonNormalizedQuaternion { index: usize },
129}
130
131impl std::fmt::Display for HostPoseError {
132 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
133 match self {
134 HostPoseError::BoneCountMismatch { expected, got } => {
135 write!(f, "bone count mismatch: expected {expected}, got {got}")
136 }
137 HostPoseError::MorphCountMismatch { expected, got } => {
138 write!(f, "morph count mismatch: expected {expected}, got {got}")
139 }
140 HostPoseError::IkCountMismatch { expected, got } => {
141 write!(f, "ik count mismatch: expected {expected}, got {got}")
142 }
143 HostPoseError::NonFiniteValue { field, index } => {
144 write!(f, "non-finite value in {field} at index {index}")
145 }
146 Self::NonNormalizedQuaternion { index } => {
147 write!(f, "non-normalized quaternion at local_rotations[{index}]")
148 }
149 }
150 }
151}
152
153impl std::error::Error for HostPoseError {}
154
155#[derive(Debug)]
156pub struct RuntimeInstance {
157 model: Arc<ModelArena>,
158 pose: PoseArena,
159 physics_mode: PhysicsMode,
160 physics_tick_config: PhysicsTickConfig,
161 physics_accumulator_seconds: f32,
162 ik_scratch: IkScratch,
163 morph_scratch: MorphScratch,
164 ik_stats: Vec<IkSolverRuntimeStats>,
165 ik_link_change_update_bones: Vec<Option<Vec<crate::BoneIndex>>>,
166 #[cfg(test)]
167 world_matrix_bone_update_count: usize,
168 #[cfg(test)]
169 world_matrix_bone_update_category: WorldMatrixBoneUpdateCategory,
170 #[cfg(test)]
171 world_matrix_bone_update_leading_bookend_count: usize,
172 #[cfg(test)]
173 world_matrix_bone_update_phase_loop_count: usize,
174 #[cfg(test)]
175 world_matrix_bone_update_trailing_bookend_count: usize,
176 #[cfg(test)]
177 world_matrix_bone_update_ik_link_change_count: usize,
178 #[cfg(test)]
179 world_matrix_bone_update_other_count: usize,
180}
181
182impl RuntimeInstance {
183 pub fn new(model: Arc<ModelArena>) -> Self {
184 let morph_count = model.morph_count() as usize;
185 Self::new_with_morph_count(model, morph_count)
186 }
187
188 pub fn new_with_morph_count(model: Arc<ModelArena>, morph_count: usize) -> Self {
189 let ik_count = model.ik_count();
190 Self::new_with_counts(model, morph_count, ik_count)
191 }
192
193 pub fn new_with_counts(model: Arc<ModelArena>, morph_count: usize, ik_count: usize) -> Self {
194 let morph_count = morph_count.max(model.morph_count() as usize);
195 let ik_count = ik_count.max(model.ik_count());
196 let pose = PoseArena::new_with_counts(model.bone_count(), morph_count, ik_count);
197 let ik_scratch = IkScratch::new(&model);
198 let morph_scratch = MorphScratch::new(morph_count);
199 let ik_stats = vec![IkSolverRuntimeStats::default(); model.ik_count()];
200 let ik_link_change_update_bones = vec![None; model.ik_count()];
201 Self {
202 model,
203 pose,
204 physics_mode: PhysicsMode::default(),
205 physics_tick_config: PhysicsTickConfig::default(),
206 physics_accumulator_seconds: 0.0,
207 ik_scratch,
208 morph_scratch,
209 ik_stats,
210 ik_link_change_update_bones,
211 #[cfg(test)]
212 world_matrix_bone_update_count: 0,
213 #[cfg(test)]
214 world_matrix_bone_update_category: WorldMatrixBoneUpdateCategory::default(),
215 #[cfg(test)]
216 world_matrix_bone_update_leading_bookend_count: 0,
217 #[cfg(test)]
218 world_matrix_bone_update_phase_loop_count: 0,
219 #[cfg(test)]
220 world_matrix_bone_update_trailing_bookend_count: 0,
221 #[cfg(test)]
222 world_matrix_bone_update_ik_link_change_count: 0,
223 #[cfg(test)]
224 world_matrix_bone_update_other_count: 0,
225 }
226 }
227
228 #[inline]
229 pub fn model(&self) -> &ModelArena {
230 &self.model
231 }
232
233 #[inline]
234 pub fn pose(&self) -> &PoseArena {
235 &self.pose
236 }
237
238 #[inline]
239 pub fn pose_mut(&mut self) -> &mut PoseArena {
240 &mut self.pose
241 }
242
243 pub fn evaluate_current_pose(&mut self) {
244 self.pose.reset_ik_rotations();
245 self.evaluate_current_pose_ordered(IkSolveOptions::default());
246 }
247
248 pub fn evaluate_current_pose_with_ik_options(&mut self, options: IkSolveOptions) {
249 self.pose.reset_ik_rotations();
250 self.evaluate_current_pose_ordered(options);
251 }
252
253 pub fn evaluate_current_pose_without_ik(&mut self) {
257 self.pose.reset_ik_rotations();
258 self.update_world_matrices();
259 }
260
261 fn evaluate_current_pose_ordered(&mut self, options: IkSolveOptions) {
262 self.begin_current_pose_evaluation();
263 let mut earliest_after_physics_eval_order_position = None;
264 self.evaluate_current_pose_phase(
265 false,
266 options,
267 &mut earliest_after_physics_eval_order_position,
268 );
269 self.evaluate_current_pose_phase(
270 true,
271 options,
272 &mut earliest_after_physics_eval_order_position,
273 );
274 self.finish_current_pose_evaluation(earliest_after_physics_eval_order_position);
275 }
276
277 pub fn evaluate_current_pose_before_physics(&mut self) {
278 self.evaluate_current_pose_before_physics_with_ik_options(IkSolveOptions::default());
279 }
280
281 pub fn evaluate_current_pose_before_physics_with_ik_options(
282 &mut self,
283 options: IkSolveOptions,
284 ) {
285 self.pose.reset_ik_rotations();
286 self.begin_current_pose_evaluation();
287 let mut earliest_after_physics_eval_order_position = None;
288 self.evaluate_current_pose_phase(
289 false,
290 options,
291 &mut earliest_after_physics_eval_order_position,
292 );
293 }
294
295 pub fn evaluate_current_pose_after_physics(&mut self) {
296 self.evaluate_current_pose_after_physics_with_ik_options(IkSolveOptions::default());
297 }
298
299 pub fn evaluate_current_pose_after_physics_with_ik_options(&mut self, options: IkSolveOptions) {
300 let mut earliest_after_physics_eval_order_position = None;
301 self.evaluate_current_pose_phase(
302 true,
303 options,
304 &mut earliest_after_physics_eval_order_position,
305 );
306 self.finish_current_pose_evaluation(earliest_after_physics_eval_order_position);
307 }
308
309 fn begin_current_pose_evaluation(&mut self) {
310 self.pose.reset_append_transforms();
311 #[cfg(test)]
312 self.set_world_matrix_bone_update_category(WorldMatrixBoneUpdateCategory::LeadingBookend);
313 self.update_world_matrices_using_current_append_from_eval_order_position(0);
314 }
315
316 fn evaluate_current_pose_phase(
317 &mut self,
318 after_physics: bool,
319 options: IkSolveOptions,
320 earliest_after_physics_eval_order_position: &mut Option<usize>,
321 ) {
322 let phase_bone_count = self.model.eval_order_for_phase(after_physics).len();
323 for phase_index in 0..phase_bone_count {
324 let bone = self.model.eval_order_for_phase(after_physics)[phase_index];
325 if after_physics {
326 let position = self.model.eval_order_position(bone);
327 *earliest_after_physics_eval_order_position = Some(
328 (*earliest_after_physics_eval_order_position)
329 .map_or(position, |earliest| earliest.min(position)),
330 );
331 }
332 if self.model.append_transform_index(bone).is_some() {
333 self.pose.reset_append_transform(bone);
334 self.update_append_transform_for_bone(bone);
335 }
336 #[cfg(test)]
337 self.set_world_matrix_bone_update_category(WorldMatrixBoneUpdateCategory::PhaseLoop);
338 self.update_world_matrix_for_bone(bone);
339
340 let ik_solver_count = self.model.ik_solver_count_for_bone(bone);
341 for local_index in 0..ik_solver_count {
342 let ik_index = self.model.ik_solver_index_for_bone(bone, local_index);
343 self.solve_ik_solver(ik_index, options, after_physics);
344 }
345 }
346 }
347
348 fn finish_current_pose_evaluation(
349 &mut self,
350 earliest_after_physics_eval_order_position: Option<usize>,
351 ) {
352 let mut trailing_refresh_start = earliest_after_physics_eval_order_position;
353 for append in self.model.append_transforms() {
354 let source_position = self.model.eval_order_position(append.source_bone);
355 let target_position = self.model.eval_order_position(append.target_bone);
356 if target_position < source_position {
357 trailing_refresh_start = Some(
358 trailing_refresh_start
359 .map_or(target_position, |start| start.min(target_position)),
360 );
361 }
362 }
363
364 if let Some(start_position) = trailing_refresh_start {
365 let start_position =
366 self.expand_update_start_for_append_dependencies(start_position, None);
367 #[cfg(test)]
368 self.set_world_matrix_bone_update_category(
369 WorldMatrixBoneUpdateCategory::TrailingBookend,
370 );
371 self.update_world_matrices_from_eval_order_position(start_position);
372 }
373 }
374
375 pub fn evaluate_rest_pose(&mut self) {
376 self.pose.reset_local_pose();
377 self.evaluate_current_pose();
378 }
379
380 pub fn apply_host_pose(&mut self, view: &HostPoseView) -> Result<(), HostPoseError> {
381 let bone_count = self.model.bone_count();
382 let morph_count = self.pose.morph_weights().len();
383 let ik_count = self.pose.ik_enabled().len();
384
385 if view.local_position_offsets.len() != bone_count {
386 return Err(HostPoseError::BoneCountMismatch {
387 expected: bone_count,
388 got: view.local_position_offsets.len(),
389 });
390 }
391 if view.local_rotations.len() != bone_count {
392 return Err(HostPoseError::BoneCountMismatch {
393 expected: bone_count,
394 got: view.local_rotations.len(),
395 });
396 }
397 if view.local_scales.len() != bone_count {
398 return Err(HostPoseError::BoneCountMismatch {
399 expected: bone_count,
400 got: view.local_scales.len(),
401 });
402 }
403 if view.morph_weights.len() != morph_count {
404 return Err(HostPoseError::MorphCountMismatch {
405 expected: morph_count,
406 got: view.morph_weights.len(),
407 });
408 }
409 if view.ik_enabled.len() != ik_count {
410 return Err(HostPoseError::IkCountMismatch {
411 expected: ik_count,
412 got: view.ik_enabled.len(),
413 });
414 }
415
416 for (i, v) in view.local_position_offsets.iter().enumerate() {
417 if !v.is_finite() {
418 return Err(HostPoseError::NonFiniteValue {
419 field: "local_position_offsets",
420 index: i,
421 });
422 }
423 }
424 for (i, q) in view.local_rotations.iter().enumerate() {
425 if !q.is_finite() {
426 return Err(HostPoseError::NonFiniteValue {
427 field: "local_rotations",
428 index: i,
429 });
430 }
431 if (q.length_squared() - 1.0).abs() > 1e-3 {
432 return Err(HostPoseError::NonNormalizedQuaternion { index: i });
433 }
434 }
435 for (i, v) in view.local_scales.iter().enumerate() {
436 if !v.is_finite() {
437 return Err(HostPoseError::NonFiniteValue {
438 field: "local_scales",
439 index: i,
440 });
441 }
442 }
443 for (i, w) in view.morph_weights.iter().enumerate() {
444 if !w.is_finite() {
445 return Err(HostPoseError::NonFiniteValue {
446 field: "morph_weights",
447 index: i,
448 });
449 }
450 }
451
452 self.pose
453 .set_local_position_offsets_from_slice(view.local_position_offsets);
454 self.pose
455 .set_local_rotations_from_slice(view.local_rotations);
456 self.pose.set_local_scales_from_slice(view.local_scales);
457 self.pose.set_morph_weights_from_slice(view.morph_weights);
458 self.pose.set_ik_enabled_from_slice(view.ik_enabled);
459
460 Ok(())
461 }
462
463 pub fn evaluate_clip_frame(&mut self, clip: &AnimationClip, frame: f32) {
464 clip.apply_to_pose(frame, &mut self.pose);
465 self.expand_morphs();
466 self.evaluate_current_pose();
467 }
468
469 pub fn evaluate_clip_frame_with_ik_options(
470 &mut self,
471 clip: &AnimationClip,
472 frame: f32,
473 options: IkSolveOptions,
474 ) {
475 clip.apply_to_pose(frame, &mut self.pose);
476 self.expand_morphs();
477 self.evaluate_current_pose_with_ik_options(options);
478 }
479
480 pub fn evaluate_clip_frame_before_physics(&mut self, clip: &AnimationClip, frame: f32) {
481 self.evaluate_clip_frame_before_physics_with_ik_options(
482 clip,
483 frame,
484 IkSolveOptions::default(),
485 );
486 }
487
488 pub fn evaluate_clip_frame_before_physics_with_ik_options(
489 &mut self,
490 clip: &AnimationClip,
491 frame: f32,
492 options: IkSolveOptions,
493 ) {
494 clip.apply_to_pose(frame, &mut self.pose);
495 self.expand_morphs();
496 self.evaluate_current_pose_before_physics_with_ik_options(options);
497 }
498
499 pub fn evaluate_clip_frame_without_ik(&mut self, clip: &AnimationClip, frame: f32) {
504 clip.apply_to_pose(frame, &mut self.pose);
505 self.expand_morphs();
506 self.pose.reset_ik_rotations();
507 self.update_world_matrices();
508 }
509
510 pub fn reset_ik_runtime_stats(&mut self) {
511 for stats in &mut self.ik_stats {
512 stats.reset();
513 }
514 }
515
516 pub fn ik_runtime_stats(&self) -> &[IkSolverRuntimeStats] {
517 &self.ik_stats
518 }
519
520 #[inline]
521 pub fn append_position_offset(&self, bone: crate::BoneIndex) -> glam::Vec3A {
522 self.pose.append_position_offset(bone)
523 }
524
525 #[inline]
526 pub fn append_rotation(&self, bone: crate::BoneIndex) -> glam::Quat {
527 self.pose.append_rotation(bone)
528 }
529
530 #[inline]
531 pub fn ik_enabled(&self) -> &[u8] {
532 self.pose.ik_enabled()
533 }
534}
535
536#[cfg(test)]
537mod tests;