1use glam::{Affine3A, Vec3A};
2
3use crate::{DHChain, DHJoint, HPChain, HPJoint, URDFChain, URDFJoint};
4use crate::{DekeError, DekeResult, FKChain, 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>::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>::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> for $name {
101 type Error = DekeError;
102
103 fn base_tf(&self) -> Affine3A {
104 match self {
105 Self::$variant(chain) => FKChain::<$n>::base_tf(chain),
106 _ => Affine3A::IDENTITY,
107 }
108 }
109
110 fn fk(&self, q: &SRobotQ<$n>) -> Result<[Affine3A; $n], Self::Error> {
111 match self {
112 Self::$variant(chain) => FKChain::<$n>::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>) -> Result<Affine3A, Self::Error> {
121 match self {
122 Self::$variant(chain) => FKChain::<$n>::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 joint_axes_positions(&self, q: &SRobotQ<$n>) -> Result<([Vec3A; $n], [Vec3A; $n], Vec3A), Self::Error> {
131 match self {
132 Self::$variant(chain) => FKChain::<$n>::joint_axes_positions(chain, q).map_err(Into::into),
133 _ => Err(DekeError::ShapeMismatch {
134 expected: self.dof(),
135 found: $n,
136 }),
137 }
138 }
139 }
140
141 impl From<$chain<$n>> for $name {
142 fn from(chain: $chain<$n>) -> Self {
143 Self::$variant(chain)
144 }
145 }
146 )+
147 };
148}
149
150dynamic_fk!(DynamicDHChain, DHChain, DHJoint);
151dynamic_fk!(DynamicHPChain, HPChain, HPJoint);
152dynamic_fk!(DynamicURDFChain, URDFChain, URDFJoint);
153
154impl DynamicDHChain {
155 pub fn from_chain(chain: impl Into<Self>) -> Self {
156 chain.into()
157 }
158}
159
160impl DynamicHPChain {
161 pub fn from_chain(chain: impl Into<Self>) -> Self {
162 chain.into()
163 }
164}
165
166impl DynamicURDFChain {
167 pub fn from_chain(chain: impl Into<Self>) -> Self {
168 chain.into()
169 }
170}
171
172trait ErasedFK<const N: usize>: Send + Sync {
173 fn base_tf(&self) -> Affine3A;
174 fn fk(&self, q: &SRobotQ<N>) -> Result<[Affine3A; N], DekeError>;
175 fn fk_end(&self, q: &SRobotQ<N>) -> Result<Affine3A, DekeError>;
176 fn joint_axes_positions(
177 &self,
178 q: &SRobotQ<N>,
179 ) -> Result<([Vec3A; N], [Vec3A; N], Vec3A), DekeError>;
180 fn clone_box(&self) -> Box<dyn ErasedFK<N>>;
181}
182
183impl<const N: usize, FK: FKChain<N> + 'static> ErasedFK<N> for FK {
184 fn base_tf(&self) -> Affine3A {
185 FKChain::base_tf(self)
186 }
187
188 fn fk(&self, q: &SRobotQ<N>) -> Result<[Affine3A; N], DekeError> {
189 FKChain::fk(self, q).map_err(Into::into)
190 }
191
192 fn fk_end(&self, q: &SRobotQ<N>) -> Result<Affine3A, DekeError> {
193 FKChain::fk_end(self, q).map_err(Into::into)
194 }
195
196 fn joint_axes_positions(
197 &self,
198 q: &SRobotQ<N>,
199 ) -> Result<([Vec3A; N], [Vec3A; N], Vec3A), DekeError> {
200 FKChain::joint_axes_positions(self, q).map_err(Into::into)
201 }
202
203 fn clone_box(&self) -> Box<dyn ErasedFK<N>> {
204 Box::new(self.clone())
205 }
206}
207
208pub struct BoxFK<const N: usize>(Box<dyn ErasedFK<N>>);
209
210impl<const N: usize> BoxFK<N> {
211 pub fn new(fk: impl FKChain<N> + 'static) -> Self {
212 Self(Box::new(fk))
213 }
214}
215
216impl<const N: usize> Clone for BoxFK<N> {
217 fn clone(&self) -> Self {
218 Self(self.0.clone_box())
219 }
220}
221
222impl<const N: usize> FKChain<N> for BoxFK<N> {
223 type Error = DekeError;
224
225 fn base_tf(&self) -> Affine3A {
226 self.0.base_tf()
227 }
228
229 fn fk(&self, q: &SRobotQ<N>) -> Result<[Affine3A; N], DekeError> {
230 self.0.fk(q)
231 }
232
233 fn fk_end(&self, q: &SRobotQ<N>) -> Result<Affine3A, DekeError> {
234 self.0.fk_end(q)
235 }
236
237 fn joint_axes_positions(
238 &self,
239 q: &SRobotQ<N>,
240 ) -> Result<([Vec3A; N], [Vec3A; N], Vec3A), DekeError> {
241 self.0.joint_axes_positions(q)
242 }
243}