Extract memory allocation functions and core utilities

PiperOrigin-RevId: 801745499
Change-Id: Iaf05c3430769d3115743d8ab020d13148cb2eb59

src/engine/CMakeLists.txt

@@ -29,6 +29,8 @@ set(MUJOCO_ENGINE_SRCS
     engine_collision_sdf.h
     engine_core_constraint.c
     engine_core_constraint.h
+    engine_core_util.c
+    engine_core_util.h
     engine_core_smooth.c
     engine_core_smooth.h
     engine_crossplatform.cc
@@ -47,6 +49,8 @@ set(MUJOCO_ENGINE_SRCS
     engine_io.c
     engine_io.h
     engine_macro.h
+    engine_memory.c
+    engine_memory.h
     engine_name.c
     engine_name.h
     engine_passive.c

src/engine/engine_collision_convex.c

@@ -25,7 +25,7 @@
 #include <mujoco/mjmodel.h>
 #include "engine/engine_collision_gjk.h"
 #include "engine/engine_collision_primitive.h"
-#include "engine/engine_io.h"
+#include "engine/engine_memory.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"

src/engine/engine_collision_driver.c

@@ -26,10 +26,11 @@
 #include "engine/engine_collision_primitive.h"
 #include "engine/engine_collision_sdf.h"
 #include "engine/engine_core_constraint.h"
+#include "engine/engine_core_util.h"
 #include "engine/engine_io.h"
 #include "engine/engine_macro.h"
+#include "engine/engine_memory.h"
 #include "engine/engine_sort.h"
-#include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"

src/engine/engine_core_constraint.c

@@ -22,9 +22,9 @@
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
 #include <mujoco/mjxmacro.h>
+#include "engine/engine_core_util.h"
 #include "engine/engine_core_smooth.h"
-#include "engine/engine_io.h"
+#include "engine/engine_memory.h"
-#include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"
@@ -84,29 +84,6 @@ static int arenaAllocEfc(const mjModel* m, mjData* d) {
 
 
 
-// determine type of friction cone
-int mj_isPyramidal(const mjModel* m) {
-  if (m->opt.cone == mjCONE_PYRAMIDAL) {
-    return 1;
-  } else {
-    return 0;
-  }
-}
-
-
-
-// determine type of constraint Jacobian
-int mj_isSparse(const mjModel* m) {
-  if (m->opt.jacobian == mjJAC_SPARSE ||
-      (m->opt.jacobian == mjJAC_AUTO && m->nv >= 60)) {
-    return 1;
-  } else {
-    return 0;
-  }
-}
-
-
-
 // determine type of solver
 int mj_isDual(const mjModel* m) {
   if (m->opt.solver == mjSOL_PGS || m->opt.noslip_iterations > 0) {
@@ -133,6 +110,7 @@ void mj_assignFriction(const mjModel* m, mjtNum* target, const mjtNum* source) {
 
 
 
+
 // assign/override contact reference parameters
 void mj_assignRef(const mjModel* m, mjtNum* target, const mjtNum* source) {
   if (mjENABLED(mjENBL_OVERRIDE)) {

src/engine/engine_core_constraint.h

@@ -27,12 +27,6 @@ extern "C" {
 
 //-------------------------- Jacobian-related ------------------------------------------------------
 
-// determine type of friction cone
-MJAPI int mj_isPyramidal(const mjModel* m);
-
-// determine type of constraint Jacobian
-MJAPI int mj_isSparse(const mjModel* m);
-
 // determine type of solver
 MJAPI int mj_isDual(const mjModel* m);
 
@@ -60,7 +54,6 @@ mjtNum mj_assignMargin(const mjModel* m, mjtNum source);
 // add contact to d->contact list; return 0 if success; 1 if buffer full
 MJAPI int mj_addContact(const mjModel* m, mjData* d, const mjContact* con);
 
-
 //-------------------------- constraint instantiation ----------------------------------------------
 
 // equality constraints

src/engine/engine_core_smooth.c

@@ -22,10 +22,10 @@
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
 #include "engine/engine_core_constraint.h"
+#include "engine/engine_core_util.h"
 #include "engine/engine_crossplatform.h"
-#include "engine/engine_io.h"
 #include "engine/engine_macro.h"
-#include "engine/engine_support.h"
+#include "engine/engine_memory.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"

src/engine/engine_core_util.c

@@ -0,0 +1,961 @@
+// Copyright 2025 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_core_util.h"
+
+#include <stddef.h>
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+#include "engine/engine_crossplatform.h"
+#include "engine/engine_memory.h"
+#include "engine/engine_util_blas.h"
+#include "engine/engine_util_errmem.h"
+#include "engine/engine_util_misc.h"
+#include "engine/engine_util_sparse.h"
+#include "engine/engine_util_spatial.h"
+
+
+
+// determine type of constraint Jacobian
+int mj_isSparse(const mjModel* m) {
+  if (m->opt.jacobian == mjJAC_SPARSE ||
+      (m->opt.jacobian == mjJAC_AUTO && m->nv >= 60)) {
+    return 1;
+  } else {
+    return 0;
+  }
+}
+
+
+// determine type of friction cone
+int mj_isPyramidal(const mjModel* m) {
+  if (m->opt.cone == mjCONE_PYRAMIDAL) {
+    return 1;
+  } else {
+    return 0;
+  }
+}
+
+
+
+//-------------------------- sparse chains ---------------------------------------------------------
+
+// merge dof chains for two bodies
+int mj_mergeChain(const mjModel* m, int* chain, int b1, int b2) {
+  int da1, da2, NV = 0;
+
+  // skip fixed bodies
+  while (b1 && !m->body_dofnum[b1]) {
+    b1 = m->body_parentid[b1];
+  }
+  while (b2 && !m->body_dofnum[b2]) {
+    b2 = m->body_parentid[b2];
+  }
+
+  // neither body is movable: empty chain
+  if (b1 == 0 && b2 == 0) {
+    return 0;
+  }
+
+  // initialize last dof address for each body
+  da1 = m->body_dofadr[b1] + m->body_dofnum[b1] - 1;
+  da2 = m->body_dofadr[b2] + m->body_dofnum[b2] - 1;
+
+  // merge chains
+  while (da1 >= 0 || da2 >= 0) {
+    chain[NV] = mjMAX(da1, da2);
+    if (da1 == chain[NV]) {
+      da1 = m->dof_parentid[da1];
+    }
+    if (da2 == chain[NV]) {
+      da2 = m->dof_parentid[da2];
+    }
+    NV++;
+  }
+
+  // reverse order of chain: make it increasing
+  for (int i=0; i < NV/2; i++) {
+    int tmp = chain[i];
+    chain[i] = chain[NV-i-1];
+    chain[NV-i-1] = tmp;
+  }
+
+  return NV;
+}
+
+
+
+// merge dof chains for two simple bodies
+int mj_mergeChainSimple(const mjModel* m, int* chain, int b1, int b2) {
+  // swap bodies if wrong order
+  if (b1 > b2) {
+    int tmp = b1;
+    b1 = b2;
+    b2 = tmp;
+  }
+
+  // init
+  int n1 = m->body_dofnum[b1], n2 = m->body_dofnum[b2];
+
+  // both fixed: nothing to do
+  if (n1 == 0 && n2 == 0) {
+    return 0;
+  }
+
+  // copy b1 dofs
+  for (int i=0; i < n1; i++) {
+    chain[i] = m->body_dofadr[b1] + i;
+  }
+
+  // copy b2 dofs
+  for (int i=0; i < n2; i++) {
+    chain[n1+i] = m->body_dofadr[b2] + i;
+  }
+
+  return (n1+n2);
+}
+
+
+
+// get body chain
+int mj_bodyChain(const mjModel* m, int body, int* chain) {
+  // simple body
+  if (m->body_simple[body]) {
+    int dofnum = m->body_dofnum[body];
+    for (int i=0; i < dofnum; i++) {
+      chain[i] = m->body_dofadr[body] + i;
+    }
+    return dofnum;
+  }
+
+  // general case
+  else {
+    // skip fixed bodies
+    while (body && !m->body_dofnum[body]) {
+      body = m->body_parentid[body];
+    }
+
+    // not movable: empty chain
+    if (body == 0) {
+      return 0;
+    }
+
+    // initialize last dof
+    int da = m->body_dofadr[body] + m->body_dofnum[body] - 1;
+    int NV = 0;
+
+    // construct chain from child to parent
+    while (da >= 0) {
+      chain[NV++] = da;
+      da = m->dof_parentid[da];
+    }
+
+    // reverse order of chain: make it increasing
+    for (int i=0; i < NV/2; i++) {
+      int tmp = chain[i];
+      chain[i] = chain[NV-i-1];
+      chain[NV-i-1] = tmp;
+    }
+
+    return NV;
+  }
+}
+
+
+
+//-------------------------- Jacobians -------------------------------------------------------------
+
+// compute 3/6-by-nv Jacobian of global point attached to given body
+void mj_jac(const mjModel* m, const mjData* d,
+            mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body) {
+  int nv = m->nv;
+  mjtNum offset[3];
+
+  // clear jacobians, compute offset if required
+  if (jacp) {
+    mju_zero(jacp, 3*nv);
+    mju_sub3(offset, point, d->subtree_com+3*m->body_rootid[body]);
+  }
+  if (jacr) {
+    mju_zero(jacr, 3*nv);
+  }
+
+  // skip fixed bodies
+  while (body && !m->body_dofnum[body]) {
+    body = m->body_parentid[body];
+  }
+
+  // no movable body found: nothing to do
+  if (!body) {
+    return;
+  }
+
+  // get last dof that affects this (as well as the original) body
+  int i = m->body_dofadr[body] + m->body_dofnum[body] - 1;
+
+  // backward pass over dof ancestor chain
+  while (i >= 0) {
+    mjtNum* cdof = d->cdof+6*i;
+
+    // construct rotation jacobian
+    if (jacr) {
+      jacr[i+0*nv] = cdof[0];
+      jacr[i+1*nv] = cdof[1];
+      jacr[i+2*nv] = cdof[2];
+    }
+
+    // construct translation jacobian (correct for rotation)
+    if (jacp) {
+      mjtNum tmp[3];
+      mju_cross(tmp, cdof, offset);
+      jacp[i+0*nv] = cdof[3] + tmp[0];
+      jacp[i+1*nv] = cdof[4] + tmp[1];
+      jacp[i+2*nv] = cdof[5] + tmp[2];
+    }
+
+    // advance to parent dof
+    i = m->dof_parentid[i];
+  }
+}
+
+
+
+// compute body Jacobian
+void mj_jacBody(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int body) {
+  mj_jac(m, d, jacp, jacr, d->xpos+3*body, body);
+}
+
+
+
+// compute body-com Jacobian
+void mj_jacBodyCom(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int body) {
+  mj_jac(m, d, jacp, jacr, d->xipos+3*body, body);
+}
+
+
+
+// compute subtree-com Jacobian
+void mj_jacSubtreeCom(const mjModel* m, mjData* d, mjtNum* jacp, int body) {
+  int nv = m->nv;
+  mj_markStack(d);
+  mjtNum* jacp_b = mjSTACKALLOC(d, 3*nv, mjtNum);
+
+  // clear output
+  mju_zero(jacp, 3*nv);
+
+  // forward pass starting from body
+  for (int b=body; b < m->nbody; b++) {
+    // end of body subtree, break from the loop
+    if (b > body && m->body_parentid[b] < body) {
+      break;
+    }
+
+    // b is in the body subtree, add mass-weighted Jacobian into jacp
+    mj_jac(m, d, jacp_b, NULL, d->xipos+3*b, b);
+    mju_addToScl(jacp, jacp_b, m->body_mass[b], 3*nv);
+  }
+
+  // normalize by subtree mass
+  mju_scl(jacp, jacp, 1/m->body_subtreemass[body], 3*nv);
+
+  mj_freeStack(d);
+}
+
+
+
+// compute geom Jacobian
+void mj_jacGeom(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int geom) {
+  mj_jac(m, d, jacp, jacr, d->geom_xpos + 3*geom, m->geom_bodyid[geom]);
+}
+
+
+
+// compute site Jacobian
+void mj_jacSite(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int site) {
+  mj_jac(m, d, jacp, jacr, d->site_xpos + 3*site, m->site_bodyid[site]);
+}
+
+
+
+// compute translation Jacobian of point, and rotation Jacobian of axis
+void mj_jacPointAxis(const mjModel* m, mjData* d, mjtNum* jacPoint, mjtNum* jacAxis,
+                     const mjtNum point[3], const mjtNum axis[3], int body) {
+  int nv = m->nv;
+
+  // get full Jacobian of point
+  mj_markStack(d);
+  mjtNum* jacp = (jacPoint ? jacPoint : mjSTACKALLOC(d, 3*nv, mjtNum));
+  mjtNum* jacr = mjSTACKALLOC(d, 3*nv, mjtNum);
+  mj_jac(m, d, jacp, jacr, point, body);
+
+  // jacAxis_col = cross(jacr_col, axis)
+  if (jacAxis) {
+    for (int i=0; i < nv; i++) {
+      jacAxis[     i] = jacr[  nv+i]*axis[2] - jacr[2*nv+i]*axis[1];
+      jacAxis[  nv+i] = jacr[2*nv+i]*axis[0] - jacr[     i]*axis[2];
+      jacAxis[2*nv+i] = jacr[     i]*axis[1] - jacr[  nv+i]*axis[0];
+    }
+  }
+
+  mj_freeStack(d);
+}
+
+
+
+// compute 3/6-by-nv sparse Jacobian of global point attached to given body
+void mj_jacSparse(const mjModel* m, const mjData* d,
+                  mjtNum* jacp, mjtNum* jacr, const mjtNum* point, int body,
+                  int NV, const int* chain) {
+  // clear jacobians
+  if (jacp) {
+    mju_zero(jacp, 3*NV);
+  }
+  if (jacr) {
+    mju_zero(jacr, 3*NV);
+  }
+
+  // compute point-com offset
+  mjtNum offset[3];
+  mju_sub3(offset, point, d->subtree_com+3*m->body_rootid[body]);
+
+  // skip fixed bodies
+  while (body && !m->body_dofnum[body]) {
+    body = m->body_parentid[body];
+  }
+
+  // no movable body found: nothing to do
+  if (!body) {
+    return;
+  }
+
+  // get last dof that affects this (as well as the original) body
+  int da = m->body_dofadr[body] + m->body_dofnum[body] - 1;
+
+  // start and the end of the chain (chain is in increasing order)
+  int ci = NV-1;
+
+  // backward pass over dof ancestor chain
+  while (da >= 0) {
+    // find chain index for this dof
+    while (ci >= 0 && chain[ci] > da) {
+      ci--;
+    }
+
+    // make sure we found it; SHOULD NOT OCCUR
+    if (ci < 0 || chain[ci] != da) {
+      mjERROR("dof index %d not found in chain", da);
+    }
+
+    const mjtNum* cdof = d->cdof + 6*da;
+
+    // construct rotation jacobian
+    if (jacr) {
+      jacr[ci+0*NV] = cdof[0];
+      jacr[ci+1*NV] = cdof[1];
+      jacr[ci+2*NV] = cdof[2];
+    }
+
+    // construct translation jacobian (correct for rotation)
+    if (jacp) {
+      mjtNum tmp[3];
+      mju_cross(tmp, cdof, offset);
+
+      jacp[ci+0*NV] = cdof[3] + tmp[0];
+      jacp[ci+1*NV] = cdof[4] + tmp[1];
+      jacp[ci+2*NV] = cdof[5] + tmp[2];
+    }
+
+    // advance to parent dof
+    da = m->dof_parentid[da];
+  }
+}
+
+
+
+// sparse Jacobian difference for simple body contacts
+void mj_jacSparseSimple(const mjModel* m, const mjData* d,
+                        mjtNum* jacdifp, mjtNum* jacdifr, const mjtNum* point,
+                        int body, int flg_second, int NV, int start) {
+  // compute point-com offset
+  mjtNum offset[3];
+  mju_sub3(offset, point, d->subtree_com+3*m->body_rootid[body]);
+
+  // skip fixed body
+  if (!m->body_dofnum[body]) {
+    return;
+  }
+
+  // process dofs
+  int ci = start;
+  int end = m->body_dofadr[body] + m->body_dofnum[body];
+  for (int da=m->body_dofadr[body]; da < end; da++) {
+    mjtNum *cdof = d->cdof+6*da;
+
+    // construct rotation jacobian
+    if (jacdifr) {
+      // plus sign
+      if (flg_second) {
+        jacdifr[ci+0*NV] = cdof[0];
+        jacdifr[ci+1*NV] = cdof[1];
+        jacdifr[ci+2*NV] = cdof[2];
+      }
+
+      // minus sign
+      else {
+        jacdifr[ci+0*NV] = -cdof[0];
+        jacdifr[ci+1*NV] = -cdof[1];
+        jacdifr[ci+2*NV] = -cdof[2];
+      }
+    }
+
+    // construct translation jacobian (correct for rotation)
+    if (jacdifp) {
+      mjtNum tmp[3];
+      mju_cross(tmp, cdof, offset);
+
+      // plus sign
+      if (flg_second) {
+        jacdifp[ci+0*NV] = (cdof[3] + tmp[0]);
+        jacdifp[ci+1*NV] = (cdof[4] + tmp[1]);
+        jacdifp[ci+2*NV] = (cdof[5] + tmp[2]);
+      }
+
+      // minus sign
+      else {
+        jacdifp[ci+0*NV] = -(cdof[3] + tmp[0]);
+        jacdifp[ci+1*NV] = -(cdof[4] + tmp[1]);
+        jacdifp[ci+2*NV] = -(cdof[5] + tmp[2]);
+      }
+    }
+
+    // advance jacdif counter
+    ci++;
+  }
+}
+
+
+
+// dense or sparse Jacobian difference for two body points: pos2 - pos1, global
+int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain,
+                  int b1, int b2, const mjtNum pos1[3], const mjtNum pos2[3],
+                  mjtNum* jac1p, mjtNum* jac2p, mjtNum* jacdifp,
+                  mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr) {
+  int issimple = (m->body_simple[b1] && m->body_simple[b2]);
+  int issparse = mj_isSparse(m);
+  int NV = m->nv;
+
+  // skip if no DOFs
+  if (!NV) {
+    return 0;
+  }
+
+  // construct merged chain of body dofs
+  if (issparse) {
+    if (issimple) {
+      NV = mj_mergeChainSimple(m, chain, b1, b2);
+    } else {
+      NV = mj_mergeChain(m, chain, b1, b2);
+    }
+  }
+
+  // skip if empty chain
+  if (!NV) {
+    return 0;
+  }
+
+  // sparse case
+  if (issparse) {
+    // simple: fast processing
+    if (issimple) {
+      // first body
+      mj_jacSparseSimple(m, d, jacdifp, jacdifr, pos1, b1, 0, NV,
+                         b1 < b2 ? 0 : m->body_dofnum[b2]);
+
+      // second body
+      mj_jacSparseSimple(m, d, jacdifp, jacdifr, pos2, b2, 1, NV,
+                         b2 < b1 ? 0 : m->body_dofnum[b1]);
+    }
+
+    // regular processing
+    else {
+      // Jacobians
+      mj_jacSparse(m, d, jac1p, jac1r, pos1, b1, NV, chain);
+      mj_jacSparse(m, d, jac2p, jac2r, pos2, b2, NV, chain);
+
+      // differences
+      if (jacdifp) {
+        mju_sub(jacdifp, jac2p, jac1p, 3*NV);
+      }
+      if (jacdifr) {
+        mju_sub(jacdifr, jac2r, jac1r, 3*NV);
+      }
+    }
+  }
+
+  // dense case
+  else {
+    // Jacobians
+    mj_jac(m, d, jac1p, jac1r, pos1, b1);
+    mj_jac(m, d, jac2p, jac2r, pos2, b2);
+
+    // differences
+    if (jacdifp) {
+      mju_sub(jacdifp, jac2p, jac1p, 3*NV);
+    }
+    if (jacdifr) {
+      mju_sub(jacdifr, jac2r, jac1r, 3*NV);
+    }
+  }
+
+  return NV;
+}
+
+
+
+// dense or sparse weighted sum of multiple body Jacobians at same point
+int mj_jacSum(const mjModel* m, mjData* d, int* chain,
+              int n, const int* body, const mjtNum* weight,
+              const mjtNum point[3], mjtNum* jac, int flg_rot) {
+  int nv = m->nv, NV;
+  mjtNum* jacp = jac;
+  mjtNum* jacr = flg_rot ? jac + 3*nv : NULL;
+
+  mj_markStack(d);
+  mjtNum* jtmp = mjSTACKALLOC(d, flg_rot ? 6*nv : 3*nv, mjtNum);
+  mjtNum* jp = jtmp;
+  mjtNum* jr = flg_rot ? jtmp + 3*nv : NULL;
+
+  // sparse
+  if (mj_isSparse(m)) {
+    mjtNum* buf = mjSTACKALLOC(d, flg_rot ? 6*nv : 3*nv, mjtNum);
+    int* buf_ind = mjSTACKALLOC(d, nv, int);
+    int* bodychain = mjSTACKALLOC(d, nv, int);
+
+    // set first
+    NV = mj_bodyChain(m, body[0], chain);
+    if (NV) {
+      // get Jacobian
+      if (m->body_simple[body[0]]) {
+        mj_jacSparseSimple(m, d, jacp, jacr, point, body[0], 1, NV, 0);
+      } else {
+        mj_jacSparse(m, d, jacp, jacr, point, body[0], NV, chain);
+      }
+
+      // apply weight
+      mju_scl(jac, jac, weight[0], flg_rot ? 6*NV : 3*NV);
+    }
+
+    // accumulate remaining
+    for (int i=1; i < n; i++) {
+      // get body chain and Jacobian
+      int bodyNV = mj_bodyChain(m, body[i], bodychain);
+      if (!bodyNV) {
+        continue;
+      }
+      if (m->body_simple[body[i]]) {
+        mj_jacSparseSimple(m, d, jp, jr, point, body[i], 1, bodyNV, 0);
+      } else {
+        mj_jacSparse(m, d, jp, jr, point, body[i], bodyNV, bodychain);
+      }
+
+      // combine sparse matrices
+      NV = mju_addToSparseMat(jac, jtmp, nv, flg_rot ? 6 : 3, weight[i],
+                              NV, bodyNV, chain, bodychain, buf, buf_ind);
+    }
+  }
+
+  // dense
+  else {
+    // set first
+    mj_jac(m, d, jacp, jacr, point, body[0]);
+    mju_scl(jac, jac, weight[0], flg_rot ? 6*nv : 3*nv);
+
+    // accumulate remaining
+    for (int i=1; i < n; i++) {
+      mj_jac(m, d, jp, jr, point, body[i]);
+      mju_addToScl(jac, jtmp, weight[i], flg_rot ? 6*nv : 3*nv);
+    }
+
+    NV = nv;
+  }
+
+  mj_freeStack(d);
+
+  return NV;
+}
+
+
+
+// compute 3/6-by-nv Jacobian time derivative of global point attached to given body
+void mj_jacDot(const mjModel* m, const mjData* d,
+               mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body) {
+  int nv = m->nv;
+  mjtNum offset[3];
+  mjtNum pvel[6];  // point velocity (rot:lin order)
+
+  // clear jacobians, compute offset and pvel if required
+  if (jacp) {
+    mju_zero(jacp, 3*nv);
+    const mjtNum* com = d->subtree_com+3*m->body_rootid[body];
+    mju_sub3(offset, point, com);
+    mju_transformSpatial(pvel, d->cvel+6*body, 0, point, com, 0);
+  }
+  if (jacr) {
+    mju_zero(jacr, 3*nv);
+  }
+
+  // skip fixed bodies
+  while (body && !m->body_dofnum[body]) {
+    body = m->body_parentid[body];
+  }
+
+  // no movable body found: nothing to do
+  if (!body) {
+    return;
+  }
+
+  // get last dof that affects this (as well as the original) body
+  int i = m->body_dofadr[body] + m->body_dofnum[body] - 1;
+
+  // backward pass over dof ancestor chain
+  while (i >= 0) {
+    mjtNum cdof_dot[6];
+    mju_copy(cdof_dot, d->cdof_dot+6*i, 6);
+    mjtNum* cdof = d->cdof+6*i;
+
+    // check for quaternion
+    mjtJoint type = m->jnt_type[m->dof_jntid[i]];
+    int dofadr = m->jnt_dofadr[m->dof_jntid[i]];
+    int is_quat = type == mjJNT_BALL || (type == mjJNT_FREE && i >= dofadr + 3);
+
+    // compute cdof_dot for quaternion (use current body cvel)
+    if (is_quat) {
+      mju_crossMotion(cdof_dot, d->cvel+6*m->dof_bodyid[i], cdof);
+    }
+
+    // construct rotation jacobian
+    if (jacr) {
+      jacr[i+0*nv] += cdof_dot[0];
+      jacr[i+1*nv] += cdof_dot[1];
+      jacr[i+2*nv] += cdof_dot[2];
+    }
+
+    // construct translation jacobian (correct for rotation)
+    if (jacp) {
+      // first correction term, account for varying cdof
+      mjtNum tmp1[3];
+      mju_cross(tmp1, cdof_dot, offset);
+
+      // second correction term, account for point translational velocity
+      mjtNum tmp2[3];
+      mju_cross(tmp2, cdof, pvel + 3);
+
+      jacp[i+0*nv] += cdof_dot[3] + tmp1[0] + tmp2[0];
+      jacp[i+1*nv] += cdof_dot[4] + tmp1[1] + tmp2[1];
+      jacp[i+2*nv] += cdof_dot[5] + tmp1[2] + tmp2[2];
+    }
+
+    // advance to parent dof
+    i = m->dof_parentid[i];
+  }
+}
+
+
+
+// compute subtree angular momentum matrix
+void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body) {
+  int nv = m->nv;
+  mj_markStack(d);
+
+  // stack allocations
+  mjtNum* jacp = mjSTACKALLOC(d, 3*nv, mjtNum);
+  mjtNum* jacr = mjSTACKALLOC(d, 3*nv, mjtNum);
+  mjtNum* term1 = mjSTACKALLOC(d, 3*nv, mjtNum);
+  mjtNum* term2 = mjSTACKALLOC(d, 3*nv, mjtNum);
+
+  // clear output
+  mju_zero(mat, 3*nv);
+
+  // save the location of the subtree COM
+  mjtNum subtree_com[3];
+  mju_copy3(subtree_com, d->subtree_com+3*body);
+
+  for (int b=body; b < m->nbody; b++) {
+    // end of body subtree, break from the loop
+    if (b > body && m->body_parentid[b] < body) {
+      break;
+    }
+
+    // linear and angular velocity Jacobian of the body COM (inertial frame)
+    mj_jacBodyCom(m, d, jacp, jacr, b);
+
+    // orientation of the COM (inertial) frame of b-th body
+    mjtNum ximat[9];
+    mju_copy(ximat, d->ximat+9*b, 9);
+
+    // save the inertia matrix of b-th body
+    mjtNum inertia[9] = {0};
+    inertia[0] = m->body_inertia[3*b];   // inertia(1,1)
+    inertia[4] = m->body_inertia[3*b+1]; // inertia(2,2)
+    inertia[8] = m->body_inertia[3*b+2]; // inertia(3,3)
+
+    // term1 = body angular momentum about self COM in world frame
+    mjtNum tmp1[9], tmp2[9];
+    mju_mulMatMat3(tmp1, ximat, inertia);          // tmp1  = ximat * inertia
+    mju_mulMatMatT3(tmp2, tmp1, ximat);            // tmp2  = ximat * inertia * ximat^T
+    mju_mulMatMat(term1, tmp2, jacr, 3, 3, nv);    // term1 = ximat * inertia * ximat^T * jacr
+
+    // location of body COM w.r.t subtree COM
+    mjtNum com[3];
+    mju_sub3(com, d->xipos+3*b, subtree_com);
+
+    // skew symmetric matrix representing body_com vector
+    mjtNum com_mat[9] = {0};
+    com_mat[1] = -com[2];
+    com_mat[2] = com[1];
+    com_mat[3] = com[2];
+    com_mat[5] = -com[0];
+    com_mat[6] = -com[1];
+    com_mat[7] = com[0];
+
+    // term2 = moment of linear momentum
+    mju_mulMatMat(term2, com_mat, jacp, 3, 3, nv);   // term2 = com_mat * jacp
+    mju_scl(term2, term2, m->body_mass[b], 3 * nv);  // term2 = com_mat * jacp * mass
+
+    // mat += term1 + term2
+    mju_addTo(mat, term1, 3*nv);
+    mju_addTo(mat, term2, 3*nv);
+  }
+
+  mj_freeStack(d);
+}
+
+
+
+// count warnings, print only the first time
+void mj_warning(mjData* d, int warning, int info) {
+  // check type
+  if (warning < 0 || warning >= mjNWARNING) {
+    mjERROR("invalid warning type %d", warning);
+  }
+
+  // save info (override previous)
+  d->warning[warning].lastinfo = info;
+
+  // print message only the first time this warning is encountered
+  if (!d->warning[warning].number) {
+    mju_warning("%s Time = %.4f.", mju_warningText(warning, info), d->time);
+  }
+
+  // increase counter
+  d->warning[warning].number++;
+}
+
+
+
+
+
+// compute object 6D velocity in object-centered frame, world/local orientation
+void mj_objectVelocity(const mjModel* m, const mjData* d,
+                       int objtype, int objid, mjtNum res[6], int flg_local) {
+  int bodyid = 0;
+  const mjtNum *pos = 0, *rot = 0;
+
+  // body-inertial
+  if (objtype == mjOBJ_BODY) {
+    bodyid = objid;
+    pos = d->xipos+3*objid;
+    rot = (flg_local ? d->ximat+9*objid : 0);
+  }
+
+  // body-regular
+  else if (objtype == mjOBJ_XBODY) {
+    bodyid = objid;
+    pos = d->xpos+3*objid;
+    rot = (flg_local ? d->xmat+9*objid : 0);
+  }
+
+  // geom
+  else if (objtype == mjOBJ_GEOM) {
+    bodyid = m->geom_bodyid[objid];
+    pos = d->geom_xpos+3*objid;
+    rot = (flg_local ? d->geom_xmat+9*objid : 0);
+  }
+
+  // site
+  else if (objtype == mjOBJ_SITE) {
+    bodyid = m->site_bodyid[objid];
+    pos = d->site_xpos+3*objid;
+    rot = (flg_local ? d->site_xmat+9*objid : 0);
+  }
+
+  // camera
+  else if (objtype == mjOBJ_CAMERA) {
+    bodyid = m->cam_bodyid[objid];
+    pos = d->cam_xpos+3*objid;
+    rot = (flg_local ? d->cam_xmat+9*objid : 0);
+  }
+
+  // object without spatial frame
+  else {
+    mjERROR("invalid object type %d", objtype);
+  }
+
+  // transform velocity
+  mju_transformSpatial(res, d->cvel+6*bodyid, 0, pos, d->subtree_com+3*m->body_rootid[bodyid], rot);
+}
+
+
+
+// compute object 6D acceleration in object-centered frame, world/local orientation
+void mj_objectAcceleration(const mjModel* m, const mjData* d,
+                           int objtype, int objid, mjtNum res[6], int flg_local) {
+  int bodyid = 0;
+  const mjtNum *pos = 0, *rot = 0;
+  mjtNum correction[3], vel[6];
+
+  // body-inertial
+  if (objtype == mjOBJ_BODY) {
+    bodyid = objid;
+    pos = d->xipos+3*objid;
+    rot = (flg_local ? d->ximat+9*objid : 0);
+  }
+
+  // body-regular
+  else if (objtype == mjOBJ_XBODY) {
+    bodyid = objid;
+    pos = d->xpos+3*objid;
+    rot = (flg_local ? d->xmat+9*objid : 0);
+  }
+
+  // geom
+  else if (objtype == mjOBJ_GEOM) {
+    bodyid = m->geom_bodyid[objid];
+    pos = d->geom_xpos+3*objid;
+    rot = (flg_local ? d->geom_xmat+9*objid : 0);
+  }
+
+  // site
+  else if (objtype == mjOBJ_SITE) {
+    bodyid = m->site_bodyid[objid];
+    pos = d->site_xpos+3*objid;
+    rot = (flg_local ? d->site_xmat+9*objid : 0);
+  }
+
+  // camera
+  else if (objtype == mjOBJ_CAMERA) {
+    bodyid = m->cam_bodyid[objid];
+    pos = d->cam_xpos+3*objid;
+    rot = (flg_local ? d->cam_xmat+9*objid : 0);
+  }
+
+  // object without spatial frame
+  else {
+    mjERROR("invalid object type %d", objtype);
+  }
+
+  // transform com-based velocity to local frame
+  mju_transformSpatial(vel, d->cvel+6*bodyid, 0, pos, d->subtree_com+3*m->body_rootid[bodyid], rot);
+
+  // transform com-based acceleration to local frame
+  mju_transformSpatial(res, d->cacc+6*bodyid, 0, pos, d->subtree_com+3*m->body_rootid[bodyid], rot);
+
+  // acc_tran += vel_rot x vel_tran
+  mju_cross(correction, vel, vel+3);
+  mju_addTo3(res+3, correction);
+}
+
+
+// extract 6D force:torque for one contact, in contact frame
+void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]) {
+  mjContact* con;
+
+  // clear result
+  mju_zero(result, 6);
+
+  // make sure contact is valid
+  if (id >= 0 && id < d->ncon && d->contact[id].efc_address >= 0) {
+    // get contact pointer
+    con = d->contact + id;
+
+    if (mj_isPyramidal(m)) {
+      mju_decodePyramid(result, d->efc_force + con->efc_address, con->friction, con->dim);
+    } else {
+      mju_copy(result, d->efc_force + con->efc_address, con->dim);
+    }
+  }
+}
+
+
+// map from body local to global Cartesian coordinates
+void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9],
+                     const mjtNum pos[3], const mjtNum quat[4],
+                     int body, mjtByte sameframe) {
+  mjtSameFrame sf = sameframe;
+
+  // position
+  if (xpos && pos) {
+    switch (sf) {
+    case mjSAMEFRAME_NONE:
+    case mjSAMEFRAME_BODYROT:
+    case mjSAMEFRAME_INERTIAROT:
+      mju_mulMatVec3(xpos, d->xmat+9*body, pos);
+      mju_addTo3(xpos, d->xpos+3*body);
+      break;
+    case mjSAMEFRAME_BODY:
+      mju_copy3(xpos, d->xpos+3*body);
+      break;
+    case mjSAMEFRAME_INERTIA:
+      mju_copy3(xpos, d->xipos+3*body);
+      break;
+    }
+  }
+
+  // orientation
+  if (xmat && quat) {
+    mjtNum tmp[4];
+    switch (sf) {
+    case mjSAMEFRAME_NONE:
+      mju_mulQuat(tmp, d->xquat+4*body, quat);
+      mju_quat2Mat(xmat, tmp);
+      break;
+    case mjSAMEFRAME_BODY:
+    case mjSAMEFRAME_BODYROT:
+      mju_copy(xmat, d->xmat+9*body, 9);
+      break;
+    case mjSAMEFRAME_INERTIA:
+    case mjSAMEFRAME_INERTIAROT:
+      mju_copy(xmat, d->ximat+9*body, 9);
+      break;
+    }
+  }
+}
+
+
+// set default solver parameters
+void mj_defaultSolRefImp(mjtNum* solref, mjtNum* solimp) {
+  if (solref) {
+    solref[0] = 0.02;       // timeconst
+    solref[1] = 1;          // dampratio
+  }
+
+  if (solimp) {
+    solimp[0] = 0.9;        // dmin
+    solimp[1] = 0.95;       // dmax
+    solimp[2] = 0.001;      // width
+    solimp[3] = 0.5;        // midpoint
+    solimp[4] = 2;          // power
+  }
+}

src/engine/engine_core_util.h

@@ -0,0 +1,139 @@
+// Copyright 2025 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_CORE_UTIL_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_CORE_UTIL_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+#include "engine/engine_memory.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// determine type of friction cone
+MJAPI int mj_isPyramidal(const mjModel* m);
+
+// determine type of constraint Jacobian
+MJAPI int mj_isSparse(const mjModel* m);
+
+
+//-------------------------- sparse chains ---------------------------------------------------------
+
+// merge dof chains for two bodies
+int mj_mergeChain(const mjModel* m, int* chain, int b1, int b2);
+
+// merge dof chains for two simple bodies
+int mj_mergeChainSimple(const mjModel* m, int* chain, int b1, int b2);
+
+// get body chain
+int mj_bodyChain(const mjModel* m, int body, int* chain);
+
+
+//-------------------------- Jacobians -------------------------------------------------------------
+
+// compute 3/6-by-nv Jacobian of global point attached to given body
+MJAPI void mj_jac(const mjModel* m, const mjData* d,
+                  mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body);
+
+// compute body frame Jacobian
+MJAPI void mj_jacBody(const mjModel* m, const mjData* d,
+                      mjtNum* jacp, mjtNum* jacr, int body);
+
+// compute body center-of-mass Jacobian
+MJAPI void mj_jacBodyCom(const mjModel* m, const mjData* d,
+                         mjtNum* jacp, mjtNum* jacr, int body);
+
+// compute subtree center-of-mass Jacobian
+MJAPI void mj_jacSubtreeCom(const mjModel* m, mjData* d, mjtNum* jacp, int body);
+
+// compute geom Jacobian
+MJAPI void mj_jacGeom(const mjModel* m, const mjData* d,
+                      mjtNum* jacp, mjtNum* jacr, int geom);
+
+// compute site Jacobian
+MJAPI void mj_jacSite(const mjModel* m, const mjData* d,
+                      mjtNum* jacp, mjtNum* jacr, int site);
+
+// compute translation Jacobian of point, and rotation Jacobian of axis
+MJAPI void mj_jacPointAxis(const mjModel* m, mjData* d,
+                           mjtNum* jacPoint, mjtNum* jacAxis,
+                           const mjtNum point[3], const mjtNum axis[3], int body);
+
+// compute 3/6-by-nv sparse Jacobian of global point attached to given body
+void mj_jacSparse(const mjModel* m, const mjData* d,
+                  mjtNum* jacp, mjtNum* jacr, const mjtNum* point, int body,
+                  int NV, const int* chain);
+
+// sparse Jacobian difference for simple body contacts
+void mj_jacSparseSimple(const mjModel* m, const mjData* d,
+                        mjtNum* jacdifp, mjtNum* jacdifr, const mjtNum* point,
+                        int body, int flg_second, int NV, int start);
+
+// dense or sparse Jacobian difference for two body points: pos2 - pos1, global
+MJAPI int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain,
+                        int b1, int b2, const mjtNum pos1[3], const mjtNum pos2[3],
+                        mjtNum* jac1p, mjtNum* jac2p, mjtNum* jacdifp,
+                        mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr);
+
+// dense or sparse weighted sum of multiple body Jacobians at same point
+int mj_jacSum(const mjModel* m, mjData* d, int* chain,
+              int n, const int* body, const mjtNum* weight,
+              const mjtNum point[3], mjtNum* jac, int flg_rot);
+
+// compute 3/6-by-nv Jacobian time derivative of global point attached to given body
+MJAPI void mj_jacDot(const mjModel* m, const mjData* d,
+                     mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body);
+
+// compute subtree angular momentum matrix
+MJAPI void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body);
+
+
+//-------------------------- coordinate transformation ---------------------------------------------
+
+// compute object 6D velocity in object-centered frame, world/local orientation
+MJAPI void mj_objectVelocity(const mjModel* m, const mjData* d,
+                             int objtype, int objid, mjtNum res[6], int flg_local);
+
+// compute object 6D acceleration in object-centered frame, world/local orientation
+MJAPI void mj_objectAcceleration(const mjModel* m, const mjData* d,
+                                 int objtype, int objid, mjtNum res[6], int flg_local);
+
+
+//-------------------------- miscellaneous ---------------------------------------------------------
+
+// map from body local to global Cartesian coordinates
+MJAPI void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9],
+                           const mjtNum pos[3], const mjtNum quat[4],
+                           int body, mjtByte sameframe);
+
+// extract 6D force:torque for one contact, in contact frame
+MJAPI void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]);
+
+// high-level warning function: count warnings in mjData, print only the first time
+MJAPI void mj_warning(mjData* d, int warning, int info);
+
+// extract 6D force:torque for one contact, in contact frame
+MJAPI void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]);
+
+// set default solver parameters
+MJAPI void mj_defaultSolRefImp(mjtNum* solref, mjtNum* solimp);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_CORE_UTIL_H_

src/engine/engine_derivative.c

@@ -18,8 +18,9 @@
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
 #include "engine/engine_core_constraint.h"
+#include "engine/engine_core_util.h"
 #include "engine/engine_crossplatform.h"
-#include "engine/engine_io.h"
+#include "engine/engine_memory.h"
 #include "engine/engine_passive.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"

src/engine/engine_derivative_fd.c

@@ -24,6 +24,7 @@
 #include "engine/engine_forward.h"
 #include "engine/engine_io.h"
 #include "engine/engine_inverse.h"
+#include "engine/engine_memory.h"
 #include "engine/engine_macro.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"

src/engine/engine_forward.c

@@ -31,6 +31,7 @@
 #include "engine/engine_island.h"
 #include "engine/engine_io.h"
 #include "engine/engine_macro.h"
+#include "engine/engine_memory.h"
 #include "engine/engine_passive.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_sensor.h"

src/engine/engine_inverse.c

@@ -25,6 +25,7 @@
 #include "engine/engine_core_smooth.h"
 #include "engine/engine_derivative.h"
 #include "engine/engine_io.h"
+#include "engine/engine_memory.h"
 #include "engine/engine_macro.h"
 #include "engine/engine_forward.h"
 #include "engine/engine_sensor.h"

src/engine/engine_io.c

@@ -27,13 +27,14 @@
 #include <mujoco/mjplugin.h>
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
 #include <mujoco/mjxmacro.h>
+#include "engine/engine_core_util.h"
 #include "engine/engine_crossplatform.h"
 #include "engine/engine_macro.h"
+#include "engine/engine_memory.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"
-#include "thread/thread_pool.h"
 
 #ifdef ADDRESS_SANITIZER
   #include <sanitizer/asan_interface.h>
@@ -48,29 +49,8 @@
   #pragma warning (disable: 4305)  // disable MSVC warning: truncation from 'double' to 'float'
 #endif
 
-// add red zone padding when built with asan, to detect out-of-bound accesses
-#ifdef ADDRESS_SANITIZER
-  #define mjREDZONE 32
-#else
-  #define mjREDZONE 0
-#endif
-
 static const int MAX_ARRAY_SIZE = INT_MAX / 4;
 
-// compute a % b with a fast code path if the second argument is a power of 2
-static inline size_t fastmod(size_t a, size_t b) {
-  // (b & (b - 1)) == 0 implies that b is a power of 2
-  if (mjLIKELY((b & (b - 1)) == 0)) {
-    return a & (b - 1);
-  }
-  return a % b;
-}
-
-typedef struct {
-  size_t pbase;   // value of d->pbase immediately before mj_markStack
-  size_t pstack;  // value of d->pstack immediately before mj_markStack
-  void* pc;       // program counter of the call site of mj_markStack (only set when under asan)
-} mjStackFrame;
 
 //------------------------------ mjLROpt -----------------------------------------------------------
 
@@ -93,22 +73,6 @@ void mj_defaultLROpt(mjLROpt* opt) {
 
 //------------------------------- mjOption ---------------------------------------------------------
 
-// set default solver parameters
-void mj_defaultSolRefImp(mjtNum* solref, mjtNum* solimp) {
-  if (solref) {
-    solref[0] = 0.02;       // timeconst
-    solref[1] = 1;          // dampratio
-  }
-
-  if (solimp) {
-    solimp[0] = 0.9;        // dmin
-    solimp[1] = 0.95;       // dmax
-    solimp[2] = 0.001;      // width
-    solimp[3] = 0.5;        // midpoint
-    solimp[4] = 2;          // power
-  }
-}
-
 
 
 // set model options to default values
@@ -1513,329 +1477,6 @@ mjData* mjv_copyData(mjData* dest, const mjModel* m, const mjData* src) {
   return mj_copyDataVisual(dest, m, src, /*flg_all=*/0);
 }
 
-static void maybe_lock_alloc_mutex(mjData* d) {
-  if (d->threadpool != 0) {
-    mju_threadPoolLockAllocMutex((mjThreadPool*)d->threadpool);
-  }
-}
-
-static void maybe_unlock_alloc_mutex(mjData* d) {
-  if (d->threadpool != 0) {
-    mju_threadPoolUnlockAllocMutex((mjThreadPool*)d->threadpool);
-  }
-}
-
-
-
-static inline mjStackInfo get_stack_info_from_data(const mjData* d) {
-  mjStackInfo stack_info;
-  stack_info.bottom = (uintptr_t)d->arena + (uintptr_t)d->narena;
-  stack_info.top = stack_info.bottom - d->pstack;
-  stack_info.limit = (uintptr_t)d->arena + (uintptr_t)d->parena;
-  stack_info.stack_base = d->pbase;
-
-  return stack_info;
-}
-
-
-#ifdef ADDRESS_SANITIZER
-// get stack usage from red-zone (under ASAN)
-static size_t stack_usage_redzone(const mjStackInfo* stack_info) {
-  size_t usage = 0;
-
-  // actual stack usage (without red zone bytes) is stored in the red zone
-  if (stack_info->top != stack_info->bottom) {
-    char* prev_pstack_ptr = (char*)(stack_info->top);
-    size_t prev_misalign = (uintptr_t)prev_pstack_ptr % _Alignof(size_t);
-    size_t* prev_usage_ptr =
-      (size_t*)(prev_pstack_ptr +
-                (prev_misalign ? _Alignof(size_t) - prev_misalign : 0));
-    ASAN_UNPOISON_MEMORY_REGION(prev_usage_ptr, sizeof(size_t));
-    usage = *prev_usage_ptr;
-    ASAN_POISON_MEMORY_REGION(prev_usage_ptr, sizeof(size_t));
-  }
-
-  return usage;
-}
-#endif
-
-// allocate memory from the mjData arena
-void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment) {
-  maybe_lock_alloc_mutex(d);
-  size_t misalignment = fastmod(d->parena, alignment);
-  size_t padding = misalignment ? alignment - misalignment : 0;
-
-  // check size
-  size_t bytes_available = d->narena - d->pstack;
-  if (mjUNLIKELY(d->parena + padding + bytes > bytes_available)) {
-    maybe_unlock_alloc_mutex(d);
-    return NULL;
-  }
-
-  size_t stack_usage = d->pstack;
-
-  // under ASAN, get stack usage from red zone
-#ifdef ADDRESS_SANITIZER
-  mjStackInfo stack_info;
-  mjStackInfo* stack_info_ptr;
-  if (!d->threadpool) {
-    stack_info = get_stack_info_from_data(d);
-    stack_info_ptr = &stack_info;
-  } else {
-    size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
-    stack_info_ptr = mju_getStackInfoForThread(d, thread_id);
-  }
-  stack_usage = stack_usage_redzone(stack_info_ptr);
-#endif
-
-  // 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, stack_usage + d->parena);
-
-#ifdef ADDRESS_SANITIZER
-  ASAN_UNPOISON_MEMORY_REGION(result, bytes);
-#endif
-
-#ifdef MEMORY_SANITIZER
-  __msan_allocated_memory(result, bytes);
-#endif
-
-  maybe_unlock_alloc_mutex(d);
-  return result;
-}
-
-
-// internal: allocate size bytes on the provided stack shard
-// declared inline so that modular arithmetic with specific alignments can be optimized out
-static inline void* stackallocinternal(mjData* d, mjStackInfo* stack_info, size_t size,
-    size_t alignment, const char* caller, int line) {
-  // return NULL if empty
-  if (mjUNLIKELY(!size)) {
-    return NULL;
-  }
-
-  // start of the memory to be allocated to the buffer
-  uintptr_t start_ptr = stack_info->top - (size + mjREDZONE);
-
-  // align the pointer
-  start_ptr -= fastmod(start_ptr, alignment);
-
-  // new top of the stack
-  uintptr_t new_top_ptr = start_ptr - mjREDZONE;
-
-  // exclude red zone from stack usage statistics
-  size_t current_alloc_usage = stack_info->top - new_top_ptr - 2 * mjREDZONE;
-  size_t usage = current_alloc_usage + (stack_info->bottom - stack_info->top);
-
-  // check size
-  size_t stack_available_bytes = stack_info->top - stack_info->limit;
-  size_t stack_required_bytes = stack_info->top - new_top_ptr;
-  if (mjUNLIKELY(stack_required_bytes > stack_available_bytes)) {
-    char info[1024];
-    if (caller) {
-      snprintf(info, sizeof(info), " at %s, line %d", caller, line);
-    } else {
-      info[0] = '\0';
-    }
-    mju_error(
-        "mj_stackAlloc: out of memory, stack overflow%s\n"
-        "  max = %" PRIuPTR ", available = %" PRIuPTR ", requested = %" PRIuPTR
-        "\n nefc = %d, ncon = %d",
-        info, stack_info->bottom - stack_info->limit, stack_available_bytes,
-        stack_required_bytes, d->nefc, d->ncon);
-  }
-
-#ifdef ADDRESS_SANITIZER
-  usage = current_alloc_usage + stack_usage_redzone(stack_info);
-
-  // store new stack usage in the red zone
-  size_t misalign = new_top_ptr % _Alignof(size_t);
-  size_t* usage_ptr =
-    (size_t*)(new_top_ptr + (misalign ? _Alignof(size_t) - misalign : 0));
-  ASAN_UNPOISON_MEMORY_REGION(usage_ptr, sizeof(size_t));
-  *usage_ptr = usage;
-  ASAN_POISON_MEMORY_REGION(usage_ptr, sizeof(size_t));
-
-  // unpoison the actual usable allocation
-  ASAN_UNPOISON_MEMORY_REGION((void*)start_ptr, size);
-#endif
-
-  // update max usage statistics
-  stack_info->top = new_top_ptr;
-  if (!d->threadpool) {
-    d->maxuse_stack = mjMAX(d->maxuse_stack, usage);
-    d->maxuse_arena = mjMAX(d->maxuse_arena, usage + d->parena);
-  } else {
-    size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
-    d->maxuse_threadstack[thread_id] = mjMAX(d->maxuse_threadstack[thread_id], usage);
-  }
-
-  return (void*)start_ptr;
-}
-
-
-
-// internal: allocate size bytes in mjData
-// declared inline so that modular arithmetic with specific alignments can be optimized out
-static inline void* stackalloc(mjData* d, size_t size, size_t alignment,
-                               const char* caller, int line) {
-  // single threaded allocation
-  if (!d->threadpool) {
-    mjStackInfo stack_info = get_stack_info_from_data(d);
-    void* result = stackallocinternal(d, &stack_info, size, alignment, caller, line);
-    d->pstack = stack_info.bottom - stack_info.top;
-    return result;
-  }
-
-  // multi threaded allocation
-  size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
-  mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
-  return stackallocinternal(d, stack_info, size, alignment, caller, line);
-}
-
-
-
-// mjStackInfo mark stack frame, inline so ASAN errors point to correct code unit
-#ifdef ADDRESS_SANITIZER
-__attribute__((always_inline))
-#endif
-static inline void markstackinternal(mjData* d, mjStackInfo* stack_info) {
-  size_t top_old = stack_info->top;
-  mjStackFrame* s =
-    (mjStackFrame*) stackallocinternal(d, stack_info, sizeof(mjStackFrame), _Alignof(mjStackFrame), NULL, 0);
-  s->pbase = stack_info->stack_base;
-  s->pstack = top_old;
-#ifdef ADDRESS_SANITIZER
-  // store the program counter to the caller so that we can compare against mj_freeStack later
-  s->pc = __sanitizer_return_address();
-#endif
-  stack_info->stack_base = (uintptr_t) s;
-}
-
-
-
-// mjData mark stack frame
-#ifndef ADDRESS_SANITIZER
-void mj_markStack(mjData* d)
-#else
-void mj__markStack(mjData* d)
-#endif
-{
-  if (!d->threadpool) {
-    mjStackInfo stack_info = get_stack_info_from_data(d);
-    markstackinternal(d, &stack_info);
-    d->pstack = stack_info.bottom - stack_info.top;
-    d->pbase = stack_info.stack_base;
-    return;
-  }
-
-  size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
-  mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
-  markstackinternal(d, stack_info);
-}
-
-
-
-#ifdef ADDRESS_SANITIZER
-__attribute__((always_inline))
-#endif
-static inline void freestackinternal(mjStackInfo* stack_info) {
-  if (mjUNLIKELY(!stack_info->stack_base)) {
-    return;
-  }
-
-  mjStackFrame* s = (mjStackFrame*) stack_info->stack_base;
-#ifdef ADDRESS_SANITIZER
-  // raise an error if caller function name doesn't match the most recent caller of mj_markStack
-  if (!mj__comparePcFuncName(s->pc, __sanitizer_return_address())) {
-    mjERROR("mj_markStack %s has no corresponding mj_freeStack (detected %s)",
-            mj__getPcDebugInfo(s->pc),
-            mj__getPcDebugInfo(__sanitizer_return_address()));
-  }
-#endif
-
-  // restore pbase and pstack
-  stack_info->stack_base = s->pbase;
-  stack_info->top = s->pstack;
-
-  // if running under asan, poison the newly freed memory region
-#ifdef ADDRESS_SANITIZER
-  ASAN_POISON_MEMORY_REGION((char*)stack_info->limit, stack_info->top - stack_info->limit);
-#endif
-}
-
-
-
-// mjData free stack frame
-#ifndef ADDRESS_SANITIZER
-void mj_freeStack(mjData* d)
-#else
-void mj__freeStack(mjData* d)
-#endif
-{
-  if (!d->threadpool) {
-    mjStackInfo stack_info = get_stack_info_from_data(d);
-    freestackinternal(&stack_info);
-    d->pstack = stack_info.bottom - stack_info.top;
-    d->pbase = stack_info.stack_base;
-    return;
-  }
-
-  size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
-  mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
-  freestackinternal(stack_info);
-}
-
-
-
-// returns the number of bytes available on the stack
-size_t mj_stackBytesAvailable(mjData* d) {
-  if (!d->threadpool) {
-    mjStackInfo stack_info = get_stack_info_from_data(d);
-    return stack_info.top - stack_info.limit;
-  } else {
-    size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
-    mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
-    return stack_info->top - stack_info->limit;
-  }
-}
-
-
-
-// allocate bytes on the stack
-void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment) {
-  return stackalloc(d, bytes, alignment, NULL, 0);
-}
-
-
-
-// allocate bytes on the stack, with caller information
-void* mj_stackAllocInfo(mjData* d, size_t bytes, size_t alignment,
-                        const char* caller, int line) {
-  return stackalloc(d, bytes, alignment, caller, line);
-}
-
-
-
-// allocate mjtNums on the stack
-mjtNum* mj_stackAllocNum(mjData* d, size_t size) {
-  if (mjUNLIKELY(size >= SIZE_MAX / sizeof(mjtNum))) {
-    mjERROR("requested size is too large (more than 2^64 bytes).");
-  }
-  return (mjtNum*) stackalloc(d, size * sizeof(mjtNum), _Alignof(mjtNum), NULL, 0);
-}
-
-
-
-// allocate ints on the stack
-int* mj_stackAllocInt(mjData* d, size_t size) {
-  if (mjUNLIKELY(size >= SIZE_MAX / sizeof(int))) {
-    mjERROR("requested size is too large (more than 2^64 bytes).");
-  }
-  return (int*) stackalloc(d, size * sizeof(int), _Alignof(int), NULL, 0);
-}
-
 
 
 // clear data, set defaults

src/engine/engine_io.h

@@ -35,9 +35,6 @@ extern "C" {
 // Set default options for length range computation.
 MJAPI void mj_defaultLROpt(mjLROpt* opt);
 
-// set default solver parameters
-MJAPI void mj_defaultSolRefImp(mjtNum* solref, mjtNum* solimp);
-
 // set options to default values
 MJAPI void mj_defaultOption(mjOption* opt);
 
@@ -130,60 +127,12 @@ MJAPI void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_va
 // Reset data. If 0 <= key < nkey, set fields from specified keyframe.
 MJAPI void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
 
-// mjData arena allocate
-MJAPI void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment);
-
 // init plugins
 MJAPI void mj_initPlugin(const mjModel* m, mjData* d);
 
-#ifndef ADDRESS_SANITIZER
-
-// mjData mark stack frame
-MJAPI void mj_markStack(mjData* d);
-
-// mjData free stack frame
-MJAPI void mj_freeStack(mjData* d);
-
-#else
-
-void mj__markStack(mjData* d) __attribute__((noinline));
-void mj__freeStack(mjData* d) __attribute__((noinline));
-
-#endif  // ADDRESS_SANITIZER
-
-// returns the number of bytes available on the stack
-MJAPI size_t mj_stackBytesAvailable(mjData* d);
-
-// allocate bytes on the stack
-MJAPI void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment);
-
-// allocate bytes on the stack, with added caller information
-MJAPI void* mj_stackAllocInfo(mjData* d, size_t bytes, size_t alignment,
-                              const char* caller, int line);
-
-// macro to allocate a stack array of given type, adds caller information
-#define mjSTACKALLOC(d, num, type) \
-(type*) mj_stackAllocInfo(d, (num) * sizeof(type), _Alignof(type), __func__, __LINE__)
-
-// mjData stack allocate for array of mjtNums
-MJAPI mjtNum* mj_stackAllocNum(mjData* d, size_t size);
-
-// mjData stack allocate for array of ints
-MJAPI int* mj_stackAllocInt(mjData* d, size_t size);
-
 // deallocate data
 MJAPI void mj_deleteData(mjData* d);
 
-// clear arena pointers in mjData
-static inline void mj_clearEfc(mjData* d) {
-#define X(type, name, nr, nc) d->name = NULL;
-  MJDATA_ARENA_POINTERS
-#undef X
-  d->nefc = 0;
-  d->nisland = 0;
-  d->contact = (mjContact*) d->arena;
-}
-
 #ifdef __cplusplus
 }
 #endif

src/engine/engine_island.c

@@ -23,7 +23,9 @@
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
 #include <mujoco/mjxmacro.h>
 #include "engine/engine_core_constraint.h"
+#include "engine/engine_core_util.h"
 #include "engine/engine_io.h"
+#include "engine/engine_memory.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"

src/engine/engine_memory.c

@@ -0,0 +1,394 @@
+// Copyright 2025 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_memory.h"
+
+#include <inttypes.h>  // NOLINT required for PRIu64, PRIuPTR
+#include <limits.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include <mujoco/mjmacro.h>
+#include <mujoco/mjmodel.h>
+#include <mujoco/mjplugin.h>
+#include <mujoco/mjsan.h>  // IWYU pragma: keep
+#include <mujoco/mjxmacro.h>
+#include "engine/engine_crossplatform.h"
+#include "engine/engine_macro.h"
+#include "engine/engine_plugin.h"
+#include "engine/engine_util_blas.h"
+#include "engine/engine_util_errmem.h"
+#include "engine/engine_util_misc.h"
+#include "thread/thread_pool.h"
+
+#ifdef ADDRESS_SANITIZER
+  #include <sanitizer/asan_interface.h>
+  #include <sanitizer/common_interface_defs.h>
+#endif
+
+#ifdef MEMORY_SANITIZER
+  #include <sanitizer/msan_interface.h>
+#endif
+
+#ifdef _MSC_VER
+  #pragma warning (disable: 4305)  // disable MSVC warning: truncation from 'double' to 'float'
+#endif
+
+// add red zone padding when built with asan, to detect out-of-bound accesses
+#ifdef ADDRESS_SANITIZER
+  #define mjREDZONE 32
+#else
+  #define mjREDZONE 0
+#endif
+
+// compute a % b with a fast code path if the second argument is a power of 2
+static inline size_t fastmod(size_t a, size_t b) {
+  // (b & (b - 1)) == 0 implies that b is a power of 2
+  if (mjLIKELY((b & (b - 1)) == 0)) {
+    return a & (b - 1);
+  }
+  return a % b;
+}
+
+typedef struct {
+  size_t pbase;   // value of d->pbase immediately before mj_markStack
+  size_t pstack;  // value of d->pstack immediately before mj_markStack
+  void* pc;       // program counter of the call site of mj_markStack (only set when under asan)
+} mjStackFrame;
+
+static void maybe_lock_alloc_mutex(mjData* d) {
+  if (d->threadpool != 0) {
+    mju_threadPoolLockAllocMutex((mjThreadPool*)d->threadpool);
+  }
+}
+
+static void maybe_unlock_alloc_mutex(mjData* d) {
+  if (d->threadpool != 0) {
+    mju_threadPoolUnlockAllocMutex((mjThreadPool*)d->threadpool);
+  }
+}
+
+
+
+static inline mjStackInfo get_stack_info_from_data(const mjData* d) {
+  mjStackInfo stack_info;
+  stack_info.bottom = (uintptr_t)d->arena + (uintptr_t)d->narena;
+  stack_info.top = stack_info.bottom - d->pstack;
+  stack_info.limit = (uintptr_t)d->arena + (uintptr_t)d->parena;
+  stack_info.stack_base = d->pbase;
+
+  return stack_info;
+}
+
+
+#ifdef ADDRESS_SANITIZER
+// get stack usage from red-zone (under ASAN)
+static size_t stack_usage_redzone(const mjStackInfo* stack_info) {
+  size_t usage = 0;
+
+  // actual stack usage (without red zone bytes) is stored in the red zone
+  if (stack_info->top != stack_info->bottom) {
+    char* prev_pstack_ptr = (char*)(stack_info->top);
+    size_t prev_misalign = (uintptr_t)prev_pstack_ptr % _Alignof(size_t);
+    size_t* prev_usage_ptr =
+      (size_t*)(prev_pstack_ptr +
+                (prev_misalign ? _Alignof(size_t) - prev_misalign : 0));
+    ASAN_UNPOISON_MEMORY_REGION(prev_usage_ptr, sizeof(size_t));
+    usage = *prev_usage_ptr;
+    ASAN_POISON_MEMORY_REGION(prev_usage_ptr, sizeof(size_t));
+  }
+
+  return usage;
+}
+#endif
+
+// allocate memory from the mjData arena
+void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment) {
+  maybe_lock_alloc_mutex(d);
+  size_t misalignment = fastmod(d->parena, alignment);
+  size_t padding = misalignment ? alignment - misalignment : 0;
+
+  // check size
+  size_t bytes_available = d->narena - d->pstack;
+  if (mjUNLIKELY(d->parena + padding + bytes > bytes_available)) {
+    maybe_unlock_alloc_mutex(d);
+    return NULL;
+  }
+
+  size_t stack_usage = d->pstack;
+
+  // under ASAN, get stack usage from red zone
+#ifdef ADDRESS_SANITIZER
+  mjStackInfo stack_info;
+  mjStackInfo* stack_info_ptr;
+  if (!d->threadpool) {
+    stack_info = get_stack_info_from_data(d);
+    stack_info_ptr = &stack_info;
+  } else {
+    size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
+    stack_info_ptr = mju_getStackInfoForThread(d, thread_id);
+  }
+  stack_usage = stack_usage_redzone(stack_info_ptr);
+#endif
+
+  // 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, stack_usage + d->parena);
+
+#ifdef ADDRESS_SANITIZER
+  ASAN_UNPOISON_MEMORY_REGION(result, bytes);
+#endif
+
+#ifdef MEMORY_SANITIZER
+  __msan_allocated_memory(result, bytes);
+#endif
+
+  maybe_unlock_alloc_mutex(d);
+  return result;
+}
+
+
+// internal: allocate size bytes on the provided stack shard
+// declared inline so that modular arithmetic with specific alignments can be optimized out
+static inline void* stackallocinternal(mjData* d, mjStackInfo* stack_info, size_t size,
+    size_t alignment, const char* caller, int line) {
+  // return NULL if empty
+  if (mjUNLIKELY(!size)) {
+    return NULL;
+  }
+
+  // start of the memory to be allocated to the buffer
+  uintptr_t start_ptr = stack_info->top - (size + mjREDZONE);
+
+  // align the pointer
+  start_ptr -= fastmod(start_ptr, alignment);
+
+  // new top of the stack
+  uintptr_t new_top_ptr = start_ptr - mjREDZONE;
+
+  // exclude red zone from stack usage statistics
+  size_t current_alloc_usage = stack_info->top - new_top_ptr - 2 * mjREDZONE;
+  size_t usage = current_alloc_usage + (stack_info->bottom - stack_info->top);
+
+  // check size
+  size_t stack_available_bytes = stack_info->top - stack_info->limit;
+  size_t stack_required_bytes = stack_info->top - new_top_ptr;
+  if (mjUNLIKELY(stack_required_bytes > stack_available_bytes)) {
+    char info[1024];
+    if (caller) {
+      snprintf(info, sizeof(info), " at %s, line %d", caller, line);
+    } else {
+      info[0] = '\0';
+    }
+    mju_error(
+        "mj_stackAlloc: out of memory, stack overflow%s\n"
+        "  max = %" PRIuPTR ", available = %" PRIuPTR ", requested = %" PRIuPTR
+        "\n nefc = %d, ncon = %d",
+        info, stack_info->bottom - stack_info->limit, stack_available_bytes,
+        stack_required_bytes, d->nefc, d->ncon);
+  }
+
+#ifdef ADDRESS_SANITIZER
+  usage = current_alloc_usage + stack_usage_redzone(stack_info);
+
+  // store new stack usage in the red zone
+  size_t misalign = new_top_ptr % _Alignof(size_t);
+  size_t* usage_ptr =
+    (size_t*)(new_top_ptr + (misalign ? _Alignof(size_t) - misalign : 0));
+  ASAN_UNPOISON_MEMORY_REGION(usage_ptr, sizeof(size_t));
+  *usage_ptr = usage;
+  ASAN_POISON_MEMORY_REGION(usage_ptr, sizeof(size_t));
+
+  // unpoison the actual usable allocation
+  ASAN_UNPOISON_MEMORY_REGION((void*)start_ptr, size);
+#endif
+
+  // update max usage statistics
+  stack_info->top = new_top_ptr;
+  if (!d->threadpool) {
+    d->maxuse_stack = mjMAX(d->maxuse_stack, usage);
+    d->maxuse_arena = mjMAX(d->maxuse_arena, usage + d->parena);
+  } else {
+    size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
+    d->maxuse_threadstack[thread_id] = mjMAX(d->maxuse_threadstack[thread_id], usage);
+  }
+
+  return (void*)start_ptr;
+}
+
+
+
+// internal: allocate size bytes in mjData
+// declared inline so that modular arithmetic with specific alignments can be optimized out
+static inline void* stackalloc(mjData* d, size_t size, size_t alignment,
+                               const char* caller, int line) {
+  // single threaded allocation
+  if (!d->threadpool) {
+    mjStackInfo stack_info = get_stack_info_from_data(d);
+    void* result = stackallocinternal(d, &stack_info, size, alignment, caller, line);
+    d->pstack = stack_info.bottom - stack_info.top;
+    return result;
+  }
+
+  // multi threaded allocation
+  size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
+  mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
+  return stackallocinternal(d, stack_info, size, alignment, caller, line);
+}
+
+
+
+// mjStackInfo mark stack frame, inline so ASAN errors point to correct code unit
+#ifdef ADDRESS_SANITIZER
+__attribute__((always_inline))
+#endif
+static inline void markstackinternal(mjData* d, mjStackInfo* stack_info) {
+  size_t top_old = stack_info->top;
+  mjStackFrame* s =
+    (mjStackFrame*) stackallocinternal(d, stack_info, sizeof(mjStackFrame), _Alignof(mjStackFrame), NULL, 0);
+  s->pbase = stack_info->stack_base;
+  s->pstack = top_old;
+#ifdef ADDRESS_SANITIZER
+  // store the program counter to the caller so that we can compare against mj_freeStack later
+  s->pc = __sanitizer_return_address();
+#endif
+  stack_info->stack_base = (uintptr_t) s;
+}
+
+
+
+// mjData mark stack frame
+#ifndef ADDRESS_SANITIZER
+void mj_markStack(mjData* d)
+#else
+void mj__markStack(mjData* d)
+#endif
+{
+  if (!d->threadpool) {
+    mjStackInfo stack_info = get_stack_info_from_data(d);
+    markstackinternal(d, &stack_info);
+    d->pstack = stack_info.bottom - stack_info.top;
+    d->pbase = stack_info.stack_base;
+    return;
+  }
+
+  size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
+  mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
+  markstackinternal(d, stack_info);
+}
+
+
+
+#ifdef ADDRESS_SANITIZER
+__attribute__((always_inline))
+#endif
+static inline void freestackinternal(mjStackInfo* stack_info) {
+  if (mjUNLIKELY(!stack_info->stack_base)) {
+    return;
+  }
+
+  mjStackFrame* s = (mjStackFrame*) stack_info->stack_base;
+#ifdef ADDRESS_SANITIZER
+  // raise an error if caller function name doesn't match the most recent caller of mj_markStack
+  if (!mj__comparePcFuncName(s->pc, __sanitizer_return_address())) {
+    mjERROR("mj_markStack %s has no corresponding mj_freeStack (detected %s)",
+            mj__getPcDebugInfo(s->pc),
+            mj__getPcDebugInfo(__sanitizer_return_address()));
+  }
+#endif
+
+  // restore pbase and pstack
+  stack_info->stack_base = s->pbase;
+  stack_info->top = s->pstack;
+
+  // if running under asan, poison the newly freed memory region
+#ifdef ADDRESS_SANITIZER
+  ASAN_POISON_MEMORY_REGION((char*)stack_info->limit, stack_info->top - stack_info->limit);
+#endif
+}
+
+
+
+// mjData free stack frame
+#ifndef ADDRESS_SANITIZER
+void mj_freeStack(mjData* d)
+#else
+void mj__freeStack(mjData* d)
+#endif
+{
+  if (!d->threadpool) {
+    mjStackInfo stack_info = get_stack_info_from_data(d);
+    freestackinternal(&stack_info);
+    d->pstack = stack_info.bottom - stack_info.top;
+    d->pbase = stack_info.stack_base;
+    return;
+  }
+
+  size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
+  mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
+  freestackinternal(stack_info);
+}
+
+
+
+// returns the number of bytes available on the stack
+size_t mj_stackBytesAvailable(mjData* d) {
+  if (!d->threadpool) {
+    mjStackInfo stack_info = get_stack_info_from_data(d);
+    return stack_info.top - stack_info.limit;
+  } else {
+    size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
+    mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
+    return stack_info->top - stack_info->limit;
+  }
+}
+
+
+
+// allocate bytes on the stack
+void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment) {
+  return stackalloc(d, bytes, alignment, NULL, 0);
+}
+
+
+
+// allocate bytes on the stack, with caller information
+void* mj_stackAllocInfo(mjData* d, size_t bytes, size_t alignment,
+                        const char* caller, int line) {
+  return stackalloc(d, bytes, alignment, caller, line);
+}
+
+
+
+// allocate mjtNums on the stack
+mjtNum* mj_stackAllocNum(mjData* d, size_t size) {
+  if (mjUNLIKELY(size >= SIZE_MAX / sizeof(mjtNum))) {
+    mjERROR("requested size is too large (more than 2^64 bytes).");
+  }
+  return (mjtNum*) stackalloc(d, size * sizeof(mjtNum), _Alignof(mjtNum), NULL, 0);
+}
+
+
+
+// allocate ints on the stack
+int* mj_stackAllocInt(mjData* d, size_t size) {
+  if (mjUNLIKELY(size >= SIZE_MAX / sizeof(int))) {
+    mjERROR("requested size is too large (more than 2^64 bytes).");
+  }
+  return (int*) stackalloc(d, size * sizeof(int), _Alignof(int), NULL, 0);
+}

src/engine/engine_memory.h

@@ -0,0 +1,85 @@
+// Copyright 2025 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_MEMORY_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_MEMORY_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+#include <mujoco/mjxmacro.h>
+
+#ifdef __cplusplus
+#include <cstddef>
+extern "C" {
+#else
+#include <stddef.h>
+#endif
+
+// internal hash map size factor (2 corresponds to a load factor of 0.5)
+#define mjLOAD_MULTIPLE 2
+
+// mjData arena allocate
+MJAPI void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment);
+
+#ifndef ADDRESS_SANITIZER
+
+// mjData mark stack frame
+MJAPI void mj_markStack(mjData* d);
+
+// mjData free stack frame
+MJAPI void mj_freeStack(mjData* d);
+
+#else
+
+void mj__markStack(mjData* d) __attribute__((noinline));
+void mj__freeStack(mjData* d) __attribute__((noinline));
+
+#endif  // ADDRESS_SANITIZER
+
+// returns the number of bytes available on the stack
+MJAPI size_t mj_stackBytesAvailable(mjData* d);
+
+// allocate bytes on the stack
+MJAPI void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment);
+
+// allocate bytes on the stack, with added caller information
+MJAPI void* mj_stackAllocInfo(mjData* d, size_t bytes, size_t alignment,
+                              const char* caller, int line);
+
+// macro to allocate a stack array of given type, adds caller information
+#define mjSTACKALLOC(d, num, type) \
+(type*) mj_stackAllocInfo(d, (num) * sizeof(type), _Alignof(type), __func__, __LINE__)
+
+// mjData stack allocate for array of mjtNums
+MJAPI mjtNum* mj_stackAllocNum(mjData* d, size_t size);
+
+// mjData stack allocate for array of ints
+MJAPI int* mj_stackAllocInt(mjData* d, size_t size);
+
+// clear arena pointers in mjData
+static inline void mj_clearEfc(mjData* d) {
+#define X(type, name, nr, nc) d->name = NULL;
+  MJDATA_ARENA_POINTERS
+#undef X
+  d->nefc = 0;
+  d->nisland = 0;
+  d->contact = (mjContact*) d->arena;
+}
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_MEMORY_H_

src/engine/engine_passive.c

@@ -22,8 +22,10 @@
 #include <mujoco/mjmodel.h>
 #include "engine/engine_callback.h"
 #include "engine/engine_core_constraint.h"
+#include "engine/engine_core_util.h"
 #include "engine/engine_crossplatform.h"
 #include "engine/engine_io.h"
+#include "engine/engine_memory.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"

src/engine/engine_print.c

@@ -25,6 +25,7 @@
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
 #include <mujoco/mjxmacro.h>
 #include "engine/engine_core_constraint.h"
+#include "engine/engine_core_util.h"
 #include "engine/engine_io.h"
 #include "engine/engine_name.h"
 #include "engine/engine_macro.h"

src/engine/engine_ray.c

@@ -24,6 +24,7 @@
 #include <mujoco/mjvisualize.h>
 #include "engine/engine_collision_sdf.h"
 #include "engine/engine_io.h"
+#include "engine/engine_memory.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"

src/engine/engine_sensor.c

@@ -23,8 +23,9 @@
 #include "engine/engine_callback.h"
 #include "engine/engine_collision_sdf.h"
 #include "engine/engine_core_smooth.h"
+#include "engine/engine_core_util.h"
 #include "engine/engine_crossplatform.h"
-#include "engine/engine_io.h"
+#include "engine/engine_memory.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_ray.h"
 #include "engine/engine_sort.h"

src/engine/engine_setconst.c

@@ -23,9 +23,10 @@
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_core_smooth.h"
+#include "engine/engine_core_util.h"
 #include "engine/engine_forward.h"
 #include "engine/engine_io.h"
-#include "engine/engine_support.h"
+#include "engine/engine_memory.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"

src/engine/engine_solver.c

@@ -23,7 +23,9 @@
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_core_smooth.h"
+#include "engine/engine_core_util.h"
 #include "engine/engine_io.h"
+#include "engine/engine_memory.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"

src/engine/engine_support.c

@@ -23,9 +23,9 @@
 #include "engine/engine_collision_driver.h"
 #include "engine/engine_collision_gjk.h"
 #include "engine/engine_collision_primitive.h"
-#include "engine/engine_core_constraint.h"
+#include "engine/engine_core_util.h"
 #include "engine/engine_crossplatform.h"
-#include "engine/engine_io.h"
+#include "engine/engine_memory.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"
@@ -270,700 +270,6 @@ void mj_setKeyframe(mjModel* m, const mjData* d, int k) {
 
 
 
-//-------------------------- sparse chains ---------------------------------------------------------
-
-// merge dof chains for two bodies
-int mj_mergeChain(const mjModel* m, int* chain, int b1, int b2) {
-  int da1, da2, NV = 0;
-
-  // skip fixed bodies
-  while (b1 && !m->body_dofnum[b1]) {
-    b1 = m->body_parentid[b1];
-  }
-  while (b2 && !m->body_dofnum[b2]) {
-    b2 = m->body_parentid[b2];
-  }
-
-  // neither body is movable: empty chain
-  if (b1 == 0 && b2 == 0) {
-    return 0;
-  }
-
-  // initialize last dof address for each body
-  da1 = m->body_dofadr[b1] + m->body_dofnum[b1] - 1;
-  da2 = m->body_dofadr[b2] + m->body_dofnum[b2] - 1;
-
-  // merge chains
-  while (da1 >= 0 || da2 >= 0) {
-    chain[NV] = mjMAX(da1, da2);
-    if (da1 == chain[NV]) {
-      da1 = m->dof_parentid[da1];
-    }
-    if (da2 == chain[NV]) {
-      da2 = m->dof_parentid[da2];
-    }
-    NV++;
-  }
-
-  // reverse order of chain: make it increasing
-  for (int i=0; i < NV/2; i++) {
-    int tmp = chain[i];
-    chain[i] = chain[NV-i-1];
-    chain[NV-i-1] = tmp;
-  }
-
-  return NV;
-}
-
-
-
-// merge dof chains for two simple bodies
-int mj_mergeChainSimple(const mjModel* m, int* chain, int b1, int b2) {
-  // swap bodies if wrong order
-  if (b1 > b2) {
-    int tmp = b1;
-    b1 = b2;
-    b2 = tmp;
-  }
-
-  // init
-  int n1 = m->body_dofnum[b1], n2 = m->body_dofnum[b2];
-
-  // both fixed: nothing to do
-  if (n1 == 0 && n2 == 0) {
-    return 0;
-  }
-
-  // copy b1 dofs
-  for (int i=0; i < n1; i++) {
-    chain[i] = m->body_dofadr[b1] + i;
-  }
-
-  // copy b2 dofs
-  for (int i=0; i < n2; i++) {
-    chain[n1+i] = m->body_dofadr[b2] + i;
-  }
-
-  return (n1+n2);
-}
-
-
-
-// get body chain
-int mj_bodyChain(const mjModel* m, int body, int* chain) {
-  // simple body
-  if (m->body_simple[body]) {
-    int dofnum = m->body_dofnum[body];
-    for (int i=0; i < dofnum; i++) {
-      chain[i] = m->body_dofadr[body] + i;
-    }
-    return dofnum;
-  }
-
-  // general case
-  else {
-    // skip fixed bodies
-    while (body && !m->body_dofnum[body]) {
-      body = m->body_parentid[body];
-    }
-
-    // not movable: empty chain
-    if (body == 0) {
-      return 0;
-    }
-
-    // initialize last dof
-    int da = m->body_dofadr[body] + m->body_dofnum[body] - 1;
-    int NV = 0;
-
-    // construct chain from child to parent
-    while (da >= 0) {
-      chain[NV++] = da;
-      da = m->dof_parentid[da];
-    }
-
-    // reverse order of chain: make it increasing
-    for (int i=0; i < NV/2; i++) {
-      int tmp = chain[i];
-      chain[i] = chain[NV-i-1];
-      chain[NV-i-1] = tmp;
-    }
-
-    return NV;
-  }
-}
-
-
-
-//-------------------------- Jacobians -------------------------------------------------------------
-
-// compute 3/6-by-nv Jacobian of global point attached to given body
-void mj_jac(const mjModel* m, const mjData* d,
-            mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body) {
-  int nv = m->nv;
-  mjtNum offset[3];
-
-  // clear jacobians, compute offset if required
-  if (jacp) {
-    mju_zero(jacp, 3*nv);
-    mju_sub3(offset, point, d->subtree_com+3*m->body_rootid[body]);
-  }
-  if (jacr) {
-    mju_zero(jacr, 3*nv);
-  }
-
-  // skip fixed bodies
-  while (body && !m->body_dofnum[body]) {
-    body = m->body_parentid[body];
-  }
-
-  // no movable body found: nothing to do
-  if (!body) {
-    return;
-  }
-
-  // get last dof that affects this (as well as the original) body
-  int i = m->body_dofadr[body] + m->body_dofnum[body] - 1;
-
-  // backward pass over dof ancestor chain
-  while (i >= 0) {
-    mjtNum* cdof = d->cdof+6*i;
-
-    // construct rotation jacobian
-    if (jacr) {
-      jacr[i+0*nv] = cdof[0];
-      jacr[i+1*nv] = cdof[1];
-      jacr[i+2*nv] = cdof[2];
-    }
-
-    // construct translation jacobian (correct for rotation)
-    if (jacp) {
-      mjtNum tmp[3];
-      mju_cross(tmp, cdof, offset);
-      jacp[i+0*nv] = cdof[3] + tmp[0];
-      jacp[i+1*nv] = cdof[4] + tmp[1];
-      jacp[i+2*nv] = cdof[5] + tmp[2];
-    }
-
-    // advance to parent dof
-    i = m->dof_parentid[i];
-  }
-}
-
-
-
-// compute body Jacobian
-void mj_jacBody(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int body) {
-  mj_jac(m, d, jacp, jacr, d->xpos+3*body, body);
-}
-
-
-
-// compute body-com Jacobian
-void mj_jacBodyCom(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int body) {
-  mj_jac(m, d, jacp, jacr, d->xipos+3*body, body);
-}
-
-
-
-// compute subtree-com Jacobian
-void mj_jacSubtreeCom(const mjModel* m, mjData* d, mjtNum* jacp, int body) {
-  int nv = m->nv;
-  mj_markStack(d);
-  mjtNum* jacp_b = mjSTACKALLOC(d, 3*nv, mjtNum);
-
-  // clear output
-  mju_zero(jacp, 3*nv);
-
-  // forward pass starting from body
-  for (int b=body; b < m->nbody; b++) {
-    // end of body subtree, break from the loop
-    if (b > body && m->body_parentid[b] < body) {
-      break;
-    }
-
-    // b is in the body subtree, add mass-weighted Jacobian into jacp
-    mj_jac(m, d, jacp_b, NULL, d->xipos+3*b, b);
-    mju_addToScl(jacp, jacp_b, m->body_mass[b], 3*nv);
-  }
-
-  // normalize by subtree mass
-  mju_scl(jacp, jacp, 1/m->body_subtreemass[body], 3*nv);
-
-  mj_freeStack(d);
-}
-
-
-
-// compute geom Jacobian
-void mj_jacGeom(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int geom) {
-  mj_jac(m, d, jacp, jacr, d->geom_xpos + 3*geom, m->geom_bodyid[geom]);
-}
-
-
-
-// compute site Jacobian
-void mj_jacSite(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int site) {
-  mj_jac(m, d, jacp, jacr, d->site_xpos + 3*site, m->site_bodyid[site]);
-}
-
-
-
-// compute translation Jacobian of point, and rotation Jacobian of axis
-void mj_jacPointAxis(const mjModel* m, mjData* d, mjtNum* jacPoint, mjtNum* jacAxis,
-                     const mjtNum point[3], const mjtNum axis[3], int body) {
-  int nv = m->nv;
-
-  // get full Jacobian of point
-  mj_markStack(d);
-  mjtNum* jacp = (jacPoint ? jacPoint : mjSTACKALLOC(d, 3*nv, mjtNum));
-  mjtNum* jacr = mjSTACKALLOC(d, 3*nv, mjtNum);
-  mj_jac(m, d, jacp, jacr, point, body);
-
-  // jacAxis_col = cross(jacr_col, axis)
-  if (jacAxis) {
-    for (int i=0; i < nv; i++) {
-      jacAxis[     i] = jacr[  nv+i]*axis[2] - jacr[2*nv+i]*axis[1];
-      jacAxis[  nv+i] = jacr[2*nv+i]*axis[0] - jacr[     i]*axis[2];
-      jacAxis[2*nv+i] = jacr[     i]*axis[1] - jacr[  nv+i]*axis[0];
-    }
-  }
-
-  mj_freeStack(d);
-}
-
-
-
-// compute 3/6-by-nv sparse Jacobian of global point attached to given body
-void mj_jacSparse(const mjModel* m, const mjData* d,
-                  mjtNum* jacp, mjtNum* jacr, const mjtNum* point, int body,
-                  int NV, const int* chain) {
-  // clear jacobians
-  if (jacp) {
-    mju_zero(jacp, 3*NV);
-  }
-  if (jacr) {
-    mju_zero(jacr, 3*NV);
-  }
-
-  // compute point-com offset
-  mjtNum offset[3];
-  mju_sub3(offset, point, d->subtree_com+3*m->body_rootid[body]);
-
-  // skip fixed bodies
-  while (body && !m->body_dofnum[body]) {
-    body = m->body_parentid[body];
-  }
-
-  // no movable body found: nothing to do
-  if (!body) {
-    return;
-  }
-
-  // get last dof that affects this (as well as the original) body
-  int da = m->body_dofadr[body] + m->body_dofnum[body] - 1;
-
-  // start and the end of the chain (chain is in increasing order)
-  int ci = NV-1;
-
-  // backward pass over dof ancestor chain
-  while (da >= 0) {
-    // find chain index for this dof
-    while (ci >= 0 && chain[ci] > da) {
-      ci--;
-    }
-
-    // make sure we found it; SHOULD NOT OCCUR
-    if (chain[ci] != da) {
-      mjERROR("dof index %d not found in chain", da);
-    }
-
-    const mjtNum* cdof = d->cdof + 6*da;
-
-    // construct rotation jacobian
-    if (jacr) {
-      jacr[ci+0*NV] = cdof[0];
-      jacr[ci+1*NV] = cdof[1];
-      jacr[ci+2*NV] = cdof[2];
-    }
-
-    // construct translation jacobian (correct for rotation)
-    if (jacp) {
-      mjtNum tmp[3];
-      mju_cross(tmp, cdof, offset);
-
-      jacp[ci+0*NV] = cdof[3] + tmp[0];
-      jacp[ci+1*NV] = cdof[4] + tmp[1];
-      jacp[ci+2*NV] = cdof[5] + tmp[2];
-    }
-
-    // advance to parent dof
-    da = m->dof_parentid[da];
-  }
-}
-
-
-
-// sparse Jacobian difference for simple body contacts
-void mj_jacSparseSimple(const mjModel* m, const mjData* d,
-                        mjtNum* jacdifp, mjtNum* jacdifr, const mjtNum* point,
-                        int body, int flg_second, int NV, int start) {
-  // compute point-com offset
-  mjtNum offset[3];
-  mju_sub3(offset, point, d->subtree_com+3*m->body_rootid[body]);
-
-  // skip fixed body
-  if (!m->body_dofnum[body]) {
-    return;
-  }
-
-  // process dofs
-  int ci = start;
-  int end = m->body_dofadr[body] + m->body_dofnum[body];
-  for (int da=m->body_dofadr[body]; da < end; da++) {
-    mjtNum *cdof = d->cdof+6*da;
-
-    // construct rotation jacobian
-    if (jacdifr) {
-      // plus sign
-      if (flg_second) {
-        jacdifr[ci+0*NV] = cdof[0];
-        jacdifr[ci+1*NV] = cdof[1];
-        jacdifr[ci+2*NV] = cdof[2];
-      }
-
-      // minus sign
-      else {
-        jacdifr[ci+0*NV] = -cdof[0];
-        jacdifr[ci+1*NV] = -cdof[1];
-        jacdifr[ci+2*NV] = -cdof[2];
-      }
-    }
-
-    // construct translation jacobian (correct for rotation)
-    if (jacdifp) {
-      mjtNum tmp[3];
-      mju_cross(tmp, cdof, offset);
-
-      // plus sign
-      if (flg_second) {
-        jacdifp[ci+0*NV] = (cdof[3] + tmp[0]);
-        jacdifp[ci+1*NV] = (cdof[4] + tmp[1]);
-        jacdifp[ci+2*NV] = (cdof[5] + tmp[2]);
-      }
-
-      // plus sign
-      else {
-        jacdifp[ci+0*NV] = -(cdof[3] + tmp[0]);
-        jacdifp[ci+1*NV] = -(cdof[4] + tmp[1]);
-        jacdifp[ci+2*NV] = -(cdof[5] + tmp[2]);
-      }
-    }
-
-    // advance jacdif counter
-    ci++;
-  }
-}
-
-
-
-// dense or sparse Jacobian difference for two body points: pos2 - pos1, global
-int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain,
-                  int b1, int b2, const mjtNum pos1[3], const mjtNum pos2[3],
-                  mjtNum* jac1p, mjtNum* jac2p, mjtNum* jacdifp,
-                  mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr) {
-  int issimple = (m->body_simple[b1] && m->body_simple[b2]);
-  int issparse = mj_isSparse(m);
-  int NV = m->nv;
-
-  // skip if no DOFs
-  if (!NV) {
-    return 0;
-  }
-
-  // construct merged chain of body dofs
-  if (issparse) {
-    if (issimple) {
-      NV = mj_mergeChainSimple(m, chain, b1, b2);
-    } else {
-      NV = mj_mergeChain(m, chain, b1, b2);
-    }
-  }
-
-  // skip if empty chain
-  if (!NV) {
-    return 0;
-  }
-
-  // sparse case
-  if (issparse) {
-    // simple: fast processing
-    if (issimple) {
-      // first body
-      mj_jacSparseSimple(m, d, jacdifp, jacdifr, pos1, b1, 0, NV,
-                         b1 < b2 ? 0 : m->body_dofnum[b2]);
-
-      // second body
-      mj_jacSparseSimple(m, d, jacdifp, jacdifr, pos2, b2, 1, NV,
-                         b2 < b1 ? 0 : m->body_dofnum[b1]);
-    }
-
-    // regular processing
-    else {
-      // Jacobians
-      mj_jacSparse(m, d, jac1p, jac1r, pos1, b1, NV, chain);
-      mj_jacSparse(m, d, jac2p, jac2r, pos2, b2, NV, chain);
-
-      // differences
-      if (jacdifp) {
-        mju_sub(jacdifp, jac2p, jac1p, 3*NV);
-      }
-      if (jacdifr) {
-        mju_sub(jacdifr, jac2r, jac1r, 3*NV);
-      }
-    }
-  }
-
-  // dense case
-  else {
-    // Jacobians
-    mj_jac(m, d, jac1p, jac1r, pos1, b1);
-    mj_jac(m, d, jac2p, jac2r, pos2, b2);
-
-    // differences
-    if (jacdifp) {
-      mju_sub(jacdifp, jac2p, jac1p, 3*NV);
-    }
-    if (jacdifr) {
-      mju_sub(jacdifr, jac2r, jac1r, 3*NV);
-    }
-  }
-
-  return NV;
-}
-
-
-
-// dense or sparse weighted sum of multiple body Jacobians at same point
-int mj_jacSum(const mjModel* m, mjData* d, int* chain,
-              int n, const int* body, const mjtNum* weight,
-              const mjtNum point[3], mjtNum* jac, int flg_rot) {
-  int nv = m->nv, NV;
-  mjtNum* jacp = jac;
-  mjtNum* jacr = flg_rot ? jac + 3*nv : NULL;
-
-  mj_markStack(d);
-  mjtNum* jtmp = mjSTACKALLOC(d, flg_rot ? 6*nv : 3*nv, mjtNum);
-  mjtNum* jp = jtmp;
-  mjtNum* jr = flg_rot ? jtmp + 3*nv : NULL;
-
-  // sparse
-  if (mj_isSparse(m)) {
-    mjtNum* buf = mjSTACKALLOC(d, flg_rot ? 6*nv : 3*nv, mjtNum);
-    int* buf_ind = mjSTACKALLOC(d, nv, int);
-    int* bodychain = mjSTACKALLOC(d, nv, int);
-
-    // set first
-    NV = mj_bodyChain(m, body[0], chain);
-    if (NV) {
-      // get Jacobian
-      if (m->body_simple[body[0]]) {
-        mj_jacSparseSimple(m, d, jacp, jacr, point, body[0], 1, NV, 0);
-      } else {
-        mj_jacSparse(m, d, jacp, jacr, point, body[0], NV, chain);
-      }
-
-      // apply weight
-      mju_scl(jac, jac, weight[0], flg_rot ? 6*NV : 3*NV);
-    }
-
-    // accumulate remaining
-    for (int i=1; i < n; i++) {
-      // get body chain and Jacobian
-      int bodyNV = mj_bodyChain(m, body[i], bodychain);
-      if (!bodyNV) {
-        continue;
-      }
-      if (m->body_simple[body[i]]) {
-        mj_jacSparseSimple(m, d, jp, jr, point, body[i], 1, bodyNV, 0);
-      } else {
-        mj_jacSparse(m, d, jp, jr, point, body[i], bodyNV, bodychain);
-      }
-
-      // combine sparse matrices
-      NV = mju_addToSparseMat(jac, jtmp, nv, flg_rot ? 6 : 3, weight[i],
-                              NV, bodyNV, chain, bodychain, buf, buf_ind);
-    }
-  }
-
-  // dense
-  else {
-    // set first
-    mj_jac(m, d, jacp, jacr, point, body[0]);
-    mju_scl(jac, jac, weight[0], flg_rot ? 6*nv : 3*nv);
-
-    // accumulate remaining
-    for (int i=1; i < n; i++) {
-      mj_jac(m, d, jp, jr, point, body[i]);
-      mju_addToScl(jac, jtmp, weight[i], flg_rot ? 6*nv : 3*nv);
-    }
-
-    NV = nv;
-  }
-
-  mj_freeStack(d);
-
-  return NV;
-}
-
-
-
-// compute 3/6-by-nv Jacobian time derivative of global point attached to given body
-void mj_jacDot(const mjModel* m, const mjData* d,
-               mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body) {
-  int nv = m->nv;
-  mjtNum offset[3];
-  mjtNum pvel[6];  // point velocity (rot:lin order)
-
-  // clear jacobians, compute offset and pvel if required
-  if (jacp) {
-    mju_zero(jacp, 3*nv);
-    const mjtNum* com = d->subtree_com+3*m->body_rootid[body];
-    mju_sub3(offset, point, com);
-    mju_transformSpatial(pvel, d->cvel+6*body, 0, point, com, 0);
-  }
-  if (jacr) {
-    mju_zero(jacr, 3*nv);
-  }
-
-  // skip fixed bodies
-  while (body && !m->body_dofnum[body]) {
-    body = m->body_parentid[body];
-  }
-
-  // no movable body found: nothing to do
-  if (!body) {
-    return;
-  }
-
-  // get last dof that affects this (as well as the original) body
-  int i = m->body_dofadr[body] + m->body_dofnum[body] - 1;
-
-  // backward pass over dof ancestor chain
-  while (i >= 0) {
-    mjtNum cdof_dot[6];
-    mju_copy(cdof_dot, d->cdof_dot+6*i, 6);
-    mjtNum* cdof = d->cdof+6*i;
-
-    // check for quaternion
-    mjtJoint type = m->jnt_type[m->dof_jntid[i]];
-    int dofadr = m->jnt_dofadr[m->dof_jntid[i]];
-    int is_quat = type == mjJNT_BALL || (type == mjJNT_FREE && i >= dofadr + 3);
-
-    // compute cdof_dot for quaternion (use current body cvel)
-    if (is_quat) {
-      mju_crossMotion(cdof_dot, d->cvel+6*m->dof_bodyid[i], cdof);
-    }
-
-    // construct rotation jacobian
-    if (jacr) {
-      jacr[i+0*nv] += cdof_dot[0];
-      jacr[i+1*nv] += cdof_dot[1];
-      jacr[i+2*nv] += cdof_dot[2];
-    }
-
-    // construct translation jacobian (correct for rotation)
-    if (jacp) {
-      // first correction term, account for varying cdof
-      mjtNum tmp1[3];
-      mju_cross(tmp1, cdof_dot, offset);
-
-      // second correction term, account for point translational velocity
-      mjtNum tmp2[3];
-      mju_cross(tmp2, cdof, pvel + 3);
-
-      jacp[i+0*nv] += cdof_dot[3] + tmp1[0] + tmp2[0];
-      jacp[i+1*nv] += cdof_dot[4] + tmp1[1] + tmp2[1];
-      jacp[i+2*nv] += cdof_dot[5] + tmp1[2] + tmp2[2];
-    }
-
-    // advance to parent dof
-    i = m->dof_parentid[i];
-  }
-}
-
-
-
-// compute subtree angular momentum matrix
-void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body) {
-  int nv = m->nv;
-  mj_markStack(d);
-
-  // stack allocations
-  mjtNum* jacp = mjSTACKALLOC(d, 3*nv, mjtNum);
-  mjtNum* jacr = mjSTACKALLOC(d, 3*nv, mjtNum);
-  mjtNum* term1 = mjSTACKALLOC(d, 3*nv, mjtNum);
-  mjtNum* term2 = mjSTACKALLOC(d, 3*nv, mjtNum);
-
-  // clear output
-  mju_zero(mat, 3*nv);
-
-  // save the location of the subtree COM
-  mjtNum subtree_com[3];
-  mju_copy3(subtree_com, d->subtree_com+3*body);
-
-  for (int b=body; b < m->nbody; b++) {
-    // end of body subtree, break from the loop
-    if (b > body && m->body_parentid[b] < body) {
-      break;
-    }
-
-    // linear and angular velocity Jacobian of the body COM (inertial frame)
-    mj_jacBodyCom(m, d, jacp, jacr, b);
-
-    // orientation of the COM (inertial) frame of b-th body
-    mjtNum ximat[9];
-    mju_copy(ximat, d->ximat+9*b, 9);
-
-    // save the inertia matrix of b-th body
-    mjtNum inertia[9] = {0};
-    inertia[0] = m->body_inertia[3*b];   // inertia(1,1)
-    inertia[4] = m->body_inertia[3*b+1]; // inertia(2,2)
-    inertia[8] = m->body_inertia[3*b+2]; // inertia(3,3)
-
-    // term1 = body angular momentum about self COM in world frame
-    mjtNum tmp1[9], tmp2[9];
-    mju_mulMatMat3(tmp1, ximat, inertia);          // tmp1  = ximat * inertia
-    mju_mulMatMatT3(tmp2, tmp1, ximat);            // tmp2  = ximat * inertia * ximat^T
-    mju_mulMatMat(term1, tmp2, jacr, 3, 3, nv);    // term1 = ximat * inertia * ximat^T * jacr
-
-    // location of body COM w.r.t subtree COM
-    mjtNum com[3];
-    mju_sub3(com, d->xipos+3*b, subtree_com);
-
-    // skew symmetric matrix representing body_com vector
-    mjtNum com_mat[9] = {0};
-    com_mat[1] = -com[2];
-    com_mat[2] = com[1];
-    com_mat[3] = com[2];
-    com_mat[5] = -com[0];
-    com_mat[6] = -com[1];
-    com_mat[7] = com[0];
-
-    // term2 = moment of linear momentum
-    mju_mulMatMat(term2, com_mat, jacp, 3, 3, nv);   // term2 = com_mat * jacp
-    mju_scl(term2, term2, m->body_mass[b], 3 * nv);  // term2 = com_mat * jacp * mass
-
-    // mat += term1 + term2
-    mju_addTo(mat, term1, 3*nv);
-    mju_addTo(mat, term2, 3*nv);
-  }
-
-  mj_freeStack(d);
-}
-
-
-
 //-------------------------- inertia functions -----------------------------------------------------
 
 // convert sparse inertia matrix M into full matrix
@@ -1118,117 +424,6 @@ void mj_xfrcAccumulate(const mjModel* m, mjData* d, mjtNum* qfrc) {
 
 
 
-// compute object 6D velocity in object-centered frame, world/local orientation
-void mj_objectVelocity(const mjModel* m, const mjData* d,
-                       int objtype, int objid, mjtNum res[6], int flg_local) {
-  int bodyid = 0;
-  const mjtNum *pos = 0, *rot = 0;
-
-  // body-inertial
-  if (objtype == mjOBJ_BODY) {
-    bodyid = objid;
-    pos = d->xipos+3*objid;
-    rot = (flg_local ? d->ximat+9*objid : 0);
-  }
-
-  // body-regular
-  else if (objtype == mjOBJ_XBODY) {
-    bodyid = objid;
-    pos = d->xpos+3*objid;
-    rot = (flg_local ? d->xmat+9*objid : 0);
-  }
-
-  // geom
-  else if (objtype == mjOBJ_GEOM) {
-    bodyid = m->geom_bodyid[objid];
-    pos = d->geom_xpos+3*objid;
-    rot = (flg_local ? d->geom_xmat+9*objid : 0);
-  }
-
-  // site
-  else if (objtype == mjOBJ_SITE) {
-    bodyid = m->site_bodyid[objid];
-    pos = d->site_xpos+3*objid;
-    rot = (flg_local ? d->site_xmat+9*objid : 0);
-  }
-
-  // camera
-  else if (objtype == mjOBJ_CAMERA) {
-    bodyid = m->cam_bodyid[objid];
-    pos = d->cam_xpos+3*objid;
-    rot = (flg_local ? d->cam_xmat+9*objid : 0);
-  }
-
-  // object without spatial frame
-  else {
-    mjERROR("invalid object type %d", objtype);
-  }
-
-  // transform velocity
-  mju_transformSpatial(res, d->cvel+6*bodyid, 0, pos, d->subtree_com+3*m->body_rootid[bodyid], rot);
-}
-
-
-
-// compute object 6D acceleration in object-centered frame, world/local orientation
-void mj_objectAcceleration(const mjModel* m, const mjData* d,
-                           int objtype, int objid, mjtNum res[6], int flg_local) {
-  int bodyid = 0;
-  const mjtNum *pos = 0, *rot = 0;
-  mjtNum correction[3], vel[6];
-
-  // body-inertial
-  if (objtype == mjOBJ_BODY) {
-    bodyid = objid;
-    pos = d->xipos+3*objid;
-    rot = (flg_local ? d->ximat+9*objid : 0);
-  }
-
-  // body-regular
-  else if (objtype == mjOBJ_XBODY) {
-    bodyid = objid;
-    pos = d->xpos+3*objid;
-    rot = (flg_local ? d->xmat+9*objid : 0);
-  }
-
-  // geom
-  else if (objtype == mjOBJ_GEOM) {
-    bodyid = m->geom_bodyid[objid];
-    pos = d->geom_xpos+3*objid;
-    rot = (flg_local ? d->geom_xmat+9*objid : 0);
-  }
-
-  // site
-  else if (objtype == mjOBJ_SITE) {
-    bodyid = m->site_bodyid[objid];
-    pos = d->site_xpos+3*objid;
-    rot = (flg_local ? d->site_xmat+9*objid : 0);
-  }
-
-  // camera
-  else if (objtype == mjOBJ_CAMERA) {
-    bodyid = m->cam_bodyid[objid];
-    pos = d->cam_xpos+3*objid;
-    rot = (flg_local ? d->cam_xmat+9*objid : 0);
-  }
-
-  // object without spatial frame
-  else {
-    mjERROR("invalid object type %d", objtype);
-  }
-
-  // transform com-based velocity to local frame
-  mju_transformSpatial(vel, d->cvel+6*bodyid, 0, pos, d->subtree_com+3*m->body_rootid[bodyid], rot);
-
-  // transform com-based acceleration to local frame
-  mju_transformSpatial(res, d->cacc+6*bodyid, 0, pos, d->subtree_com+3*m->body_rootid[bodyid], rot);
-
-  // acc_tran += vel_rot x vel_tran
-  mju_cross(correction, vel, vel+3);
-  mju_addTo3(res+3, correction);
-}
-
-
 
 //-------------------------- miscellaneous ---------------------------------------------------------
 
@@ -1314,28 +509,6 @@ mjtNum mj_geomDistance(const mjModel* m, const mjData* d, int geom1, int geom2,
 
 
 
-// extract 6D force:torque for one contact, in contact frame
-void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]) {
-  mjContact* con;
-
-  // clear result
-  mju_zero(result, 6);
-
-  // make sure contact is valid
-  if (id >= 0 && id < d->ncon && d->contact[id].efc_address >= 0) {
-    // get contact pointer
-    con = d->contact + id;
-
-    if (mj_isPyramidal(m)) {
-      mju_decodePyramid(result, d->efc_force + con->efc_address, con->friction, con->dim);
-    } else {
-      mju_copy(result, d->efc_force + con->efc_address, con->dim);
-    }
-  }
-}
-
-
-
 // compute velocity by finite-differencing two positions
 void mj_differentiatePos(const mjModel* m, mjtNum* qvel, mjtNum dt,
                          const mjtNum* qpos1, const mjtNum* qpos2) {
@@ -1418,50 +591,6 @@ void mj_normalizeQuat(const mjModel* m, mjtNum* qpos) {
 
 
 
-// map from body local to global Cartesian coordinates
-void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9],
-                     const mjtNum pos[3], const mjtNum quat[4],
-                     int body, mjtByte sameframe) {
-  mjtSameFrame sf = sameframe;
-
-  // position
-  if (xpos && pos) {
-    switch (sf) {
-    case mjSAMEFRAME_NONE:
-    case mjSAMEFRAME_BODYROT:
-    case mjSAMEFRAME_INERTIAROT:
-      mju_mulMatVec3(xpos, d->xmat+9*body, pos);
-      mju_addTo3(xpos, d->xpos+3*body);
-      break;
-    case mjSAMEFRAME_BODY:
-      mju_copy3(xpos, d->xpos+3*body);
-      break;
-    case mjSAMEFRAME_INERTIA:
-      mju_copy3(xpos, d->xipos+3*body);
-      break;
-    }
-  }
-
-  // orientation
-  if (xmat && quat) {
-    mjtNum tmp[4];
-    switch (sf) {
-    case mjSAMEFRAME_NONE:
-      mju_mulQuat(tmp, d->xquat+4*body, quat);
-      mju_quat2Mat(xmat, tmp);
-      break;
-    case mjSAMEFRAME_BODY:
-    case mjSAMEFRAME_BODYROT:
-      mju_copy(xmat, d->xmat+9*body, 9);
-      break;
-    case mjSAMEFRAME_INERTIA:
-    case mjSAMEFRAME_INERTIAROT:
-      mju_copy(xmat, d->ximat+9*body, 9);
-      break;
-    }
-  }
-}
-
 // return 1 if actuator i is disabled, 0 otherwise
 int mj_actuatorDisabled(const mjModel* m, int i) {
   int group = m->actuator_group[i];
@@ -1503,27 +632,6 @@ void mj_setTotalmass(mjModel* m, mjtNum newmass) {
 
 
 
-// count warnings, print only the first time
-void mj_warning(mjData* d, int warning, int info) {
-  // check type
-  if (warning < 0 || warning >= mjNWARNING) {
-    mjERROR("invalid warning type %d", warning);
-  }
-
-  // save info (override previous)
-  d->warning[warning].lastinfo = info;
-
-  // print message only the first time this warning is encountered
-  if (!d->warning[warning].number) {
-    mju_warning("%s Time = %.4f.", mju_warningText(warning, info), d->time);
-  }
-
-  // increase counter
-  d->warning[warning].number++;
-}
-
-
-
 // version number
 int mj_version(void) {
   return mjVERSION;

src/engine/engine_support.h

@@ -47,77 +47,6 @@ MJAPI void mj_setState(const mjModel* m, mjData* d, const mjtNum* state, unsigne
 // copy current state to the k-th model keyframe
 MJAPI void mj_setKeyframe(mjModel* m, const mjData* d, int k);
 
-//-------------------------- sparse chains ---------------------------------------------------------
-
-// merge dof chains for two bodies
-int mj_mergeChain(const mjModel* m, int* chain, int b1, int b2);
-
-// merge dof chains for two simple bodies
-int mj_mergeChainSimple(const mjModel* m, int* chain, int b1, int b2);
-
-// get body chain
-int mj_bodyChain(const mjModel* m, int body, int* chain);
-
-
-//-------------------------- Jacobians -------------------------------------------------------------
-
-// compute 3/6-by-nv Jacobian of global point attached to given body
-MJAPI void mj_jac(const mjModel* m, const mjData* d,
-                  mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body);
-
-// compute body frame Jacobian
-MJAPI void mj_jacBody(const mjModel* m, const mjData* d,
-                      mjtNum* jacp, mjtNum* jacr, int body);
-
-// compute body center-of-mass Jacobian
-MJAPI void mj_jacBodyCom(const mjModel* m, const mjData* d,
-                         mjtNum* jacp, mjtNum* jacr, int body);
-
-// compute subtree center-of-mass Jacobian
-MJAPI void mj_jacSubtreeCom(const mjModel* m, mjData* d, mjtNum* jacp, int body);
-
-// compute geom Jacobian
-MJAPI void mj_jacGeom(const mjModel* m, const mjData* d,
-                      mjtNum* jacp, mjtNum* jacr, int geom);
-
-// compute site Jacobian
-MJAPI void mj_jacSite(const mjModel* m, const mjData* d,
-                      mjtNum* jacp, mjtNum* jacr, int site);
-
-// compute translation Jacobian of point, and rotation Jacobian of axis
-MJAPI void mj_jacPointAxis(const mjModel* m, mjData* d,
-                           mjtNum* jacPoint, mjtNum* jacAxis,
-                           const mjtNum point[3], const mjtNum axis[3], int body);
-
-// compute 3/6-by-nv sparse Jacobian of global point attached to given body
-void mj_jacSparse(const mjModel* m, const mjData* d,
-                  mjtNum* jacp, mjtNum* jacr, const mjtNum* point, int body,
-                  int NV, const int* chain);
-
-// sparse Jacobian difference for simple body contacts
-void mj_jacSparseSimple(const mjModel* m, const mjData* d,
-                        mjtNum* jacdifp, mjtNum* jacdifr, const mjtNum* point,
-                        int body, int flg_second, int NV, int start);
-
-// dense or sparse Jacobian difference for two body points: pos2 - pos1, global
-MJAPI int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain,
-                        int b1, int b2, const mjtNum pos1[3], const mjtNum pos2[3],
-                        mjtNum* jac1p, mjtNum* jac2p, mjtNum* jacdifp,
-                        mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr);
-
-// dense or sparse weighted sum of multiple body Jacobians at same point
-int mj_jacSum(const mjModel* m, mjData* d, int* chain,
-              int n, const int* body, const mjtNum* weight,
-              const mjtNum point[3], mjtNum* jac, int flg_rot);
-
-// compute 3/6-by-nv Jacobian time derivative of global point attached to given body
-MJAPI void mj_jacDot(const mjModel* m, const mjData* d,
-                     mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body);
-
-// compute subtree angular momentum matrix
-MJAPI void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body);
-
-
 //-------------------------- inertia functions -----------------------------------------------------
 
 // convert sparse inertia matrix M into full matrix
@@ -146,26 +75,12 @@ MJAPI void mj_applyFT(const mjModel* m, mjData* d,
 void mj_xfrcAccumulate(const mjModel* m, mjData* d, mjtNum* qfrc);
 
 
-//-------------------------- coordinate transformation ---------------------------------------------
-
-// compute object 6D velocity in object-centered frame, world/local orientation
-MJAPI void mj_objectVelocity(const mjModel* m, const mjData* d,
-                             int objtype, int objid, mjtNum res[6], int flg_local);
-
-// compute object 6D acceleration in object-centered frame, world/local orientation
-MJAPI void mj_objectAcceleration(const mjModel* m, const mjData* d,
-                                 int objtype, int objid, mjtNum res[6], int flg_local);
-
-
 //-------------------------- miscellaneous ---------------------------------------------------------
 
 // returns the smallest distance between two geoms
 MJAPI mjtNum mj_geomDistance(const mjModel* m, const mjData* d, int geom1, int geom2,
                              mjtNum distmax, mjtNum fromto[6]);
 
-// extract 6D force:torque for one contact, in contact frame
-MJAPI void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]);
-
 // compute velocity by finite-differencing two positions
 MJAPI void mj_differentiatePos(const mjModel* m, mjtNum* qvel, mjtNum dt,
                                const mjtNum* qpos1, const mjtNum* qpos2);
@@ -176,11 +91,6 @@ MJAPI void mj_integratePos(const mjModel* m, mjtNum* qpos, const mjtNum* qvel, m
 // normalize all quaternions in qpos-type vector
 MJAPI void mj_normalizeQuat(const mjModel* m, mjtNum* qpos);
 
-// map from body local to global Cartesian coordinates
-MJAPI void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9],
-                           const mjtNum pos[3], const mjtNum quat[4],
-                           int body, mjtByte sameframe);
-
 // return 1 if actuator i is disabled, 0 otherwise
 MJAPI int mj_actuatorDisabled(const mjModel* m, int i);
 
@@ -190,9 +100,6 @@ MJAPI mjtNum mj_getTotalmass(const mjModel* m);
 // scale body masses and inertias to achieve specified total mass
 MJAPI void mj_setTotalmass(mjModel* m, mjtNum newmass);
 
-// high-level warning function: count warnings in mjData, print only the first time
-MJAPI void mj_warning(mjData* d, int warning, int info);
-
 // version number
 MJAPI int mj_version(void);
 

src/engine/engine_util_container.c

@@ -20,7 +20,7 @@
 
 #include <mujoco/mujoco.h>
 #include "engine/engine_crossplatform.h"
-#include "engine/engine_io.h"
+#include "engine/engine_memory.h"
 
 // stack allocate and initialize new mjArrayList
 mjArrayList* mju_arrayListCreate(mjData* d, size_t element_size, size_t initial_capacity) {

src/engine/engine_util_solve.c

@@ -19,11 +19,11 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmacro.h>
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
-#include "engine/engine_io.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"
 #include "engine/engine_util_sparse.h"
+#include "engine/engine_memory.h"
 #include "engine/engine_util_spatial.h"
 
 //---------------------------- dense Cholesky ------------------------------------------------------

src/engine/engine_util_sparse.c

@@ -21,7 +21,7 @@
 #include <mujoco/mjmacro.h>
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
 #include <mujoco/mjtnum.h>
-#include "engine/engine_io.h"
+#include "engine/engine_memory.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_misc.h"
 

src/engine/engine_vis_interact.c

@@ -22,7 +22,9 @@
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
 #include <mujoco/mjvisualize.h>
 #include "engine/engine_core_smooth.h"
+#include "engine/engine_core_util.h"
 #include "engine/engine_io.h"
+#include "engine/engine_memory.h"
 #include "engine/engine_ray.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"

src/engine/engine_vis_visualize.c

@@ -23,7 +23,8 @@
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
 #include <mujoco/mjvisualize.h>
 #include "engine/engine_array_safety.h"
-#include "engine/engine_io.h"
+#include "engine/engine_core_util.h"
+#include "engine/engine_memory.h"
 #include "engine/engine_name.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_support.h"

src/user/user_init.c

@@ -16,6 +16,7 @@
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjvisualize.h>
 #include <mujoco/mjspec.h>
+#include "engine/engine_core_util.h"
 #include "engine/engine_io.h"
 #include "user/user_api.h"
 

test/engine/engine_core_constraint_test.cc

@@ -24,6 +24,7 @@
 #include <mujoco/mjmodel.h>
 #include <mujoco/mujoco.h>
 #include "src/engine/engine_core_constraint.h"
+#include "src/engine/engine_core_util.h"
 #include "src/engine/engine_support.h"
 #include "src/engine/engine_util_misc.h"
 #include "test/fixture.h"
@@ -197,7 +198,9 @@ TEST_F(CoreConstraintTest, RestPenetration) {
       // simulate for 50 seconds
       mj_resetData(model, data);
       while (data->time < 50) {
+        mjtNum time = data->time;
         mj_step(model, data);
+        ASSERT_GT(data->time, time) << "Divergence detected";
       }
 
       mjtNum depth = -data->contact[0].dist;
@@ -263,7 +266,9 @@ TEST_F(CoreConstraintTest, EqualityBodySite) {
 
   // simulate, get diag(A)
   while (data->time < 0.1) {
+    mjtNum time = data->time;
     mj_step(model, data);
+    ASSERT_GT(data->time, time) << "Divergence detected";
   }
   int nefc_site = data->nefc;
   std::vector<mjtNum> dA = AsVector(data->efc_diagApprox, nefc_site);
@@ -276,7 +281,9 @@ TEST_F(CoreConstraintTest, EqualityBodySite) {
 
   // simulate again, get diag(A)
   while (data->time < 0.1) {
+    mjtNum time = data->time;
     mj_step(model, data);
+    ASSERT_GT(data->time, time) << "Divergence detected";
   }
 
   // compare
@@ -310,8 +317,12 @@ TEST_F(CoreConstraintTest, ConstraintUpdateImpl) {
       mj_resetData(model, d1);
       mj_resetData(model, d2);
       while (d1->time < 0.2) {
+        mjtNum time1 = d1->time;
+        mjtNum time2 = d2->time;
         mj_step(model, d1);
+        ASSERT_GT(d1->time, time1) << "Divergence detected";
         mj_step(model, d2);
+        ASSERT_GT(d2->time, time2) << "Divergence detected";
       }
       mj_forward(model, d1);
       mj_forward(model, d2);

test/engine/engine_core_smooth_test.cc

@@ -187,7 +187,9 @@ TEST_F(CoreSmoothTest, TendonJdot) {
       } else {
         mj_resetData(m, d);
         while (d->time < 1) {
+          mjtNum time = d->time;
           mj_step(m, d);
+          ASSERT_GT(d->time, time) << "Divergence detected";
         }
       }
 
@@ -307,7 +309,9 @@ TEST_F(CoreSmoothTest, TendonArmatureConservesEnergy) {
 
       double eps = std::max(energy_0, 1.0) * 1e-5;
       while (d->time < 1) {
+        mjtNum time = d->time;
         mj_step(m, d);
+        ASSERT_GT(d->time, time) << "Divergence detected";
         double energy_t = d->energy[0] + d->energy[1];
         EXPECT_THAT(energy_t, DoubleNear(energy_0, eps));
       }
@@ -336,7 +340,9 @@ TEST_F(CoreSmoothTest, TendonArmatureConservesMomentum) {
 
     double eps = 1e-5;
     while (d->time < 1) {
+      mjtNum time = d->time;
       mj_step(m, d);
+      ASSERT_GT(d->time, time) << "Divergence detected";
       vector<mjtNum> sdata_t = AsVector(d->sensordata, m->nsensordata);
       EXPECT_THAT(sdata_t, Pointwise(DoubleNear(eps), sdata_0));
     }
@@ -380,10 +386,14 @@ TEST_F(CoreSmoothTest, TendonInertiaEquivalent) {
     double eps = lpath == kTen_i0 ? 1e-6 : 1e-3;
 
     while (d->time < 1) {
+      mjtNum time = d->time;
       mj_step(m, d);
+      ASSERT_GT(d->time, time) << "Divergence detected";
       vector<mjtNum> xpos = AsVector(d->geom_xpos + 3*gid, 3);
 
+      time = d_e->time;
       mj_step(m_e, d_e);
+      ASSERT_GT(d_e->time, time) << "Divergence detected";
       vector<mjtNum> xpos_e = AsVector(d_e->geom_xpos + 3*gid_e, 3);
 
       EXPECT_THAT(xpos, Pointwise(DoubleNear(eps), xpos_e));
@@ -560,7 +570,9 @@ TEST_F(CoreSmoothTest, EqualityBodySite) {
 
   // simulate again, get sensordata
   while (data->time < 0.1) {
+    mjtNum time = data->time;
     mj_step(model, data);
+    ASSERT_GT(data->time, time) << "Divergence detected";
   }
 
   // compare
@@ -589,7 +601,9 @@ TEST_F(CoreSmoothTest, RefsiteBringsToPose) {
 
   // step for 5 seconds
   while (data->time < 10) {
+    mjtNum time = data->time;
     mj_step(model, data);
+    ASSERT_GT(data->time, time) << "Divergence detected";
   }
 
   // get site IDs
@@ -628,7 +642,9 @@ TEST_F(CoreSmoothTest, RefsiteConservesMomentum) {
   // simulate, assert that momentum is conserved
   mjtNum eps = 1e-9;
   while (data->time < 1) {
+    mjtNum time = data->time;
     mj_step(model, data);
+    ASSERT_GT(data->time, time) << "Divergence detected";
     for (int i=0; i < 6; i++) {
       EXPECT_LT(mju_abs(data->sensordata[i]), eps);
     }

test/engine/engine_io_test.cc

@@ -12,15 +12,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Tests for engine/engine_io.c.
+// Tests for engine/{engine_io.c and engine_memory.c}.
 
 #include "src/engine/engine_io.h"
+#include "src/engine/engine_memory.h"
 
 #include <array>
 #include <cstdint>
 #include <cstdio>
 #include <cstring>
-#include <filesystem>
+#include <filesystem>  // NOLINT
 #include <string>
 #include <vector>
 

test/engine/engine_support_test.cc

@@ -12,8 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Tests for engine/engine_support.c.
+// Tests for engine/{engine_support.c and engine_core_util.c}
 
+#include "src/engine/engine_core_util.h"
 #include "src/engine/engine_support.h"
 
 #include <limits>