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 Ok(match joints.len() {
23 1 => Self::J1($chain::new(joints.try_into().ok().unwrap())),
24 2 => Self::J2($chain::new(joints.try_into().ok().unwrap())),
25 3 => Self::J3($chain::new(joints.try_into().ok().unwrap())),
26 4 => Self::J4($chain::new(joints.try_into().ok().unwrap())),
27 5 => Self::J5($chain::new(joints.try_into().ok().unwrap())),
28 6 => Self::J6($chain::new(joints.try_into().ok().unwrap())),
29 7 => Self::J7($chain::new(joints.try_into().ok().unwrap())),
30 8 => Self::J8($chain::new(joints.try_into().ok().unwrap())),
31 n => return Err(DekeError::ShapeMismatch { expected: 8, found: n }),
32 })
33 }
34
35 pub fn dof(&self) -> usize {
36 match self {
37 Self::J1(_) => 1,
38 Self::J2(_) => 2,
39 Self::J3(_) => 3,
40 Self::J4(_) => 4,
41 Self::J5(_) => 5,
42 Self::J6(_) => 6,
43 Self::J7(_) => 7,
44 Self::J8(_) => 8,
45 }
46 }
47
48 pub fn fk_dyn(&self, q: &[f32]) -> DekeResult<Vec<Affine3A>> {
49 dynamic_fk!(@dispatch_fk self, q,
50 J1 1, J2 2, J3 3, J4 4, J5 5, J6 6, J7 7, J8 8
51 )
52 }
53
54 pub fn fk_end_dyn(&self, q: &[f32]) -> DekeResult<Affine3A> {
55 dynamic_fk!(@dispatch_fk_end self, q,
56 J1 1, J2 2, J3 3, J4 4, J5 5, J6 6, J7 7, J8 8
57 )
58 }
59 }
60
61 dynamic_fk!(@impl_fkchain $name, $chain, 1 J1, 2 J2, 3 J3, 4 J4, 5 J5, 6 J6, 7 J7, 8 J8);
62 };
63
64 (@dispatch_fk $self:ident, $q:ident, $($variant:ident $n:literal),+) => {
65 match $self {
66 $(Self::$variant(chain) => {
67 let arr: &[f32; $n] = $q.try_into().map_err(|_| DekeError::ShapeMismatch {
68 expected: $n,
69 found: $q.len(),
70 })?;
71 Ok(FKChain::<$n>::fk(chain, &SRobotQ(*arr)).map_err(|e| -> DekeError { e.into() })?.to_vec())
72 }),+
73 }
74 };
75
76 (@dispatch_fk_end $self:ident, $q:ident, $($variant:ident $n:literal),+) => {
77 match $self {
78 $(Self::$variant(chain) => {
79 let arr: &[f32; $n] = $q.try_into().map_err(|_| DekeError::ShapeMismatch {
80 expected: $n,
81 found: $q.len(),
82 })?;
83 FKChain::<$n>::fk_end(chain, &SRobotQ(*arr)).map_err(|e| -> DekeError { e.into() })
84 }),+
85 }
86 };
87
88 (@impl_fkchain $name:ident, $chain:ident, $($n:literal $variant:ident),+) => {
89 $(
90 impl FKChain<$n> for $name {
91 type Error = DekeError;
92
93 fn fk(&self, q: &SRobotQ<$n>) -> Result<[Affine3A; $n], Self::Error> {
94 match self {
95 Self::$variant(chain) => FKChain::<$n>::fk(chain, q).map_err(Into::into),
96 _ => Err(DekeError::ShapeMismatch {
97 expected: self.dof(),
98 found: $n,
99 }),
100 }
101 }
102
103 fn fk_end(&self, q: &SRobotQ<$n>) -> Result<Affine3A, Self::Error> {
104 match self {
105 Self::$variant(chain) => FKChain::<$n>::fk_end(chain, q).map_err(Into::into),
106 _ => Err(DekeError::ShapeMismatch {
107 expected: self.dof(),
108 found: $n,
109 }),
110 }
111 }
112
113 fn joint_axes_positions(&self, q: &SRobotQ<$n>) -> Result<([Vec3A; $n], [Vec3A; $n], Vec3A), Self::Error> {
114 match self {
115 Self::$variant(chain) => FKChain::<$n>::joint_axes_positions(chain, q).map_err(Into::into),
116 _ => Err(DekeError::ShapeMismatch {
117 expected: self.dof(),
118 found: $n,
119 }),
120 }
121 }
122 }
123
124 impl From<$chain<$n>> for $name {
125 fn from(chain: $chain<$n>) -> Self {
126 Self::$variant(chain)
127 }
128 }
129 )+
130 };
131}
132
133dynamic_fk!(DynamicDHChain, DHChain, DHJoint);
134dynamic_fk!(DynamicHPChain, HPChain, HPJoint);
135dynamic_fk!(DynamicURDFChain, URDFChain, URDFJoint);
136
137impl DynamicDHChain {
138 pub fn from_chain(chain: impl Into<Self>) -> Self {
139 chain.into()
140 }
141}
142
143impl DynamicHPChain {
144 pub fn from_chain(chain: impl Into<Self>) -> Self {
145 chain.into()
146 }
147}
148
149impl DynamicURDFChain {
150 pub fn from_chain(chain: impl Into<Self>) -> Self {
151 chain.into()
152 }
153}
154
155trait ErasedFK<const N: usize>: Send + Sync {
156 fn fk(&self, q: &SRobotQ<N>) -> Result<[Affine3A; N], DekeError>;
157 fn fk_end(&self, q: &SRobotQ<N>) -> Result<Affine3A, DekeError>;
158 fn joint_axes_positions(
159 &self,
160 q: &SRobotQ<N>,
161 ) -> Result<([Vec3A; N], [Vec3A; N], Vec3A), DekeError>;
162 fn clone_box(&self) -> Box<dyn ErasedFK<N>>;
163}
164
165impl<const N: usize, FK: FKChain<N> + 'static> ErasedFK<N> for FK {
166 fn fk(&self, q: &SRobotQ<N>) -> Result<[Affine3A; N], DekeError> {
167 FKChain::fk(self, q).map_err(Into::into)
168 }
169
170 fn fk_end(&self, q: &SRobotQ<N>) -> Result<Affine3A, DekeError> {
171 FKChain::fk_end(self, q).map_err(Into::into)
172 }
173
174 fn joint_axes_positions(
175 &self,
176 q: &SRobotQ<N>,
177 ) -> Result<([Vec3A; N], [Vec3A; N], Vec3A), DekeError> {
178 FKChain::joint_axes_positions(self, q).map_err(Into::into)
179 }
180
181 fn clone_box(&self) -> Box<dyn ErasedFK<N>> {
182 Box::new(self.clone())
183 }
184}
185
186pub struct BoxFK<const N: usize>(Box<dyn ErasedFK<N>>);
187
188impl<const N: usize> BoxFK<N> {
189 pub fn new(fk: impl FKChain<N> + 'static) -> Self {
190 Self(Box::new(fk))
191 }
192}
193
194impl<const N: usize> Clone for BoxFK<N> {
195 fn clone(&self) -> Self {
196 Self(self.0.clone_box())
197 }
198}
199
200impl<const N: usize> FKChain<N> for BoxFK<N> {
201 type Error = DekeError;
202
203 fn fk(&self, q: &SRobotQ<N>) -> Result<[Affine3A; N], DekeError> {
204 self.0.fk(q)
205 }
206
207 fn fk_end(&self, q: &SRobotQ<N>) -> Result<Affine3A, DekeError> {
208 self.0.fk_end(q)
209 }
210
211 fn joint_axes_positions(
212 &self,
213 q: &SRobotQ<N>,
214 ) -> Result<([Vec3A; N], [Vec3A; N], Vec3A), DekeError> {
215 self.0.joint_axes_positions(q)
216 }
217}