mujoco_rs/
util.rs

1//! Utility related data
2use std::{marker::PhantomData, ops::{Deref, DerefMut}};
3
4
5/// Creates a (start, length) tuple based on
6/// lookup variables that define the mapping in MuJoCo's mjModel struct.
7/// The tuple is used to create views to correct addresses in corresponding structs.
8/// Format: item id, map from item id to index inside the array of all items' values,
9///         number of items, maximum number of elements inside the array of all items' values
10#[macro_export]
11#[doc(hidden)]
12macro_rules! mj_view_indices {
13    ($id:expr, $addr_map:expr, $njnt:expr, $max_n:expr) => {
14        {
15            let start_addr = *$addr_map.add($id) as isize;
16            if start_addr == -1 {
17                (0, 0)
18            }
19            else
20            {
21                let end_addr = if $id + 1 < $njnt as usize {*$addr_map.add($id as usize + 1) as usize} else {$max_n as usize};
22                let n = end_addr - start_addr as usize;
23                (start_addr as usize, n)
24            }
25        }
26    };
27}
28
29/// Returns the correct address mapping based on the X in nX (nq, nv, nu, ...).
30#[macro_export]
31#[doc(hidden)]
32macro_rules! mj_model_nx_to_mapping {
33    ($model_ffi:ident, nq) => {
34        $model_ffi.jnt_qposadr
35    };
36
37    ($model_ffi:ident, nv) => {
38        $model_ffi.jnt_dofadr
39    };
40
41    ($model_ffi:ident, nsensordata) => {
42        $model_ffi.sensor_adr
43    }
44}
45
46
47/// Returns the correct number of items based on the X in nX (nq, nv, nu, ...).
48#[macro_export]
49#[doc(hidden)]
50macro_rules! mj_model_nx_to_nitem {
51    ($model_ffi:ident, nq) => {
52        $model_ffi.njnt
53    };
54
55    ($model_ffi:ident, nv) => {
56        $model_ffi.njnt
57    };
58
59    ($model_ffi:ident, nsensordata) => {
60        $model_ffi.nsensor
61    };
62}
63
64/// Provides a more direct view to a C array.
65/// # Safety
66/// This does not check if the data is valid. It is assumed
67/// the correct data is given and that it doesn't get dropped before this struct.
68/// This does not break Rust's checks as we create the view each
69/// time from the saved pointers.
70/// This should ONLY be used within a wrapper who fully encapsulates the underlying data.
71#[derive(Debug)]
72pub struct PointerViewMut<'d, T> {
73    ptr: *mut T,
74    len: usize,
75    phantom: PhantomData<&'d mut ()>
76}
77
78impl<'d, T> PointerViewMut<'d, T> {
79    pub(crate) fn new(ptr: *mut T, len: usize) -> Self {
80        Self {ptr, len, phantom: PhantomData}
81    }
82
83    #[allow(unused)]
84    pub(crate) unsafe fn set_len(&mut self, len: usize) {
85        self.len = len;
86    }
87}
88
89/// Compares if the two views point to the same data.
90impl<T> PartialEq for PointerViewMut<'_, T> {
91    fn eq(&self, other: &Self) -> bool {
92        self.ptr == other.ptr  // if the pointer differs, this isn't a view to the same data
93    }
94}
95
96impl<T> Deref for PointerViewMut<'_, T> {
97    type Target = [T];
98    fn deref(&self) -> &Self::Target {
99        unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
100    }
101}
102
103impl<T> DerefMut for PointerViewMut<'_, T> {
104    fn deref_mut(&mut self) -> &mut Self::Target {
105        unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
106    }
107}
108
109/// Provides a more direct view to a C array.
110/// # Safety
111/// This does not check if the data is valid. It is assumed
112/// the correct data is given and that it doesn't get dropped before this struct.
113/// This does not break Rust's checks as we create the view each
114/// time from the saved pointers.
115/// This should ONLY be used within a wrapper who fully encapsulates the underlying data.
116#[derive(Debug)]
117pub struct PointerView<'d, T> {
118    ptr: *const T,
119    len: usize,
120    phantom: PhantomData<&'d ()>
121}
122
123impl<'d, T> PointerView<'d, T> {
124    pub(crate) fn new(ptr: *const T, len: usize) -> Self {
125        Self {ptr, len, phantom: PhantomData}
126    }
127    
128    #[allow(unused)]
129    pub(crate) unsafe fn set_len(&mut self, len: usize) {
130        self.len = len;
131    }
132}
133
134/// Compares if the two views point to the same data.
135impl<T> PartialEq for PointerView<'_, T> {
136    fn eq(&self, other: &Self) -> bool {
137        self.ptr == other.ptr  // if the pointer differs, this isn't a view to the same data
138    }
139}
140
141impl<T> Deref for PointerView<'_, T> {
142    type Target = [T];
143    fn deref(&self) -> &Self::Target {
144        unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
145    }
146}
147
148
149/**************************************************************************************************/
150// View creation for MjData and MjModel
151/**************************************************************************************************/
152
153/// Creates a $view struct, mapping $field and $opt_field to the same location as in $data.
154#[macro_export]
155#[doc(hidden)]
156macro_rules! view_creator {
157    /* Pointer view */
158    ($self:expr, $view:ident, $data:expr, [$($field:ident),*], [$($opt_field:ident),*], $ptr_view:expr) => {
159        unsafe {
160            $view {
161                $(
162                    $field: $ptr_view($data.$field.add($self.$field.0).cast(), $self.$field.1),
163                )*
164                $(
165                    $opt_field: if $self.$opt_field.1 > 0 {
166                        Some($ptr_view($data.$opt_field.add($self.$opt_field.0).cast(), $self.$opt_field.1))
167                    } else {None},
168                )*
169            }
170        }
171    };
172
173    ($self:expr, $view:ident, $data:expr, $prefix:ident, [$($field:ident),*], [$($opt_field:ident),*], $ptr_view:expr) => {
174        paste::paste! {
175            unsafe {
176                $view {
177                    $(
178                        $field: $ptr_view($data.[<$prefix $field>].add($self.$field.0).cast(), $self.$field.1),
179                    )*
180                    $(
181                        $opt_field: if $self.$opt_field.1 > 0 {
182                            Some($ptr_view($data.[<$prefix $opt_field>].add($self.$opt_field.0).cast(), $self.$opt_field.1))
183                        } else {None},
184                    )*
185                }
186            }
187        }
188    };
189}
190
191
192/// Macro for reducing duplicated code when creating info structs to
193/// items that have fixed size arrays in [`MjData`](crate::prelude::MjData) or [`MjModel`](crate::prelude::MjModel).
194/// This creates a method `X(self, name; &str) -> XInfo`.
195#[doc(hidden)]
196#[macro_export]
197macro_rules! fixed_size_info_method {
198    ($info_type:ident, $ffi:expr, $type_:ident, [$($attr:ident: $len:expr),*]) => {
199        paste::paste! {
200            #[doc = concat!(
201                "Obtains a [`", stringify!([<Mj $type_:camel $info_type Info>]), "`] struct containing information about the name, id, and ",
202                "indices required for obtaining references to the correct locations in [`Mj", stringify!($info_type), "`]. ",
203                "The actual view can be obtained via [`", stringify!([<Mj $type_:camel $info_type Info>]), "::view`]."
204            )]
205            pub fn $type_(&self, name: &str) -> Option<[<Mj $type_:camel $info_type Info>]> {
206                let c_name = CString::new(name).unwrap();
207                let id = unsafe { mj_name2id(self.$ffi, MjtObj::[<mjOBJ_ $type_:upper>] as i32, c_name.as_ptr())};
208                if id == -1 {  // not found
209                    return None;
210                }
211
212                let id = id as usize;
213                $(
214                    let $attr = (id * $len, $len);
215                )*
216
217                Some([<Mj $type_:camel $info_type Info>] {name: name.to_string(), id, $($attr),*})
218            }
219        }
220    }
221}
222
223
224/// Creates the xInfo struct along with corresponding xView and xViewMut structs.
225#[doc(hidden)]
226#[macro_export]
227macro_rules! info_with_view {
228    /* PointerView */
229
230    /* name of the view/info, attribute prefix in MjData/MjModel, [attributes always present], [attributes that can be None] */
231    ($info_type:ident, $name:ident, $prefix:ident, [$($attr:ident: $type_:ty),*], [$($opt_attr:ident: $type_opt:ty),*]) => {
232        paste::paste! {
233            #[doc = "Stores information required to create views to [`Mj" $info_type "`] arrays corresponding to a " $name "."]
234            #[allow(non_snake_case)]
235            pub struct [<Mj $name:camel $info_type Info>] {
236                pub name: String,
237                pub id: usize,
238                $(
239                    $attr: (usize, usize),
240                )*
241                $(
242                    $opt_attr: (usize, usize),
243                )*
244            }
245
246            impl [<Mj $name:camel $info_type Info>] {
247                #[doc = "Returns a mutable view to the correct fields in [`Mj" $info_type "`]"]
248                pub fn view_mut<'d>(&self, [<$info_type:lower>]: &'d mut [<Mj $info_type>]) -> [<Mj $name:camel $info_type ViewMut>]<'d> {
249                    view_creator!(self, [<Mj $name:camel $info_type ViewMut>], [<$info_type:lower>].ffi(), $prefix, [$($attr),*], [$($opt_attr),*], crate::util::PointerViewMut::new)
250                }
251
252                #[doc = "Returns a view to the correct fields in [`Mj" $info_type "`]"]
253                pub fn view<'d>(&self, [<$info_type:lower>]: &'d [<Mj $info_type>]) -> [<Mj $name:camel $info_type View>]<'d> {
254                    view_creator!(self, [<Mj $name:camel $info_type View>], [<$info_type:lower>].ffi(), $prefix, [$($attr),*], [$($opt_attr),*], crate::util::PointerView::new)
255                }
256            }
257
258            #[doc = "A mutable view to " $name " variables of [`Mj" $info_type "`]."]
259            #[allow(non_snake_case)]
260            pub struct [<Mj $name:camel $info_type ViewMut>]<'d> {
261                $(
262                    pub $attr: crate::util::PointerViewMut<'d, $type_>,
263                )*
264                $(
265                    pub $opt_attr: Option<crate::util::PointerViewMut<'d, $type_opt>>,
266                )*
267            }
268
269            impl [<Mj $name:camel $info_type ViewMut>]<'_> {
270                /// Resets the internal variables to 0.0.
271                pub fn zero(&mut self) {
272                    $(
273                        self.$attr.fill(unsafe { std::mem::zeroed() });
274                    )*
275                    $(
276                        if let Some(x) = &mut self.$opt_attr {
277                            x.fill(unsafe { std::mem::zeroed() });
278                        }
279                    )*
280                }
281            }
282
283            #[doc = "An immutable view to " $name " variables of [`Mj" $info_type "`]."]
284            #[allow(non_snake_case)]
285            pub struct [<Mj $name:camel $info_type View>]<'d> {
286                $(
287                    pub $attr: crate::util::PointerView<'d, $type_>,
288                )*
289                $(
290                    pub $opt_attr: Option<crate::util::PointerView<'d, $type_opt>>,
291                )*
292            }
293        }
294    };
295
296    /* name of the view/info, [attributes always present], [attributes that can be None] */
297    ($info_type:ident, $name:ident, [$($attr:ident: $type_:ty),*], [$($opt_attr:ident: $type_opt:ty),*]) => {
298        paste::paste! {
299            #[doc = "Stores information required to create views to [`Mj" $info_type "`] arrays corresponding to a " $name "."]
300            #[allow(non_snake_case)]
301            pub struct [<Mj $name:camel $info_type Info>] {
302                pub name: String,
303                pub id: usize,
304                $(
305                    $attr: (usize, usize),
306                )*
307                $(
308                    $opt_attr: (usize, usize),
309                )*
310            }
311
312            impl [<Mj $name:camel $info_type Info>] {
313                #[doc = "Returns a mutable view to the correct fields in [`Mj" $info_type "`]"]
314                pub fn view_mut<'d>(&self, [<$info_type:lower>]: &'d mut [<Mj $info_type>]) -> [<Mj $name:camel $info_type ViewMut>]<'d> {
315                    view_creator!(self, [<Mj $name:camel $info_type ViewMut>], [<$info_type:lower>].ffi(), [$($attr),*], [$($opt_attr),*], crate::util::PointerViewMut::new)
316                }
317
318                #[doc = "Returns a view to the correct fields in [`Mj" $info_type "`]"]
319                pub fn view<'d>(&self, [<$info_type:lower>]: &'d [<Mj $info_type>]) -> [<Mj $name:camel $info_type View>]<'d> {
320                    view_creator!(self, [<Mj $name:camel $info_type View>], [<$info_type:lower>].ffi(), [$($attr),*], [$($opt_attr),*], crate::util::PointerView::new)
321                }
322            }
323
324            #[doc = "A mutable view to " $name " variables of [`Mj" $info_type "`]."]
325            #[allow(non_snake_case)]
326            pub struct [<Mj $name:camel $info_type ViewMut>]<'d> {
327                $(
328                    pub $attr: crate::util::PointerViewMut<'d, $type_>,
329                )*
330                $(
331                    pub $opt_attr: Option<crate::util::PointerViewMut<'d, $type_opt>>,
332                )*
333            }
334
335            impl [<Mj $name:camel $info_type ViewMut>]<'_> {
336                /// Resets the internal variables to 0.0.
337                pub fn zero(&mut self) {
338                    $(
339                        self.$attr.fill(unsafe { std::mem::zeroed() });
340                    )*
341                    $(
342                        if let Some(x) = &mut self.$opt_attr {
343                            x.fill(unsafe { std::mem::zeroed() });
344                        }
345                    )*
346                }
347            }
348
349            #[doc = "An immutable view to " $name " variables of [`Mj" $info_type "`]."]
350            #[allow(non_snake_case)]
351            pub struct [<Mj $name:camel $info_type View>]<'d> {
352                $(
353                    pub $attr: crate::util::PointerView<'d, $type_>,
354                )*
355                $(
356                    pub $opt_attr: Option<crate::util::PointerView<'d, $type_opt>>,
357                )*
358            }
359        }
360    };
361}
362
363
364/// assert_eq!, but with tolerance for floating point rounding.
365#[doc(hidden)]
366#[macro_export]
367macro_rules! assert_relative_eq {
368    ($a:expr, $b:expr, epsilon = $eps:expr) => {{
369        let (a, b, eps) = ($a as f64, $b as f64, $eps as f64);
370        assert!((a - b).abs() <= eps, "left={:?} right={:?} eps={:?}", a, b, eps);
371    }};
372}