Add mj_markStack and mj_freeStack as public API functions.

Also add asan instrumentation to detect stack frame leakages (i.e. `mj_markStack` without a corresponding `mj_freeStack` in the same caller function).

PiperOrigin-RevId: 562625645
Change-Id: I4e3ff66ca0b9d08ed0a95cef45393db8e3053e22

include/mujoco/mjdata.h

@@ -157,6 +157,7 @@ struct mjData_ {
 
   // stack pointer
   size_t  pstack;            // first available mjtNum address in stack
+  size_t  pbase;             // value of pstack when mj_markStack was last called
 
   // arena pointer
   size_t  parena;            // first available byte in arena

include/mujoco/mjmacro.h

@@ -15,8 +15,6 @@
 #ifndef MUJOCO_MJMACRO_H_
 #define MUJOCO_MJMACRO_H_
 
-#include <stddef.h>
-
 // include asan interface header, or provide stubs for poison/unpoison macros when not using asan
 #ifdef ADDRESS_SANITIZER
   #include <sanitizer/asan_interface.h>
@@ -33,8 +31,8 @@
 #define mjMIN(a, b) (((a) < (b)) ? (a) : (b))
 
 // mjData stack frame management
-#define mjMARKSTACK   size_t _mark = d->pstack;
-#define mjFREESTACK   d->pstack = _mark;
+#define mjMARKSTACK   mj_markStack(d);
+#define mjFREESTACK   mj_freeStack(d);
 
 // return current value of mjOption enable/disable flags
 #define mjDISABLED(x) (m->opt.disableflags & (x))
@@ -49,15 +47,4 @@
   #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->narena - d->pstack - d->parena);                      \
-  }
-#endif
-
 #endif  // MUJOCO_MJMACRO_H_

include/mujoco/mjxmacro.h

@@ -625,6 +625,7 @@
     X( size_t,    nbuffer            ) \
     X( int,       nplugin            ) \
     X( size_t,    pstack             ) \
+    X( size_t,    pbase              ) \
     X( size_t,    parena             ) \
     X( size_t,    maxuse_stack       ) \
     X( size_t,    maxuse_arena       ) \

include/mujoco/mujoco.h

@@ -30,10 +30,6 @@ 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>
@@ -188,6 +184,13 @@ MJAPI void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_va
 // Reset data, set fields from specified keyframe.
 MJAPI void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
 
+// Mark a new frame on the mjData stack.
+MJAPI void mj_markStack(mjData* d);
+
+// Free the current mjData stack frame. All pointers returned by mj_stackAlloc since the last call
+// to mj_markStack must no longer be used afterwards.
+MJAPI void mj_freeStack(mjData* d);
+
 // Allocate a number of bytes on mjData stack at a specific alignment.
 // Call mju_error on stack overflow.
 MJAPI void* mj_stackAlloc(mjData* d, size_t bytes, size_t alignment);