Fix memory allocation error when memory size is larger than 2G. Fixes #1606

PiperOrigin-RevId: 627058355
Change-Id: Iebef647a0107d1f112a876d705f5b3ea888b4a41

doc/changelog.rst

@@ -15,34 +15,37 @@ General
    - :ref:`material/metallic<asset-material-metallic>` attribute and corresponding ``mjModel.material_metallic`` field.
    - :ref:`material/roughness<asset-material-roughness>` attribute and corresponding ``mjModel.material_roughness``
      field.
+3. The type of the ``size`` argument of :ref:`mj_stackAllocNum` and :ref:`mj_stackAllocInt` was changed from ``int``
+   to ``size_t``.
 
 MJX
 ^^^
 .. admonition:: Breaking API changes
    :class: attention
 
-   3. Removed deprecated ``mjx.device_get_into`` and ``mjx.device_put`` functions as they lack critical new
+   4. Removed deprecated ``mjx.device_get_into`` and ``mjx.device_put`` functions as they lack critical new
       functionality.
 
       **Migration:** Use ``mjx.get_data_into`` instead of ``mjx.device_get_into``, and ``mjx.put_data`` instead of
       ``mjx.device_put``.
 
-4. Added cylinder plane collisions.
-5. Added ``efc_type`` to ``mjx.Data`` and ``dim``, ``efc_address`` to ``mjx.Contact``.
-6. Added ``geom`` to ``mjx.Contact`` and marked ``geom1``, ``geom2`` deprecated.
-7. Added ``ne``, ``nf``, ``nl``, ``nefc``, and ``ncon`` to ``mjx.Data`` to match ``mujoco.MjData``.
-8. Given the above added fields, removed ``mjx.get_params``, ``mjx.ncon``, and ``mjx.count_constraints``.
-9. Changed the way meshes are organized on device to speed up collision detection when a mesh is replicated for many
-   geoms.
-10. Fixed a bug where capsules might be ignored in broadphase colliision checking.
-11. Added cylinder collisions using SDFs.
+5. Added cylinder plane collisions.
+6. Added ``efc_type`` to ``mjx.Data`` and ``dim``, ``efc_address`` to ``mjx.Contact``.
+7. Added ``geom`` to ``mjx.Contact`` and marked ``geom1``, ``geom2`` deprecated.
+8. Added ``ne``, ``nf``, ``nl``, ``nefc``, and ``ncon`` to ``mjx.Data`` to match ``mujoco.MjData``.
+9. Given the above added fields, removed ``mjx.get_params``, ``mjx.ncon``, and ``mjx.count_constraints``.
+10. Changed the way meshes are organized on device to speed up collision detection when a mesh is replicated for many
+    geoms.
+11. Fixed a bug where capsules might be ignored in broadphase colliision checking.
+12. Added cylinder collisions using SDFs.
 
 Bug fixes
 ^^^^^^^^^
-12. Defaults of lights were not being saved, now fixed.
-13. Prevent overwriting of frame names by body names when saving an XML. Bug introduced in 3.1.4.
-14. Fixed bug in Python binding of :ref:`mj_saveModel`: ``buffer`` argument was documented as optional but was actually
+13. Defaults of lights were not being saved, now fixed.
+14. Prevent overwriting of frame names by body names when saving an XML. Bug introduced in 3.1.4.
+15. Fixed bug in Python binding of :ref:`mj_saveModel`: ``buffer`` argument was documented as optional but was actually
     not optional.
+16. Fixed bug that prevented memory allocations larger than 2.15 GB.
 
 
 Version 3.1.4 (April 10th, 2024)

doc/includes/references.h

@@ -2453,8 +2453,8 @@ void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
 void mj_markStack(mjData* d);
 void mj_freeStack(mjData* d);
 void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment);
-mjtNum* mj_stackAllocNum(mjData* d, int size);
-int* mj_stackAllocInt(mjData* d, int size);
+mjtNum* mj_stackAllocNum(mjData* d, size_t size);
+int* mj_stackAllocInt(mjData* d, size_t size);
 void mj_deleteData(mjData* d);
 void mj_resetCallbacks(void);
 void mj_setConst(mjModel* m, mjData* d);

include/mujoco/mujoco.h

@@ -203,10 +203,10 @@ MJAPI void mj_freeStack(mjData* d);
 MJAPI void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment);
 
 // Allocate array of mjtNums on mjData stack. Call mju_error on stack overflow.
-MJAPI mjtNum* mj_stackAllocNum(mjData* d, int size);
+MJAPI mjtNum* mj_stackAllocNum(mjData* d, size_t size);
 
 // Allocate array of ints on mjData stack. Call mju_error on stack overflow.
-MJAPI int* mj_stackAllocInt(mjData* d, int size);
+MJAPI int* mj_stackAllocInt(mjData* d, size_t size);
 
 // Free memory allocation in mjData.
 MJAPI void mj_deleteData(mjData* d);

introspect/functions.py

@@ -774,7 +774,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
              ),
              FunctionParameterDecl(
                  name='size',
-                 type=ValueType(name='int'),
+                 type=ValueType(name='size_t'),
              ),
          ),
          doc='Allocate array of mjtNums on mjData stack. Call mju_error on stack overflow.',  # pylint: disable=line-too-long
@@ -794,7 +794,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
              ),
              FunctionParameterDecl(
                  name='size',
-                 type=ValueType(name='int'),
+                 type=ValueType(name='size_t'),
              ),
          ),
          doc='Allocate array of ints on mjData stack. Call mju_error on stack overflow.',  # pylint: disable=line-too-long

src/engine/engine_io.c

@@ -1524,14 +1524,20 @@ void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment) {
 
 
 // allocate mjtNums on the stack
-mjtNum* mj_stackAllocNum(mjData* d, int size) {
+mjtNum* mj_stackAllocNum(mjData* d, size_t size) {
+  if (mjUNLIKELY(size >= SIZE_MAX / sizeof(mjtNum))) {
+    mjERROR("requested size is too large.");
+  }
   return (mjtNum*) stackalloc(d, size * sizeof(mjtNum), _Alignof(mjtNum));
 }
 
 
 
 // allocate ints on the stack
-int* mj_stackAllocInt(mjData* d, int size) {
+int* mj_stackAllocInt(mjData* d, size_t size) {
+  if (mjUNLIKELY(size >= SIZE_MAX / sizeof(int))) {
+    mjERROR("requested size is too large.");
+  }
   return (int*) stackalloc(d, size * sizeof(int), _Alignof(int));
 }
 

src/engine/engine_io.h

@@ -129,10 +129,10 @@ void mj__freeStack(mjData* d) __attribute__((noinline));
 MJAPI void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment);
 
 // mjData stack allocate for array of mjtNums
-MJAPI mjtNum* mj_stackAllocNum(mjData* d, int size);
+MJAPI mjtNum* mj_stackAllocNum(mjData* d, size_t size);
 
 // mjData stack allocate for array of ints
-MJAPI int* mj_stackAllocInt(mjData* d, int size);
+MJAPI int* mj_stackAllocInt(mjData* d, size_t size);
 
 // de-allocate data
 MJAPI void mj_deleteData(mjData* d);

src/xml/xml_native_reader.cc

@@ -1186,7 +1186,7 @@ void mjXReader::Size(XMLElement* section, mjSpec* mod) {
     }();
 
     if (memory.has_value()) {
-      if (*memory / sizeof(mjtNum) > std::numeric_limits<int>::max()) {
+      if (*memory / sizeof(mjtNum) > std::numeric_limits<std::size_t>::max()) {
         throw mjXError(section, "%s", err_msg);
       }
       mod->memory = *memory;

test/engine/engine_io_test.cc

@@ -667,6 +667,73 @@ TEST_F(EngineIoTest, CanMarkAndFreeStack) {
   mj_deleteModel(model);
 }
 
+TEST_F(EngineIoTest, LargeMemory) {
+  constexpr char xml[] = R"(
+  <mujoco>
+    <size memory="2400M"/>
+  </mujoco>
+  )";
+
+  std::array<char, 1024> error;
+  mjModel* model = LoadModelFromString(xml, error.data(), error.size());
+  ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data();
+  mjData* data = mj_makeData(model);
+  ASSERT_THAT(data, NotNull());
+
+  // allocate 2.3G of mjtNums
+  mj_markStack(data);
+  size_t num = 2300000000 / sizeof(mjtNum);
+  mjtNum* testNum = mj_stackAllocNum(data, num);
+  testNum[num-1] = 1;
+  mj_freeStack(data);
+
+  // allocate 2.3G of bytes
+  mj_markStack(data);
+  num = 2300000000;
+  char* testByte = (char*) mj_stackAllocByte(data, num, alignof(char));
+  testByte[num-1] = 1;
+  mj_freeStack(data);
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
+TEST_F(EngineIoTest, VeryLargeMemory) {
+  constexpr char xml[] = R"(
+  <mujoco>
+    <size memory="64G"/>
+  </mujoco>
+  )";
+
+  std::array<char, 1024> error;
+  mjModel* model = LoadModelFromString(xml, error.data(), error.size());
+  if (!model) {
+    // in some test environments, 64GB is too large
+    EXPECT_THAT(error.data(), HasSubstr("Could not allocate memory"));
+  } else {
+    ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data();
+    mjData* data = mj_makeData(model);
+    ASSERT_THAT(data, NotNull());
+
+    // allocate 63G of mjtNums
+    mj_markStack(data);
+    size_t num = 63000000000 / sizeof(mjtNum);
+    mjtNum* testNum = mj_stackAllocNum(data, num);
+    testNum[num-1] = 1;
+    mj_freeStack(data);
+
+    // allocate 63G of bytes
+    mj_markStack(data);
+    num = 63000000000;
+    char* testByte = (char*) mj_stackAllocByte(data, num, alignof(char));
+    testByte[num-1] = 1;
+    mj_freeStack(data);
+
+    mj_deleteData(data);
+    mj_deleteModel(model);
+  }
+}
+
 struct TestFunctionArgs_ {
   mjData* d;
   int input;

unity/Runtime/Bindings/MjBindings.cs

@@ -6430,10 +6430,10 @@ public static unsafe extern void mj_freeStack(mjData_* d);
 public static unsafe extern void* mj_stackAllocByte(mjData_* d, UIntPtr bytes, UIntPtr alignment);
 
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
-public static unsafe extern double* mj_stackAllocNum(mjData_* d, int size);
+public static unsafe extern double* mj_stackAllocNum(mjData_* d, UIntPtr size);
 
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
-public static unsafe extern int* mj_stackAllocInt(mjData_* d, int size);
+public static unsafe extern int* mj_stackAllocInt(mjData_* d, UIntPtr size);
 
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern void mj_deleteData(mjData_* d);