1use glam::{Affine3A, Vec3A};
2
3use super::{DHChain, DHJoint, FKChain, FKScalar, HPChain, HPJoint, URDFChain, URDFJoint};
4use crate::{DekeError, DekeResult, SRobotQ};
5
6macro_rules! dynamic_fk {
7 ($name:ident, $chain:ident, $joint:ident) => {
8 #[derive(Debug, Clone)]
9 pub enum $name {
10 J1($chain<1>),
11 J2($chain<2>),
12 J3($chain<3>),
13 J4($chain<4>),
14 J5($chain<5>),
15 J6($chain<6>),
16 J7($chain<7>),
17 J8($chain<8>),
18 }
19
20 impl $name {
21 pub fn try_new(joints: Vec<$joint>) -> DekeResult<Self> {
22 let n = joints.len();
23 let err = || DekeError::ShapeMismatch { expected: n, found: n };
24 Ok(match n {
25 1 => Self::J1(dynamic_fk!(@ctor $chain, joints, err, 1)),
26 2 => Self::J2(dynamic_fk!(@ctor $chain, joints, err, 2)),
27 3 => Self::J3(dynamic_fk!(@ctor $chain, joints, err, 3)),
28 4 => Self::J4(dynamic_fk!(@ctor $chain, joints, err, 4)),
29 5 => Self::J5(dynamic_fk!(@ctor $chain, joints, err, 5)),
30 6 => Self::J6(dynamic_fk!(@ctor $chain, joints, err, 6)),
31 7 => Self::J7(dynamic_fk!(@ctor $chain, joints, err, 7)),
32 8 => Self::J8(dynamic_fk!(@ctor $chain, joints, err, 8)),
33 _ => return Err(DekeError::ShapeMismatch { expected: 8, found: n }),
34 })
35 }
36
37 pub fn dof(&self) -> usize {
38 match self {
39 Self::J1(_) => 1,
40 Self::J2(_) => 2,
41 Self::J3(_) => 3,
42 Self::J4(_) => 4,
43 Self::J5(_) => 5,
44 Self::J6(_) => 6,
45 Self::J7(_) => 7,
46 Self::J8(_) => 8,
47 }
48 }
49
50 pub fn fk_dyn(&self, q: &[f32]) -> DekeResult<Vec<Affine3A>> {
51 dynamic_fk!(@dispatch_fk self, q,
52 J1 1, J2 2, J3 3, J4 4, J5 5, J6 6, J7 7, J8 8
53 )
54 }
55
56 pub fn fk_end_dyn(&self, q: &[f32]) -> DekeResult<Affine3A> {
57 dynamic_fk!(@dispatch_fk_end self, q,
58 J1 1, J2 2, J3 3, J4 4, J5 5, J6 6, J7 7, J8 8
59 )
60 }
61 }
62
63 dynamic_fk!(@impl_fkchain $name, $chain, 1 J1, 2 J2, 3 J3, 4 J4, 5 J5, 6 J6, 7 J7, 8 J8);
64 };
65
66 (@ctor URDFChain, $joints:ident, $err:ident, $n:literal) => {
67 URDFChain::<$n>::new($joints.try_into().map_err(|_| $err())?)?
68 };
69
70 (@ctor $chain:ident, $joints:ident, $err:ident, $n:literal) => {
71 $chain::<$n>::new($joints.try_into().map_err(|_| $err())?)
72 };
73
74 (@dispatch_fk $self:ident, $q:ident, $($variant:ident $n:literal),+) => {
75 match $self {
76 $(Self::$variant(chain) => {
77 let arr: &[f32; $n] = $q.try_into().map_err(|_| DekeError::ShapeMismatch {
78 expected: $n,
79 found: $q.len(),
80 })?;
81 Ok(FKChain::<$n, f32>::fk(chain, &SRobotQ(*arr)).map_err(|e| -> DekeError { e.into() })?.to_vec())
82 }),+
83 }
84 };
85
86 (@dispatch_fk_end $self:ident, $q:ident, $($variant:ident $n:literal),+) => {
87 match $self {
88 $(Self::$variant(chain) => {
89 let arr: &[f32; $n] = $q.try_into().map_err(|_| DekeError::ShapeMismatch {
90 expected: $n,
91 found: $q.len(),
92 })?;
93 FKChain::<$n, f32>::fk_end(chain, &SRobotQ(*arr)).map_err(|e| -> DekeError { e.into() })
94 }),+
95 }
96 };
97
98 (@impl_fkchain $name:ident, $chain:ident, $($n:literal $variant:ident),+) => {
99 $(
100 impl FKChain<$n, f32> for $name {
101 type Error = DekeError;
102
103 fn base_tf(&self) -> Affine3A {
104 match self {
105 Self::$variant(chain) => FKChain::<$n, f32>::base_tf(chain),
106 _ => Affine3A::IDENTITY,
107 }
108 }
109
110 fn fk(&self, q: &SRobotQ<$n, f32>) -> Result<[Affine3A; $n], Self::Error> {
111 match self {
112 Self::$variant(chain) => FKChain::<$n, f32>::fk(chain, q).map_err(Into::into),
113 _ => Err(DekeError::ShapeMismatch {
114 expected: self.dof(),
115 found: $n,
116 }),
117 }
118 }
119
120 fn fk_end(&self, q: &SRobotQ<$n, f32>) -> Result<Affine3A, Self::Error> {
121 match self {
122 Self::$variant(chain) => FKChain::<$n, f32>::fk_end(chain, q).map_err(Into::into),
123 _ => Err(DekeError::ShapeMismatch {
124 expected: self.dof(),
125 found: $n,
126 }),
127 }
128 }
129
130 fn all_fk(&self, q: &SRobotQ<$n, f32>) -> Result<(Affine3A, [Affine3A; $n], Affine3A), Self::Error> {
131 match self {
132 Self::$variant(chain) => FKChain::<$n, f32>::all_fk(chain, q).map_err(Into::into),
133 _ => Err(DekeError::ShapeMismatch {
134 expected: self.dof(),
135 found: $n,
136 }),
137 }
138 }
139
140 fn joint_axes_positions(&self, q: &SRobotQ<$n, f32>) -> Result<([Vec3A; $n], [Vec3A; $n], Vec3A), Self::Error> {
141 match self {
142 Self::$variant(chain) => FKChain::<$n, f32>::joint_axes_positions(chain, q).map_err(Into::into),
143 _ => Err(DekeError::ShapeMismatch {
144 expected: self.dof(),
145 found: $n,
146 }),
147 }
148 }
149 }
150
151 impl From<$chain<$n>> for $name {
152 fn from(chain: $chain<$n>) -> Self {
153 Self::$variant(chain)
154 }
155 }
156 )+
157 };
158}
159
160dynamic_fk!(DynamicDHChain, DHChain, DHJoint);
161dynamic_fk!(DynamicHPChain, HPChain, HPJoint);
162dynamic_fk!(DynamicURDFChain, URDFChain, URDFJoint);
163
164impl DynamicDHChain {
165 pub fn from_chain(chain: impl Into<Self>) -> Self {
166 chain.into()
167 }
168}
169
170impl DynamicHPChain {
171 pub fn from_chain(chain: impl Into<Self>) -> Self {
172 chain.into()
173 }
174}
175
176impl DynamicURDFChain {
177 pub fn from_chain(chain: impl Into<Self>) -> Self {
178 chain.into()
179 }
180}
181
182trait ErasedFK<const N: usize, F: FKScalar>: Send + Sync {
183 fn base_tf(&self) -> super::AAffine3<F>;
184 fn fk(&self, q: &SRobotQ<N, F>) -> Result<[super::AAffine3<F>; N], DekeError>;
185 fn fk_end(&self, q: &SRobotQ<N, F>) -> Result<super::AAffine3<F>, DekeError>;
186 fn all_fk(
187 &self,
188 q: &SRobotQ<N, F>,
189 ) -> Result<(super::AAffine3<F>, [super::AAffine3<F>; N], super::AAffine3<F>), DekeError>;
190 fn joint_axes_positions(
191 &self,
192 q: &SRobotQ<N, F>,
193 ) -> Result<([super::AVec3<F>; N], [super::AVec3<F>; N], super::AVec3<F>), DekeError>;
194 fn clone_box(&self) -> Box<dyn ErasedFK<N, F>>;
195}
196
197impl<const N: usize, F: FKScalar, FK: FKChain<N, F> + 'static> ErasedFK<N, F> for FK {
198 fn base_tf(&self) -> super::AAffine3<F> {
199 FKChain::base_tf(self)
200 }
201
202 fn fk(&self, q: &SRobotQ<N, F>) -> Result<[super::AAffine3<F>; N], DekeError> {
203 FKChain::fk(self, q).map_err(Into::into)
204 }
205
206 fn fk_end(&self, q: &SRobotQ<N, F>) -> Result<super::AAffine3<F>, DekeError> {
207 FKChain::fk_end(self, q).map_err(Into::into)
208 }
209
210 fn all_fk(
211 &self,
212 q: &SRobotQ<N, F>,
213 ) -> Result<(super::AAffine3<F>, [super::AAffine3<F>; N], super::AAffine3<F>), DekeError> {
214 FKChain::all_fk(self, q).map_err(Into::into)
215 }
216
217 fn joint_axes_positions(
218 &self,
219 q: &SRobotQ<N, F>,
220 ) -> Result<([super::AVec3<F>; N], [super::AVec3<F>; N], super::AVec3<F>), DekeError> {
221 FKChain::joint_axes_positions(self, q).map_err(Into::into)
222 }
223
224 fn clone_box(&self) -> Box<dyn ErasedFK<N, F>> {
225 Box::new(self.clone())
226 }
227}
228
229pub struct BoxFK<const N: usize, F: FKScalar = f32>(Box<dyn ErasedFK<N, F>>);
230
231impl<const N: usize, F: FKScalar> BoxFK<N, F> {
232 pub fn new(fk: impl FKChain<N, F> + 'static) -> Self {
233 Self(Box::new(fk))
234 }
235}
236
237impl<const N: usize, F: FKScalar> Clone for BoxFK<N, F> {
238 fn clone(&self) -> Self {
239 Self(self.0.clone_box())
240 }
241}
242
243impl<const N: usize, F: FKScalar> FKChain<N, F> for BoxFK<N, F> {
244 type Error = DekeError;
245
246 fn base_tf(&self) -> super::AAffine3<F> {
247 self.0.base_tf()
248 }
249
250 fn fk(&self, q: &SRobotQ<N, F>) -> Result<[super::AAffine3<F>; N], DekeError> {
251 self.0.fk(q)
252 }
253
254 fn fk_end(&self, q: &SRobotQ<N, F>) -> Result<super::AAffine3<F>, DekeError> {
255 self.0.fk_end(q)
256 }
257
258 fn all_fk(
259 &self,
260 q: &SRobotQ<N, F>,
261 ) -> Result<(super::AAffine3<F>, [super::AAffine3<F>; N], super::AAffine3<F>), DekeError> {
262 self.0.all_fk(q)
263 }
264
265 fn joint_axes_positions(
266 &self,
267 q: &SRobotQ<N, F>,
268 ) -> Result<([super::AVec3<F>; N], [super::AVec3<F>; N], super::AVec3<F>), DekeError> {
269 self.0.joint_axes_positions(q)
270 }
271}