Rename nstack to narena, change it and pstack to count bytes.

PiperOrigin-RevId: 561066311
Change-Id: I4d0e645e6c254dd906bfc42159e0e99149a7c6c8

doc/XMLreference.rst

@@ -450,7 +450,7 @@ compilation.
 
 :at:`nstack`: :at-val:`int, "-1"` |nbsp| |nbsp| |nbsp| (legacy)
    This is a deprecated legacy attribute. In versions prior to 2.3.0, it determined the maximum size of the
-   :ref:`stack <siStack>`. After version 2.3.0, if :at:`nstack` is specified, then the size of ``mjData.arena`` is
+   :ref:`stack <siStack>`. After version 2.3.0, if :at:`nstack` is specified, then the size of ``mjData.narena`` is
    ``nstack * sizeof(mjtNum)`` bytes, plus an additional space for the constraint solver. Specifying both :at:`nstack`
    and :at:`memory` leads to an error.
 

doc/changelog.rst

@@ -39,6 +39,8 @@ General
 #. Added the flag :ref:`invdiscrete<option-flag-invdiscrete>`, which enables discrete-time inverse dynamics for all
    :ref:`integrators<option-integrator>` other than ``RK4``. See the flag documentation for more details.
 #. Renamed the function ``mj_stackAlloc`` to ``mj_stackAllocNum``.
+#. Renamed the ``nstack`` field in ``mjModel`` and ``mjData`` to ``narena``. Changed ``narena``, ``pstack``, and
+   ``maxuse_stack`` to count number of bytes rather than number of ``mjtNum``s.
 
 Python bindings
 ^^^^^^^^^^^^^^^

doc/includes/references.h

@@ -124,8 +124,8 @@ struct mjSolverStat_ {       // per-iteration solver statistics
 typedef struct mjSolverStat_ mjSolverStat;
 struct mjData_ {
   // constant sizes
-  int     nstack;            // number of mjtNums that can fit in the arena+stack space
+  size_t  narena;            // size of the arena in bytes (inclusive of the stack)
-  int     nbuffer;           // size of main buffer in bytes
+  size_t  nbuffer;           // size of main buffer in bytes
   int     nplugin;           // number of plugin instances
 
   // stack pointer
@@ -135,7 +135,7 @@ struct mjData_ {
   size_t  parena;            // first available byte in arena
 
   // memory utilization stats
-  int     maxuse_stack;      // maximum stack allocation
+  size_t  maxuse_stack;      // maximum stack allocation
   size_t  maxuse_arena;      // maximum arena allocation
   int     maxuse_con;        // maximum number of contacts
   int     maxuse_efc;        // maximum number of scalar constraints
@@ -881,12 +881,12 @@ struct mjModel_ {
   int nemax;                      // number of potential equality-constraint rows
   int njmax;                      // number of available rows in constraint Jacobian
   int nconmax;                    // number of potential contacts in contact list
-  int nstack;                     // number of fields in mjData stack
   int nuserdata;                  // number of extra fields in mjData
   int nsensordata;                // number of fields in sensor data vector
   int npluginstate;               // number of fields in plugin state vector
 
-  int nbuffer;                    // number of bytes in buffer
+  size_t narena;                  // number of bytes in the mjData arena (inclusive of stack)
+  size_t nbuffer;                 // number of bytes in buffer
 
   // ------------------------------- options and statistics
 

include/mujoco/mjdata.h

@@ -151,8 +151,8 @@ typedef struct mjSolverStat_ mjSolverStat;
 
 struct mjData_ {
   // constant sizes
-  int     nstack;            // number of mjtNums that can fit in the arena+stack space
+  size_t  narena;            // size of the arena in bytes (inclusive of the stack)
-  int     nbuffer;           // size of main buffer in bytes
+  size_t  nbuffer;           // size of main buffer in bytes
   int     nplugin;           // number of plugin instances
 
   // stack pointer
@@ -162,7 +162,7 @@ struct mjData_ {
   size_t  parena;            // first available byte in arena
 
   // memory utilization stats
-  int     maxuse_stack;      // maximum stack allocation
+  size_t  maxuse_stack;      // maximum stack allocation
   size_t  maxuse_arena;      // maximum arena allocation
   int     maxuse_con;        // maximum number of contacts
   int     maxuse_efc;        // maximum number of scalar constraints

include/mujoco/mjmacro.h

@@ -54,7 +54,7 @@
     d->pstack = _mark;                                           \
     ASAN_POISON_MEMORY_REGION(                                   \
         (char*)d->arena + d->parena,                             \
-        (d->nstack - d->pstack) * sizeof(mjtNum) - d->parena );  \
+        d->narena - d->pstack - d->parena);                      \
   }
 #endif
 

include/mujoco/mjmodel.h

@@ -15,6 +15,8 @@
 #ifndef MUJOCO_MJMODEL_H_
 #define MUJOCO_MJMODEL_H_
 
+#include <stddef.h>
+
 #include <mujoco/mjtnum.h>
 
 // global constants
@@ -623,12 +625,12 @@ struct mjModel_ {
   int nemax;                      // number of potential equality-constraint rows
   int njmax;                      // number of available rows in constraint Jacobian
   int nconmax;                    // number of potential contacts in contact list
-  int nstack;                     // number of fields in mjData stack
   int nuserdata;                  // number of extra fields in mjData
   int nsensordata;                // number of fields in sensor data vector
   int npluginstate;               // number of fields in plugin state vector
 
-  int nbuffer;                    // number of bytes in buffer
+  size_t narena;                  // number of bytes in the mjData arena (inclusive of stack)
+  size_t nbuffer;                 // number of bytes in buffer
 
   // ------------------------------- options and statistics
 

include/mujoco/mjxmacro.h

@@ -125,10 +125,10 @@
     X   ( njmax )              \
     X   ( nconmax )            \
     XMJV( ntree )              \
-    X   ( nstack )             \
     X   ( nuserdata )          \
     XMJV( nsensordata )        \
     X   ( npluginstate )       \
+    X   ( narena )             \
     X   ( nbuffer )
 
 
@@ -620,13 +620,13 @@
 
 // scalar fields of mjData
 #define MJDATA_SCALAR                  \
-    X( int,       nstack             ) \
+    X( size_t,    narena             ) \
-    X( int,       nbuffer            ) \
+    X( size_t,    nbuffer            ) \
     X( int,       nplugin            ) \
-    X( int,       pstack             ) \
+    X( size_t,    pstack             ) \
-    X( int,       parena             ) \
+    X( size_t,    parena             ) \
-    X( int,       maxuse_stack       ) \
+    X( size_t,    maxuse_stack       ) \
-    X( int,       maxuse_arena       ) \
+    X( size_t,    maxuse_arena       ) \
     X( int,       maxuse_con         ) \
     X( int,       maxuse_efc         ) \
     X( int,       solver_iter        ) \

introspect/structs.py

@@ -1100,11 +1100,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  type=ValueType(name='int'),
                  doc='number of potential contacts in contact list',
              ),
-             StructFieldDecl(
-                 name='nstack',
-                 type=ValueType(name='int'),
-                 doc='number of fields in mjData stack',
-             ),
              StructFieldDecl(
                  name='nuserdata',
                  type=ValueType(name='int'),
@@ -1120,9 +1115,14 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  type=ValueType(name='int'),
                  doc='number of fields in plugin state vector',
              ),
+             StructFieldDecl(
+                 name='narena',
+                 type=ValueType(name='size_t'),
+                 doc='number of bytes in the mjData arena (inclusive of stack)',
+             ),
              StructFieldDecl(
                  name='nbuffer',
-                 type=ValueType(name='int'),
+                 type=ValueType(name='size_t'),
                  doc='number of bytes in buffer',
              ),
              StructFieldDecl(
@@ -3483,13 +3483,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([
          declname='struct mjData_',
          fields=(
              StructFieldDecl(
-                 name='nstack',
+                 name='narena',
-                 type=ValueType(name='int'),
+                 type=ValueType(name='size_t'),
-                 doc='number of mjtNums that can fit in the arena+stack space',
+                 doc='size of the arena in bytes (inclusive of the stack)',
              ),
              StructFieldDecl(
                  name='nbuffer',
-                 type=ValueType(name='int'),
+                 type=ValueType(name='size_t'),
                  doc='size of main buffer in bytes',
              ),
              StructFieldDecl(
@@ -3509,7 +3509,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              ),
              StructFieldDecl(
                  name='maxuse_stack',
-                 type=ValueType(name='int'),
+                 type=ValueType(name='size_t'),
                  doc='maximum stack allocation',
              ),
              StructFieldDecl(

python/mujoco/bindings_test.py

@@ -946,7 +946,7 @@ Euler integrator, semi-implicit in velocity.
                      mujoco.mjtDisableBit.mjDSBL_CONTACT)
 
   def test_can_raise_error(self):
-    self.data.pstack = self.data.nstack
+    self.data.pstack = self.data.narena
     with self.assertRaisesRegex(mujoco.FatalError,
                                 r'\Amj_stackAllocByte: stack overflow'):
       mujoco.mj_forward(self.model, self.data)

python/mujoco/functions.cc

@@ -1341,7 +1341,7 @@ PYBIND11_MODULE(_functions, pymodule) {
 #ifdef ADDRESS_SANITIZER
         ASAN_POISON_MEMORY_REGION(
             static_cast<char*>(data->arena),
-            (data->nstack - data->pstack) * sizeof(mjtNum));
+            data->narena - data->pstack);
 #endif
           data->ncon = 0;
           data->nefc = 0;

python/mujoco/mjdata_meta.h

@@ -88,7 +88,7 @@ struct MjDataMetadata {
         is_dual(mj_isDual(m)) {
   }
 
-#define X(var) int var;
+#define X(var) decltype(raw::MjModel::var) var;
   MJMODEL_INTS
 #undef X
 #define X(type, var, n) std::shared_ptr<type[]> var;

python/mujoco/structs.cc

@@ -78,7 +78,7 @@ constexpr auto XArrayShapeImpl(const std::string_view dim1_str) {
 }
 
 inline std::size_t NConMax(const mjData* d) {
-  return d->nstack * sizeof(mjtNum) / sizeof(mjContact);
+  return d->narena / sizeof(mjContact);
 }
 }  // namespace
 
@@ -459,7 +459,7 @@ MjModelWrapper MjModelWrapper::Deserialize(std::istream& input) {
     throw py::value_error("Incompatible serialization version.");
   }
 
-  int model_size = ReadInt(input);
+  std::size_t model_size = ReadInt(input);
   CheckInput(input, "mjModel");
   if (model_size < 0) {
     throw py::value_error("Invalid serialized mjModel.");

simulate/simulate.cc

@@ -525,8 +525,8 @@ void UpdateInfoText(mj::Simulate* sim, const mjModel* m, const mjData* d,
                    d->timer[mjTIMER_FORWARD].duration / mjMAX(1, d->timer[mjTIMER_FORWARD].number),
                    solerr, d->solver_iter,
                    fps,
-                   d->maxuse_arena/(double)(d->nstack * sizeof(mjtNum)),
+                   d->maxuse_arena/(double)(d->narena),
-                   mju_writeNumBytes(d->nstack * sizeof(mjtNum)));
+                   mju_writeNumBytes(d->narena));
 
   // add Energy if enabled
   {

src/engine/engine_core_constraint.c

@@ -71,13 +71,13 @@ static int arenaAllocEfc(const mjModel* m, mjData* d) {
   // poison remaining memory
 #ifdef ADDRESS_SANITIZER
   ASAN_POISON_MEMORY_REGION(
-    (char*)d->arena + d->parena, (d->nstack - d->pstack) * sizeof(mjtNum) - d->parena);
+    (char*)d->arena + d->parena, d->narena - d->pstack - d->parena);
 #endif
 
 #define X(type, name, nr, nc)                                             \
   d->name = mj_arenaAlloc(d, sizeof(type) * (nr) * (nc), _Alignof(type)); \
   if (!d->name) {                                                         \
-    mj_warning(d, mjWARN_CNSTRFULL, d->nstack * sizeof(mjtNum));          \
+    mj_warning(d, mjWARN_CNSTRFULL, d->narena);                           \
     clearEfc(d);                                                          \
     d->parena = d->ncon * sizeof(mjContact);                              \
     return 0;                                                             \
@@ -179,7 +179,7 @@ int mj_addContact(const mjModel* m, mjData* d, const mjContact* con) {
   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);
+    (char*)d->arena + d->parena, d->narena - d->pstack - d->parena);
 #endif
   clearEfc(d);
 

src/engine/engine_io.c

@@ -16,6 +16,7 @@
 #include "engine/engine_io.h"
 
 #include <limits.h>
+#include <stddef.h>
 #include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -251,16 +252,36 @@ void mj_defaultStatistic(mjStatistic* stat) {
 
 //----------------------------------- static utility functions -------------------------------------
 
-// id used to indentify binary mjModel file/buffer
+// id used to identify binary mjModel file/buffer
 static const int ID = 54321;
 
 
+// number of ints in the mjb header
+#define NHEADER 5
+
+
+// macro for referring to a mjModel member in generic expressions
+#define MJMODEL_MEMBER(name) (((mjModel*) NULL)->name)
+
 
 // count ints in mjModel
 static int getnint(void) {
   int cnt = 0;
 
-#define X(name) cnt++;
+#define X(name) cnt += _Generic(MJMODEL_MEMBER(name), int: 1, default: 0);
+  MJMODEL_INTS
+#undef X
+
+  return cnt;
+}
+
+
+
+// count size_t members in mjModel
+static int getnsize(void) {
+  int cnt = 0;
+
+#define X(name) cnt += _Generic(MJMODEL_MEMBER(name), size_t: 1, default: 0);
   MJMODEL_INTS
 #undef X
 
@@ -334,7 +355,7 @@ static inline unsigned int SKIP(intptr_t offset) {
 // set pointers in mjModel buffer
 static void mj_setPtrModel(mjModel* m) {
   char* ptr = (char*)m->buffer;
-  int sz;
+  ptrdiff_t sz;
 
   // prepare symbols needed by xmacro
   MJMODEL_POINTERS_PREAMBLE(m);
@@ -349,9 +370,9 @@ static void mj_setPtrModel(mjModel* m) {
 #undef X
 
   // check size
-  sz = (int)(ptr - (char*)m->buffer);
+  sz = ptr - (char*)m->buffer;
   if (m->nbuffer != sz) {
-    printf("expected size: %d,  actual size: %d\n", m->nbuffer, sz);
+    printf("expected size: %zu,  actual size: %zu\n", m->nbuffer, sz);
     mjERROR("mjModel buffer size mismatch");
   }
 }
@@ -362,7 +383,8 @@ static void mj_setPtrModel(mjModel* m) {
 // performs the following operations:
 // *nbuffer += SKIP(*offset) + type_size*nr*nc;
 // *offset += SKIP(*offset) + type_size*nr*nc;
-static int safeAddToBufferSize(intptr_t* offset, int* nbuffer, size_t type_size, int nr, int nc) {
+static int safeAddToBufferSize(intptr_t* offset, size_t* nbuffer,
+                               size_t type_size, int nr, int nc) {
   if (type_size < 0 || nr < 0 || nc < 0) {
     return 0;
   }
@@ -583,7 +605,7 @@ void mj_saveModel(const mjModel* m, const char* filename, void* buffer, int buff
   int ptrbuf = 0;
 
   // standard header
-  int header[4] = {ID, sizeof(mjtNum), getnint(), getnptr()};
+  int header[NHEADER] = {ID, sizeof(mjtNum), getnint(), getnsize(), getnptr()};
 
   // open file for writing if no buffer
   if (!buffer) {
@@ -596,8 +618,10 @@ void mj_saveModel(const mjModel* m, const char* filename, void* buffer, int buff
 
   // write standard header, info, options, buffer (omit pointers)
   if (fp) {
-    fwrite(header, sizeof(int), 4, fp);
+    fwrite(header, sizeof(int), NHEADER, fp);
-    fwrite(m, sizeof(int), getnint(), fp);
+    #define X(name) fwrite(&m->name, sizeof(m->name), 1, fp);
+    MJMODEL_INTS
+    #undef X
     fwrite((void*)&m->opt, sizeof(mjOption), 1, fp);
     fwrite((void*)&m->vis, sizeof(mjVisual), 1, fp);
     fwrite((void*)&m->stat, sizeof(mjStatistic), 1, fp);
@@ -609,8 +633,10 @@ void mj_saveModel(const mjModel* m, const char* filename, void* buffer, int buff
       #undef X
     }
   } else {
-    bufwrite(header, sizeof(int)*4, buffer_sz, buffer, &ptrbuf);
+    bufwrite(header, sizeof(int)*sizeof(header) / sizeof(int), buffer_sz, buffer, &ptrbuf);
-    bufwrite(m, sizeof(int)*getnint(), buffer_sz, buffer, &ptrbuf);
+    #define X(name) bufwrite(&m->name, sizeof(m->name), buffer_sz, buffer, &ptrbuf);
+    MJMODEL_INTS
+    #undef X
     bufwrite((void*)&m->opt, sizeof(mjOption), buffer_sz, buffer, &ptrbuf);
     bufwrite((void*)&m->vis, sizeof(mjVisual), buffer_sz, buffer, &ptrbuf);
     bufwrite((void*)&m->stat, sizeof(mjStatistic), buffer_sz, buffer, &ptrbuf);
@@ -632,9 +658,10 @@ void mj_saveModel(const mjModel* m, const char* filename, void* buffer, int buff
 
 // load model from binary MJB resource
 static mjModel* _mj_loadModel(const char* filename, int vfs_provider) {
-  int header[4] = {0};
+  int header[NHEADER] = {0};
-  int expected_header[4] = {ID, sizeof(mjtNum), getnint(), getnptr()};
+  int expected_header[NHEADER] = {ID, sizeof(mjtNum), getnint(), getnsize(), getnptr()};
-  int info[2000];
+  int ints[256];
+  size_t sizes[8];
   int ptrbuf = 0;
   mjModel *m = 0;
   mjResource* r = NULL;
@@ -650,16 +677,16 @@ static mjModel* _mj_loadModel(const char* filename, int vfs_provider) {
     return NULL;
   }
 
-  if (buffer_sz < 4*sizeof(int)) {
+  if (buffer_sz < NHEADER*sizeof(int)) {
     mju_warning("Model file has an incomplete header");
     mju_closeResource(r);
     return NULL;
   }
 
-  bufread(header, 4*sizeof(int), buffer_sz, buffer, &ptrbuf);
+  bufread(header, NHEADER*sizeof(int), buffer_sz, buffer, &ptrbuf);
 
   // check header
-  for (int i=0; i < 4; i++) {
+  for (int i=0; i < NHEADER; i++) {
     if (header[i] != expected_header[i]) {
       switch (i) {
       case 0:
@@ -677,6 +704,11 @@ static mjModel* _mj_loadModel(const char* filename, int vfs_provider) {
         mju_closeResource(r);
         return NULL;
 
+      case 3:
+        mju_warning("Model and executable have different number of size_t members in mjModel");
+        mju_closeResource(r);
+        return NULL;
+
       default:
         mju_warning("Model and executable have different number of pointers in mjModel");
         mju_closeResource(r);
@@ -686,26 +718,34 @@ static mjModel* _mj_loadModel(const char* filename, int vfs_provider) {
   }
 
   // read mjModel structure: info only
-  bufread(info, sizeof(int)*getnint(), buffer_sz, buffer, &ptrbuf);
+  bufread(ints, sizeof(int)*getnint(), buffer_sz, buffer, &ptrbuf);
+  bufread(sizes, sizeof(size_t)*getnsize(), buffer_sz, buffer, &ptrbuf);
 
   // allocate new mjModel, check sizes
-  m = mj_makeModel(info[0],  info[1],  info[2],  info[3],  info[4],  info[5],  info[6],
+  m = mj_makeModel(ints[0],  ints[1],  ints[2],  ints[3],  ints[4],  ints[5],  ints[6],
-                   info[7],  info[8],  info[9],  info[10], info[11], info[12], info[13],
+                   ints[7],  ints[8],  ints[9],  ints[10], ints[11], ints[12], ints[13],
-                   info[14], info[15], info[16], info[17], info[18], info[19], info[20],
+                   ints[14], ints[15], ints[16], ints[17], ints[18], ints[19], ints[20],
-                   info[21], info[22], info[23], info[24], info[25], info[26], info[27],
+                   ints[21], ints[22], ints[23], ints[24], ints[25], ints[26], ints[27],
-                   info[28], info[29], info[30], info[31], info[32], info[33], info[34],
+                   ints[28], ints[29], ints[30], ints[31], ints[32], ints[33], ints[34],
-                   info[35], info[36], info[37], info[38], info[39], info[40], info[41],
+                   ints[35], ints[36], ints[37], ints[38], ints[39], ints[40], ints[41],
-                   info[42], info[43], info[44], info[45], info[46], info[47], info[48],
+                   ints[42], ints[43], ints[44], ints[45], ints[46], ints[47], ints[48],
-                   info[49], info[50], info[51], info[52]);
+                   ints[49], ints[50], ints[51], ints[52]);
-  if (!m || m->nbuffer != info[getnint()-1]) {
+  if (!m || m->nbuffer != sizes[getnsize()-1]) {
     mju_closeResource(r);
     mju_warning("Corrupted model, wrong size parameters");
     mj_deleteModel(m);
     return NULL;
   }
 
-  // set info fields
+  // set integer fields
-  memcpy(m, info, sizeof(int)*getnint());
+  {
+    int int_idx = 0;
+    int size_idx = 0;
+    #define X(name) \
+        m->name = _Generic(m->name, size_t: sizes[size_idx++], default: ints[int_idx++]);
+    MJMODEL_INTS
+    #undef X
+  }
 
   // read options and buffer
   bufread((void*)&m->opt, sizeof(mjOption), buffer_sz, buffer, &ptrbuf);
@@ -775,7 +815,8 @@ void mj_deleteModel(mjModel* m) {
 // size of buffer needed to hold model
 int mj_sizeModel(const mjModel* m) {
   int size = (
-    sizeof(int)*(4+getnint())
+    sizeof(int)*(NHEADER+getnint())
+    + sizeof(size_t)*getnsize()
     + sizeof(mjOption)
     + sizeof(mjVisual)
     + sizeof(mjStatistic));
@@ -1030,7 +1071,7 @@ static mjData* _makeData(const mjModel* m) {
 #undef X
 
   // copy stack size from model
-  d->nstack = m->nstack;
+  d->narena = m->narena;
 
   // allocate buffer
   d->buffer = mju_malloc(d->nbuffer);
@@ -1040,7 +1081,7 @@ static mjData* _makeData(const mjModel* m) {
   }
 
   // allocate arena
-  d->arena = mju_malloc(d->nstack * sizeof(mjtNum));
+  d->arena = mju_malloc(d->narena);
   if (!d->arena) {
     mju_free(d->buffer);
     mju_free(d);
@@ -1091,7 +1132,7 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) {
   if (dest->nbuffer != src->nbuffer) {
     mjERROR("dest and src data buffers have different size");
   }
-  if (dest->nstack != src->nstack) {
+  if (dest->narena != src->narena) {
     mjERROR("dest and src stacks have different size");
   }
 
@@ -1164,7 +1205,7 @@ void* mj_arenaAlloc(mjData* d, int bytes, int alignment) {
   int padding = misalignment ? alignment - misalignment : 0;
 
   // check size
-  size_t bytes_available = (d->nstack - d->pstack) * sizeof(mjtNum);
+  size_t bytes_available = d->narena - d->pstack;
   if (d->parena + padding + bytes > bytes_available) {
     return NULL;
   }
@@ -1172,7 +1213,7 @@ void* mj_arenaAlloc(mjData* d, int bytes, int alignment) {
   // allocate, update max, return pointer to buffer
   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);
+  d->maxuse_arena = mjMAX(d->maxuse_arena, d->pstack + d->parena);
 
 #ifdef ADDRESS_SANITIZER
   ASAN_UNPOISON_MEMORY_REGION(result, bytes);
@@ -1202,26 +1243,26 @@ void* mj_stackAllocByte(mjData* d, size_t size) {
 #endif
 
   // size of entire arena/stack in bytes
-  size_t stack_size_bytes = d->nstack * sizeof(mjtNum);
+  size_t stack_size_bytes = d->narena;
 
   // end of the arena
   uintptr_t end_of_arena_ptr = (uintptr_t)d->arena + stack_size_bytes;
 
   // current top of the stack
-  uintptr_t end_ptr = end_of_arena_ptr - (d->pstack * sizeof(mjtNum));
+  uintptr_t end_ptr = end_of_arena_ptr - d->pstack;
 
   // start of the memory to be allocated to the buffer
   uintptr_t start_ptr = end_ptr - (size + mjREDZONE);
 
-  // move start_ptr back to align to mjtNum
+  // move start_ptr back to align to max_align_t
   start_ptr -= start_ptr % _Alignof(max_align_t);  // NOLINT
 
   // new top of the stack
   uintptr_t new_pstack_ptr = start_ptr - mjREDZONE;
-  size_t new_pstack = (end_of_arena_ptr - new_pstack_ptr) / sizeof(mjtNum);
+  size_t new_pstack = end_of_arena_ptr - new_pstack_ptr;
 
   // exclude red zone from stack usage statistics
-  size_t current_alloc_usage = (end_ptr - new_pstack_ptr - 2 * mjREDZONE) / sizeof(mjtNum);
+  size_t current_alloc_usage = end_ptr - new_pstack_ptr - 2 * mjREDZONE;
   size_t usage = current_alloc_usage + d->pstack;
 
   // check size
@@ -1235,13 +1276,9 @@ void* mj_stackAllocByte(mjData* d, size_t size) {
   }
 
 #ifdef ADDRESS_SANITIZER
-  if ((uintptr_t)start_ptr % sizeof(mjtNum)) {
-    mjERROR("mj_stackAlloc failed to align to sizeof(mjtNum)");
-  }
-
   // actual stack usage (without red zone bytes) is stored in the red zone
   if (d->pstack) {
-    size_t* prev_usage_ptr = (size_t*)(end_of_arena_ptr - d->pstack*sizeof(mjtNum));
+    size_t* prev_usage_ptr = (size_t*)(end_of_arena_ptr - d->pstack);
     ASAN_UNPOISON_MEMORY_REGION(prev_usage_ptr, sizeof(size_t));
     usage = current_alloc_usage + *prev_usage_ptr;
     ASAN_POISON_MEMORY_REGION(prev_usage_ptr, sizeof(size_t));
@@ -1261,7 +1298,7 @@ void* mj_stackAllocByte(mjData* d, size_t size) {
   // 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);
+  d->maxuse_arena = mjMAX(d->maxuse_arena, usage + d->parena);
 
   return (void*)start_ptr;
 }
@@ -1294,11 +1331,11 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
 
   // poison the entire arena+stack memory region when built with asan
 #ifdef ADDRESS_SANITIZER
-  ASAN_POISON_MEMORY_REGION(d->arena, d->nstack * sizeof(mjtNum));
+  ASAN_POISON_MEMORY_REGION(d->arena, d->narena);
 #endif
 
 #ifdef MEMORY_SANITIZER
-  __msan_allocated_memory(d->arena, d->nstack * sizeof(mjtNum));
+  __msan_allocated_memory(d->arena, d->narena);
 #endif
 
 #define X(type, name, nr, nc) d->name = NULL;

src/engine/engine_island.c

@@ -100,7 +100,7 @@ static void clearIsland(mjData* d, size_t parena) {
   // poison remaining memory
 #ifdef ADDRESS_SANITIZER
   ASAN_POISON_MEMORY_REGION(
-    (char*)d->arena + d->parena, (d->nstack - d->pstack) * sizeof(mjtNum) - d->parena);
+    (char*)d->arena + d->parena, d->narena - d->pstack - d->parena);
 #endif
 }
 
@@ -118,7 +118,7 @@ static int arenaAllocIsland(const mjModel* m, mjData* d) {
 #define X(type, name, nr, nc)                                             \
   d->name = mj_arenaAlloc(d, sizeof(type) * (nr) * (nc), _Alignof(type)); \
   if (!d->name) {                                                         \
-    mj_warning(d, mjWARN_CNSTRFULL, d->nstack * sizeof(mjtNum));          \
+    mj_warning(d, mjWARN_CNSTRFULL, d->narena);                           \
     clearIsland(d, parena_old);                                           \
     return 0;                                                             \
   }

src/engine/engine_print.c

@@ -242,8 +242,13 @@ void mj_printFormattedModel(const mjModel* m, const char* filename, const char*
   // sizes
 #define X( name )                                 \
   if (m->name) {                                  \
+    const char* format = _Generic(                \
+        m->name,                                  \
+        size_t : SIZE_T_FORMAT,                   \
+        default : INT_FORMAT);                    \
     fprintf(fp, NAME_FORMAT, #name);              \
-    fprintf(fp, INT_FORMAT "\n", m->name);        \
+    fprintf(fp, format, m->name);                 \
+    fprintf(fp, "\n");                            \
   }
 
   MJMODEL_INTS

src/user/user_model.cc

@@ -2803,23 +2803,23 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, int vfs_provider) {
     mj_setTotalmass(m, settotalmass);
   }
 
-  // set arena size into m->nstack
+  // set arena size into m->narena
   if (memory != -1) {
-    // memory size is user-specified in bytes, round down to nearest sizeof(mjtNum)
+    // memory size is user-specified in bytes
-    m->nstack = memory / sizeof(mjtNum);
+    m->narena = memory;
   } else {
     const int nconmax = m->nconmax == -1 ? 100 : m->nconmax;
     const int njmax = m->njmax == -1 ? 500 : m->njmax;
     if (nstack != -1) {
-      // (legacy) stack size is user-specified, already as multiple of sizeof(mjtNum)
+      // (legacy) stack size is user-specified as multiple of sizeof(mjtNum)
-      m->nstack = nstack;
+      m->narena = sizeof(mjtNum) * nstack;
     } else {
       // use a conservative heuristic if neither memory nor nstack is specified in XML
-      m->nstack = mjMAX(
+      m->narena = sizeof(mjtNum) * static_cast<size_t>(mjMAX(
           1000,
           5*(njmax + m->neq + m->nv)*(njmax + m->neq + m->nv) +
           20*(m->nq + m->nv + m->nu + m->na + m->nbody + m->njnt +
-              m->ngeom + m->nsite + m->neq + m->ntendon +  m->nwrap));
+              m->ngeom + m->nsite + m->neq + m->ntendon +  m->nwrap)));
     }
 
     // add an arena space equal to memory footprint prior to the introduction of the arena
@@ -2829,13 +2829,13 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, int vfs_provider) {
         m->nv * (3 * sizeof(int)) +
         njmax * m->nv * (2 * sizeof(int) + 2 * sizeof(mjtNum)) +
         njmax * njmax * (sizeof(int) + sizeof(mjtNum)));
-    m->nstack += (arena_bytes / sizeof(mjtNum)) + (arena_bytes % sizeof(mjtNum) ? 1 : 0);
+    m->narena += arena_bytes;
 
     // round up to the nearest megabyte
-    constexpr int kMegabyte = (1 << 20) / sizeof(mjtNum);  // number of mjtNum's in 1 Mb
+    constexpr std::size_t kMegabyte = 1 << 20;
-    int nstack_mb = m->nstack / kMegabyte;
+    std::size_t nstack_mb = m->narena / kMegabyte;
-    int residual_mb = m->nstack % kMegabyte ? 1 : 0;
+    std::size_t residual_mb = m->narena % kMegabyte ? 1 : 0;
-    m->nstack = kMegabyte * (nstack_mb + residual_mb);
+    m->narena = kMegabyte * (nstack_mb + residual_mb);
   }
 
   // create data

src/user/user_model.h

@@ -151,7 +151,7 @@ class mjCModel {
   int nemax;                      // max number of equality constraints
   int njmax;                      // max number of constraints (Jacobian rows)
   int nconmax;                    // max number of detected contacts (mjContact array size)
-  int nstack;                     // (deprecated) number of fields in mjData stack
+  size_t nstack;                  // (deprecated) number of fields in mjData stack
   int nuserdata;                  // number extra fields in mjData
   int nuser_body;                 // number of mjtNums in body_user
   int nuser_jnt;                  // number of mjtNums in jnt_user

test/engine/engine_util_container_test.cc

@@ -61,8 +61,7 @@ TEST(TestMjArrayList, TestMjArrayListSingleThreaded) {
   EXPECT_EQ(mju_arrayListSize(array_list), kNumElements);
 
   constexpr int kExpectedMaxUseStack =
-      GetExpectedStackUsageBytes<DataType, kNumElements, kInitialCapacity>() /
+      GetExpectedStackUsageBytes<DataType, kNumElements, kInitialCapacity>();
-      sizeof(mjtNum);
   EXPECT_EQ(d->maxuse_stack, kExpectedMaxUseStack);
 
   for (int i = 0; i < kNumElements; ++i) {

test/xml/xml_native_reader_test.cc

@@ -48,7 +48,7 @@ TEST_F(XMLReaderTest, MemorySize) {
     )";
     mjModel* model = LoadModelFromString(xml, error.data(), error.size());
     ASSERT_THAT(model, NotNull()) << error.data();
-    EXPECT_EQ(model->nstack, 256 / sizeof(mjtNum));
+    EXPECT_EQ(model->narena, 256);
     mj_deleteModel(model);
   }
   {
@@ -59,7 +59,7 @@ TEST_F(XMLReaderTest, MemorySize) {
     )";
     mjModel* model = LoadModelFromString(xml, error.data(), error.size());
     ASSERT_THAT(model, NotNull()) << error.data();
-    EXPECT_EQ(model->nstack, 1024 / sizeof(mjtNum));
+    EXPECT_EQ(model->narena, 1024);
     mj_deleteModel(model);
   }
   {
@@ -70,7 +70,7 @@ TEST_F(XMLReaderTest, MemorySize) {
     )";
     mjModel* model = LoadModelFromString(xml, error.data(), error.size());
     ASSERT_THAT(model, NotNull()) << error.data();
-    EXPECT_EQ(model->nstack, 10240 / sizeof(mjtNum));
+    EXPECT_EQ(model->narena, 10240);
     mj_deleteModel(model);
   }
   {
@@ -81,7 +81,7 @@ TEST_F(XMLReaderTest, MemorySize) {
     )";
     mjModel* model = LoadModelFromString(xml, error.data(), error.size());
     ASSERT_THAT(model, NotNull()) << error.data();
-    EXPECT_EQ(model->nstack, 4*1024*1024 / sizeof(mjtNum));
+    EXPECT_EQ(model->narena, 4*1024*1024);
     mj_deleteModel(model);
   }
   {
@@ -92,7 +92,7 @@ TEST_F(XMLReaderTest, MemorySize) {
     )";
     mjModel* model = LoadModelFromString(xml, error.data(), error.size());
     ASSERT_THAT(model, NotNull()) << error.data();
-    EXPECT_EQ(model->nstack, 1024*1024*1024 / sizeof(mjtNum));
+    EXPECT_EQ(model->narena, 1024*1024*1024);
     mj_deleteModel(model);
   }
   {
@@ -103,7 +103,7 @@ TEST_F(XMLReaderTest, MemorySize) {
     )";
     mjModel* model = LoadModelFromString(xml, error.data(), error.size());
     ASSERT_THAT(model, NotNull()) << error.data();
-    EXPECT_EQ(model->nstack, 1024*1024*1024 / sizeof(mjtNum));
+    EXPECT_EQ(model->narena, 1024*1024*1024);
     mj_deleteModel(model);
   }
 }

unity/Runtime/Bindings/MjBindings.cs

@@ -570,12 +570,12 @@ public unsafe struct mjSolverStat_ {
 
 [StructLayout(LayoutKind.Sequential)]
 public unsafe struct mjData_ {
-  public int nstack;
+  public UIntPtr narena;
-  public int nbuffer;
+  public UIntPtr nbuffer;
   public int nplugin;
   public UIntPtr pstack;
   public UIntPtr parena;
-  public int maxuse_stack;
+  public UIntPtr maxuse_stack;
   public UIntPtr maxuse_arena;
   public int maxuse_con;
   public int maxuse_efc;
@@ -1965,11 +1965,11 @@ public unsafe struct mjModel_ {
   public int nemax;
   public int njmax;
   public int nconmax;
-  public int nstack;
   public int nuserdata;
   public int nsensordata;
   public int npluginstate;
-  public int nbuffer;
+  public UIntPtr narena;
+  public UIntPtr nbuffer;
   public mjOption_ opt;
   public mjVisual_ vis;
   public mjStatistic_ stat;