Add ASAN and MSAN instrumentation for arena and mjData stack.

PiperOrigin-RevId: 535635490
Change-Id: Ia1e58771bae25bf72c861f57452fb8faa2b918fb

CMakeLists.txt

@@ -55,6 +55,7 @@ include(MujocoDependencies)
 set(MUJOCO_HEADERS
     include/mujoco/mjdata.h
     include/mujoco/mjexport.h
+    include/mujoco/mjmacro.h
     include/mujoco/mjmodel.h
     include/mujoco/mjplugin.h
     include/mujoco/mjrender.h

doc/changelog.rst

@@ -56,6 +56,14 @@ General
   between the two extremal values given by the `Millard et al. (2013) <https://doi.org/10.1115/1.4023390>`__ muscle
   model, within a range of width tausmooth.  See :ref:`Muscle actuators<CMuscle>` for more details.
   Relatedly, :ref:`mju_muscleDynamics` now takes 3 parameters instead of 2, adding the new smoothing-width parameter.
+- Moved public C macro definitions out of mujoco.h into a new public header file called
+  `mjmacro.h <https://github.com/deepmind/mujoco/blob/main/include/mujoco/mjmacro.h>`__. The new file
+  is included by mujoco.h so this change does not break existing user code.
+- Added instrumentation for the `Address Sanitizer (ASAN) <https://clang.llvm.org/docs/AddressSanitizer.html>`__ and
+  `Memory Sanitizer (MSAN) <https://clang.llvm.org/docs/MemorySanitizer.html>`__ to detect memory bugs when allocating
+  from the ``mjData`` stack and arena.
+- Removed ``pstack`` and ``parena`` from the output of ``mj_printData``, since these are implementation details of the
+  ``mjData`` allocators that are affected by diagnostic paddings in instrumented builds.
 
 .. youtube:: hqIMTNGaLF4
    :align: right
@@ -553,7 +561,8 @@ General
 #. ``mju_rotVecMat`` and ``mju_rotVecMatT`` now support in-place multiplication.
 #. ``mjData.ctrl`` values are no longer clamped in-place, remain untouched by the engine.
 #. Arrays in mjData's buffer now align to 64-byte boundaries rather than 8-byte.
-#. Added memory poisoning when building with Address Sanitizer (ASAN) and Memory Sanitizer (MSAN). This allows ASAN to
+#. Added memory poisoning when building with `Address Sanitizer (ASAN) <https://clang.llvm.org/docs/AddressSanitizer.html>`__
+   and `Memory Sanitizer (MSAN) <https://clang.llvm.org/docs/MemorySanitizer.html>`__. This allows ASAN to
    detect reads and writes to regions in ``mjModel.buffer`` and ``mjData.buffer`` that do not lie within an array, and
    for MSAN to detect reads from uninitialised fields in ``mjData`` following ``mj_resetData``.
 #. Added a `slider-crank example model <https://github.com/deepmind/mujoco/tree/main/model/slider_crank>`_.

doc/programming/index.rst

@@ -136,7 +136,7 @@ links below, to make this documentation self-contained.
 
 mujoco.h   `(source) <https://github.com/deepmind/mujoco/blob/main/include/mujoco/mujoco.h>`__
    This is the main header file and must be included in all programs using MuJoCo. It defines all API functions and
-   global variables, and includes the next 5 files which provide the necessary type definitions.
+   global variables, and includes the all other header files except mjxmacro.h.
 mjmodel.h   `(source) <https://github.com/deepmind/mujoco/blob/main/include/mujoco/mjmodel.h>`__
    Defines the C structure :ref:`mjModel` which is the runtime representation of the
    model being simulated. It also defines a number of primitive types and other structures needed to define mjModel.
@@ -152,6 +152,8 @@ mjui.h   `(source) <https://github.com/deepmind/mujoco/blob/main/include/mujoc
    Defines the primitive types and structures needed by the UI framework.
 mjtnum.h   `(source) <https://github.com/deepmind/mujoco/blob/main/include/mujoco/mjtnum.h>`__
    Defines MuJoCo's ``mjtNum`` floating-point type to be either ``double`` or ``float``. See :ref:`mjtNum`.
+mjmacro.h   `(source) <https://github.com/deepmind/mujoco/blob/main/include/mujoco/mjmacro.h>`__
+   Defines C macros that are useful in user code.
 mjxmacro.h   `(source) <https://github.com/deepmind/mujoco/blob/main/include/mujoco/mjxmacro.h>`__
    This file is optional and is not included by mujoco.h. It defines :ref:`X Macros <tyXMacro>` that can
    automate the mapping of mjModel and mjData into scripting languages, as well as other operations that require

include/mujoco/mjmacro.h

@@ -0,0 +1,61 @@
+// Copyright 2023 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_MJMACRO_H_
+#define MUJOCO_MJMACRO_H_
+
+// include asan interface header, or provide stubs for poison/unpoison macros when not using asan
+#ifdef ADDRESS_SANITIZER
+  #include <sanitizer/asan_interface.h>
+#elif defined(_MSC_VER)
+  #define ASAN_POISON_MEMORY_REGION(addr, size)
+  #define ASAN_UNPOISON_MEMORY_REGION(addr, size)
+#else
+  #define ASAN_POISON_MEMORY_REGION(addr, size) ((void)(addr), (void)(size))
+  #define ASAN_UNPOISON_MEMORY_REGION(addr, size) ((void)(addr), (void)(size))
+#endif
+
+// max and min (use only for primitive types)
+#define mjMAX(a, b) (((a) > (b)) ? (a) : (b))
+#define mjMIN(a, b) (((a) < (b)) ? (a) : (b))
+
+// mjData stack frame management
+#define mjMARKSTACK   int _mark = d->pstack;
+#define mjFREESTACK   d->pstack = _mark;
+
+// return current value of mjOption enable/disable flags
+#define mjDISABLED(x) (m->opt.disableflags & (x))
+#define mjENABLED(x)  (m->opt.enableflags & (x))
+
+// annotation for functions that accept printf-like variadic arguments
+#ifndef mjPRINTFLIKE
+  #if defined(__GNUC__)
+    #define mjPRINTFLIKE(n, m) __attribute__((format(printf, n, m)))
+  #else
+    #define mjPRINTFLIKE(n, m)
+  #endif
+#endif
+
+// implementation of mjFREESTACK when using the address sanitizer
+#ifdef ADDRESS_SANITIZER
+  #undef mjFREESTACK
+  #define mjFREESTACK {                                          \
+    d->pstack = _mark;                                           \
+    ASAN_POISON_MEMORY_REGION(                                   \
+        (char*)d->arena + d->parena,                             \
+        (d->nstack - d->pstack) * sizeof(mjtNum) - d->parena );  \
+  }
+#endif
+
+#endif  // MUJOCO_MJMACRO_H_

include/mujoco/mujoco.h

@@ -30,10 +30,14 @@ extern "C" {
 #include <stdlib.h>
 #include <math.h>
 
+#ifdef ADDRESS_SANITIZER
+#include <sanitizer/asan_interface.h>
+#endif
 
 // type definitions
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjplugin.h>
 #include <mujoco/mjrender.h>
 #include <mujoco/mjtnum.h>
@@ -41,21 +45,6 @@ extern "C" {
 #include <mujoco/mjvisualize.h>
 
 
-// macros
-#define mjMARKSTACK   int _mark = d->pstack;
-#define mjFREESTACK   d->pstack = _mark;
-#define mjDISABLED(x) (m->opt.disableflags & (x))
-#define mjENABLED(x)  (m->opt.enableflags & (x))
-
-#ifndef mjPRINTFLIKE
-  #if defined(__GNUC__)
-    #define mjPRINTFLIKE(n, m) __attribute__((format(printf, n, m)))
-  #else
-    #define mjPRINTFLIKE(n, m)
-  #endif // __GNUC__
-#endif // mjPRINTFLIKE
-
-
 // user error and memory handlers
 MJAPI extern void  (*mju_user_error)(const char*);
 MJAPI extern void  (*mju_user_warning)(const char*);
@@ -822,9 +811,6 @@ MJAPI void mj_deactivate(void);
 
 //---------------------------------- Standard math -------------------------------------------------
 
-#define mjMAX(a, b) (((a) > (b)) ? (a) : (b))
-#define mjMIN(a, b) (((a) < (b)) ? (a) : (b))
-
 #ifdef mjUSEDOUBLE
   #define mju_sqrt    sqrt
   #define mju_exp     exp

python/mujoco/CMakeLists.txt

@@ -203,6 +203,11 @@ add_subdirectory(simulate)
 
 add_subdirectory(util)
 
+add_library(mujoco_private_api_header INTERFACE)
+target_sources(mujoco_private_api_header INTERFACE private.h)
+set_target_properties(mujoco_private_api_header PROPERTIES PUBLIC_HEADER private.h)
+target_link_libraries(mujoco_private_api_header INTERFACE mujoco)
+
 if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/enum_traits.h)
   add_library(enum_traits INTERFACE)
   target_sources(enum_traits INTERFACE enum_traits.h)
@@ -242,7 +247,13 @@ target_link_libraries(function_traits INTERFACE mujoco absl::core_headers)
 add_library(errors_header INTERFACE)
 target_sources(errors_header INTERFACE errors.h)
 set_target_properties(errors_header PROPERTIES PUBLIC_HEADER errors.h)
-target_link_libraries(errors_header INTERFACE crossplatform func_wrap mujoco)
+target_link_libraries(
+  errors_header
+  INTERFACE crossplatform
+            func_wrap
+            mujoco
+            mujoco_private_api_header
+)
 
 add_library(raw INTERFACE)
 target_sources(raw INTERFACE raw.h)
@@ -381,6 +392,7 @@ target_link_libraries(
   PRIVATE absl::flat_hash_map
           crossplatform
           mujoco
+          mujoco_private_api_header
           raw
           errors_header
           func_wrap

python/mujoco/bindings_test.py

@@ -515,8 +515,8 @@ class MuJoCoBindingsTest(parameterized.TestCase):
   def test_mj_contact_list(self):
     self.assertEmpty(self.data.contact)
 
-    expected_ncon = 1234
-    self.data.ncon = expected_ncon
+    expected_ncon = 4
+    mujoco.mj_forward(self.model, self.data)
     self.assertLen(self.data.contact, expected_ncon)
 
     expected_pos = []

python/mujoco/errors.h

@@ -21,16 +21,11 @@
 #include <type_traits>
 
 #include <mujoco/mjexport.h>
+#include "private.h"
 #include "util/crossplatform.h"
 #include "util/func_wrap.h"
 #include <pybind11/pybind11.h>
 
-// DO NOT USE THIS FUNCTION ELSEWHERE.
-// It should be regarded as part of MuJoCo's internal implementation detail.
-extern "C" {
-MJAPI void _mjPRIVATE__set_tls_error_fn(void (*h)(const char*));
-}
-
 // When building for Linux and statically linking against a "hermetic" libc++abi
 // (i.e. where libc++/libc++abi symbols all have "hidden" visibility), exception
 // types do not propagate correctly across shared library boundaries.

python/mujoco/private.h

@@ -0,0 +1,28 @@
+// Copyright 2023 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_PYTHON_PRIVATE_H_
+#define MUJOCO_PYTHON_PRIVATE_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+
+// DO NOT USE THESE FUNCTIONS ELSEWHERE.
+// They should be regarded as part of MuJoCo's internal implementation detail.
+extern "C" {
+MJAPI void _mjPRIVATE__set_tls_error_fn(void (*h)(const char*));
+MJAPI void* mj_arenaAlloc(mjData* d, int bytes, int alignment);
+}
+
+#endif  // MUJOCO_PYTHON_PRIVATE_H_

python/mujoco/structs.cc

@@ -38,6 +38,7 @@
 #include "function_traits.h"
 #include "indexers.h"
 #include "mjdata_meta.h"
+#include "private.h"
 #include "raw.h"
 #include "serialization.h"
 #include <pybind11/numpy.h>
@@ -695,7 +696,6 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
 
   // Write struct and scalar fields
 #define X(var) WriteBytes(output, &ptr_->var, sizeof(ptr_->var))
-  X(parena);
   X(maxuse_stack);
   X(maxuse_arena);
   X(maxuse_con);
@@ -727,7 +727,6 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
 #define MJ_D(x) this->ptr_->x
 #define X(type, name, nr, nc)                                   \
   if ((nr) * (nc)) {                                            \
-    WriteInt(output, PTRDIFF(ptr_->name, ptr_->arena));         \
     WriteBytes(output, ptr_->name, sizeof(type) * (nr) * (nc)); \
   }
 
@@ -786,7 +785,6 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
   ReadBytes(input, (void*) &d->var, sizeof(d->var)); \
   CheckInput(input, "mjData");
 
-  X(parena);
   X(maxuse_stack);
   X(maxuse_arena);
   X(maxuse_con);
@@ -816,11 +814,11 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
 #define MJ_M(x) m.x
 #undef MJ_D
 #define MJ_D(x) d->x
-#define X(type, name, nr, nc)                              \
-  if ((nr) * (nc)) {                                       \
-    d->name = reinterpret_cast<decltype(d->name)>(         \
-        static_cast<char*>(d->arena) + ReadInt(input));    \
-    ReadBytes(input, d->name, sizeof(type) * (nr) * (nc)); \
+#define X(type, name, nr, nc)                                          \
+  if ((nr) * (nc)) {                                                   \
+    d->name = static_cast<decltype(d->name)>(                          \
+        mj_arenaAlloc(d, sizeof(type) * (nr) * (nc), alignof(type)));  \
+    ReadBytes(input, d->name, sizeof(type) * (nr) * (nc));             \
   }
 
     MJDATA_ARENA_POINTERS_CONTACT

src/engine/engine_collision_box.c

@@ -15,6 +15,7 @@
 #include <math.h>
 #include <string.h>
 
+#include <mujoco/mjmacro.h>
 #include "engine/engine_collision_primitive.h"
 #include "engine/engine_util_blas.h"
 

src/engine/engine_collision_convex.c

@@ -20,9 +20,9 @@
 #include <ccd/vec3.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_collision_primitive.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"

src/engine/engine_collision_driver.c

@@ -18,6 +18,7 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_callback.h"
 #include "engine/engine_collision_convex.h"
@@ -25,7 +26,6 @@
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_crossplatform.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"

src/engine/engine_collision_primitive.c

@@ -17,6 +17,7 @@
 #include <math.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_spatial.h"

src/engine/engine_core_constraint.c

@@ -19,12 +19,12 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjxmacro.h>
 #include "engine/engine_array_safety.h"
 #include "engine/engine_core_smooth.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
@@ -32,6 +32,10 @@
 #include "engine/engine_util_sparse.h"
 #include "engine/engine_util_spatial.h"
 
+#ifdef MEMORY_SANITIZER
+  #include <sanitizer/msan_interface.h>
+#endif
+
 #ifdef mjUSEPLATFORMSIMD
   #if defined(__AVX__) && defined(mjUSEDOUBLE)
     #define mjUSEAVX
@@ -127,6 +131,10 @@ int mj_addContact(const mjModel* m, mjData* d, const mjContact* con) {
 
   // move arena pointer back to the end of the existing contact array and invalidate efc_ arrays
   d->parena = d->ncon * sizeof(mjContact);
+#ifdef ADDRESS_SANITIZER
+  ASAN_POISON_MEMORY_REGION(
+      (char*)d->arena + d->parena, (d->nstack - d->pstack) * sizeof(mjtNum) - d->parena);
+#endif
   clearEfc(d);
 
   // copy contact
@@ -1580,6 +1588,10 @@ void mj_makeConstraint(const mjModel* m, mjData* d) {
 
   // move arena pointer to end of contact array
   d->parena = d->ncon * sizeof(mjContact);
+#ifdef ADDRESS_SANITIZER
+  ASAN_POISON_MEMORY_REGION(
+      (char*)d->arena + d->parena, (d->nstack - d->pstack) * sizeof(mjtNum) - d->parena);
+#endif
 
 #define X(type, name, nr, nc)                                             \
   d->name = mj_arenaAlloc(d, sizeof(type) * (nr) * (nc), _Alignof(type)); \

src/engine/engine_core_smooth.c

@@ -18,11 +18,11 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_crossplatform.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"

src/engine/engine_derivative.c

@@ -15,11 +15,11 @@
 #include "engine/engine_derivative.h"
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_crossplatform.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_passive.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"

src/engine/engine_derivative_fd.c

@@ -18,6 +18,7 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_forward.h"
 #include "engine/engine_io.h"

src/engine/engine_forward.c

@@ -18,6 +18,7 @@
 #include <stdio.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjplugin.h>
 #include "engine/engine_callback.h"

src/engine/engine_inverse.c

@@ -17,6 +17,7 @@
 #include <stddef.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_collision_driver.h"
 #include "engine/engine_core_constraint.h"

src/engine/engine_io.c

@@ -22,6 +22,7 @@
 #include <string.h>
 
 #include <mujoco/mjmodel.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjplugin.h>
 #include <mujoco/mjxmacro.h>
 #include "engine/engine_array_safety.h"
@@ -33,6 +34,10 @@
 #include "engine/engine_util_misc.h"
 #include "engine/engine_vfs.h"
 
+#ifdef MEMORY_SANITIZER
+  #include <sanitizer/msan_interface.h>
+#endif
+
 #ifdef _MSC_VER
   #pragma warning (disable: 4305)  // disable MSVC warning: truncation from 'double' to 'float'
 #endif
@@ -1163,6 +1168,15 @@ void* mj_arenaAlloc(mjData* d, int bytes, int alignment) {
   void* result = (char*)d->arena + d->parena + padding;
   d->parena += padding + bytes;
   d->maxuse_arena = mjMAX(d->maxuse_arena, d->pstack*sizeof(mjtNum) + d->parena);
+
+#ifdef ADDRESS_SANITIZER
+  ASAN_UNPOISON_MEMORY_REGION(result, bytes);
+#endif
+
+#ifdef MEMORY_SANITIZER
+  __msan_allocated_memory(result, bytes);
+#endif
+
   return result;
 }
 
@@ -1172,12 +1186,19 @@ void* mj_arenaAlloc(mjData* d, int bytes, int alignment) {
 mjtNum* mj_stackAlloc(mjData* d, int size) {
   // return NULL if empty
   if (!size) {
-    return 0;
+    return NULL;
   }
 
+  // add red zone padding when built with asan, to detect out-of-bound accesses
+#ifdef ADDRESS_SANITIZER
+  #define mjREDZONE 4
+#else
+  #define mjREDZONE 0
+#endif
+
   // check size
   size_t stack_available_bytes = d->nstack * sizeof(mjtNum) - d->parena;
-  size_t stack_required_bytes = (d->pstack + size) * sizeof(mjtNum);
+  size_t stack_required_bytes = (d->pstack + size + 2*mjREDZONE) * sizeof(mjtNum);
   if (stack_required_bytes > stack_available_bytes) {
     mju_error("stack overflow: max = %zu, available = %zu, requested = %zu "
               "(ne = %d, nf = %d, nefc = %d, ncon = %d)",
@@ -1187,18 +1208,45 @@ mjtNum* mj_stackAlloc(mjData* d, int size) {
 
   // allocate at end of arena
   char* end_ptr = (char*)d->arena + d->nstack * sizeof(mjtNum);
-  char* result = end_ptr - (d->pstack + size + 1) * sizeof(mjtNum);
+  char* result = end_ptr - (d->pstack + size + mjREDZONE) * sizeof(mjtNum);
+  size_t new_pstack = d->pstack + size + 2*mjREDZONE;
+  #undef mjREDZONE
+
+  // new stack usage level
+  size_t usage;
 
 #ifdef ADDRESS_SANITIZER
   if ((uintptr_t)result % sizeof(mjtNum)) {
     mju_error("mj_stackAlloc fails to align to sizeof(mjtNum)");
   }
+
+  // actual stack usage (without red zone bytes) is stored in the red zone
+  if (d->pstack) {
+    size_t* prev_ptr = (size_t*)(end_ptr - d->pstack*sizeof(mjtNum));
+    ASAN_UNPOISON_MEMORY_REGION(prev_ptr, sizeof(size_t));
+    usage = *prev_ptr + size;
+    ASAN_POISON_MEMORY_REGION(prev_ptr, sizeof(size_t));
+  } else {
+    usage = size;
+  }
+
+  // store new stack usage in the red zone
+  size_t* cur_ptr = (size_t*)(end_ptr - new_pstack*sizeof(mjtNum));
+  ASAN_UNPOISON_MEMORY_REGION(cur_ptr, sizeof(size_t));
+  *cur_ptr = usage;
+  ASAN_POISON_MEMORY_REGION(cur_ptr, sizeof(size_t));
+
+  // unpoison the actual usable allocation
+  ASAN_UNPOISON_MEMORY_REGION(result, size*sizeof(mjtNum));
+#else
+  usage = d->pstack + size;
 #endif
 
-  // update max, return pointer to buffer
-  d->pstack += size;
-  d->maxuse_stack = mjMAX(d->maxuse_stack, d->pstack);
-  d->maxuse_arena = mjMAX(d->maxuse_arena, d->pstack*sizeof(mjtNum) + d->parena);
+  // update pstack and max usage statistics
+  d->pstack = new_pstack;
+  d->maxuse_stack = mjMAX(d->maxuse_stack, usage);
+  d->maxuse_arena = mjMAX(d->maxuse_arena, usage*sizeof(mjtNum) + d->parena);
+
   return (mjtNum*)result;
 }
 
@@ -1232,6 +1280,16 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
 
   // clear arena pointers
   d->parena = 0;
+
+  // poison the entire arena+stack memory region when built with asan
+#ifdef ADDRESS_SANITIZER
+  ASAN_POISON_MEMORY_REGION(d->arena, d->nstack * sizeof(mjtNum));
+#endif
+
+#ifdef MEMORY_SANITIZER
+  __msan_allocated_memory(d->arena, d->nstack * sizeof(mjtNum));
+#endif
+
 #define X(type, name, nr, nc) d->name = NULL;
   MJDATA_ARENA_POINTERS
 #undef X

src/engine/engine_io.h

@@ -100,7 +100,7 @@ MJAPI void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_va
 MJAPI void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
 
 // mjData arena allocate
-void* mj_arenaAlloc(mjData* d, int bytes, int alignment);
+MJAPI void* mj_arenaAlloc(mjData* d, int bytes, int alignment);
 
 // mjData stack allocate for array of mjtNums
 MJAPI mjtNum* mj_stackAlloc(mjData* d, int size);

src/engine/engine_macro.h

@@ -17,22 +17,11 @@
 
 #include <stdint.h>
 
+#include <mujoco/mjmacro.h>
 #include "engine/engine_callback.h"  // IWYU pragma: export
 
 //-------------------------------- utility macros --------------------------------------------------
 
-// mark and free stack
-#define mjMARKSTACK int _mark = d->pstack;
-#define mjFREESTACK d->pstack = _mark;
-
-// check bitflag
-#define mjDISABLED(x) (m->opt.disableflags & (x))
-#define mjENABLED(x)  (m->opt.enableflags & (x))
-
-// max and min macros
-#define mjMAX(a, b)    (((a) > (b)) ? (a) : (b))
-#define mjMIN(a, b)    (((a) < (b)) ? (a) : (b))
-
 // thread local macro
 #ifdef _MSC_VER
 #define mjTHREADLOCAL __declspec(thread)
@@ -49,21 +38,7 @@
 #define TM_START1 mjtNum _tm1 = (mjcb_time ? mjcb_time() : 0);
 #define TM_END1(i) {d->timer[i].duration += ((mjcb_time ? mjcb_time() : 0) - _tm1); d->timer[i].number++;}
 
-//-------------------------- sanitizer macros ------------------------------------------------------
-
-#ifdef ADDRESS_SANITIZER
-  #include <sanitizer/asan_interface.h>
-#elif defined(_MSC_VER)
-  #define ASAN_POISON_MEMORY_REGION(addr, size)
-  #define ASAN_UNPOISON_MEMORY_REGION(addr, size)
-#else
-  #define ASAN_POISON_MEMORY_REGION(addr, size) ((void)(addr), (void)(size))
-  #define ASAN_UNPOISON_MEMORY_REGION(addr, size) ((void)(addr), (void)(size))
-#endif
-
-#ifdef MEMORY_SANITIZER
-  #include <sanitizer/msan_interface.h>
-#endif
+//-------------------------- compiler builtin ------------------------------------------------------
 
 #ifndef __has_builtin
 #define __has_builtin(x) 0

src/engine/engine_passive.c

@@ -18,11 +18,11 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_callback.h"
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_crossplatform.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"

src/engine/engine_print.c

@@ -20,11 +20,11 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjxmacro.h>
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"
@@ -770,7 +770,7 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
 
   fprintf(fp, "SIZES\n");
 #define X(type, name)                                                         \
-  {                                                                           \
+  if (strcmp(#name, "pstack") != 0 && strcmp(#name, "parena") != 0) {         \
     const char* format = _Generic(                                            \
         d->name,                                                              \
         int : INT_FORMAT,                                                     \

src/engine/engine_ray.c

@@ -18,10 +18,10 @@
 #include <stddef.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjvisualize.h>
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"

src/engine/engine_sensor.c

@@ -17,13 +17,13 @@
 #include <stddef.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjplugin.h>
 #include "engine/engine_callback.h"
 #include "engine/engine_core_smooth.h"
 #include "engine/engine_crossplatform.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_ray.h"
 #include "engine/engine_support.h"

src/engine/engine_setconst.c

@@ -17,12 +17,12 @@
 #include <stdio.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_core_smooth.h"
 #include "engine/engine_forward.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"

src/engine/engine_solver.c

@@ -18,11 +18,11 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_core_smooth.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"

src/engine/engine_support.c

@@ -17,12 +17,12 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_array_safety.h"
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_crossplatform.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"

src/engine/engine_util_blas.c

@@ -16,6 +16,7 @@
 
 #include <string.h>
 
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjtnum.h>
 
 #ifdef mjUSEPLATFORMSIMD

src/engine/engine_util_errmem.c

@@ -27,6 +27,7 @@
 
 #include "engine/engine_array_safety.h"
 #include "engine/engine_macro.h"
+
 //------------------------- cross-platform aligned malloc/free -------------------------------------
 
 static inline void* mju_alignedMalloc(size_t size, size_t align) {

src/engine/engine_util_errmem.h

@@ -18,6 +18,7 @@
 #include <stddef.h>
 
 #include <mujoco/mjexport.h>
+#include <mujoco/mjmacro.h>
 
 #ifdef __cplusplus
 extern "C" {

src/engine/engine_util_misc.c

@@ -20,6 +20,7 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_array_safety.h"
 #include "engine/engine_macro.h"

src/engine/engine_util_solve.c

@@ -19,8 +19,8 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"

src/engine/engine_util_sparse.c

@@ -18,9 +18,9 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjtnum.h>
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_util_blas.h"
 
 

src/engine/engine_vis_init.c

@@ -17,10 +17,10 @@
 #include <math.h>
 #include <string.h>
 
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjvisualize.h>
 #include "engine/engine_array_safety.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"
 

src/engine/engine_vis_interact.c

@@ -19,11 +19,11 @@
 
 #include <mujoco/mjdata.h>
 #include <mujoco/mjexport.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjvisualize.h>
 #include "engine/engine_core_smooth.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_ray.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"

src/engine/engine_vis_state.c

@@ -22,7 +22,6 @@
 #include <mujoco/mjvisualize.h>
 #include <mujoco/mjxmacro.h>
 #include "engine/engine_core_constraint.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_errmem.h"
@@ -30,6 +29,10 @@
 #include "engine/engine_vis_interact.h"
 #include "engine/engine_vis_visualize.h"
 
+#ifdef MEMORY_SANITIZER
+  #include <sanitizer/msan_interface.h>
+#endif
+
 // this source file needs to treat XMJV differently from other X macros
 #undef XMJV
 

src/engine/engine_vis_visualize.c

@@ -18,10 +18,10 @@
 #include <string.h>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjvisualize.h>
 #include "engine/engine_array_safety.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"

src/render/render_context.c

@@ -19,6 +19,7 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjvisualize.h>
 #include <mujoco/mujoco.h>
 #include "render/render_util.h"

src/render/render_gl2.c

@@ -17,6 +17,7 @@
 #include <string.h>
 #include <stdio.h>
 
+#include <mujoco/mjmacro.h>
 #include <mujoco/mujoco.h>
 #include "engine/engine_array_safety.h"
 #include "render/glad/glad.h"

src/render/render_gl3.c

@@ -19,6 +19,7 @@
 #include <string.h>
 #include <stdlib.h>
 
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjvisualize.h>
 #include <mujoco/mujoco.h>
 #include "engine/engine_array_safety.h"

src/ui/ui_main.c

@@ -17,6 +17,7 @@
 #include <stdio.h>
 #include <string.h>
 
+#include <mujoco/mjmacro.h>
 #include <mujoco/mujoco.h>
 #include "engine/engine_array_safety.h"
 #include "render/glad/glad.h"

src/user/user_composite.cc

@@ -23,6 +23,7 @@
 #include <utility>
 #include <vector>
 
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "cc/array_safety.h"
 #include "engine/engine_io.h"

src/user/user_mesh.cc

@@ -26,11 +26,11 @@
 #define TINYOBJLOADER_IMPLEMENTATION
 #endif
 
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "cc/array_safety.h"
 #include "engine/engine_crossplatform.h"
 #include "engine/engine_resource.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_solve.h"

src/user/user_model.cc

@@ -23,13 +23,13 @@
 #include <vector>
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjplugin.h>
 #include <mujoco/mjvisualize.h>
 #include "cc/array_safety.h"
 #include "engine/engine_forward.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_setconst.h"
 #include "engine/engine_resource.h"

src/user/user_objects.cc

@@ -24,6 +24,7 @@
 #include <vector>
 
 #include "lodepng.h"
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjplugin.h>
 #include "cc/array_safety.h"
@@ -31,7 +32,6 @@
 #include "engine/engine_crossplatform.h"
 #include "engine/engine_resource.h"
 #include "engine/engine_io.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_passive.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_util_blas.h"

src/user/user_util.cc

@@ -21,9 +21,9 @@
 #include <string>
 #include <string_view>
 
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjtnum.h>
 #include <mujoco/mujoco.h>
-#include "engine/engine_macro.h"
 #include "engine/engine_util_spatial.h"
 
 using std::isnan;

src/xml/xml_native_reader.cc

@@ -30,9 +30,9 @@
 #include <utility>
 #include <vector>
 
+#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjvisualize.h>
-#include "engine/engine_macro.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"

test/benchmark/engine_util_sparse_benchmark_test.cc

@@ -422,6 +422,7 @@ BENCHMARK(BM_MatVecSparse_1);
 
 static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
   static mjModel* m = LoadModelFromPath("humanoid/humanoid.xml");
+  m->opt.jacobian = mjJAC_SPARSE;
 
   mjData* d = mj_makeData(m);
 

test/sample/testspeed_test.sh

@@ -25,10 +25,15 @@ test_model() {
   local model="$1"
   echo "Testing $model" >&2
 
+  local iterations=10
+  if [[ "$model" == */composite/particle.xml && ${TESTSPEED_ASAN:-0} != 0 ]]; then
+    iterations=2
+  fi
+
   # run testspeed, writing its output to stderr.
   # die if testspeed returns a failure code, or if it doesn't have the string
   # "Simulation time" in the output.
-  ("$TARGET_BINARY" "$model" 10 || die "testspeed failed") \
+  ("$TARGET_BINARY" "$model" "$iterations" || die "testspeed failed") \
       | tee >(cat 1>&2) | grep -q "$EXPECTED_STR"
 
   if [ "$?" != 0 ]; then

test/xml/xml_native_reader_test.cc

@@ -43,12 +43,12 @@ TEST_F(XMLReaderTest, MemorySize) {
   {
     static constexpr char xml[] = R"(
     <mujoco>
-      <size memory="128"/>
+      <size memory="256"/>
     </mujoco>
     )";
     mjModel* model = LoadModelFromString(xml, error.data(), error.size());
-    ASSERT_THAT(model, NotNull());
-    EXPECT_EQ(model->nstack, 128 / sizeof(mjtNum));
+    ASSERT_THAT(model, NotNull()) << error.data();
+    EXPECT_EQ(model->nstack, 256 / sizeof(mjtNum));
     mj_deleteModel(model);
   }
   {
@@ -58,7 +58,7 @@ TEST_F(XMLReaderTest, MemorySize) {
     </mujoco>
     )";
     mjModel* model = LoadModelFromString(xml, error.data(), error.size());
-    ASSERT_THAT(model, NotNull());
+    ASSERT_THAT(model, NotNull()) << error.data();
     EXPECT_EQ(model->nstack, 1024 / sizeof(mjtNum));
     mj_deleteModel(model);
   }
@@ -69,7 +69,7 @@ TEST_F(XMLReaderTest, MemorySize) {
     </mujoco>
     )";
     mjModel* model = LoadModelFromString(xml, error.data(), error.size());
-    ASSERT_THAT(model, NotNull());
+    ASSERT_THAT(model, NotNull()) << error.data();
     EXPECT_EQ(model->nstack, 10240 / sizeof(mjtNum));
     mj_deleteModel(model);
   }
@@ -80,7 +80,7 @@ TEST_F(XMLReaderTest, MemorySize) {
     </mujoco>
     )";
     mjModel* model = LoadModelFromString(xml, error.data(), error.size());
-    ASSERT_THAT(model, NotNull());
+    ASSERT_THAT(model, NotNull()) << error.data();
     EXPECT_EQ(model->nstack, 4*1024*1024 / sizeof(mjtNum));
     mj_deleteModel(model);
   }
@@ -91,7 +91,7 @@ TEST_F(XMLReaderTest, MemorySize) {
     </mujoco>
     )";
     mjModel* model = LoadModelFromString(xml, error.data(), error.size());
-    ASSERT_THAT(model, NotNull());
+    ASSERT_THAT(model, NotNull()) << error.data();
     EXPECT_EQ(model->nstack, 1024*1024*1024 / sizeof(mjtNum));
     mj_deleteModel(model);
   }
@@ -102,7 +102,7 @@ TEST_F(XMLReaderTest, MemorySize) {
     </mujoco>
     )";
     mjModel* model = LoadModelFromString(xml, error.data(), error.size());
-    ASSERT_THAT(model, NotNull());
+    ASSERT_THAT(model, NotNull()) << error.data();
     EXPECT_EQ(model->nstack, 1024*1024*1024 / sizeof(mjtNum));
     mj_deleteModel(model);
   }

unity/Runtime/Bindings/MjBindings.cs

@@ -32,6 +32,7 @@ public const bool MUJOCO_HELPER_DLL_IMPORT = true;
 public const bool MUJOCO_HELPER_DLL_EXPORT = true;
 public const bool MJAPI = true;
 public const bool MJLOCAL = true;
+public const bool THIRD_PARTY_MUJOCO_MJMACRO_H_ = true;
 public const bool THIRD_PARTY_MUJOCO_MJMODEL_H_ = true;
 public const double mjPI = 3.141592653589793;
 public const double mjMAXVAL = 10000000000.0;