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    /// Returns joint rotation axes and axis-origin positions in world frame at
117    /// configuration `q`, plus the end-effector position.
118    fn joint_axes_positions(
119        &self,
120        q: &SRobotQ<N, F>,
121    ) -> Result<([AVec3<F>; N], [AVec3<F>; N], AVec3<F>), Self::Error>;
122
123    /// Geometric Jacobian (6×N) at configuration `q`.
124    /// Rows 0–2: linear velocity, rows 3–5: angular velocity.
125    fn jacobian(&self, q: &SRobotQ<N, F>) -> Result<[[F; N]; 6], Self::Error> {
126        let (z, p, p_ee) = self.joint_axes_positions(q)?;
127        let mut j = [[F::zero(); N]; 6];
128        for i in 0..N {
129            let dp = p_ee - p[i];
130            let c = z[i].cross(dp);
131            j[0][i] = c.x();
132            j[1][i] = c.y();
133            j[2][i] = c.z();
134            j[3][i] = z[i].x();
135            j[4][i] = z[i].y();
136            j[5][i] = z[i].z();
137        }
138        Ok(j)
139    }
140
141    /// First time-derivative of the geometric Jacobian.
142    fn jacobian_dot(
143        &self,
144        q: &SRobotQ<N, F>,
145        qdot: &SRobotQ<N, F>,
146    ) -> Result<[[F; N]; 6], Self::Error> {
147        let (z, p, p_ee) = self.joint_axes_positions(q)?;
148
149        let mut omega = AVec3::<F>::ZERO;
150        let mut z_dot = [AVec3::<F>::ZERO; N];
151        let mut p_dot = [AVec3::<F>::ZERO; N];
152        let mut pdot_acc = AVec3::<F>::ZERO;
153
154        for i in 0..N {
155            p_dot[i] = pdot_acc;
156            z_dot[i] = omega.cross(z[i]);
157            omega += z[i] * qdot.0[i];
158            let next_p = if i + 1 < N { p[i + 1] } else { p_ee };
159            pdot_acc += omega.cross(next_p - p[i]);
160        }
161        let p_ee_dot = pdot_acc;
162
163        let mut jd = [[F::zero(); N]; 6];
164        for i in 0..N {
165            let dp = p_ee - p[i];
166            let dp_dot = p_ee_dot - p_dot[i];
167            let c1 = z_dot[i].cross(dp);
168            let c2 = z[i].cross(dp_dot);
169            jd[0][i] = c1.x() + c2.x();
170            jd[1][i] = c1.y() + c2.y();
171            jd[2][i] = c1.z() + c2.z();
172            jd[3][i] = z_dot[i].x();
173            jd[4][i] = z_dot[i].y();
174            jd[5][i] = z_dot[i].z();
175        }
176        Ok(jd)
177    }
178
179    /// Second time-derivative of the geometric Jacobian.
180    fn jacobian_ddot(
181        &self,
182        q: &SRobotQ<N, F>,
183        qdot: &SRobotQ<N, F>,
184        qddot: &SRobotQ<N, F>,
185    ) -> Result<[[F; N]; 6], Self::Error> {
186        let (z, p, p_ee) = self.joint_axes_positions(q)?;
187
188        let mut omega = AVec3::<F>::ZERO;
189        let mut omega_dot = AVec3::<F>::ZERO;
190        let mut z_dot = [AVec3::<F>::ZERO; N];
191        let mut z_ddot = [AVec3::<F>::ZERO; N];
192        let mut p_dot = [AVec3::<F>::ZERO; N];
193        let mut p_ddot = [AVec3::<F>::ZERO; N];
194        let mut pdot_acc = AVec3::<F>::ZERO;
195        let mut pddot_acc = AVec3::<F>::ZERO;
196
197        for i in 0..N {
198            p_dot[i] = pdot_acc;
199            p_ddot[i] = pddot_acc;
200            let zd = omega.cross(z[i]);
201            z_dot[i] = zd;
202            z_ddot[i] = omega_dot.cross(z[i]) + omega.cross(zd);
203            omega_dot += z[i] * qddot.0[i] + zd * qdot.0[i];
204            omega += z[i] * qdot.0[i];
205            let next_p = if i + 1 < N { p[i + 1] } else { p_ee };
206            let delta = next_p - p[i];
207            let delta_dot = omega.cross(delta);
208            pdot_acc += delta_dot;
209            pddot_acc += omega_dot.cross(delta) + omega.cross(delta_dot);
210        }
211        let p_ee_dot = pdot_acc;
212        let p_ee_ddot = pddot_acc;
213
214        let mut jdd = [[F::zero(); N]; 6];
215        for i in 0..N {
216            let dp = p_ee - p[i];
217            let dp_dot = p_ee_dot - p_dot[i];
218            let dp_ddot = p_ee_ddot - p_ddot[i];
219            let c1 = z_ddot[i].cross(dp);
220            let c2 = z_dot[i].cross(dp_dot);
221            let c3 = z[i].cross(dp_ddot);
222            let two = F::one() + F::one();
223            jdd[0][i] = c1.x() + two * c2.x() + c3.x();
224            jdd[1][i] = c1.y() + two * c2.y() + c3.y();
225            jdd[2][i] = c1.z() + two * c2.z() + c3.z();
226            jdd[3][i] = z_ddot[i].x();
227            jdd[4][i] = z_ddot[i].y();
228            jdd[5][i] = z_ddot[i].z();
229        }
230        Ok(jdd)
231    }
232}
233
234#[inline(always)]
235#[cfg(debug_assertions)]
236fn check_finite<const N: usize, F: FloatScalar>(q: &SRobotQ<N, F>) -> Result<(), DekeError> {
237    if q.any_non_finite() {
238        return Err(DekeError::JointsNonFinite);
239    }
240    Ok(())
241}
242
243#[inline(always)]
244#[cfg(not(debug_assertions))]
245fn check_finite<const N: usize, F: FloatScalar>(_: &SRobotQ<N, F>) -> Result<(), std::convert::Infallible> {
246    Ok(())
247}
248
249#[inline(always)]
250const fn abs_f32(x: f32) -> f32 {
251    if x < 0.0 { -x } else { x }
252}
253
254/// Const-friendly affine transform backed by plain f32 arrays. `glam`'s
255/// `Vec3A`/`Mat3A` types use SIMD and expose components via a non-const
256/// `Deref`, so `const fn` code that needs per-component arithmetic (compose,
257/// identity check) routes through this type and only converts to
258/// `Affine3A` at the boundary.
259#[derive(Debug, Clone, Copy)]
260struct AffineRaw {
261    c0: [f32; 3],
262    c1: [f32; 3],
263    c2: [f32; 3],
264    t: [f32; 3],
265}
266
267impl AffineRaw {
268    const IDENTITY: Self = Self {
269        c0: [1.0, 0.0, 0.0],
270        c1: [0.0, 1.0, 0.0],
271        c2: [0.0, 0.0, 1.0],
272        t: [0.0, 0.0, 0.0],
273    };
274
275    /// `self * other` — applies `other` first, then `self`.
276    #[inline(always)]
277    const fn mul(self, other: Self) -> Self {
278        let nc0 = [
279            self.c0[0] * other.c0[0] + self.c1[0] * other.c0[1] + self.c2[0] * other.c0[2],
280            self.c0[1] * other.c0[0] + self.c1[1] * other.c0[1] + self.c2[1] * other.c0[2],
281            self.c0[2] * other.c0[0] + self.c1[2] * other.c0[1] + self.c2[2] * other.c0[2],
282        ];
283        let nc1 = [
284            self.c0[0] * other.c1[0] + self.c1[0] * other.c1[1] + self.c2[0] * other.c1[2],
285            self.c0[1] * other.c1[0] + self.c1[1] * other.c1[1] + self.c2[1] * other.c1[2],
286            self.c0[2] * other.c1[0] + self.c1[2] * other.c1[1] + self.c2[2] * other.c1[2],
287        ];
288        let nc2 = [
289            self.c0[0] * other.c2[0] + self.c1[0] * other.c2[1] + self.c2[0] * other.c2[2],
290            self.c0[1] * other.c2[0] + self.c1[1] * other.c2[1] + self.c2[1] * other.c2[2],
291            self.c0[2] * other.c2[0] + self.c1[2] * other.c2[1] + self.c2[2] * other.c2[2],
292        ];
293        let nt = [
294            self.c0[0] * other.t[0]
295                + self.c1[0] * other.t[1]
296                + self.c2[0] * other.t[2]
297                + self.t[0],
298            self.c0[1] * other.t[0]
299                + self.c1[1] * other.t[1]
300                + self.c2[1] * other.t[2]
301                + self.t[1],
302            self.c0[2] * other.t[0]
303                + self.c1[2] * other.t[1]
304                + self.c2[2] * other.t[2]
305                + self.t[2],
306        ];
307        Self {
308            c0: nc0,
309            c1: nc1,
310            c2: nc2,
311            t: nt,
312        }
313    }
314
315    #[inline(always)]
316    const fn is_identity(&self) -> bool {
317        const EPS: f32 = 1e-6;
318        abs_f32(self.c0[0] - 1.0) <= EPS
319            && abs_f32(self.c0[1]) <= EPS
320            && abs_f32(self.c0[2]) <= EPS
321            && abs_f32(self.c1[0]) <= EPS
322            && abs_f32(self.c1[1] - 1.0) <= EPS
323            && abs_f32(self.c1[2]) <= EPS
324            && abs_f32(self.c2[0]) <= EPS
325            && abs_f32(self.c2[1]) <= EPS
326            && abs_f32(self.c2[2] - 1.0) <= EPS
327            && abs_f32(self.t[0]) <= EPS
328            && abs_f32(self.t[1]) <= EPS
329            && abs_f32(self.t[2]) <= EPS
330    }
331
332    /// Build the URDF RPY-convention rotation (`Rz(yaw)·Ry(pitch)·Rx(roll)`)
333    /// and translate by `xyz`, using [`const_sin_cos`] for const evaluation.
334    #[inline(always)]
335    const fn from_xyz_rpy(xyz: (f64, f64, f64), rpy: (f64, f64, f64)) -> Self {
336        let (ox, oy, oz) = xyz;
337        let (roll, pitch, yaw) = rpy;
338        let (sr, cr) = const_sin_cos(roll as f32);
339        let (sp, cp) = const_sin_cos(pitch as f32);
340        let (sy, cy) = const_sin_cos(yaw as f32);
341        Self {
342            c0: [cy * cp, sy * cp, -sp],
343            c1: [cy * sp * sr - sy * cr, sy * sp * sr + cy * cr, cp * sr],
344            c2: [cy * sp * cr + sy * sr, sy * sp * cr - cy * sr, cp * cr],
345            t: [ox as f32, oy as f32, oz as f32],
346        }
347    }
348
349    #[inline(always)]
350    const fn to_affine3a(self) -> Affine3A {
351        Affine3A {
352            matrix3: Mat3A::from_cols(
353                Vec3A::new(self.c0[0], self.c0[1], self.c0[2]),
354                Vec3A::new(self.c1[0], self.c1[1], self.c1[2]),
355                Vec3A::new(self.c2[0], self.c2[1], self.c2[2]),
356            ),
357            translation: Vec3A::new(self.t[0], self.t[1], self.t[2]),
358        }
359    }
360
361    #[inline(always)]
362    const fn c0_vec3a(&self) -> Vec3A {
363        Vec3A::new(self.c0[0], self.c0[1], self.c0[2])
364    }
365
366    #[inline(always)]
367    const fn c1_vec3a(&self) -> Vec3A {
368        Vec3A::new(self.c1[0], self.c1[1], self.c1[2])
369    }
370
371    #[inline(always)]
372    const fn c2_vec3a(&self) -> Vec3A {
373        Vec3A::new(self.c2[0], self.c2[1], self.c2[2])
374    }
375
376    #[inline(always)]
377    const fn t_vec3a(&self) -> Vec3A {
378        Vec3A::new(self.t[0], self.t[1], self.t[2])
379    }
380}
381
382#[inline(always)]
383const fn abs_f64(x: f64) -> f64 {
384    if x < 0.0 { -x } else { x }
385}
386
387/// `f64` analogue of [`AffineRaw`] for const-context URDF construction in
388/// `URDFChain<N, f64>`. Stores plain `[f64; 3]` columns + translation; converts
389/// to `glam::DAffine3` only at the boundary.
390#[derive(Debug, Clone, Copy)]
391pub(crate) struct AffineRaw64 {
392    pub(crate) c0: [f64; 3],
393    pub(crate) c1: [f64; 3],
394    pub(crate) c2: [f64; 3],
395    pub(crate) t: [f64; 3],
396}
397
398impl AffineRaw64 {
399    pub(crate) const IDENTITY: Self = Self {
400        c0: [1.0, 0.0, 0.0],
401        c1: [0.0, 1.0, 0.0],
402        c2: [0.0, 0.0, 1.0],
403        t: [0.0, 0.0, 0.0],
404    };
405
406    /// `self * other` — applies `other` first, then `self`.
407    #[inline(always)]
408    pub(crate) const fn mul(self, other: Self) -> Self {
409        let nc0 = [
410            self.c0[0] * other.c0[0] + self.c1[0] * other.c0[1] + self.c2[0] * other.c0[2],
411            self.c0[1] * other.c0[0] + self.c1[1] * other.c0[1] + self.c2[1] * other.c0[2],
412            self.c0[2] * other.c0[0] + self.c1[2] * other.c0[1] + self.c2[2] * other.c0[2],
413        ];
414        let nc1 = [
415            self.c0[0] * other.c1[0] + self.c1[0] * other.c1[1] + self.c2[0] * other.c1[2],
416            self.c0[1] * other.c1[0] + self.c1[1] * other.c1[1] + self.c2[1] * other.c1[2],
417            self.c0[2] * other.c1[0] + self.c1[2] * other.c1[1] + self.c2[2] * other.c1[2],
418        ];
419        let nc2 = [
420            self.c0[0] * other.c2[0] + self.c1[0] * other.c2[1] + self.c2[0] * other.c2[2],
421            self.c0[1] * other.c2[0] + self.c1[1] * other.c2[1] + self.c2[1] * other.c2[2],
422            self.c0[2] * other.c2[0] + self.c1[2] * other.c2[1] + self.c2[2] * other.c2[2],
423        ];
424        let nt = [
425            self.c0[0] * other.t[0]
426                + self.c1[0] * other.t[1]
427                + self.c2[0] * other.t[2]
428                + self.t[0],
429            self.c0[1] * other.t[0]
430                + self.c1[1] * other.t[1]
431                + self.c2[1] * other.t[2]
432                + self.t[1],
433            self.c0[2] * other.t[0]
434                + self.c1[2] * other.t[1]
435                + self.c2[2] * other.t[2]
436                + self.t[2],
437        ];
438        Self { c0: nc0, c1: nc1, c2: nc2, t: nt }
439    }
440
441    #[inline(always)]
442    pub(crate) const fn is_identity(&self) -> bool {
443        const EPS: f64 = 1e-12;
444        abs_f64(self.c0[0] - 1.0) <= EPS
445            && abs_f64(self.c0[1]) <= EPS
446            && abs_f64(self.c0[2]) <= EPS
447            && abs_f64(self.c1[0]) <= EPS
448            && abs_f64(self.c1[1] - 1.0) <= EPS
449            && abs_f64(self.c1[2]) <= EPS
450            && abs_f64(self.c2[0]) <= EPS
451            && abs_f64(self.c2[1]) <= EPS
452            && abs_f64(self.c2[2] - 1.0) <= EPS
453            && abs_f64(self.t[0]) <= EPS
454            && abs_f64(self.t[1]) <= EPS
455            && abs_f64(self.t[2]) <= EPS
456    }
457
458    /// URDF RPY-convention rotation `Rz(yaw)·Ry(pitch)·Rx(roll)` plus `xyz`
459    /// translation, evaluated in `const` via [`const_sin_cos_f64`].
460    #[inline(always)]
461    pub(crate) const fn from_xyz_rpy(xyz: (f64, f64, f64), rpy: (f64, f64, f64)) -> Self {
462        let (ox, oy, oz) = xyz;
463        let (roll, pitch, yaw) = rpy;
464        let (sr, cr) = const_sin_cos_f64(roll);
465        let (sp, cp) = const_sin_cos_f64(pitch);
466        let (sy, cy) = const_sin_cos_f64(yaw);
467        Self {
468            c0: [cy * cp, sy * cp, -sp],
469            c1: [cy * sp * sr - sy * cr, sy * sp * sr + cy * cr, cp * sr],
470            c2: [cy * sp * cr + sy * sr, sy * sp * cr - cy * sr, cp * cr],
471            t: [ox, oy, oz],
472        }
473    }
474
475    #[inline(always)]
476    pub(crate) const fn to_daffine3(self) -> DAffine3 {
477        DAffine3 {
478            matrix3: DMat3::from_cols(
479                DVec3::new(self.c0[0], self.c0[1], self.c0[2]),
480                DVec3::new(self.c1[0], self.c1[1], self.c1[2]),
481                DVec3::new(self.c2[0], self.c2[1], self.c2[2]),
482            ),
483            translation: DVec3::new(self.t[0], self.t[1], self.t[2]),
484        }
485    }
486
487    #[inline(always)]
488    pub(crate) const fn c0_dvec3(&self) -> DVec3 {
489        DVec3::new(self.c0[0], self.c0[1], self.c0[2])
490    }
491
492    #[inline(always)]
493    pub(crate) const fn c1_dvec3(&self) -> DVec3 {
494        DVec3::new(self.c1[0], self.c1[1], self.c1[2])
495    }
496
497    #[inline(always)]
498    pub(crate) const fn c2_dvec3(&self) -> DVec3 {
499        DVec3::new(self.c2[0], self.c2[1], self.c2[2])
500    }
501
502    #[inline(always)]
503    pub(crate) const fn t_dvec3(&self) -> DVec3 {
504        DVec3::new(self.t[0], self.t[1], self.t[2])
505    }
506}
507
deke_types/fk/mod.rs

deke_types/fk/
mod.rs
