1use crate::FrameIdString;
5use crate::frames::{FrameId, FramePair};
6use crate::velocity::VelocityTransform;
7use bincode::{Decode, Encode};
8use cu_spatial_payloads::Transform3D;
9#[allow(unused_imports)]
10use cu29::bevy_reflect;
11use cu29::clock::{CuTime, CuTimeRange, Tov};
12use cu29::cutask::CuStampedData;
13use cu29::prelude::{CuMsgPayload, Reflect};
14use cu29::units::si::f32::Angle as Angle32;
15use cu29::units::si::f32::Length as Length32;
16use cu29::units::si::f64::Angle as Angle64;
17use cu29::units::si::f64::Length as Length64;
18use num_traits;
19use serde::{Deserialize, Serialize};
20use std::fmt::Debug;
21
22pub type StampedFrameTransform<T> = CuStampedData<FrameTransform<T>, ()>;
24
25#[derive(Clone, Debug, Serialize, Deserialize, Default, Reflect)]
48#[reflect(opaque, from_reflect = false, no_field_bounds)]
49pub struct FrameTransform<T: Copy + Debug + Default + Serialize + 'static> {
50 pub transform: Transform3D<T>,
52 pub parent_frame: FrameIdString,
54 pub child_frame: FrameIdString,
56}
57
58impl<T: Copy + Debug + Default + Serialize + 'static> FrameTransform<T> {
59 pub fn new(
61 transform: Transform3D<T>,
62 parent_frame: impl AsRef<str>,
63 child_frame: impl AsRef<str>,
64 ) -> Self {
65 Self {
66 transform,
67 parent_frame: FrameIdString::from(parent_frame.as_ref())
68 .expect("Parent frame name too long (max 64 chars)"),
69 child_frame: FrameIdString::from(child_frame.as_ref())
70 .expect("Child frame name too long (max 64 chars)"),
71 }
72 }
73
74 pub fn from_stamped(stamped: &crate::transform::StampedTransform<T>) -> Self {
76 Self {
77 transform: stamped.transform,
78 parent_frame: FrameIdString::from(stamped.parent_frame.as_str())
79 .expect("Parent frame name too long"),
80 child_frame: FrameIdString::from(stamped.child_frame.as_str())
81 .expect("Child frame name too long"),
82 }
83 }
84}
85
86impl<T: Copy + Debug + Default + Serialize + 'static> Encode for FrameTransform<T>
88where
89 T: Encode,
90{
91 fn encode<E: bincode::enc::Encoder>(
92 &self,
93 encoder: &mut E,
94 ) -> Result<(), bincode::error::EncodeError> {
95 self.transform.encode(encoder)?;
96 self.parent_frame.encode(encoder)?;
97 self.child_frame.encode(encoder)?;
98 Ok(())
99 }
100}
101
102impl<T: Copy + Debug + Default + Serialize + 'static> Decode<()> for FrameTransform<T>
103where
104 T: Decode<()>,
105{
106 fn decode<D: bincode::de::Decoder<Context = ()>>(
107 decoder: &mut D,
108 ) -> Result<Self, bincode::error::DecodeError> {
109 let transform = <Transform3D<T> as Decode<()>>::decode(decoder)?;
110 let parent_frame_str = String::decode(decoder)?;
111 let child_frame_str = String::decode(decoder)?;
112 let parent_frame = FrameIdString::from(&parent_frame_str).map_err(|_| {
113 bincode::error::DecodeError::OtherString("Parent frame name too long".to_string())
114 })?;
115 let child_frame = FrameIdString::from(&child_frame_str).map_err(|_| {
116 bincode::error::DecodeError::OtherString("Child frame name too long".to_string())
117 })?;
118 Ok(Self {
119 transform,
120 parent_frame,
121 child_frame,
122 })
123 }
124}
125
126pub type TypedStampedFrameTransform<T> = CuStampedData<Transform3D<T>, ()>;
128
129#[derive(Debug, Clone, Copy)]
153pub struct TypedTransform3D<T, Parent, Child>
154where
155 T: Copy + Debug + 'static,
156 Parent: FrameId,
157 Child: FrameId,
158{
159 transform: Transform3D<T>,
160 pub frames: FramePair<Parent, Child>,
162}
163
164impl<T, Parent, Child> TypedTransform3D<T, Parent, Child>
165where
166 T: Copy + Debug + 'static,
167 Parent: FrameId,
168 Child: FrameId,
169{
170 pub const fn new(transform: Transform3D<T>) -> Self {
172 Self {
173 transform,
174 frames: FramePair::new(),
175 }
176 }
177
178 pub const fn transform(&self) -> &Transform3D<T> {
180 &self.transform
181 }
182
183 pub fn at(self, time: CuTime) -> TypedTransform<T, Parent, Child>
185 where
186 T: CuMsgPayload,
187 {
188 TypedTransform::new(self.transform, time)
189 }
190
191 pub const fn parent_id(&self) -> u32 {
193 Parent::ID
194 }
195
196 pub const fn child_id(&self) -> u32 {
198 Child::ID
199 }
200
201 pub const fn parent_name(&self) -> &'static str {
203 Parent::NAME
204 }
205
206 pub const fn child_name(&self) -> &'static str {
208 Child::NAME
209 }
210}
211
212macro_rules! impl_const_typed_transform {
213 ($ty:ty, $len:ty, $ang:ty) => {
214 impl<Parent, Child> TypedTransform3D<$ty, Parent, Child>
215 where
216 Parent: FrameId,
217 Child: FrameId,
218 {
219 pub const fn from_translation_euler_xyz(
223 translation: [$len; 3],
224 rotation: [$ang; 3],
225 ) -> Self {
226 Self::new(Transform3D::<$ty>::from_translation_euler_xyz(
227 translation,
228 rotation,
229 ))
230 }
231
232 pub const fn then<Next>(
253 self,
254 next: TypedTransform3D<$ty, Child, Next>,
255 ) -> TypedTransform3D<$ty, Parent, Next>
256 where
257 Next: FrameId,
258 {
259 TypedTransform3D::new(self.transform.compose(next.transform))
260 }
261 }
262 };
263}
264
265impl_const_typed_transform!(f32, Length32, Angle32);
266impl_const_typed_transform!(f64, Length64, Angle64);
267
268#[derive(Debug, Clone)]
270pub struct TypedTransform<T, Parent, Child>
271where
272 T: CuMsgPayload + Copy + Debug + 'static,
273 Parent: FrameId,
274 Child: FrameId,
275{
276 pub transform: TypedStampedFrameTransform<T>,
278 pub frames: FramePair<Parent, Child>,
280}
281
282impl<T, Parent, Child> TypedTransform<T, Parent, Child>
283where
284 T: CuMsgPayload + Copy + Debug + 'static,
285 Parent: FrameId,
286 Child: FrameId,
287{
288 pub fn new(transform: Transform3D<T>, time: CuTime) -> Self {
290 let mut transform = TypedStampedFrameTransform::new(Some(transform));
291 transform.tov = Tov::Time(time);
292
293 let frames = FramePair::new();
294
295 Self { transform, frames }
296 }
297
298 pub fn transform(&self) -> Option<&Transform3D<T>> {
300 self.transform.payload()
301 }
302
303 pub fn timestamp(&self) -> Option<CuTime> {
305 match self.transform.tov {
306 Tov::Time(time) => Some(time),
307 _ => None,
308 }
309 }
310
311 pub fn parent_id(&self) -> u32 {
313 Parent::ID
314 }
315
316 pub fn child_id(&self) -> u32 {
318 Child::ID
319 }
320
321 pub fn parent_name(&self) -> &'static str {
323 Parent::NAME
324 }
325
326 pub fn child_name(&self) -> &'static str {
328 Child::NAME
329 }
330}
331
332#[derive(Debug)]
335pub struct TypedTransformBuffer<T, Parent, Child, const N: usize>
336where
337 T: CuMsgPayload + Copy + Debug + 'static,
338 Parent: FrameId,
339 Child: FrameId,
340{
341 transforms: [Option<TypedTransform<T, Parent, Child>>; N],
343 count: usize,
345}
346
347impl<T, Parent, Child, const N: usize> TypedTransformBuffer<T, Parent, Child, N>
348where
349 T: CuMsgPayload + Copy + Debug + 'static,
350 Parent: FrameId,
351 Child: FrameId,
352{
353 pub fn new() -> Self {
355 Self {
356 transforms: std::array::from_fn(|_| None),
357 count: 0,
358 }
359 }
360
361 pub fn add_transform(&mut self, transform_msg: TypedTransform<T, Parent, Child>) {
363 if self.count < N {
364 self.transforms[self.count] = Some(transform_msg);
366 self.count += 1;
367 } else {
368 for i in 0..N - 1 {
370 self.transforms[i] = self.transforms[i + 1].take();
371 }
372 self.transforms[N - 1] = Some(transform_msg);
373 }
374
375 self.sort_by_time();
377 }
378
379 fn transform_at(&self, index: usize) -> Option<&TypedTransform<T, Parent, Child>> {
380 self.transforms.get(index)?.as_ref()
381 }
382
383 fn timed_indices(&self) -> Vec<(usize, CuTime)> {
384 (0..self.count)
385 .filter_map(|index| {
386 let transform = self.transform_at(index)?;
387 Some((index, transform.timestamp()?))
388 })
389 .collect()
390 }
391
392 #[allow(clippy::type_complexity)]
393 fn transform_pair(
394 &self,
395 first: usize,
396 second: usize,
397 ) -> Option<(
398 &TypedTransform<T, Parent, Child>,
399 &TypedTransform<T, Parent, Child>,
400 )> {
401 Some((self.transform_at(first)?, self.transform_at(second)?))
402 }
403
404 fn sort_by_time(&mut self) {
406 let mut time_indices = self.timed_indices();
407
408 time_indices.sort_by_key(|(_, time)| *time);
410
411 let mut new_transforms: [Option<TypedTransform<T, Parent, Child>>; N] =
413 std::array::from_fn(|_| None);
414
415 for (new_idx, (old_idx, _)) in time_indices.into_iter().enumerate() {
416 new_transforms[new_idx] = self.transforms[old_idx].take();
417 }
418
419 self.transforms = new_transforms;
420 }
421
422 pub fn get_latest_transform(&self) -> Option<&TypedTransform<T, Parent, Child>> {
424 self.count
425 .checked_sub(1)
426 .and_then(|index| self.transform_at(index))
427 }
428
429 pub fn get_closest_transform(&self, time: CuTime) -> Option<&TypedTransform<T, Parent, Child>> {
431 if self.count == 0 {
432 return None;
433 }
434
435 let closest_idx = self
436 .timed_indices()
437 .into_iter()
438 .min_by_key(|(_, transform_time)| time.as_nanos().abs_diff(transform_time.as_nanos()))
439 .map(|(index, _)| index)
440 .unwrap_or(0);
441
442 self.transform_at(closest_idx)
443 }
444
445 pub fn get_time_range(&self) -> Option<CuTimeRange> {
447 if self.count == 0 {
448 return None;
449 }
450
451 let end_index = self.count.checked_sub(1)?;
453 let start = self.transform_at(0)?.timestamp()?;
454 let end = self.transform_at(end_index)?.timestamp()?;
455
456 Some(CuTimeRange { start, end })
457 }
458
459 #[allow(clippy::type_complexity)]
461 pub fn get_transforms_around(
462 &self,
463 time: CuTime,
464 ) -> Option<(
465 &TypedTransform<T, Parent, Child>,
466 &TypedTransform<T, Parent, Child>,
467 )> {
468 if self.count < 2 {
469 return None;
470 }
471
472 let mut before_idx = None;
474 let mut after_idx = None;
475
476 for i in 0..self.count {
477 let Some(transform) = self.transform_at(i) else {
478 continue;
479 };
480 let Some(transform_time) = transform.timestamp() else {
481 continue;
482 };
483
484 if transform_time <= time {
485 before_idx = Some(i);
486 } else if after_idx.is_none() {
487 after_idx = Some(i);
488 break;
489 }
490 }
491
492 match (before_idx, after_idx) {
493 (Some(before), Some(after)) => self.transform_pair(before, after),
494 (Some(before), None) if before > 0 => self.transform_pair(before - 1, before),
495 (None, Some(after)) if after + 1 < self.count => self.transform_pair(after, after + 1),
496 _ => None,
497 }
498 }
499}
500
501impl<T, Parent, Child, const N: usize> Default for TypedTransformBuffer<T, Parent, Child, N>
502where
503 T: CuMsgPayload + Copy + Debug + 'static,
504 Parent: FrameId,
505 Child: FrameId,
506{
507 fn default() -> Self {
508 Self::new()
509 }
510}
511
512impl<T, Parent, Child> TypedTransform<T, Parent, Child>
514where
515 T: CuMsgPayload
516 + Copy
517 + Debug
518 + Default
519 + std::ops::Add<Output = T>
520 + std::ops::Sub<Output = T>
521 + std::ops::Mul<Output = T>
522 + std::ops::Div<Output = T>
523 + num_traits::NumCast
524 + 'static,
525 Parent: FrameId,
526 Child: FrameId,
527{
528 pub fn compute_velocity(&self, previous: &Self) -> Option<VelocityTransform<T>> {
530 let current_time = self.timestamp()?;
531 let previous_time = previous.timestamp()?;
532 let current_transform = self.transform()?;
533 let previous_transform = previous.transform()?;
534
535 let dt_nanos = current_time.as_nanos() as i64 - previous_time.as_nanos() as i64;
537 if dt_nanos <= 0 {
538 return None;
539 }
540
541 let dt = dt_nanos as f64 / 1_000_000_000.0;
543
544 let dt_t = num_traits::cast::cast::<f64, T>(dt)?;
545
546 let current_mat = current_transform.to_matrix();
548 let previous_mat = previous_transform.to_matrix();
549 let mut linear_velocity = [T::default(); 3];
550 for (i, vel) in linear_velocity.iter_mut().enumerate() {
551 let pos_diff = current_mat[i][3] - previous_mat[i][3];
552 *vel = pos_diff / dt_t;
553 }
554
555 let angular_velocity = [T::default(); 3];
557
558 Some(VelocityTransform {
559 linear: linear_velocity,
560 angular: angular_velocity,
561 })
562 }
563}
564
565#[cfg(test)]
566mod tests {
567 use super::*;
568 use crate::frames::{BaseFrame, RobotFrame, WorldFrame};
569
570 const WORLD_TO_BASE: TypedTransform3D<f32, WorldFrame, BaseFrame> =
571 TypedTransform3D::<f32, WorldFrame, BaseFrame>::from_translation_euler_xyz(
572 [
573 Length32 { value: 1.0 },
574 Length32 { value: 2.0 },
575 Length32 { value: 0.0 },
576 ],
577 [
578 Angle32 { value: 0.0 },
579 Angle32 { value: 0.0 },
580 Angle32 {
581 value: core::f32::consts::FRAC_PI_2,
582 },
583 ],
584 );
585 const BASE_TO_ROBOT: TypedTransform3D<f32, BaseFrame, RobotFrame> =
586 TypedTransform3D::<f32, BaseFrame, RobotFrame>::from_translation_euler_xyz(
587 [
588 Length32 { value: 1.0 },
589 Length32 { value: 0.0 },
590 Length32 { value: 0.0 },
591 ],
592 [Angle32 { value: 0.0 }; 3],
593 );
594 const WORLD_TO_ROBOT: TypedTransform3D<f32, WorldFrame, RobotFrame> =
595 WORLD_TO_BASE.then(BASE_TO_ROBOT);
596 fn assert_approx_eq(actual: f32, expected: f32, epsilon: f32) {
598 let diff = (actual - expected).abs();
599 assert!(
600 diff <= epsilon,
601 "expected {expected}, got {actual}, difference {diff} exceeds epsilon {epsilon}",
602 );
603 }
604 use cu29::clock::CuTime;
605
606 type WorldToRobotFrameTransform = TypedTransform<f32, WorldFrame, RobotFrame>;
607 type WorldToRobotBuffer = TypedTransformBuffer<f32, WorldFrame, RobotFrame, 10>;
608
609 #[test]
610 fn test_const_typed_transform_composition_and_stamping() {
611 assert_eq!(
612 std::mem::size_of::<TypedTransform3D<f32, WorldFrame, RobotFrame>>(),
613 std::mem::size_of::<Transform3D<f32>>(),
614 );
615
616 let position = WORLD_TO_ROBOT.transform().position();
617 assert_approx_eq(position.x.raw(), 1.0, 1e-5);
618 assert_approx_eq(position.y.raw(), 3.0, 1e-5);
619 assert_approx_eq(position.z.raw(), 0.0, 1e-5);
620 assert_eq!(WORLD_TO_ROBOT.parent_id(), WorldFrame::ID);
621 assert_eq!(WORLD_TO_ROBOT.child_id(), RobotFrame::ID);
622
623 let stamped = WORLD_TO_ROBOT.at(CuTime::from_nanos(42));
624 assert_eq!(stamped.timestamp(), Some(CuTime::from_nanos(42)));
625 assert_approx_eq(stamped.transform().unwrap().position().y.raw(), 3.0, 1e-5);
626 }
627
628 #[test]
629 fn test_typed_transform_msg_creation() {
630 let transform = Transform3D::<f32>::default();
631 let time = CuTime::from_nanos(1000);
632
633 let msg = WorldToRobotFrameTransform::new(transform, time);
634
635 assert_eq!(msg.parent_id(), WorldFrame::ID);
636 assert_eq!(msg.child_id(), RobotFrame::ID);
637 assert_eq!(msg.parent_name(), "world");
638 assert_eq!(msg.child_name(), "robot");
639 assert_eq!(msg.timestamp().unwrap().as_nanos(), 1000);
640 }
641
642 #[test]
643 fn test_typed_transform_buffer() {
644 let mut buffer = WorldToRobotBuffer::new();
645
646 let transform1 = Transform3D::<f32>::default();
647 let msg1 = WorldToRobotFrameTransform::new(transform1, CuTime::from_nanos(1000));
648
649 let transform2 = Transform3D::<f32>::default();
650 let msg2 = WorldToRobotFrameTransform::new(transform2, CuTime::from_nanos(2000));
651
652 buffer.add_transform(msg1);
653 buffer.add_transform(msg2);
654
655 let latest = buffer.get_latest_transform().unwrap();
656 assert_eq!(latest.timestamp().unwrap().as_nanos(), 2000);
657
658 let range = buffer.get_time_range().unwrap();
659 assert_eq!(range.start.as_nanos(), 1000);
660 assert_eq!(range.end.as_nanos(), 2000);
661 }
662
663 #[test]
664 fn test_closest_transform() {
665 let mut buffer = WorldToRobotBuffer::new();
666
667 let transform1 = Transform3D::<f32>::default();
668 let msg1 = WorldToRobotFrameTransform::new(transform1, CuTime::from_nanos(1000));
669
670 let transform2 = Transform3D::<f32>::default();
671 let msg2 = WorldToRobotFrameTransform::new(transform2, CuTime::from_nanos(3000));
672
673 buffer.add_transform(msg1);
674 buffer.add_transform(msg2);
675
676 let closest = buffer.get_closest_transform(CuTime::from_nanos(1500));
677 assert_eq!(closest.unwrap().timestamp().unwrap().as_nanos(), 1000);
678
679 let closest = buffer.get_closest_transform(CuTime::from_nanos(2500));
680 assert_eq!(closest.unwrap().timestamp().unwrap().as_nanos(), 3000);
681 }
682
683 #[test]
684 fn test_velocity_computation() {
685 use crate::test_utils::translation_transform;
686
687 let transform1 = translation_transform(0.0f32, 0.0, 0.0);
688 let transform2 = translation_transform(1.0f32, 2.0, 0.0);
689
690 let msg1 = WorldToRobotFrameTransform::new(transform1, CuTime::from_nanos(1_000_000_000)); let msg2 = WorldToRobotFrameTransform::new(transform2, CuTime::from_nanos(2_000_000_000)); let velocity = msg2.compute_velocity(&msg1).unwrap();
694
695 assert_approx_eq(velocity.linear[0], 1.0, 1e-5);
696 assert_approx_eq(velocity.linear[1], 2.0, 1e-5);
697 assert_approx_eq(velocity.linear[2], 0.0, 1e-5);
698 }
699}