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