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    ($model_ffi:ident, ntupledata) => {
45        $model_ffi.tuple_adr
46    };
47    ($model_ffi:ident, ntexdata) => {
48        $model_ffi.tex_adr
49    };
50    ($model_ffi:ident, nnumericdata) => {
51        $model_ffi.numeric_adr
52    };
53    ($model_ffi:ident, nhfielddata) => {
54        $model_ffi.hfield_adr
55    };
56}
57
58
59/// Returns the correct number of items based on the X in nX (nq, nv, nu, ...).
60#[macro_export]
61#[doc(hidden)]
62macro_rules! mj_model_nx_to_nitem {
63    ($model_ffi:ident, nq) => {
64        $model_ffi.njnt
65    };
66
67    ($model_ffi:ident, nv) => {
68        $model_ffi.njnt
69    };
70
71    ($model_ffi:ident, nsensordata) => {
72        $model_ffi.nsensor
73    };
74    ($model_ffi:ident, ntupledata) => {
75        $model_ffi.ntuple
76    };
77    ($model_ffi:ident, ntexdata) => {
78        $model_ffi.ntex
79    };
80    ($model_ffi:ident, nnumericdata) => {
81        $model_ffi.nnumeric
82    };
83    ($model_ffi:ident, nhfielddata) => {
84        $model_ffi.nhfield
85    };
86}
87
88/// Provides a more direct view to a C array.
89/// # Safety
90/// This does not check if the data is valid. It is assumed
91/// the correct data is given and that it doesn't get dropped before this struct.
92/// This does not break Rust's checks as we create the view each
93/// time from the saved pointers.
94/// This should ONLY be used within a wrapper who fully encapsulates the underlying data.
95#[derive(Debug)]
96pub struct PointerViewMut<'d, T> {
97    ptr: *mut T,
98    len: usize,
99    phantom: PhantomData<&'d mut ()>
100}
101
102impl<'d, T> PointerViewMut<'d, T> {
103    pub(crate) fn new(ptr: *mut T, len: usize) -> Self {
104        Self {ptr, len, phantom: PhantomData}
105    }
106
107    #[allow(unused)]
108    pub(crate) unsafe fn set_len(&mut self, len: usize) {
109        self.len = len;
110    }
111}
112
113/// Compares if the two views point to the same data.
114impl<T> PartialEq for PointerViewMut<'_, T> {
115    fn eq(&self, other: &Self) -> bool {
116        self.ptr == other.ptr  // if the pointer differs, this isn't a view to the same data
117    }
118}
119
120impl<T> Deref for PointerViewMut<'_, T> {
121    type Target = [T];
122    fn deref(&self) -> &Self::Target {
123        unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
124    }
125}
126
127impl<T> DerefMut for PointerViewMut<'_, T> {
128    fn deref_mut(&mut self) -> &mut Self::Target {
129        unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
130    }
131}
132
133/// Provides a more direct view to a C array.
134/// # Safety
135/// This does not check if the data is valid. It is assumed
136/// the correct data is given and that it doesn't get dropped before this struct.
137/// This does not break Rust's checks as we create the view each
138/// time from the saved pointers.
139/// This should ONLY be used within a wrapper who fully encapsulates the underlying data.
140#[derive(Debug)]
141pub struct PointerView<'d, T> {
142    ptr: *const T,
143    len: usize,
144    phantom: PhantomData<&'d ()>
145}
146
147impl<'d, T> PointerView<'d, T> {
148    pub(crate) fn new(ptr: *const T, len: usize) -> Self {
149        Self {ptr, len, phantom: PhantomData}
150    }
151    
152    #[allow(unused)]
153    pub(crate) unsafe fn set_len(&mut self, len: usize) {
154        self.len = len;
155    }
156}
157
158/// Compares if the two views point to the same data.
159impl<T> PartialEq for PointerView<'_, T> {
160    fn eq(&self, other: &Self) -> bool {
161        self.ptr == other.ptr  // if the pointer differs, this isn't a view to the same data
162    }
163}
164
165impl<T> Deref for PointerView<'_, T> {
166    type Target = [T];
167    fn deref(&self) -> &Self::Target {
168        unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
169    }
170}
171
172
173/**************************************************************************************************/
174// View creation for MjData and MjModel
175/**************************************************************************************************/
176
177/// Creates a $view struct, mapping $field and $opt_field to the same location as in $data.
178#[macro_export]
179#[doc(hidden)]
180macro_rules! view_creator {
181    /* Pointer view */
182    ($self:expr, $view:ident, $data:expr, [$($field:ident),*], [$($opt_field:ident),*], $ptr_view:expr) => {
183        unsafe {
184            $view {
185                $(
186                    $field: $ptr_view($data.$field.add($self.$field.0).cast(), $self.$field.1),
187                )*
188                $(
189                    $opt_field: if $self.$opt_field.1 > 0 {
190                        Some($ptr_view($data.$opt_field.add($self.$opt_field.0).cast(), $self.$opt_field.1))
191                    } else {None},
192                )*
193            }
194        }
195    };
196
197    ($self:expr, $view:ident, $data:expr, $prefix:ident, [$($field:ident),*], [$($opt_field:ident),*], $ptr_view:expr) => {
198        paste::paste! {
199            unsafe {
200                $view {
201                    $(
202                        $field: $ptr_view($data.[<$prefix $field>].add($self.$field.0).cast(), $self.$field.1),
203                    )*
204                    $(
205                        $opt_field: if $self.$opt_field.1 > 0 {
206                            Some($ptr_view($data.[<$prefix $opt_field>].add($self.$opt_field.0).cast(), $self.$opt_field.1))
207                        } else {None},
208                    )*
209                }
210            }
211        }
212    };
213}
214
215
216/// Macro for reducing duplicated code when creating info structs to
217/// items in [`MjData`](crate::prelude::MjData) or [`MjModel`](crate::prelude::MjModel).
218/// This creates a method `X(self, name; &str) -> XInfo`.
219/// Compatible entries: (..., [name: fixed number], [name: ffi().attribute (* repeats)], [length of the item's data array (e.g., hfield -> nhfielddata, texture -> ntexdata)])
220#[doc(hidden)]
221#[macro_export]
222macro_rules! info_method {
223    ($info_type:ident, $ffi:expr, $type_:ident, [$($attr:ident: $len:expr),*], [$($attr_ffi:ident: $len_ffi:ident $(* $multiplier:expr)?),*], [$($attr_dyn:ident: $ffi_len_dyn:expr),*]) => {
224        paste::paste! {
225            #[doc = concat!(
226                "Obtains a [`", stringify!([<Mj $type_:camel $info_type Info>]), "`] struct containing information about the name, id, and ",
227                "indices required for obtaining references to the correct locations in [`Mj", stringify!($info_type), "`]. ",
228                "The actual view can be obtained via [`", stringify!([<Mj $type_:camel $info_type Info>]), "::view`].\n",
229                "# Panics\n",
230                "A panic will occur if `name` contains `\\0` characters."
231            )]
232            pub fn $type_(&self, name: &str) -> Option<[<Mj $type_:camel $info_type Info>]> {
233                let c_name = CString::new(name).unwrap();
234                let ffi = self.$ffi;
235                let id = unsafe { mj_name2id(ffi, MjtObj::[<mjOBJ_ $type_:upper>] as i32, c_name.as_ptr())};
236                if id == -1 {  // not found
237                    return None;
238                }
239
240                let id = id as usize;
241                $(
242                    let $attr = (id * $len, $len);
243                )*
244
245                $(
246                    let $attr_ffi = (id * ffi.$len_ffi as usize $( * $multiplier)*, ffi.$len_ffi as usize $( * $multiplier)*);
247                )*
248
249                $(
250                    let $attr_dyn = unsafe { mj_view_indices!(
251                        id,
252                        mj_model_nx_to_mapping!(ffi, $ffi_len_dyn),
253                        mj_model_nx_to_nitem!(ffi, $ffi_len_dyn),
254                        ffi.$ffi_len_dyn
255                    ) };
256                )*
257
258                Some([<Mj $type_:camel $info_type Info>] {name: name.to_string(), id, $($attr,)* $($attr_ffi,)* $($attr_dyn),*})
259            }
260        }
261    }
262}
263
264
265/// This creates a method `X(self, name: &str) -> XInfo`.
266/// 
267/// # Compatible Entry Types
268/// The macro supports the following types of entries for struct fields:
269/// 
270/// - **Fixed number**:  
271///   `[field_name: fixed_length]`.  
272///   Example: `[id: 1]` (for a field with a fixed length of 1)
273/// 
274/// - **FFI attribute (possibly with multiplier)**:  
275///   `[field_name: ffi_attribute (* multiplier)]`.  
276///   Example: `[matid: nmatid * 2]` (where `nmatid` is an attribute from the FFI struct, and the field length is `nmatid * 2`)
277/// 
278/// - **Dynamic length (from item's data array)**:  
279///   `[field_name: data_array_length]`.  
280///   Example: `[hfielddata: nhfielddata]` (where `nhfielddata` is the length of the hfield data array).
281///   Note: nhfielddata is the major index (i.e., [nhfielddata x something] in the C array.)
282///
283#[doc(hidden)]
284#[macro_export]
285macro_rules! info_with_view {
286    /* PointerView */
287
288    /* name of the view/info, attribute prefix in MjData/MjModel, [attributes always present], [attributes that can be None] */
289    ($info_type:ident, $name:ident, $prefix:ident, [$($attr:ident: $type_:ty),*], [$($opt_attr:ident: $type_opt:ty),*]) => {
290        paste::paste! {
291            #[doc = "Stores information required to create views to [`Mj" $info_type "`] arrays corresponding to a " $name "."]
292            #[allow(non_snake_case)]
293            pub struct [<Mj $name:camel $info_type Info>] {
294                pub name: String,
295                pub id: usize,
296                $(
297                    $attr: (usize, usize),
298                )*
299                $(
300                    $opt_attr: (usize, usize),
301                )*
302            }
303
304            impl [<Mj $name:camel $info_type Info>] {
305                #[doc = "Returns a mutable view to the correct fields in [`Mj" $info_type "`]"]
306                pub fn view_mut<'d>(&self, [<$info_type:lower>]: &'d mut [<Mj $info_type>]) -> [<Mj $name:camel $info_type ViewMut>]<'d> {
307                    view_creator!(self, [<Mj $name:camel $info_type ViewMut>], [<$info_type:lower>].ffi(), $prefix, [$($attr),*], [$($opt_attr),*], crate::util::PointerViewMut::new)
308                }
309
310                #[doc = "Returns a view to the correct fields in [`Mj" $info_type "`]"]
311                pub fn view<'d>(&self, [<$info_type:lower>]: &'d [<Mj $info_type>]) -> [<Mj $name:camel $info_type View>]<'d> {
312                    view_creator!(self, [<Mj $name:camel $info_type View>], [<$info_type:lower>].ffi(), $prefix, [$($attr),*], [$($opt_attr),*], crate::util::PointerView::new)
313                }
314            }
315
316            #[doc = "A mutable view to " $name " variables of [`Mj" $info_type "`]."]
317            #[allow(non_snake_case)]
318            pub struct [<Mj $name:camel $info_type ViewMut>]<'d> {
319                $(
320                    pub $attr: crate::util::PointerViewMut<'d, $type_>,
321                )*
322                $(
323                    pub $opt_attr: Option<crate::util::PointerViewMut<'d, $type_opt>>,
324                )*
325            }
326
327            impl [<Mj $name:camel $info_type ViewMut>]<'_> {
328                /// Resets the internal variables to 0.0.
329                pub fn zero(&mut self) {
330                    $(
331                        self.$attr.fill(unsafe { std::mem::zeroed() });
332                    )*
333                    $(
334                        if let Some(x) = &mut self.$opt_attr {
335                            x.fill(unsafe { std::mem::zeroed() });
336                        }
337                    )*
338                }
339            }
340
341            #[doc = "An immutable view to " $name " variables of [`Mj" $info_type "`]."]
342            #[allow(non_snake_case)]
343            pub struct [<Mj $name:camel $info_type View>]<'d> {
344                $(
345                    pub $attr: crate::util::PointerView<'d, $type_>,
346                )*
347                $(
348                    pub $opt_attr: Option<crate::util::PointerView<'d, $type_opt>>,
349                )*
350            }
351        }
352    };
353
354    /* name of the view/info, [attributes always present], [attributes that can be None] */
355    ($info_type:ident, $name:ident, [$($attr:ident: $type_:ty),*], [$($opt_attr:ident: $type_opt:ty),*]) => {
356        paste::paste! {
357            #[doc = "Stores information required to create views to [`Mj" $info_type "`] arrays corresponding to a " $name "."]
358            #[allow(non_snake_case)]
359            pub struct [<Mj $name:camel $info_type Info>] {
360                pub name: String,
361                pub id: usize,
362                $(
363                    $attr: (usize, usize),
364                )*
365                $(
366                    $opt_attr: (usize, usize),
367                )*
368            }
369
370            impl [<Mj $name:camel $info_type Info>] {
371                #[doc = "Returns a mutable view to the correct fields in [`Mj" $info_type "`]"]
372                pub fn view_mut<'d>(&self, [<$info_type:lower>]: &'d mut [<Mj $info_type>]) -> [<Mj $name:camel $info_type ViewMut>]<'d> {
373                    view_creator!(self, [<Mj $name:camel $info_type ViewMut>], [<$info_type:lower>].ffi(), [$($attr),*], [$($opt_attr),*], crate::util::PointerViewMut::new)
374                }
375
376                #[doc = "Returns a view to the correct fields in [`Mj" $info_type "`]"]
377                pub fn view<'d>(&self, [<$info_type:lower>]: &'d [<Mj $info_type>]) -> [<Mj $name:camel $info_type View>]<'d> {
378                    view_creator!(self, [<Mj $name:camel $info_type View>], [<$info_type:lower>].ffi(), [$($attr),*], [$($opt_attr),*], crate::util::PointerView::new)
379                }
380            }
381
382            #[doc = "A mutable view to " $name " variables of [`Mj" $info_type "`]."]
383            #[allow(non_snake_case)]
384            pub struct [<Mj $name:camel $info_type ViewMut>]<'d> {
385                $(
386                    pub $attr: crate::util::PointerViewMut<'d, $type_>,
387                )*
388                $(
389                    pub $opt_attr: Option<crate::util::PointerViewMut<'d, $type_opt>>,
390                )*
391            }
392
393            impl [<Mj $name:camel $info_type ViewMut>]<'_> {
394                /// Resets the internal variables to 0.0.
395                pub fn zero(&mut self) {
396                    $(
397                        self.$attr.fill(unsafe { std::mem::zeroed() });
398                    )*
399                    $(
400                        if let Some(x) = &mut self.$opt_attr {
401                            x.fill(unsafe { std::mem::zeroed() });
402                        }
403                    )*
404                }
405            }
406
407            #[doc = "An immutable view to " $name " variables of [`Mj" $info_type "`]."]
408            #[allow(non_snake_case)]
409            pub struct [<Mj $name:camel $info_type View>]<'d> {
410                $(
411                    pub $attr: crate::util::PointerView<'d, $type_>,
412                )*
413                $(
414                    pub $opt_attr: Option<crate::util::PointerView<'d, $type_opt>>,
415                )*
416            }
417        }
418    };
419}
420
421
422#[doc(hidden)]
423#[macro_export]
424macro_rules! getter_setter {
425    (get, [$($name:ident: bool; $comment:expr);* $(;)?]) => {paste::paste!{
426        $(
427            #[doc = concat!("Check ", $comment)]
428            pub fn [<$name:lower>](&self) -> bool {
429                self.ffi().$name == 1
430            }
431        )*
432    }};
433
434    (get, [$($name:ident: & $type:ty; $comment:expr);* $(;)?]) => {paste::paste!{
435        $(
436            #[doc = concat!("Return an immutable reference to ", $comment)]
437            pub fn [<$name:lower>](&self) -> &$type {
438                &self.ffi().$name
439            }
440
441            #[doc = concat!("Return a mutable reference to ", $comment)]
442            pub fn [<$name:camel:snake _mut>](&mut self) -> &mut $type {
443                unsafe { &mut self.ffi_mut().$name }
444            }
445        )*
446    }};
447
448    (get, [$($name:ident: $type:ty; $comment:expr);* $(;)?]) => {paste::paste!{
449        $(
450            #[doc = concat!("Return value of ", $comment)]
451            pub fn [<$name:lower>](&self) -> $type {
452                self.ffi().$name.into()
453            }
454        )*
455    }};
456
457    (set, [$($name:ident: $type:ty; $comment:expr);* $(;)?]) => {
458        paste::paste!{ 
459            $(
460                #[doc = concat!("Set ", $comment)]
461                pub fn [<set_ $name:camel:snake>](&mut self, value: $type) {
462                    unsafe { self.ffi_mut().$name = value.into() };
463                }
464            )*
465        }
466    };
467
468    /* Enum conversion */
469    (force!, get, [$($name:ident: $type:ty; $comment:expr);* $(;)?]) => {paste::paste!{
470        $(
471            #[doc = concat!("Return value of ", $comment)]
472            pub fn [<$name:lower>](&self) -> $type {
473                unsafe { std::mem::transmute(self.ffi().$name) }
474            }
475        )*
476    }};
477
478    (force!, set, [$($name:ident: $type:ty; $comment:expr);* $(;)?]) => {
479        paste::paste!{ 
480            $(
481                #[doc = concat!("Set ", $comment)]
482                pub fn [<set_ $name:camel:snake>](&mut self, value: $type) {
483                    #[allow(unnecessary_transmutes)]
484                    unsafe { self.ffi_mut().$name = std::mem::transmute(value) };
485                }
486            )*
487        }
488    };
489
490    /* Builder pattern */
491    (force!, with, [$($name:ident: $type:ty; $comment:expr);* $(;)?]) => {
492        paste::paste!{ 
493            $(
494                #[doc = concat!("Builder method for setting ", $comment)]
495                pub fn [<with_ $name:camel:snake>](mut self, value: $type) -> Self {
496                    #[allow(unnecessary_transmutes)]
497                    unsafe { self.ffi_mut().$name = std::mem::transmute(value) };
498                    self
499                }
500            )*
501        }
502    };
503
504    (with, [$($name:ident: $type:ty; $comment:expr);* $(;)?]) => {
505        paste::paste!{ 
506            $(
507                #[doc = concat!("Builder method for setting ", $comment)]
508                pub fn [<with_ $name:camel:snake>](mut self, value: $type) -> Self {
509                    unsafe { self.ffi_mut().$name = value.into() };
510                    self
511                }
512            )*
513        }
514    };
515    
516    /* Handling of optional arguments */
517    /* Enum pass */
518    (force!, get, set, [ $( $name:ident : $type:ty ; $comment:expr );* $(;)?]) => {
519        $crate::getter_setter!(force!, get, [ $( $name : $type ; $comment );* ]);
520        $crate::getter_setter!(force!, set, [ $( $name : $type ; $comment );* ]);
521    };
522
523    (get, set, [ $( $name:ident : bool ; $comment:expr );* $(;)?]) => {
524        $crate::getter_setter!(get, [ $( $name : bool ; $comment );* ]);
525        $crate::getter_setter!(set, [ $( $name : bool ; $comment );* ]);
526    };
527
528    (get, set, [ $( $name:ident : $type:ty ; $comment:expr );* $(;)?]) => {
529        $crate::getter_setter!(get, [ $( $name : $type ; $comment );* ]);
530        $crate::getter_setter!(set, [ $( $name : $type ; $comment );* ]);
531    };
532
533    /* Builder pattern */
534    (with, get, set, [ $( $name:ident : bool ; $comment:expr );* $(;)?]) => {
535        $crate::getter_setter!(get, [ $( $name : bool ; $comment );* ]);
536        $crate::getter_setter!(set, [ $( $name : bool ; $comment );* ]);
537        $crate::getter_setter!(with, [ $( $name : bool ; $comment );* ]);
538    };
539
540    (with, get, set, [ $( $name:ident : $type:ty ; $comment:expr );* $(;)?]) => {
541        $crate::getter_setter!(get, [ $( $name : $type ; $comment );* ]);
542        $crate::getter_setter!(set, [ $( $name : $type ; $comment );* ]);
543        $crate::getter_setter!(with, [ $( $name : $type ; $comment );* ]);
544    };
545
546    (force!, with, get, set, [ $( $name:ident : $type:ty ; $comment:expr );* $(;)?]) => {
547        $crate::getter_setter!(force!, get, [ $( $name : $type ; $comment );* ]);
548        $crate::getter_setter!(force!, set, [ $( $name : $type ; $comment );* ]);
549        $crate::getter_setter!(force!, with, [ $( $name : $type ; $comment );* ]);
550    };
551
552    (with, get, [ $( $name:ident : & $type:ty ; $comment:expr );* $(;)?]) => {
553        $crate::getter_setter!(get, [ $( $name : & $type ; $comment );* ]);
554        $crate::getter_setter!(with, [ $( $name : $type ; $comment );* ]);
555    };
556}
557
558
559#[doc(hidden)]
560#[macro_export]
561/// Constructs builder methods.
562macro_rules! builder_setters {
563    ($($name:ident: $type:ty; $comment:expr);* $(;)?) => {
564        $(
565            #[doc = concat!("Set ", $comment)]
566            pub fn $name(mut self, value: $type) -> Self {
567                self.$name = value;
568                self
569            }
570        )*
571    };
572}
573
574/// assert_eq!, but with tolerance for floating point rounding.
575#[doc(hidden)]
576#[macro_export]
577macro_rules! assert_relative_eq {
578    ($a:expr, $b:expr, epsilon = $eps:expr) => {{
579        let (a, b, eps) = ($a as f64, $b as f64, $eps as f64);
580        assert!((a - b).abs() <= eps, "left={:?} right={:?} eps={:?}", a, b, eps);
581    }};
582}