From ad96acfb16ec397e3719ddedd02b82a90bd1bf81 Mon Sep 17 00:00:00 2001
From: Kyle Bayes <kylebayes@google.com>
Date: Thu, 21 May 2026 05:15:50 -0700
Subject: [PATCH] Add internal mj_narrowphase for batch processing of potential
 geom pair collisions.

PiperOrigin-RevId: 918988042
Change-Id: I871dcdf42f8105e2c8a57e275a3775ea46c69e65
---
 src/engine/engine_collision_convex.c |  31 +-
 src/engine/engine_collision_convex.h |   3 +
 src/engine/engine_collision_driver.c | 437 ++++++++++++++++++---------
 src/engine/engine_collision_driver.h |  11 -
 4 files changed, 315 insertions(+), 167 deletions(-)

diff --git a/src/engine/engine_collision_convex.c b/src/engine/engine_collision_convex.c
index 2c52a789..30dd4b08 100644
--- a/src/engine/engine_collision_convex.c
+++ b/src/engine/engine_collision_convex.c
@@ -24,7 +24,7 @@
 #include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_collision_gjk.h"
-#include "engine/engine_collision_primitive.h"
+#include "engine/engine_macro.h"
 #include "engine/engine_memory.h"
 #include "engine/engine_inline.h"
 #include "engine/engine_util_blas.h"
@@ -34,6 +34,15 @@
 
 #define mjMINVAL2 (mjMINVAL * mjMINVAL)
 
+// CCD internal buffer used for batched processing; if NULL, stack memory allocated on each
+// mjc_penetration call
+static mjTHREADLOCAL void* ccd_buffer = NULL;
+
+// set CCD internal buffer
+void mjc_setCCDBuffer(void* buffer) {
+  ccd_buffer = buffer;
+}
+
 // ccd prism first dir
 static void prism_firstdir(const void* o1, const void* o2, ccd_vec3_t *vec) {
   ccdVec3Set(vec, 0, 0, 1);
@@ -85,18 +94,23 @@ static int mjc_penetration(const mjModel* m, mjData* d, mjCCDObj* obj1, mjCCDObj
   mjCCDConfig config;
   mjCCDStatus status;
   mjtNum dist;
+  int nwitness = 0;
+  void* buffer = ccd_buffer;
 
   // set config
-  mj_markStack(d);
   config.max_iterations = m->opt.ccd_iterations;
   config.tolerance = m->opt.ccd_tolerance;
   config.max_contacts = ncon;
   config.dist_cutoff = 0;  // no geom distances needed
-  config.buffer = mj_stackAllocByte(d, mjc_ccdSize(config.max_iterations), sizeof(mjtNum));
+  if (buffer) {
+    config.buffer = buffer;
+  } else {
+    mj_markStack(d);
+    config.buffer = mj_stackAllocByte(d, mjc_ccdSize(config.max_iterations), sizeof(mjtNum));
+  }
 
   if ((dist = mjc_ccd(&config, &status, obj1, obj2)) < 0) {
-    mj_freeStack(d);
-    int nwitness = status.nx;
+    nwitness = status.nx;
     for (int i = 0; i < nwitness; i++) {
       con[i].dist = margin + dist;
       con[i].pos[0] = 0.5*(status.x1[3*i + 0] + status.x2[3*i + 0]);
@@ -106,10 +120,11 @@ static int mjc_penetration(const mjModel* m, mjData* d, mjCCDObj* obj1, mjCCDObj
       mju_normalize3(con[i].normal);
       mji_zero3(con[i].tangent);
     }
-    return nwitness;
   }
-  mj_freeStack(d);
-  return 0;
+  if (!buffer) {
+    mj_freeStack(d);
+  }
+  return nwitness;
 }
 
 
diff --git a/src/engine/engine_collision_convex.h b/src/engine/engine_collision_convex.h
index 88fda2ea..eae1e72d 100644
--- a/src/engine/engine_collision_convex.h
+++ b/src/engine/engine_collision_convex.h
@@ -124,6 +124,9 @@ int mjc_HFieldElem(const mjModel* m, mjData* d, mjContact* con, int g, int f, in
 // fix contact frame normal
 void mjc_fixNormal(const mjModel* m, const mjData* d, mjPreContact* con, int g1, int g2);
 
+// set CCD internal buffer
+void mjc_setCCDBuffer(void* buffer);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/src/engine/engine_collision_driver.c b/src/engine/engine_collision_driver.c
index a598602f..677a457d 100644
--- a/src/engine/engine_collision_driver.c
+++ b/src/engine/engine_collision_driver.c
@@ -23,6 +23,7 @@
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
 #include "engine/engine_callback.h"
 #include "engine/engine_collision_convex.h"
+#include "engine/engine_collision_gjk.h"
 #include "engine/engine_collision_primitive.h"
 #include "engine/engine_collision_sdf.h"
 #include "engine/engine_core_constraint.h"
@@ -155,6 +156,24 @@ int mj_maxContact(const mjModel* m, int g1, int g2, int has_margin) {
 }
 
 
+// return the margin for a given geom pair
+static inline mjtNum getMargin(const mjModel* m, int g1, int g2, int ipair) {
+  if (ipair >= 0) {
+    return mj_assignMargin(m, m->pair_margin[ipair]);
+  }
+  return mj_assignMargin(m, m->geom_margin[g1] + m->geom_margin[g2]);
+}
+
+
+// return the gap for a given geom pair
+static inline mjtNum getGap(const mjModel* m, int g1, int g2, int ipair) {
+  if (ipair >= 0) {
+    return m->pair_gap[ipair];
+  }
+  return m->geom_gap[g1] + m->geom_gap[g2];
+}
+
+
 // move arena pointer back to the end of the contact array
 static inline void resetArena(mjData* d) {
   d->parena = d->ncon * sizeof(mjContact);
@@ -167,6 +186,16 @@ static inline void resetArena(mjData* d) {
 }
 
 
+// realign the arena pointer to the specified alignment
+static inline size_t alignArena(mjData* d, size_t alignment) {
+  size_t misalignment = d->parena % alignment;
+  if (misalignment) {
+    d->parena += alignment - misalignment;
+  }
+  return d->parena;
+}
+
+
 // plane to geom_center squared distance, g1 is a plane
 static mjtNum planeGeomDist(const mjModel* m, mjData* d, int g1, int g2) {
   mjtNum* mat1 = d->geom_xmat + 9*g1;
@@ -322,6 +351,16 @@ int mj_isElemActive(const mjModel* m, int f, int e) {
 
 //----------------------------- collision detection entry point ------------------------------------
 
+// binary search between two bodyflex trees
+void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2,
+                    int merged, int startadr, int pairadr);
+
+// compute contacts for a batch of collision pairs contained in a buffer of
+// stride 3 ints (g1, g2, ipair)
+// if buffer is NULL, results are read from arena starting at parena
+void mj_narrowphase(const mjModel* m, mjData* d, const int* buffer, int npair, size_t parena);
+
+
 // compare contact pairs by their geom/elem/vert IDs
 static inline int contactcompare(const mjContact* c1, const mjContact* c2, void* context) {
   const mjModel* m = (const mjModel*) context;
@@ -333,7 +372,7 @@ static inline int contactcompare(const mjContact* c1, const mjContact* c2, void*
   int con2_obj2 = c2->geom[1] >= 0 ? c2->geom[1] : (c2->elem[1] >= 0 ? c2->elem[1] : c2->vert[1]);
 
   // for geom:geom, reproduce the order of contacts without mj_collideTree
-  // normally sorted by (g1, g2), but in mj_collideGeoms, g1 and g2 are swapped based on geom_type
+  // normally sorted by (g1, g2), but g1 and g2 are swapped based on geom_type
   // here we undo this swapping for the purpose of sorting - needs to be done for each mjContact
   if (c1->geom[0] >= 0 && c1->geom[1] >= 0 &&
       c2->geom[0] >= 0 && c2->geom[1] >= 0) {
@@ -431,6 +470,72 @@ static void filterFlexContacts(mjData* d, int ncon_before) {
 }
 
 
+// push a candidate collision pair onto the arena
+static void pushPairArena(const mjModel* m, mjData* d, int g1, int g2, int ipair) {
+  // allocate geom pair on the arena
+  int* pair = (int*) mj_arenaAllocByte(d, 3 * sizeof(int), _Alignof(int));
+  if (!pair) {
+    mjERROR("arena too small to allocate geom pair");
+  }
+
+  if (m->geom_type[g1] > m->geom_type[g2]) {
+    pair[0] = g2;
+    pair[1] = g1;
+  } else {
+    pair[0] = g1;
+    pair[1] = g2;
+  }
+  pair[2] = ipair;
+}
+
+
+// filter candidate collision pair; return 0 if pair should be discarded
+static int filterCollisionPair(const mjModel* m, mjData* d, int g1, int g2, int ipair,
+                               int merged, int startadr, int pairadr) {
+  // merged, find matching pair
+  if (merged) {
+    for (int k=startadr; k < pairadr; k++) {
+      if ((m->pair_geom1[k] == g1 && m->pair_geom2[k] == g2) ||
+          (m->pair_geom1[k] == g2 && m->pair_geom2[k] == g1)) {
+        return 0;
+      }
+    }
+  }
+
+  if (ipair >= 0) {
+    if (mjENABLED(mjENBL_SLEEP)) {
+      int b1 = m->geom_bodyid[g1];
+      int b2 = m->geom_bodyid[g2];
+      if (d->body_awake[b1] != mjS_AWAKE && d->body_awake[b2] != mjS_AWAKE) {
+        return 0;
+      }
+    }
+  }
+
+  if (ipair < 0) {
+    if (mjcb_contactfilter) {
+      if (mjcb_contactfilter(m, d, g1, g2)) {
+        return 0;
+      }
+    } else if (filterBitmask(m->geom_contype[g1], m->geom_conaffinity[g1],
+                             m->geom_contype[g2], m->geom_conaffinity[g2])) {
+      return 0;
+    }
+  }
+
+  // bounding sphere filter
+  mjtNum margin = getMargin(m, g1, g2, ipair);
+  mjtNum gap = getGap(m, g1, g2, ipair);
+  if (mj_filterSphere(m, d, g1, g2, margin + gap)) {
+    return 0;
+  }
+
+  // highly unlikely, but check collision function is well-defined
+  int type1 = mjMIN(m->geom_type[g1], m->geom_type[g2]);
+  int type2 = mjMAX(m->geom_type[g1], m->geom_type[g2]);
+  return mjCOLLISIONFUNC[type1][type2] != NULL;
+}
+
 
 // main collision function
 void mj_collision(const mjModel* m, mjData* d) {
@@ -469,7 +574,11 @@ void mj_collision(const mjModel* m, mjData* d) {
   TM_RESTART;
 
   // process bodyflex pairs returned by broadphase, merge with predefined geom pairs
-  int pairadr = 0;
+  int pairadr = 0, ngeompair = 0, g1, g2;
+
+  // align the arena for candidate collision pairs for narrowphase
+  size_t parena = alignArena(d, _Alignof(int));
+
   for (int i=0; i < nbfpair; i++) {
     // reconstruct bodyflex pair ids
     int bf1 = (broadphasepair[i]>>16) & 0xFFFF;
@@ -490,9 +599,11 @@ void mj_collision(const mjModel* m, mjData* d) {
     // test all predefined pairs for which pair_signature <= signature
     for (; pairadr < npair && m->pair_signature[pairadr] <= signature; pairadr++) {
       merged = (m->pair_signature[pairadr] == signature);
-      int g1 = m->pair_geom1[pairadr];
-      int g2 = m->pair_geom2[pairadr];
-      mj_collideGeoms(m, d, pairadr, g1, g2);
+      g1 = m->pair_geom1[pairadr], g2 = m->pair_geom2[pairadr];
+      if (filterCollisionPair(m, d, g1, g2, pairadr, 0, 0, 0)) {
+        pushPairArena(m, d, g1, g2, pairadr);
+        ngeompair++;
+      }
     }
 
     // apply bitmask filtering at the bodyflex level
@@ -524,11 +635,20 @@ void mj_collision(const mjModel* m, mjData* d) {
 
     // process bodyflex pair: two single-geom bodies
     if (isbody1 && isbody2 && m->body_geomnum[bf1] == 1 && m->body_geomnum[bf2] == 1) {
-      mj_collideGeomPair(m, d, geomadr1, geomadr2, merged, startadr, pairadr);
+      if (filterCollisionPair(m, d, geomadr1, geomadr2, -1, merged, startadr, pairadr)) {
+        pushPairArena(m, d, geomadr1, geomadr2, -1);
+        ngeompair++;
+      }
     }
 
     // process bodyflex pair: midphase
     else if (!mjDISABLED(mjDSBL_MIDPHASE) && bvh1 >= 0 && bvh2 >= 0) {
+      // flush geom pairs before calling mj_collideTree as post sorting needs to happen
+      if (ngeompair > 0) {
+        mj_narrowphase(m, d, NULL, ngeompair, parena);
+        ngeompair = 0;
+      }
+
       int ncon_before = d->ncon;
       mj_collideTree(m, d, bf1, bf2, merged, startadr, pairadr);
       int ncon_after = d->ncon;
@@ -547,6 +667,9 @@ void mj_collision(const mjModel* m, mjData* d) {
         contactSort(d->contact + ncon_before, buf, n, (void*)m);
         mj_freeStack(d);
       }
+
+      // realign the arena after adding contacts
+      parena = alignArena(d, _Alignof(int));
     }
 
     // process bodyflex pair: all-to-all
@@ -556,15 +679,23 @@ void mj_collision(const mjModel* m, mjData* d) {
 
       // body : body
       if (isbody1 && isbody2) {
-        for (int g1=geomadr1; g1 < geomadr_end1; g1++) {
-          for (int g2=geomadr2; g2 < geomadr_end2; g2++) {
-            mj_collideGeomPair(m, d, g1, g2, merged, startadr, pairadr);
+        for (g1=geomadr1; g1 < geomadr_end1; g1++) {
+          for (g2=geomadr2; g2 < geomadr_end2; g2++) {
+            if (filterCollisionPair(m, d, g1, g2, -1, merged, startadr, pairadr)) {
+              pushPairArena(m, d, g1, g2, -1);
+              ngeompair++;
+            }
           }
         }
       }
 
       // body : flex
       else if (isbody1) {
+        if (ngeompair > 0) {
+          mj_narrowphase(m, d, NULL, ngeompair, parena);
+          ngeompair = 0;
+        }
+
         int f = bf2 - nbody;
 
         // process body geoms
@@ -599,10 +730,18 @@ void mj_collision(const mjModel* m, mjData* d) {
           }
           filterFlexContacts(d, ncon_before);
         }
+        // realign the arena after adding contacts
+        parena = alignArena(d, _Alignof(int));
       }
 
       // flex : flex
       else {
+        // flush accumulated geompairs before flex:flex processing
+        if (ngeompair > 0) {
+          mj_narrowphase(m, d, NULL, ngeompair, parena);
+          ngeompair = 0;
+        }
+
         int f1 = bf1 - nbody;
         int f2 = bf2 - nbody;
 
@@ -614,6 +753,9 @@ void mj_collision(const mjModel* m, mjData* d) {
           }
         }
         filterFlexContacts(d, ncon_before);
+
+        // realign the arena after adding contacts
+        parena = alignArena(d, _Alignof(int));
       }
     }
   }
@@ -621,9 +763,17 @@ void mj_collision(const mjModel* m, mjData* d) {
 
   // finish merging predefined geom pairs
   for (; pairadr < npair; pairadr++) {
-    int g1 = m->pair_geom1[pairadr];
-    int g2 = m->pair_geom2[pairadr];
-    mj_collideGeoms(m, d, pairadr, g1, g2);
+    g1 = m->pair_geom1[pairadr], g2 = m->pair_geom2[pairadr];
+    if (filterCollisionPair(m, d, g1, g2, pairadr, 0, 0, 0)) {
+      pushPairArena(m, d, g1, g2, pairadr);
+      ngeompair++;
+    }
+  }
+
+  // flush remaining collision pairs
+  if (ngeompair > 0) {
+    mj_narrowphase(m, d, NULL, ngeompair, parena);
+    ngeompair = 0;
   }
 
   // flex self-collisions
@@ -689,25 +839,6 @@ typedef struct  {
   int node2;
 } mjCollisionTree;
 
-
-// checks if the proposed collision pair is already present in pair_geom and calls narrow phase
-void mj_collideGeomPair(const mjModel* m, mjData* d, int g1, int g2, int merged,
-                        int startadr, int pairadr) {
-  // merged, find matching pair
-  if (merged) {
-    for (int k=startadr; k < pairadr; k++) {
-      if ((m->pair_geom1[k] == g1 && m->pair_geom2[k] == g2) ||
-          (m->pair_geom1[k] == g2 && m->pair_geom2[k] == g1)) {
-        return;
-      }
-    }
-  }
-
-  // not merged, always test
-  mj_collideGeoms(m, d, -1, g1, g2);
-}
-
-
 // oriented bounding boxes collision (see Gottschalk et al.)
 int mj_collideOBB(const mjtNum aabb1[6], const mjtNum aabb2[6],
                   const mjtNum xpos1[3], const mjtNum xmat1[9],
@@ -806,7 +937,6 @@ int mj_collideOBB(const mjtNum aabb1[6], const mjtNum aabb2[6],
 }
 
 
-// binary search between two bodyflex trees
 void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2,
                     int merged, int startadr, int pairadr) {
   int nbody = m->nbody, nbvhstatic = m->nbvhstatic;
@@ -894,7 +1024,16 @@ void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2,
                             d->geom_xpos + 3*nodeid1, d->geom_xmat + 9*nodeid1,
                             d->geom_xpos + 3*nodeid2, d->geom_xmat + 9*nodeid2,
                             margin + gap, NULL, NULL, &initialize)) {
-            mj_collideGeomPair(m, d, nodeid1, nodeid2, merged, startadr, pairadr);
+
+            if (filterCollisionPair(m, d, nodeid1, nodeid2, -1, merged, startadr, pairadr)) {
+              int n1 = nodeid1, n2 = nodeid2;
+              if (m->geom_type[n1] > m->geom_type[n2]) {
+                n1 = nodeid2;
+                n2 = nodeid1;
+              }
+              int pair[3] = {n1, n2, -1};
+              mj_narrowphase(m, d, pair, 1, 0);
+            }
             if (mark_active) {
               d->bvh_active[node1 + bvhadr1] = 1;
               d->bvh_active[node2 + bvhadr2] = 1;
@@ -1675,139 +1814,141 @@ static void mj_makeCapsule(const mjModel* m, mjData* d, int f, const int vid[2],
 }
 
 
-// test two geoms for collision, apply filters, add to contact list
-void mj_collideGeoms(const mjModel* m, mjData* d, int ipair, int g1, int g2) {
-  int num, type1, type2, condim;
-  mjtNum margin, gap, friction[5], solref[mjNREF], solimp[mjNIMP];
-  mjtNum solreffriction[mjNREF] = {0};
+// compute contacts for a batch of collision pairs contained in a buffer of
+// stride 3 ints (g1, g2, ipair)
+// if buffer is NULL, results are read from arena starting at parena
+void mj_narrowphase(const mjModel* m, mjData* d, const int* buffer, int npair, size_t parena) {
+  int ccd_size = mjc_ccdSize(m->opt.ccd_iterations);
+  mjtNum margin, gap;
 
-  // sleep filtering for explicit pairs
-  if (ipair >= 0) {
-    if (mjENABLED(mjENBL_SLEEP)) {
-      int b1 = m->geom_bodyid[g1];
-      int b2 = m->geom_bodyid[g2];
-      if (d->body_awake[b1] != mjS_AWAKE && d->body_awake[b2] != mjS_AWAKE) {
-        return;
-      }
-    }
-  }
-
-  // order geoms by type
-  if (m->geom_type[g1] > m->geom_type[g2]) {
-    int i = g1;
-    g1 = g2;
-    g2 = i;
-  }
-
-  // copy types and bodies
-  type1 = m->geom_type[g1];
-  type2 = m->geom_type[g2];
-
-  mjfCollision collisionFunc = mjCOLLISIONFUNC[type1][type2];
-
-  // return if no collision function
-  if (!collisionFunc) {
-    return;
-  }
-
-  // apply filters if not predefined pair
-  if (ipair < 0) {
-    // user filter if defined
-    if (mjcb_contactfilter) {
-      if (mjcb_contactfilter(m, d, g1, g2)) {
-        return;
-      }
-    }
-
-    // otherwise built-in filter
-    else if (filterBitmask(m->geom_contype[g1], m->geom_conaffinity[g1],
-                           m->geom_contype[g2], m->geom_conaffinity[g2])) {
-      return;
-    }
-  }
-
-  // set margin: dynamic or pair
-  if (ipair < 0) {
-    margin = mj_assignMargin(m, m->geom_margin[g1] + m->geom_margin[g2]);
+  // set buffer and arena pointer
+  if (!buffer) {
+    buffer = (const int*) ((char*) d->arena + parena);
   } else {
-    margin = mj_assignMargin(m, m->pair_margin[ipair]);
+    parena = d->parena;
   }
 
-  // set gap: dynamic or pair
-  if (ipair < 0) {
-    gap = m->geom_gap[g1] + m->geom_gap[g2];
-  } else {
-    gap = m->pair_gap[ipair];
+  mj_markStack(d);
+
+  // buffer store how many contacts are generated for each pair
+  int* nconbuffer = mj_stackAllocInt(d, npair);
+
+  // buffer for pair data (g1, g2, ipair, index into conbuffer)
+  int* pairbuffer = mj_stackAllocInt(d, 4 * npair);
+  int maxcon = 0;
+  for (int i = 0; i < npair; i++) {
+    int g1 = buffer[3*i + 0];
+    int g2 = buffer[3*i + 1];
+    int ipair = buffer[3*i + 2];
+
+    pairbuffer[4*i + 0] = g1;
+    pairbuffer[4*i + 1] = g2;
+    pairbuffer[4*i + 2] = ipair;
+    pairbuffer[4*i + 3] = maxcon;
+    margin = getMargin(m, g1, g2, ipair);
+    gap = getGap(m, g1, g2, ipair);
+    maxcon += mj_maxContact(m, g1, g2, margin + gap > 0);
   }
 
-  // bounding sphere filter
-  if (mj_filterSphere(m, d, g1, g2, margin + gap)) {
+  // buffer for precontact data
+  mjPreContact* conbuffer = mjSTACKALLOC(d, maxcon, mjPreContact);
+
+  // buffer data has been copied to metadata on the stack;
+  // reclaim arena space so contacts can overwrite the buffer region
+  d->parena = parena;
+
+  // set buffer for nativeccd
+  mj_markStack(d);
+  mjc_setCCDBuffer(mj_stackAllocByte(d, ccd_size, sizeof(mjtNum)));
+
+  for (int i = 0; i < npair; i++) {
+    int g1 = pairbuffer[4*i + 0];
+    int g2 = pairbuffer[4*i + 1];
+    int ipair = pairbuffer[4*i + 2];
+    int idx = pairbuffer[4*i + 3];
+    mjfCollision collision_func = mjCOLLISIONFUNC[m->geom_type[g1]][m->geom_type[g2]];
+    margin = getMargin(m, g1, g2, ipair);
+    gap = getGap(m, g1, g2, ipair);
+    nconbuffer[i] = collision_func(m, d, conbuffer + idx, g1, g2, margin + gap);
+
+    // SHOULD NOT OCCUR
+    int expected_max = (i + 1 < npair ? pairbuffer[4*(i+1) + 3] : maxcon) - idx;
+    if (nconbuffer[i] > expected_max) {
+      mjERROR("collision function returned %d contacts for geom pair (%d, %d), "
+              "expected at most %d from mj_maxContact", nconbuffer[i], g1, g2, expected_max);
+    }
+  }
+
+  // set nativeccd buffer back to NULL
+  mjc_setCCDBuffer(NULL);
+  mj_freeStack(d);
+
+  int ncon = 0;
+  for (int i = 0; i < npair; i++) {
+    ncon += nconbuffer[i];
+  }
+
+  if (ncon == 0) {
+    mj_freeStack(d);
     return;
   }
 
-
-
-  // call collision detector to generate contacts
-  mjPreContact precon[mjMAXCONPAIR];
-  if (!(num = collisionFunc(m, d, precon, g1, g2, margin + gap))) {
-    return;
-  }
-
-  // check number of contacts, SHOULD NOT OCCUR
-  if (num > mjMAXCONPAIR) {
-    mjERROR("too many contacts returned by collision function");
-  }
-
-  mjContact* con = (mjContact*) mj_arenaAllocByte(d, sizeof(mjContact) * num, _Alignof(mjContact));
+  // try allocate contact buffer in arena
+  mjContact* con =
+    (mjContact*) mj_arenaAllocByte(d, sizeof(mjContact) * ncon, _Alignof(mjContact));
   if (!con) {
     mj_warning(d, mjWARN_CONTACTFULL, d->ncon);
+    mj_freeStack(d);
     return;
   }
+  d->ncon += ncon;
 
-  // set condim, solref, solimp, friction: dynamic
-  if (ipair < 0) {
-    mj_contactParam(m, &condim, solref, solimp, friction, g1, g2, -1, -1);
-  }
+  // fill in contact data
+  int conpos = 0;
+  for (int i = 0; i < npair; i++) {
+    if (!(ncon = nconbuffer[i]))
+      continue;
 
-  // set condim, solref, solimp, friction: pair
-  else {
-    condim = m->pair_dim[ipair];
-    mju_copy(solref, m->pair_solref+mjNREF*ipair, mjNREF);
-    mju_copy(solimp, m->pair_solimp+mjNIMP*ipair, mjNIMP);
-    mju_copy(friction, m->pair_friction+5*ipair, 5);
+    int condim;
+    mjtNum friction[5], solref[mjNREF], solimp[mjNIMP];
+    mjtNum solreffriction[mjNREF] = {0};
+    int g1 = pairbuffer[4*i + 0];
+    int g2 = pairbuffer[4*i + 1];
+    int ipair = pairbuffer[4*i + 2];
 
-    // reference, friction directions
-    if (m->pair_solreffriction[mjNREF*ipair] || m->pair_solreffriction[mjNREF*ipair + 1]) {
-      mju_copy(solreffriction, m->pair_solreffriction+mjNREF*ipair, mjNREF);
+    if (ipair >= 0) {
+      condim = m->pair_dim[ipair];
+      mju_copy(solref, m->pair_solref+mjNREF*ipair, mjNREF);
+      mju_copy(solimp, m->pair_solimp+mjNIMP*ipair, mjNIMP);
+      mju_copy(friction, m->pair_friction+5*ipair, 5);
+      if (m->pair_solreffriction[mjNREF*ipair] || m->pair_solreffriction[mjNREF*ipair + 1]) {
+        mju_copy(solreffriction, m->pair_solreffriction+mjNREF*ipair, mjNREF);
+      }
+    } else {
+      mj_contactParam(m, &condim, solref, solimp, friction, g1, g2, -1, -1);
     }
+
+    mjPreContact* bc = conbuffer + pairbuffer[4*i + 3];
+    margin = getMargin(m, g1, g2, ipair);
+    for (int j=0; j < ncon; j++) {
+      mjContact* c = con + conpos + j;
+      c->dist = bc[j].dist;
+      mji_copy3(c->pos, bc[j].pos);
+      mji_copy3(c->frame, bc[j].normal);
+      mji_copy3(c->frame + 3, bc[j].tangent);
+      c->geom[0] = g1;
+      c->geom[1] = g2;
+      c->flex[0] = -1;
+      c->flex[1] = -1;
+      c->elem[0] = -1;
+      c->elem[1] = -1;
+      c->vert[0] = -1;
+      c->vert[1] = -1;
+      mj_setContact(m, c, condim, margin, solref, solreffriction, solimp, friction);
+    }
+    conpos += ncon;
   }
-
-  // add contacts returned by collision detector
-  for (int i=0; i < num; i++) {
-    // set contact parameters
-    con[i].dist = precon[i].dist;
-    mji_copy3(con[i].pos, precon[i].pos);
-    mji_copy3(con[i].frame + 0, precon[i].normal);
-    mji_copy3(con[i].frame + 3, precon[i].tangent);
-
-    con[i].geom[0] = g1;
-    con[i].geom[1] = g2;
-    con[i].flex[0] = -1;
-    con[i].flex[1] = -1;
-    con[i].elem[0] = -1;
-    con[i].elem[1] = -1;
-    con[i].vert[0] = -1;
-    con[i].vert[1] = -1;
-
-    // set remaining contact parameters
-    mj_setContact(m, con + i, condim, margin, solref, solreffriction, solimp, friction);
-  }
-
-  // add to ncon
-  d->ncon += num;
-
-  // move arena pointer back to the end of the contact array
-  resetArena(d);
+  mj_freeStack(d);
 }
 
 
diff --git a/src/engine/engine_collision_driver.h b/src/engine/engine_collision_driver.h
index e231eaaf..0e2b9edd 100644
--- a/src/engine/engine_collision_driver.h
+++ b/src/engine/engine_collision_driver.h
@@ -44,20 +44,9 @@ MJAPI int mj_collideOBB(const mjtNum aabb1[6], const mjtNum aabb2[6],
 // is element active (for collisions)
 MJAPI int mj_isElemActive(const mjModel* m, int f, int e);
 
-// checks if pair is already present in pair_geom and calls narrow phase
-void mj_collideGeomPair(const mjModel* m, mjData* d, int g1, int g2, int merged,
-                        int startadr, int pairadr);
-
-// binary search between two bodyflex trees
-void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2,
-                    int merged, int startadr, int pairadr);
-
 // broad phase collision detection; return list of bodyflex pairs
 int mj_broadphase(const mjModel* m, mjData* d, int* bfpair, int maxpair);
 
-// test two geoms for collision, apply filters, add to contact list
-void mj_collideGeoms(const mjModel* m, mjData* d, int ipair, int g1, int g2);
-
 // test a plane geom and a flex for collision, add to contact list
 void mj_collidePlaneFlex(const mjModel* m, mjData* d, int g, int f);