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

PiperOrigin-RevId: 627058355
Change-Id: Iebef647a0107d1f112a876d705f5b3ea888b4a41

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);