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