Add mjModel and mjData arguments to mj_makeModel and mj_makeRawData. If they are NULL, then the functions allocate the memory, otherwise they will use the memory allocated by the caller.

PiperOrigin-RevId: 639283021
Change-Id: I765f95fe1b5c27202ae033b05a13ae7b5fb8dbc3
This commit is contained in:
Alessio Quaglino
2024-06-01 01:07:50 -07:00
committed by Copybara-Service
parent 814deb0e9f
commit 936866be0f
3 changed files with 127 additions and 86 deletions
+100 -61
View File
@@ -447,22 +447,40 @@ static int safeAddToBufferSize(intptr_t* offset, size_t* nbuffer,
// free model memory without destroying the struct
static void freeModelBuffers(mjModel* m) {
mju_free(m->buffer);
}
// allocate and initialize mjModel structure
mjModel* mj_makeModel(
int nq, int nv, int nu, int na, int nbody, int nbvh, int nbvhstatic, int nbvhdynamic, int njnt,
int ngeom, int nsite, int ncam, int nlight, int nflex, int nflexvert, int nflexedge,
int nflexelem, int nflexelemdata, int nflexshelldata, int nflexevpair, int nflextexcoord,
int nmesh, int nmeshvert, int nmeshnormal, int nmeshtexcoord, int nmeshface, int nmeshgraph,
int nskin, int nskinvert, int nskintexvert, int nskinface, int nskinbone, int nskinbonevert,
int nhfield, int nhfielddata, int ntex, int ntexdata, int nmat, int npair, int nexclude, int neq,
int ntendon, int nwrap, int nsensor, int nnumeric, int nnumericdata, int ntext, int ntextdata,
int ntuple, int ntupledata, int nkey, int nmocap, int nplugin, int npluginattr, int nuser_body,
int nuser_jnt, int nuser_geom, int nuser_site, int nuser_cam, int nuser_tendon,
int nuser_actuator, int nuser_sensor, int nnames, int npaths) {
void mj_makeModel(mjModel** dest,
int nq, int nv, int nu, int na, int nbody, int nbvh,
int nbvhstatic, int nbvhdynamic, int njnt, int ngeom, int nsite, int ncam,
int nlight, int nflex, int nflexvert, int nflexedge, int nflexelem,
int nflexelemdata, int nflexshelldata, int nflexevpair, int nflextexcoord,
int nmesh, int nmeshvert, int nmeshnormal, int nmeshtexcoord, int nmeshface,
int nmeshgraph, int nskin, int nskinvert, int nskintexvert, int nskinface,
int nskinbone, int nskinbonevert, int nhfield, int nhfielddata, int ntex,
int ntexdata, int nmat, int npair, int nexclude, int neq, int ntendon,
int nwrap, int nsensor, int nnumeric, int nnumericdata, int ntext,
int ntextdata, int ntuple, int ntupledata, int nkey, int nmocap,
int nplugin, int npluginattr, int nuser_body, int nuser_jnt, int nuser_geom,
int nuser_site, int nuser_cam, int nuser_tendon, int nuser_actuator,
int nuser_sensor, int nnames, int npaths) {
intptr_t offset = 0;
int allocate = *dest ? 0 : 1;
mjModel* m = NULL;
// allocate mjModel
mjModel* m = (mjModel*)mju_malloc(sizeof(mjModel));
if (!allocate) {
m = *dest;
freeModelBuffers(m);
} else {
m = (mjModel*)mju_malloc(sizeof(mjModel));
}
if (!m) {
mjERROR("could not allocate mjModel");
}
@@ -536,43 +554,43 @@ mjModel* mj_makeModel(
nhfield + ntex + nmat + npair + nexclude + neq + ntendon + nu + nsensor +
nnumeric + ntext + ntuple + nkey + nplugin;
if (nnames_map >= INT_MAX / mjLOAD_MULTIPLE) {
mju_free(m);
if (allocate) mju_free(m);
mju_warning("Invalid model: size of nnames_map is larger than INT_MAX");
return 0;
return;
}
m->nnames_map = mjLOAD_MULTIPLE * nnames_map;
m->npaths = npaths;
#define X(name) \
if ((m->name) < 0) { \
mju_free(m); \
if (allocate) mju_free(m); \
mju_warning("Invalid model: negative " #name); \
return 0; \
return; \
}
MJMODEL_INTS;
#undef X
// nbody should always be positive
if (m->nbody == 0) {
mju_free(m);
if (allocate) mju_free(m);
mju_warning("Invalid model: nbody == 0");
return 0;
return;
}
// nmocap is going to get multiplied by 4, and shouldn't overflow
if (m->nmocap >= MAX_ARRAY_SIZE) {
mju_free(m);
if (allocate) mju_free(m);
mju_warning("Invalid model: nmocap too large");
return 0;
return;
}
// compute buffer size
m->nbuffer = 0;
#define X(type, name, nr, nc) \
if (!safeAddToBufferSize(&offset, &m->nbuffer, sizeof(type), m->nr, nc)) { \
mju_free(m); \
if (allocate) mju_free(m); \
mju_warning("Invalid model: " #name " too large."); \
return 0; \
return; \
}
MJMODEL_POINTERS
@@ -581,7 +599,7 @@ mjModel* mj_makeModel(
// allocate buffer
m->buffer = mju_malloc(m->nbuffer);
if (!m->buffer) {
mju_free(m);
if (allocate) mju_free(m);
mjERROR("could not allocate mjModel buffer");
}
@@ -598,18 +616,19 @@ mjModel* mj_makeModel(
mj_defaultVisual(&m->vis);
mj_defaultStatistic(&m->stat);
return m;
// copy pointer if allocated here
if (allocate) {
*dest = m;
}
}
// copy mjModel, if dest==NULL create new model
mjModel* mj_copyModel(mjModel* dest, const mjModel* src) {
void* save_bufptr;
// allocate new model if needed
if (!dest) {
dest = mj_makeModel(
mj_makeModel(&dest,
src->nq, src->nv, src->nu, src->na, src->nbody, src->nbvh,
src->nbvhstatic, src->nbvhdynamic, src->njnt, src->ngeom, src->nsite,
src->ncam, src->nlight, src->nflex, src->nflexvert, src->nflexedge,
@@ -764,16 +783,17 @@ mjModel* mj_loadModelBuffer(const void* buffer, int buffer_sz) {
bufread(sizes, sizeof(size_t)*getnsize(), buffer_sz, buffer, &ptrbuf);
// allocate new mjModel, check sizes
m = mj_makeModel(ints[0], ints[1], ints[2], ints[3], ints[4], ints[5], ints[6],
ints[7], ints[8], ints[9], ints[10], ints[11], ints[12], ints[13],
ints[14], ints[15], ints[16], ints[17], ints[18], ints[19], ints[20],
ints[21], ints[22], ints[23], ints[24], ints[25], ints[26], ints[27],
ints[28], ints[29], ints[30], ints[31], ints[32], ints[33], ints[34],
ints[35], ints[36], ints[37], ints[38], ints[39], ints[40], ints[41],
ints[42], ints[43], ints[44], ints[45], ints[46], ints[47], ints[48],
ints[49], ints[50], ints[51], ints[52], ints[53], ints[54], ints[55],
ints[56], ints[57], ints[58], ints[59], ints[60], ints[61], ints[62],
ints[63]);
mj_makeModel(&m,
ints[0], ints[1], ints[2], ints[3], ints[4], ints[5], ints[6],
ints[7], ints[8], ints[9], ints[10], ints[11], ints[12], ints[13],
ints[14], ints[15], ints[16], ints[17], ints[18], ints[19], ints[20],
ints[21], ints[22], ints[23], ints[24], ints[25], ints[26], ints[27],
ints[28], ints[29], ints[30], ints[31], ints[32], ints[33], ints[34],
ints[35], ints[36], ints[37], ints[38], ints[39], ints[40], ints[41],
ints[42], ints[43], ints[44], ints[45], ints[46], ints[47], ints[48],
ints[49], ints[50], ints[51], ints[52], ints[53], ints[54], ints[55],
ints[56], ints[57], ints[58], ints[59], ints[60], ints[61], ints[62],
ints[63]);
if (!m || m->nbuffer != sizes[getnsize()-1]) {
mju_warning("Corrupted model, wrong size parameters");
mj_deleteModel(m);
@@ -834,7 +854,7 @@ mjModel* mj_loadModelBuffer(const void* buffer, int buffer_sz) {
// de-allocate mjModel
void mj_deleteModel(mjModel* m) {
if (m) {
mju_free(m->buffer);
freeModelBuffers(m);
mju_free(m);
}
}
@@ -1089,12 +1109,40 @@ static void _initPlugin(const mjModel* m, mjData* d) {
// free mjData memory without destroying the struct
static void freeDataBuffers(mjData* d) {
#ifdef ADDRESS_SANITIZER
// raise an error if there's a dangling stack frame
mj_freeStack(d);
#endif
// destroy plugin instances
for (int i = 0; i < d->nplugin; ++i) {
const mjpPlugin* plugin = mjp_getPluginAtSlot(d->plugin[i]);
if (plugin->destroy) {
plugin->destroy(d, i);
}
}
mju_free(d->buffer);
mju_free(d->arena);
}
// allocate and initialize raw mjData structure
mjData* mj_makeRawData(const mjModel* m) {
void mj_makeRawData(mjData** dest, const mjModel* m) {
intptr_t offset = 0;
int allocate = *dest ? 0 : 1;
mjData* d = NULL;
// allocate mjData
mjData* d = (mjData*) mju_malloc(sizeof(mjData));
if (!allocate) {
d = *dest;
freeDataBuffers(d);
} else {
d = (mjData*) mju_malloc(sizeof(mjData));
}
if (!d) {
mjERROR("could not allocate mjData");
}
@@ -1107,9 +1155,9 @@ mjData* mj_makeRawData(const mjModel* m) {
d->buffer = d->arena = NULL;
#define X(type, name, nr, nc) \
if (!safeAddToBufferSize(&offset, &d->nbuffer, sizeof(type), m->nr, nc)) { \
mju_free(d); \
if (allocate) mju_free(d); \
mju_warning("Invalid data: " #name " too large."); \
return 0; \
return; \
}
MJDATA_POINTERS
@@ -1121,7 +1169,7 @@ mjData* mj_makeRawData(const mjModel* m) {
// allocate buffer
d->buffer = mju_malloc(d->nbuffer);
if (!d->buffer) {
mju_free(d);
if (allocate) mju_free(d);
mjERROR("could not allocate mjData buffer");
}
@@ -1129,7 +1177,7 @@ mjData* mj_makeRawData(const mjModel* m) {
d->arena = mju_malloc(d->narena);
if (!d->arena) {
mju_free(d->buffer);
mju_free(d);
if (allocate) mju_free(d);
mjERROR("could not allocate mjData arena");
}
@@ -1142,14 +1190,18 @@ mjData* mj_makeRawData(const mjModel* m) {
// clear nplugin (overwritten by _initPlugin)
d->nplugin = 0;
return d;
// copy pointer if allocated here
if (allocate) {
*dest = d;
}
}
// allocate and initialize mjData structure
mjData* mj_makeData(const mjModel* m) {
mjData* d = mj_makeRawData(m);
mjData* d = NULL;
mj_makeRawData(&d, m);
if (d) {
_initPlugin(m, d);
mj_resetData(m, d);
@@ -1166,7 +1218,7 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) {
// allocate new data if needed
if (!dest) {
dest = mj_makeRawData(m);
mj_makeRawData(&dest, m);
_initPlugin(m, dest);
}
@@ -1697,20 +1749,7 @@ void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key) {
// de-allocate mjData
void mj_deleteData(mjData* d) {
if (d) {
#ifdef ADDRESS_SANITIZER
// raise an error if there's a dangling stack frame
mj_freeStack(d);
#endif
// destroy plugin instances
for (int i = 0; i < d->nplugin; ++i) {
const mjpPlugin* plugin = mjp_getPluginAtSlot(d->plugin[i]);
if (plugin->destroy) {
plugin->destroy(d, i);
}
}
mju_free(d->buffer);
mju_free(d->arena);
freeDataBuffers(d);
mju_free(d);
}
}
+15 -14
View File
@@ -51,19 +51,20 @@ void mj_defaultStatistic(mjStatistic* stat);
//------------------------------- mjModel ----------------------------------------------------------
// allocate mjModel
mjModel* mj_makeModel(int nq, int nv, int nu, int na, int nbody, int nbvh, int nbvhstatic, int nbvhdynamic,
int njnt, int ngeom, int nsite, int ncam, int nlight, int nflex, int nflexvert,
int nflexedge, int nflexelem, int nflexelemdata, int nflexshelldata, int nflexevpair,
int nflextexcoord, int nmesh, int nmeshvert, int nmeshnormal, int nmeshtexcoord, int nmeshface,
int nmeshgraph, int nskin, int nskinvert, int nskintexvert, int nskinface,
int nskinbone, int nskinbonevert, int nhfield, int nhfielddata,
int ntex, int ntexdata, int nmat, int npair, int nexclude,
int neq, int ntendon, int nwrap, int nsensor,
int nnumeric, int nnumericdata, int ntext, int ntextdata,
int ntuple, int ntupledata, int nkey, int nmocap, int nplugin,
int npluginattr, int nuser_body, int nuser_jnt, int nuser_geom,
int nuser_site, int nuser_cam, int nuser_tendon, int nuser_actuator,
int nuser_sensor, int nnames, int npaths);
void mj_makeModel(mjModel** dest,
int nq, int nv, int nu, int na, int nbody, int nbvh, int nbvhstatic, int nbvhdynamic,
int njnt, int ngeom, int nsite, int ncam, int nlight, int nflex, int nflexvert,
int nflexedge, int nflexelem, int nflexelemdata, int nflexshelldata, int nflexevpair,
int nflextexcoord, int nmesh, int nmeshvert, int nmeshnormal, int nmeshtexcoord, int nmeshface,
int nmeshgraph, int nskin, int nskinvert, int nskintexvert, int nskinface,
int nskinbone, int nskinbonevert, int nhfield, int nhfielddata,
int ntex, int ntexdata, int nmat, int npair, int nexclude,
int neq, int ntendon, int nwrap, int nsensor,
int nnumeric, int nnumericdata, int ntext, int ntextdata,
int ntuple, int ntupledata, int nkey, int nmocap, int nplugin,
int npluginattr, int nuser_body, int nuser_jnt, int nuser_geom,
int nuser_site, int nuser_cam, int nuser_tendon, int nuser_actuator,
int nuser_sensor, int nnames, int npaths);
// copy mjModel; allocate new if dest is NULL
MJAPI mjModel* mj_copyModel(mjModel* dest, const mjModel* src);
@@ -91,7 +92,7 @@ MJAPI const char* mj_validateReferences(const mjModel* m);
MJAPI mjData* mj_makeData(const mjModel* m);
// allocate mjData corresponding to given model, used internally
MJAPI mjData* mj_makeRawData(const mjModel* m);
MJAPI void mj_makeRawData(mjData** dest, const mjModel* m);
// Copy mjData.
// m is only required to contain the size fields from MJMODEL_INTS.
+12 -11
View File
@@ -3303,16 +3303,17 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
}
// create low-level model
m = mj_makeModel(nq, nv, nu, na, nbody, nbvh, nbvhstatic, nbvhdynamic, njnt, ngeom, nsite,
ncam, nlight, nflex, nflexvert, nflexedge, nflexelem,
nflexelemdata, nflexshelldata, nflexevpair, nflextexcoord,
nmesh, nmeshvert, nmeshnormal, nmeshtexcoord, nmeshface, nmeshgraph,
nskin, nskinvert, nskintexvert, nskinface, nskinbone, nskinbonevert,
nhfield, nhfielddata, ntex, ntexdata, nmat, npair, nexclude,
neq, ntendon, nwrap, nsensor, nnumeric, nnumericdata, ntext, ntextdata,
ntuple, ntupledata, nkey, nmocap, nplugin, npluginattr,
nuser_body, nuser_jnt, nuser_geom, nuser_site, nuser_cam,
nuser_tendon, nuser_actuator, nuser_sensor, nnames, npaths);
mj_makeModel(&m,
nq, nv, nu, na, nbody, nbvh, nbvhstatic, nbvhdynamic, njnt, ngeom, nsite,
ncam, nlight, nflex, nflexvert, nflexedge, nflexelem,
nflexelemdata, nflexshelldata, nflexevpair, nflextexcoord,
nmesh, nmeshvert, nmeshnormal, nmeshtexcoord, nmeshface, nmeshgraph,
nskin, nskinvert, nskintexvert, nskinface, nskinbone, nskinbonevert,
nhfield, nhfielddata, ntex, ntexdata, nmat, npair, nexclude,
neq, ntendon, nwrap, nsensor, nnumeric, nnumericdata, ntext, ntextdata,
ntuple, ntupledata, nkey, nmocap, nplugin, npluginattr,
nuser_body, nuser_jnt, nuser_geom, nuser_site, nuser_cam,
nuser_tendon, nuser_actuator, nuser_sensor, nnames, npaths);
if (!m) {
throw mjCError(0, "could not create mjModel");
}
@@ -3464,7 +3465,7 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
// create data
int disableflags = m->opt.disableflags;
m->opt.disableflags |= mjDSBL_CONTACT;
d = mj_makeRawData(m);
mj_makeRawData(&d, m);
if (!d) {
mj_deleteModel(m);
throw mjCError(0, "could not create mjData");