Fix more mjData stack leakages.

PiperOrigin-RevId: 562545230
Change-Id: I2facf25b00878669dc12d5865767a44229d2ca48

src/engine/engine_collision_driver.c

@@ -387,7 +387,6 @@ void mj_collision(const mjModel* m, mjData* d) {
   int g1, g2, merged, b1 = 0, b2 = 0, exadr = 0, pairadr = 0, startadr;
   int nexclude = m->nexclude, npair = m->npair, nbodypair = ((m->nbody-1)*m->nbody)/2;
   int *broadphasepair = 0;
-  mjMARKSTACK;
 
   // reset the size of the contact array and invalidate efc arrays
   d->ncon = 0;
@@ -409,6 +408,8 @@ void mj_collision(const mjModel* m, mjData* d) {
     return;
   }
 
+  mjMARKSTACK;
+
   // predefined only; ignore exclude
   if (m->opt.collision == mjCOL_PAIR) {
     d->nbodypair_broad = npair;
@@ -677,7 +678,6 @@ int mj_broadphase(const mjModel* m, mjData* d, int* pair, int maxpair) {
   mjtNum cov[9], cen[3], dif[3], eigval[3], frame[9], quat[4];
   mjtBroadphase *sortbuf, *activebuf;
   mjtNum *aabb;
-  mjMARKSTACK;
 
   int dsbl_filterparent = mjDISABLED(mjDSBL_FILTERPARENT);
   // world with geoms, and body with plane or hfield, can collide all bodies
@@ -750,6 +750,7 @@ int mj_broadphase(const mjModel* m, mjData* d, int* pair, int maxpair) {
   mju_eig3(eigval, frame, quat, cov);
 
   // allocate AABB; clear world entry (not used)
+  mjMARKSTACK;
   aabb = mj_stackAllocNum(d, 6*nbody);
   mju_zero(aabb, 6);
 

src/engine/engine_core_constraint.c

@@ -495,13 +495,14 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
   mjtNum cpos[6], pos[2][3], ref[2], dif, deriv;
   mjtNum quat[4], quat1[4], quat2[4], quat3[4], axis[3];
   mjtNum *jac[2], *jacdif, *data, *sparse_buf = NULL;
-  mjMARKSTACK;
 
   // disabled or no equality constraints: return
   if (mjDISABLED(mjDSBL_EQUALITY) || m->nemax == 0) {
     return;
   }
 
+  mjMARKSTACK;
+
   // allocate space
   jac[0] = mj_stackAllocNum(d, 6*nv);
   jac[1] = mj_stackAllocNum(d, 6*nv);
@@ -703,13 +704,14 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
 void mj_instantiateFriction(const mjModel* m, mjData* d) {
   int nv = m->nv, issparse = mj_isSparse(m);
   mjtNum* jac;
-  mjMARKSTACK;
 
   // disabled: return
   if (mjDISABLED(mjDSBL_FRICTIONLOSS)) {
     return;
   }
 
+  mjMARKSTACK;
+
   // allocate Jacobian
   jac = mj_stackAllocNum(d, nv);
 
@@ -764,13 +766,14 @@ void mj_instantiateLimit(const mjModel* m, mjData* d) {
   int side, nv = m->nv, issparse = mj_isSparse(m);
   mjtNum margin, value, dist, angleAxis[3];
   mjtNum *jac;
-  mjMARKSTACK;
 
   // disabled: return
   if (mjDISABLED(mjDSBL_LIMIT)) {
     return;
   }
 
+  mjMARKSTACK;
+
   // allocate Jacobian
   jac = mj_stackAllocNum(d, nv);
 
@@ -863,7 +866,7 @@ void mj_instantiateLimit(const mjModel* m, mjData* d) {
   // find tendon limits
   for (int i=0; i < m->ntendon; i++) {
     if (m->tendon_limited[i]) {
-      // get value = lenth, margin
+      // get value = length, margin
       value = d->ten_length[i];
       margin = m->tendon_margin[i];
 

src/engine/engine_core_smooth.c

@@ -394,13 +394,13 @@ void mj_tendon(const mjModel* m, mjData* d) {
   mjtNum dif[3], divisor, wpnt[12], wlen;
   mjtNum *L = d->ten_length, *J = d->ten_J;
   mjtNum *jac1, *jac2, *jacdif, *tmp, *sparse_buf = NULL;
-  mjMARKSTACK;
 
   if (!nten) {
     return;
   }
 
   // allocate space
+  mjMARKSTACK;
   jac1 = mj_stackAllocNum(d, 3*nv);
   jac2 = mj_stackAllocNum(d, 3*nv);
   jacdif = mj_stackAllocNum(d, 3*nv);
@@ -622,13 +622,13 @@ void mj_transmission(const mjModel* m, mjData* d) {
   mjtNum *jac, *jacA, *jacS;
   mjtNum *length = d->actuator_length, *moment = d->actuator_moment, *gear;
   mjtNum *jacref = NULL, *moment_tmp = NULL;  // required for site actuators
-  mjMARKSTACK;
 
   if (!nu) {
     return;
   }
 
   // allocate space, clear moments
+  mjMARKSTACK;
   jac  = mj_stackAllocNum(d, 3*nv);
   jacA = mj_stackAllocNum(d, 3*nv);
   jacS = mj_stackAllocNum(d, 3*nv);

src/engine/engine_derivative_fd.c

@@ -634,7 +634,6 @@ void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_actuatio
                    mjtNum *DsDq, mjtNum *DsDv, mjtNum *DsDa,
                    mjtNum *DmDq) {
   int nq = m->nq, nv = m->nv, nM = m->nM, ns = m->nsensordata;
-  mjMARKSTACK;
 
   if (m->opt.integrator == mjINT_RK4) {
     mjERROR("RK4 integrator is not supported");
@@ -648,6 +647,7 @@ void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_actuatio
   int skipsensor = !DsDq && !DsDv && !DsDa;
 
   // local vectors
+  mjMARKSTACK;
   mjtNum *pos        = mj_stackAllocNum(d, nq);                      // position
   mjtNum *force      = mj_stackAllocNum(d, nv);                      // force
   mjtNum *force_plus = mj_stackAllocNum(d, nv);                      // nudged force

src/engine/engine_forward.c

@@ -665,7 +665,6 @@ void mj_RungeKutta(const mjModel* m, mjData* d, int N) {
   mjtNum C[9], T[9], *X[10], *F[10], *dX;
   const mjtNum* A = (N == 4 ? RK4_A : 0);
   const mjtNum* B = (N == 4 ? RK4_B : 0);
-  mjMARKSTACK;
 
   // check order
   if (!A) {
@@ -673,6 +672,7 @@ void mj_RungeKutta(const mjModel* m, mjData* d, int N) {
   }
 
   // allocate space for intermediate solutions
+  mjMARKSTACK;
   dX = mj_stackAllocNum(d, 2*nv+na);
   for (int i=0; i < N; i++) {
     X[i] = mj_stackAllocNum(d, nq+nv+na);

src/engine/engine_inverse.c

@@ -119,6 +119,7 @@ static void mj_discreteAcc(const mjModel* m, mjData* d) {
 
     // if disabled or no dof damping, nothing to do
     if (!dof_damping) {
+      mjFREESTACK;
       return;
     }
 
@@ -277,7 +278,6 @@ void mj_inverse(const mjModel* m, mjData* d) {
 void mj_compareFwdInv(const mjModel* m, mjData* d) {
   int nv = m->nv, nefc = d->nefc;
   mjtNum *qforce, *dif, *save_qfrc_constraint, *save_efc_force;
-  mjMARKSTACK;
 
   // clear result, return if no constraints
   d->solver_fwdinv[0] = d->solver_fwdinv[1] = 0;
@@ -286,6 +286,7 @@ void mj_compareFwdInv(const mjModel* m, mjData* d) {
   }
 
   // allocate
+  mjMARKSTACK;
   qforce = mj_stackAllocNum(d, nv);
   dif = mj_stackAllocNum(d, nv);
   save_qfrc_constraint = mj_stackAllocNum(d, nv);

src/engine/engine_island.c

@@ -459,7 +459,7 @@ void mj_island(const mjModel* m, mjData* d) {
   memset(d->island_dofnum, 0, nisland*sizeof(int));
   for (int i=0; i < nv; i++) {
     // dof_island
-    int island = tree_island[m->dof_treeid[i]];;
+    int island = tree_island[m->dof_treeid[i]];
     d->dof_island[i] = island;
 
     // island_dofnum
@@ -497,7 +497,7 @@ void mj_island(const mjModel* m, mjData* d) {
     mjERROR("not all islands assigned to dofs");
   }
 
-  // finalize dof_islandind: set remaning indices to -1
+  // finalize dof_islandind: set remaining indices to -1
   for (int i=num_dof_island; i < nv; i++) {
     d->island_dofind[i] = -1;
   }

src/engine/engine_sensor.c

@@ -854,13 +854,13 @@ void mj_energyPos(const mjModel* m, mjData* d) {
 // velocity-dependent energy (kinetic)
 void mj_energyVel(const mjModel* m, mjData* d) {
   mjtNum *vec;
-  mjMARKSTACK;
 
   // return if disabled (already cleared in potential)
   if (!mjENABLED(mjENBL_ENERGY)) {
     return;
   }
 
+  mjMARKSTACK;
   vec = mj_stackAllocNum(d, m->nv);
 
   // kinetic energy:  0.5 * qvel' * M * qvel

src/engine/engine_vis_interact.c

@@ -529,6 +529,7 @@ void mjv_initPerturb(const mjModel* m, mjData* d, const mjvScene* scn, mjvPertur
 
   // invalid selected body: return
   if (sel <= 0 || sel >= m->nbody) {
+    mjFREESTACK;
     return;
   }