deke_types/fk/
transformed.rs

1use glam_traits_ext::{FloatVec, TAffine3};
2
3use crate::SRobotQ;
4
5use super::{AAffine3, AVec3, FKChain, FKScalar, URDFBuildError, URDFJoint, compose_fixed_joints};
6
7/// Wraps an `FKChain` with an optional prefix (base) and/or suffix (tool) transform.
8///
9/// - `fk` applies only the prefix — intermediate frames stay in world coordinates
10///   without the tool offset.
11/// - `fk_end` and `joint_axes_positions` apply both — the end-effector includes
12///   the tool tip.
13#[derive(Debug, Clone)]
14pub struct TransformedFK<const N: usize, F: FKScalar, FK: FKChain<N, F>> {
15    inner: FK,
16    prefix: Option<AAffine3<F>>,
17    suffix: Option<AAffine3<F>>,
18}
19
20impl<const N: usize, F: FKScalar, FK: FKChain<N, F>> TransformedFK<N, F, FK> {
21    pub const fn new(inner: FK) -> Self {
22        Self {
23            inner,
24            prefix: None,
25            suffix: None,
26        }
27    }
28
29    pub fn with_prefix(mut self, prefix: AAffine3<F>) -> Self {
30        self.prefix = Some(prefix);
31        self
32    }
33
34    pub fn with_suffix(mut self, suffix: AAffine3<F>) -> Self {
35        self.suffix = Some(suffix);
36        self
37    }
38
39    /// Infallible setter for the prefix. `None` clears any previously set prefix.
40    pub fn with_prefix_opt(mut self, prefix: Option<AAffine3<F>>) -> Self {
41        self.prefix = prefix;
42        self
43    }
44
45    /// Infallible setter for the suffix. `None` clears any previously set suffix.
46    pub fn with_suffix_opt(mut self, suffix: Option<AAffine3<F>>) -> Self {
47        self.suffix = suffix;
48        self
49    }
50
51    pub fn set_prefix(&mut self, prefix: Option<AAffine3<F>>) {
52        self.prefix = prefix;
53    }
54
55    pub fn set_suffix(&mut self, suffix: Option<AAffine3<F>>) {
56        self.suffix = suffix;
57    }
58
59    pub fn prefix(&self) -> Option<&AAffine3<F>> {
60        self.prefix.as_ref()
61    }
62
63    pub fn suffix(&self) -> Option<&AAffine3<F>> {
64        self.suffix.as_ref()
65    }
66
67    pub fn inner(&self) -> &FK {
68        &self.inner
69    }
70}
71
72impl<const N: usize, FK: FKChain<N, f32>> TransformedFK<N, f32, FK> {
73    /// Compose a slice of fixed URDF joints (parent→child order) and set the
74    /// result as the base-side prefix transform, replacing any existing
75    /// prefix. Every joint in `joints` must be `URDFJointType::Fixed`. An
76    /// empty slice clears the prefix.
77    pub fn with_prefix_joints(mut self, joints: &[URDFJoint]) -> Result<Self, URDFBuildError> {
78        if joints.is_empty() {
79            self.prefix = None;
80            Ok(self)
81        } else {
82            self.prefix = Some(compose_fixed_joints(joints)?);
83            Ok(self)
84        }
85    }
86
87    /// Compose a slice of fixed URDF joints (parent→child order) and set the
88    /// result as the tool-side suffix transform, replacing any existing
89    /// suffix. Every joint in `joints` must be `URDFJointType::Fixed`. An
90    /// empty slice clears the suffix.
91    pub fn with_suffix_joints(mut self, joints: &[URDFJoint]) -> Result<Self, URDFBuildError> {
92        if joints.is_empty() {
93            self.suffix = None;
94            Ok(self)
95        } else {
96            self.suffix = Some(compose_fixed_joints(joints)?);
97            Ok(self)
98        }
99    }
100}
101
102impl<const N: usize, F: FKScalar, FK: FKChain<N, F>> FKChain<N, F> for TransformedFK<N, F, FK> {
103    type Error = FK::Error;
104
105    fn base_tf(&self) -> AAffine3<F> {
106        match &self.prefix {
107            Some(p) => *p * self.inner.base_tf(),
108            None => self.inner.base_tf(),
109        }
110    }
111
112    fn max_reach(&self) -> Result<F, Self::Error> {
113        let mut reach = self.inner.max_reach()?;
114        if let Some(suf) = &self.suffix {
115            reach = reach + suf.translation().length();
116        }
117        Ok(reach)
118    }
119
120    fn fk(&self, q: &SRobotQ<N, F>) -> Result<[AAffine3<F>; N], Self::Error> {
121        let mut frames = self.inner.fk(q)?;
122        if let Some(pre) = &self.prefix {
123            for f in &mut frames {
124                *f = *pre * *f;
125            }
126        }
127        Ok(frames)
128    }
129
130    fn fk_end(&self, q: &SRobotQ<N, F>) -> Result<AAffine3<F>, Self::Error> {
131        let mut end = self.inner.fk_end(q)?;
132        if let Some(pre) = &self.prefix {
133            end = *pre * end;
134        }
135        if let Some(suf) = &self.suffix {
136            end = end * *suf;
137        }
138        Ok(end)
139    }
140
141    fn all_fk(
142        &self,
143        q: &SRobotQ<N, F>,
144    ) -> Result<(AAffine3<F>, [AAffine3<F>; N], AAffine3<F>), Self::Error> {
145        let (inner_base, mut frames, mut end) = self.inner.all_fk(q)?;
146
147        if let Some(pre) = &self.prefix {
148            for f in &mut frames {
149                *f = *pre * *f;
150            }
151            end = *pre * end;
152        }
153        if let Some(suf) = &self.suffix {
154            end = end * *suf;
155        }
156
157        let base = match &self.prefix {
158            Some(p) => *p * inner_base,
159            None => inner_base,
160        };
161
162        Ok((base, frames, end))
163    }
164
165    fn joint_axes_positions(
166        &self,
167        q: &SRobotQ<N, F>,
168    ) -> Result<([AVec3<F>; N], [AVec3<F>; N], AVec3<F>), Self::Error> {
169        let (mut axes, mut positions, inner_p_ee) = self.inner.joint_axes_positions(q)?;
170
171        if let Some(pre) = &self.prefix {
172            let rot = pre.matrix3();
173            let t = pre.translation();
174            for i in 0..N {
175                axes[i] = rot * axes[i];
176                positions[i] = rot * positions[i] + t;
177            }
178        }
179
180        let p_ee = if self.prefix.is_some() || self.suffix.is_some() {
181            self.fk_end(q)?.translation()
182        } else {
183            inner_p_ee
184        };
185
186        Ok((axes, positions, p_ee))
187    }
188}
189
190impl<const N: usize, FK32, FK64> From<TransformedFK<N, f32, FK32>> for TransformedFK<N, f64, FK64>
191where
192    FK32: FKChain<N, f32>,
193    FK64: FKChain<N, f64> + From<FK32>,
194{
195    #[inline]
196    fn from(t: TransformedFK<N, f32, FK32>) -> Self {
197        TransformedFK {
198            inner: FK64::from(t.inner),
199            prefix: t.prefix.map(|p| glam::DAffine3 {
200                matrix3: p.matrix3.as_dmat3(),
201                translation: p.translation.as_dvec3(),
202            }),
203            suffix: t.suffix.map(|p| glam::DAffine3 {
204                matrix3: p.matrix3.as_dmat3(),
205                translation: p.translation.as_dvec3(),
206            }),
207        }
208    }
209}
210
211impl<const N: usize, FK64, FK32> From<TransformedFK<N, f64, FK64>> for TransformedFK<N, f32, FK32>
212where
213    FK64: FKChain<N, f64>,
214    FK32: FKChain<N, f32> + From<FK64>,
215{
216    #[inline]
217    fn from(t: TransformedFK<N, f64, FK64>) -> Self {
218        TransformedFK {
219            inner: FK32::from(t.inner),
220            prefix: t.prefix.map(|p| glam::Affine3A {
221                matrix3: glam::Mat3A::from(p.matrix3.as_mat3()),
222                translation: p.translation.as_vec3a(),
223            }),
224            suffix: t.suffix.map(|p| glam::Affine3A {
225                matrix3: glam::Mat3A::from(p.matrix3.as_mat3()),
226                translation: p.translation.as_vec3a(),
227            }),
228        }
229    }
230}
deke_types/fk/transformed.rs

deke_types/fk/
transformed.rs
