deke_types/fk/
mod.rs

1use std::ops::Mul;
2
3use const_soft_float::soft_f32::SoftF32;
4use const_soft_float::soft_f64::SoftF64;
5use glam::{Affine3A, DAffine3, DMat3, DVec3, Mat3A, Vec3A};
6use glam_traits_ext::{FloatAffine, FloatMat, FloatScalar, FloatVec, TAffine3, TMat3, TVec3};
7
8use crate::{DekeError, SRobotQ};
9
10mod dh;
11mod dynamic;
12mod fp_dispatch;
13mod hp;
14mod prismatic;
15mod transformed;
16mod urdf;
17
18pub use dh::{DHChain, DHJoint};
19pub use dynamic::{BoxFK, DynamicDHChain, DynamicHPChain, DynamicURDFChain};
20pub use fp_dispatch::FPDispatch;
21pub use hp::{HPChain, HPJoint};
22pub use prismatic::PrismaticFK;
23pub use transformed::TransformedFK;
24pub use urdf::{
25    URDFBuildError, URDFChain, URDFJoint, URDFJointType, compose_fixed_joints,
26    compose_fixed_joints_f64,
27};
28
29/// Const-context sine/cosine via soft-float. Use only inside `const fn`
30/// builders where the runtime intrinsic is unavailable; hot paths must call
31/// `f32::sin_cos` directly.
32#[inline(always)]
33const fn const_sin_cos(x: f32) -> (f32, f32) {
34    let sf = SoftF32::from_f32(x);
35    (sf.sin().to_f32(), sf.cos().to_f32())
36}
37
38/// `f64` analogue of [`const_sin_cos`].
39#[inline(always)]
40const fn const_sin_cos_f64(x: f64) -> (f64, f64) {
41    let sf = SoftF64::from_f64(x);
42    (sf.sin().to_f64(), sf.cos().to_f64())
43}
44
45mod sealed {
46    pub trait Sealed {}
47    impl Sealed for f32 {}
48    impl Sealed for f64 {}
49}
50
51/// Scalar bound for FK code: `FloatScalar` plus the SIMD-aligned vec/mat/affine
52/// types the FK code operates on.
53///
54/// For `f32`, the aligned types are `Vec3A`/`Mat3A`/`Affine3A` (16-byte SIMD).
55/// For `f64`, they are `DVec3`/`DMat3`/`DAffine3` (already efficient packing).
56/// Both share a uniform interface via the `T*` traits in `glam-traits-ext`.
57pub trait FKScalar: FloatScalar + Copy + std::fmt::Debug + Send + Sync + 'static + sealed::Sealed {
58    type AVec3: TVec3<Self, MaybeAligned = Self::AVec3>;
59    type AMat3: TMat3<Self, MaybeAligned = Self::AMat3>
60        + FloatMat<Self, Col = Self::AVec3>
61        + Mul<Self::AVec3, Output = Self::AVec3>;
62    type AAffine3: TAffine3<Self, MaybeAligned = Self::AAffine3>
63        + FloatAffine<Self, Vec = Self::AVec3, Mat = Self::AMat3>
64        + Mul<Self::AAffine3, Output = Self::AAffine3>;
65}
66
67impl FKScalar for f32 {
68    type AVec3 = Vec3A;
69    type AMat3 = Mat3A;
70    type AAffine3 = Affine3A;
71}
72
73impl FKScalar for f64 {
74    type AVec3 = DVec3;
75    type AMat3 = DMat3;
76    type AAffine3 = DAffine3;
77}
78
79#[allow(type_alias_bounds)]
80pub(crate) type AAffine3<F: FKScalar> = F::AAffine3;
81#[allow(type_alias_bounds)]
82pub(crate) type AMat3<F: FKScalar> = F::AMat3;
83#[allow(type_alias_bounds)]
84pub(crate) type AVec3<F: FKScalar> = F::AVec3;
85
86pub trait FKChain<const N: usize, F: FKScalar = f32>: Clone + Send + Sync {
87    type Error: Into<DekeError>;
88
89    fn dof(&self) -> usize {
90        N
91    }
92    /// Configuration-independent transform from the robot's base frame to the
93    /// world frame. Defaults to identity; wrappers that install a static
94    /// prefix (e.g. [`TransformedFK`] with a prefix set, or [`URDFChain`]
95    /// with fixed leading joints baked in) override this so downstream
96    /// consumers (collision validators, visualizers) can place the static
97    /// base body at the correct pose.
98    fn base_tf(&self) -> AAffine3<F> {
99        AAffine3::<F>::IDENTITY
100    }
101    /// Theoretical maximum reach: sum of link lengths (upper bound, ignores joint limits).
102    fn max_reach(&self) -> Result<F, Self::Error> {
103        let (_, p, p_ee) = self.joint_axes_positions(&SRobotQ::zeros())?;
104        let mut total = F::zero();
105        let mut prev = p[0];
106        for i in 1..N {
107            total = total + (p[i] - prev).length();
108            prev = p[i];
109        }
110        total = total + (p_ee - prev).length();
111        Ok(total)
112    }
113
114    fn fk(&self, q: &SRobotQ<N, F>) -> Result<[AAffine3<F>; N], Self::Error>;
115    fn fk_end(&self, q: &SRobotQ<N, F>) -> Result<AAffine3<F>, Self::Error>;
116
117    /// Compute base transform, per-link frames, and the end-effector frame
118    /// in one call. Consumers that need all three (collision validators,
119    /// visualizers) should prefer this over invoking [`base_tf`], [`fk`],
120    /// and [`fk_end`] separately.
121    ///
122    /// The default implementation simply chains the three calls; chains
123    /// where the work overlaps (e.g. `fk_end` repeats the joint
124    /// accumulation done by `fk`, or wrappers like
125    /// [`TransformedFK`]/[`PrismaticFK`] that compose an inner chain)
126    /// override this to share the inner computation.
127    ///
128    /// [`base_tf`]: FKChain::base_tf
129    /// [`fk`]: FKChain::fk
130    /// [`fk_end`]: FKChain::fk_end
131    fn all_fk(
132        &self,
133        q: &SRobotQ<N, F>,
134    ) -> Result<(AAffine3<F>, [AAffine3<F>; N], AAffine3<F>), Self::Error> {
135        let base = self.base_tf();
136        let frames = self.fk(q)?;
137        let end = self.fk_end(q)?;
138        Ok((base, frames, end))
139    }
140    /// Returns joint rotation axes and axis-origin positions in world frame at
141    /// configuration `q`, plus the end-effector position.
142    fn joint_axes_positions(
143        &self,
144        q: &SRobotQ<N, F>,
145    ) -> Result<([AVec3<F>; N], [AVec3<F>; N], AVec3<F>), Self::Error>;
146
147    /// Geometric Jacobian (6×N) at configuration `q`.
148    /// Rows 0–2: linear velocity, rows 3–5: angular velocity.
149    fn jacobian(&self, q: &SRobotQ<N, F>) -> Result<[[F; N]; 6], Self::Error> {
150        let (z, p, p_ee) = self.joint_axes_positions(q)?;
151        let mut j = [[F::zero(); N]; 6];
152        for i in 0..N {
153            let dp = p_ee - p[i];
154            let c = z[i].cross(dp);
155            j[0][i] = c.x();
156            j[1][i] = c.y();
157            j[2][i] = c.z();
158            j[3][i] = z[i].x();
159            j[4][i] = z[i].y();
160            j[5][i] = z[i].z();
161        }
162        Ok(j)
163    }
164
165    /// First time-derivative of the geometric Jacobian.
166    fn jacobian_dot(
167        &self,
168        q: &SRobotQ<N, F>,
169        qdot: &SRobotQ<N, F>,
170    ) -> Result<[[F; N]; 6], Self::Error> {
171        let (z, p, p_ee) = self.joint_axes_positions(q)?;
172
173        let mut omega = AVec3::<F>::ZERO;
174        let mut z_dot = [AVec3::<F>::ZERO; N];
175        let mut p_dot = [AVec3::<F>::ZERO; N];
176        let mut pdot_acc = AVec3::<F>::ZERO;
177
178        for i in 0..N {
179            p_dot[i] = pdot_acc;
180            z_dot[i] = omega.cross(z[i]);
181            omega += z[i] * qdot.0[i];
182            let next_p = if i + 1 < N { p[i + 1] } else { p_ee };
183            pdot_acc += omega.cross(next_p - p[i]);
184        }
185        let p_ee_dot = pdot_acc;
186
187        let mut jd = [[F::zero(); N]; 6];
188        for i in 0..N {
189            let dp = p_ee - p[i];
190            let dp_dot = p_ee_dot - p_dot[i];
191            let c1 = z_dot[i].cross(dp);
192            let c2 = z[i].cross(dp_dot);
193            jd[0][i] = c1.x() + c2.x();
194            jd[1][i] = c1.y() + c2.y();
195            jd[2][i] = c1.z() + c2.z();
196            jd[3][i] = z_dot[i].x();
197            jd[4][i] = z_dot[i].y();
198            jd[5][i] = z_dot[i].z();
199        }
200        Ok(jd)
201    }
202
203    /// Second time-derivative of the geometric Jacobian.
204    fn jacobian_ddot(
205        &self,
206        q: &SRobotQ<N, F>,
207        qdot: &SRobotQ<N, F>,
208        qddot: &SRobotQ<N, F>,
209    ) -> Result<[[F; N]; 6], Self::Error> {
210        let (z, p, p_ee) = self.joint_axes_positions(q)?;
211
212        let mut omega = AVec3::<F>::ZERO;
213        let mut omega_dot = AVec3::<F>::ZERO;
214        let mut z_dot = [AVec3::<F>::ZERO; N];
215        let mut z_ddot = [AVec3::<F>::ZERO; N];
216        let mut p_dot = [AVec3::<F>::ZERO; N];
217        let mut p_ddot = [AVec3::<F>::ZERO; N];
218        let mut pdot_acc = AVec3::<F>::ZERO;
219        let mut pddot_acc = AVec3::<F>::ZERO;
220
221        for i in 0..N {
222            p_dot[i] = pdot_acc;
223            p_ddot[i] = pddot_acc;
224            let zd = omega.cross(z[i]);
225            z_dot[i] = zd;
226            z_ddot[i] = omega_dot.cross(z[i]) + omega.cross(zd);
227            omega_dot += z[i] * qddot.0[i] + zd * qdot.0[i];
228            omega += z[i] * qdot.0[i];
229            let next_p = if i + 1 < N { p[i + 1] } else { p_ee };
230            let delta = next_p - p[i];
231            let delta_dot = omega.cross(delta);
232            pdot_acc += delta_dot;
233            pddot_acc += omega_dot.cross(delta) + omega.cross(delta_dot);
234        }
235        let p_ee_dot = pdot_acc;
236        let p_ee_ddot = pddot_acc;
237
238        let mut jdd = [[F::zero(); N]; 6];
239        for i in 0..N {
240            let dp = p_ee - p[i];
241            let dp_dot = p_ee_dot - p_dot[i];
242            let dp_ddot = p_ee_ddot - p_ddot[i];
243            let c1 = z_ddot[i].cross(dp);
244            let c2 = z_dot[i].cross(dp_dot);
245            let c3 = z[i].cross(dp_ddot);
246            let two = F::one() + F::one();
247            jdd[0][i] = c1.x() + two * c2.x() + c3.x();
248            jdd[1][i] = c1.y() + two * c2.y() + c3.y();
249            jdd[2][i] = c1.z() + two * c2.z() + c3.z();
250            jdd[3][i] = z_ddot[i].x();
251            jdd[4][i] = z_ddot[i].y();
252            jdd[5][i] = z_ddot[i].z();
253        }
254        Ok(jdd)
255    }
256}
257
258#[inline(always)]
259#[cfg(debug_assertions)]
260fn check_finite<const N: usize, F: FloatScalar>(q: &SRobotQ<N, F>) -> Result<(), DekeError> {
261    if q.any_non_finite() {
262        return Err(DekeError::JointsNonFinite);
263    }
264    Ok(())
265}
266
267#[inline(always)]
268#[cfg(not(debug_assertions))]
269fn check_finite<const N: usize, F: FloatScalar>(_: &SRobotQ<N, F>) -> Result<(), std::convert::Infallible> {
270    Ok(())
271}
272
273#[inline(always)]
274const fn abs_f32(x: f32) -> f32 {
275    if x < 0.0 { -x } else { x }
276}
277
278/// Const-friendly affine transform backed by plain f32 arrays. `glam`'s
279/// `Vec3A`/`Mat3A` types use SIMD and expose components via a non-const
280/// `Deref`, so `const fn` code that needs per-component arithmetic (compose,
281/// identity check) routes through this type and only converts to
282/// `Affine3A` at the boundary.
283#[derive(Debug, Clone, Copy)]
284struct AffineRaw {
285    c0: [f32; 3],
286    c1: [f32; 3],
287    c2: [f32; 3],
288    t: [f32; 3],
289}
290
291impl AffineRaw {
292    const IDENTITY: Self = Self {
293        c0: [1.0, 0.0, 0.0],
294        c1: [0.0, 1.0, 0.0],
295        c2: [0.0, 0.0, 1.0],
296        t: [0.0, 0.0, 0.0],
297    };
298
299    /// `self * other` — applies `other` first, then `self`.
300    #[inline(always)]
301    const fn mul(self, other: Self) -> Self {
302        let nc0 = [
303            self.c0[0] * other.c0[0] + self.c1[0] * other.c0[1] + self.c2[0] * other.c0[2],
304            self.c0[1] * other.c0[0] + self.c1[1] * other.c0[1] + self.c2[1] * other.c0[2],
305            self.c0[2] * other.c0[0] + self.c1[2] * other.c0[1] + self.c2[2] * other.c0[2],
306        ];
307        let nc1 = [
308            self.c0[0] * other.c1[0] + self.c1[0] * other.c1[1] + self.c2[0] * other.c1[2],
309            self.c0[1] * other.c1[0] + self.c1[1] * other.c1[1] + self.c2[1] * other.c1[2],
310            self.c0[2] * other.c1[0] + self.c1[2] * other.c1[1] + self.c2[2] * other.c1[2],
311        ];
312        let nc2 = [
313            self.c0[0] * other.c2[0] + self.c1[0] * other.c2[1] + self.c2[0] * other.c2[2],
314            self.c0[1] * other.c2[0] + self.c1[1] * other.c2[1] + self.c2[1] * other.c2[2],
315            self.c0[2] * other.c2[0] + self.c1[2] * other.c2[1] + self.c2[2] * other.c2[2],
316        ];
317        let nt = [
318            self.c0[0] * other.t[0]
319                + self.c1[0] * other.t[1]
320                + self.c2[0] * other.t[2]
321                + self.t[0],
322            self.c0[1] * other.t[0]
323                + self.c1[1] * other.t[1]
324                + self.c2[1] * other.t[2]
325                + self.t[1],
326            self.c0[2] * other.t[0]
327                + self.c1[2] * other.t[1]
328                + self.c2[2] * other.t[2]
329                + self.t[2],
330        ];
331        Self {
332            c0: nc0,
333            c1: nc1,
334            c2: nc2,
335            t: nt,
336        }
337    }
338
339    #[inline(always)]
340    const fn is_identity(&self) -> bool {
341        const EPS: f32 = 1e-6;
342        abs_f32(self.c0[0] - 1.0) <= EPS
343            && abs_f32(self.c0[1]) <= EPS
344            && abs_f32(self.c0[2]) <= EPS
345            && abs_f32(self.c1[0]) <= EPS
346            && abs_f32(self.c1[1] - 1.0) <= EPS
347            && abs_f32(self.c1[2]) <= EPS
348            && abs_f32(self.c2[0]) <= EPS
349            && abs_f32(self.c2[1]) <= EPS
350            && abs_f32(self.c2[2] - 1.0) <= EPS
351            && abs_f32(self.t[0]) <= EPS
352            && abs_f32(self.t[1]) <= EPS
353            && abs_f32(self.t[2]) <= EPS
354    }
355
356    /// Build the URDF RPY-convention rotation (`Rz(yaw)·Ry(pitch)·Rx(roll)`)
357    /// and translate by `xyz`, using [`const_sin_cos`] for const evaluation.
358    #[inline(always)]
359    const fn from_xyz_rpy(xyz: (f64, f64, f64), rpy: (f64, f64, f64)) -> Self {
360        let (ox, oy, oz) = xyz;
361        let (roll, pitch, yaw) = rpy;
362        let (sr, cr) = const_sin_cos(roll as f32);
363        let (sp, cp) = const_sin_cos(pitch as f32);
364        let (sy, cy) = const_sin_cos(yaw as f32);
365        Self {
366            c0: [cy * cp, sy * cp, -sp],
367            c1: [cy * sp * sr - sy * cr, sy * sp * sr + cy * cr, cp * sr],
368            c2: [cy * sp * cr + sy * sr, sy * sp * cr - cy * sr, cp * cr],
369            t: [ox as f32, oy as f32, oz as f32],
370        }
371    }
372
373    #[inline(always)]
374    const fn to_affine3a(self) -> Affine3A {
375        Affine3A {
376            matrix3: Mat3A::from_cols(
377                Vec3A::new(self.c0[0], self.c0[1], self.c0[2]),
378                Vec3A::new(self.c1[0], self.c1[1], self.c1[2]),
379                Vec3A::new(self.c2[0], self.c2[1], self.c2[2]),
380            ),
381            translation: Vec3A::new(self.t[0], self.t[1], self.t[2]),
382        }
383    }
384
385    #[inline(always)]
386    const fn c0_vec3a(&self) -> Vec3A {
387        Vec3A::new(self.c0[0], self.c0[1], self.c0[2])
388    }
389
390    #[inline(always)]
391    const fn c1_vec3a(&self) -> Vec3A {
392        Vec3A::new(self.c1[0], self.c1[1], self.c1[2])
393    }
394
395    #[inline(always)]
396    const fn c2_vec3a(&self) -> Vec3A {
397        Vec3A::new(self.c2[0], self.c2[1], self.c2[2])
398    }
399
400    #[inline(always)]
401    const fn t_vec3a(&self) -> Vec3A {
402        Vec3A::new(self.t[0], self.t[1], self.t[2])
403    }
404}
405
406#[inline(always)]
407const fn abs_f64(x: f64) -> f64 {
408    if x < 0.0 { -x } else { x }
409}
410
411/// `f64` analogue of [`AffineRaw`] for const-context URDF construction in
412/// `URDFChain<N, f64>`. Stores plain `[f64; 3]` columns + translation; converts
413/// to `glam::DAffine3` only at the boundary.
414#[derive(Debug, Clone, Copy)]
415pub(crate) struct AffineRaw64 {
416    pub(crate) c0: [f64; 3],
417    pub(crate) c1: [f64; 3],
418    pub(crate) c2: [f64; 3],
419    pub(crate) t: [f64; 3],
420}
421
422impl AffineRaw64 {
423    pub(crate) const IDENTITY: Self = Self {
424        c0: [1.0, 0.0, 0.0],
425        c1: [0.0, 1.0, 0.0],
426        c2: [0.0, 0.0, 1.0],
427        t: [0.0, 0.0, 0.0],
428    };
429
430    /// `self * other` — applies `other` first, then `self`.
431    #[inline(always)]
432    pub(crate) const fn mul(self, other: Self) -> Self {
433        let nc0 = [
434            self.c0[0] * other.c0[0] + self.c1[0] * other.c0[1] + self.c2[0] * other.c0[2],
435            self.c0[1] * other.c0[0] + self.c1[1] * other.c0[1] + self.c2[1] * other.c0[2],
436            self.c0[2] * other.c0[0] + self.c1[2] * other.c0[1] + self.c2[2] * other.c0[2],
437        ];
438        let nc1 = [
439            self.c0[0] * other.c1[0] + self.c1[0] * other.c1[1] + self.c2[0] * other.c1[2],
440            self.c0[1] * other.c1[0] + self.c1[1] * other.c1[1] + self.c2[1] * other.c1[2],
441            self.c0[2] * other.c1[0] + self.c1[2] * other.c1[1] + self.c2[2] * other.c1[2],
442        ];
443        let nc2 = [
444            self.c0[0] * other.c2[0] + self.c1[0] * other.c2[1] + self.c2[0] * other.c2[2],
445            self.c0[1] * other.c2[0] + self.c1[1] * other.c2[1] + self.c2[1] * other.c2[2],
446            self.c0[2] * other.c2[0] + self.c1[2] * other.c2[1] + self.c2[2] * other.c2[2],
447        ];
448        let nt = [
449            self.c0[0] * other.t[0]
450                + self.c1[0] * other.t[1]
451                + self.c2[0] * other.t[2]
452                + self.t[0],
453            self.c0[1] * other.t[0]
454                + self.c1[1] * other.t[1]
455                + self.c2[1] * other.t[2]
456                + self.t[1],
457            self.c0[2] * other.t[0]
458                + self.c1[2] * other.t[1]
459                + self.c2[2] * other.t[2]
460                + self.t[2],
461        ];
462        Self { c0: nc0, c1: nc1, c2: nc2, t: nt }
463    }
464
465    #[inline(always)]
466    pub(crate) const fn is_identity(&self) -> bool {
467        const EPS: f64 = 1e-12;
468        abs_f64(self.c0[0] - 1.0) <= EPS
469            && abs_f64(self.c0[1]) <= EPS
470            && abs_f64(self.c0[2]) <= EPS
471            && abs_f64(self.c1[0]) <= EPS
472            && abs_f64(self.c1[1] - 1.0) <= EPS
473            && abs_f64(self.c1[2]) <= EPS
474            && abs_f64(self.c2[0]) <= EPS
475            && abs_f64(self.c2[1]) <= EPS
476            && abs_f64(self.c2[2] - 1.0) <= EPS
477            && abs_f64(self.t[0]) <= EPS
478            && abs_f64(self.t[1]) <= EPS
479            && abs_f64(self.t[2]) <= EPS
480    }
481
482    /// URDF RPY-convention rotation `Rz(yaw)·Ry(pitch)·Rx(roll)` plus `xyz`
483    /// translation, evaluated in `const` via [`const_sin_cos_f64`].
484    #[inline(always)]
485    pub(crate) const fn from_xyz_rpy(xyz: (f64, f64, f64), rpy: (f64, f64, f64)) -> Self {
486        let (ox, oy, oz) = xyz;
487        let (roll, pitch, yaw) = rpy;
488        let (sr, cr) = const_sin_cos_f64(roll);
489        let (sp, cp) = const_sin_cos_f64(pitch);
490        let (sy, cy) = const_sin_cos_f64(yaw);
491        Self {
492            c0: [cy * cp, sy * cp, -sp],
493            c1: [cy * sp * sr - sy * cr, sy * sp * sr + cy * cr, cp * sr],
494            c2: [cy * sp * cr + sy * sr, sy * sp * cr - cy * sr, cp * cr],
495            t: [ox, oy, oz],
496        }
497    }
498
499    #[inline(always)]
500    pub(crate) const fn to_daffine3(self) -> DAffine3 {
501        DAffine3 {
502            matrix3: DMat3::from_cols(
503                DVec3::new(self.c0[0], self.c0[1], self.c0[2]),
504                DVec3::new(self.c1[0], self.c1[1], self.c1[2]),
505                DVec3::new(self.c2[0], self.c2[1], self.c2[2]),
506            ),
507            translation: DVec3::new(self.t[0], self.t[1], self.t[2]),
508        }
509    }
510
511    #[inline(always)]
512    pub(crate) const fn c0_dvec3(&self) -> DVec3 {
513        DVec3::new(self.c0[0], self.c0[1], self.c0[2])
514    }
515
516    #[inline(always)]
517    pub(crate) const fn c1_dvec3(&self) -> DVec3 {
518        DVec3::new(self.c1[0], self.c1[1], self.c1[2])
519    }
520
521    #[inline(always)]
522    pub(crate) const fn c2_dvec3(&self) -> DVec3 {
523        DVec3::new(self.c2[0], self.c2[1], self.c2[2])
524    }
525
526    #[inline(always)]
527    pub(crate) const fn t_dvec3(&self) -> DVec3 {
528        DVec3::new(self.t[0], self.t[1], self.t[2])
529    }
530}
531
deke_types/fk/mod.rs

deke_types/fk/
mod.rs
