From 49c33716a1cced63d5fc35ef514d8cb1c6c95521 Mon Sep 17 00:00:00 2001
From: Erik Frey <erikfrey@google.com>
Date: Tue, 20 May 2025 15:25:10 -0700
Subject: [PATCH] Allow _realloc_con_efc to also allocate new nJ arrays.

PiperOrigin-RevId: 761252071
Change-Id: Idce822c6f2b8f2d81a5d43228f5dc35281c5ad07
---
 python/mujoco/bindings_test.py |  4 +++-
 python/mujoco/functions.cc     | 14 ++++++++------
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py
index 3896b8ce..99ca73ab 100644
--- a/python/mujoco/bindings_test.py
+++ b/python/mujoco/bindings_test.py
@@ -623,10 +623,12 @@ class MuJoCoBindingsTest(parameterized.TestCase):
 
     ncon = 13
     nefc = 17
-    mujoco._functions._realloc_con_efc(self.data, ncon=ncon, nefc=nefc)
+    nj = 21
+    mujoco._functions._realloc_con_efc(self.data, ncon=ncon, nefc=nefc, nJ=nj)
 
     self.assertLen(self.data.contact, ncon)
     self.assertEqual(self.data.efc_id.shape, (nefc,))
+    self.assertEqual(self.data.efc_J.shape, (nj,))
     self.assertEqual(self.data.efc_KBIP.shape, (nefc, 4))
 
     expected_error = 'insufficient arena memory available'
diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc
index 7fed04f7..d4708254 100644
--- a/python/mujoco/functions.cc
+++ b/python/mujoco/functions.cc
@@ -1467,10 +1467,10 @@ PYBIND11_MODULE(_functions, pymodule) {
 
   pymodule.def(
       "_realloc_con_efc",
-      [](MjDataWrapper& d, int ncon, int nefc) {
+      [](MjDataWrapper& d, int ncon, int nefc, int nJ) {
         raw::MjData* data = d.get();
 
-        auto cleanup = [](raw::MjData* data) {
+        auto cleanup = [](raw::MjData* data, int nJ) {
 #ifdef ADDRESS_SANITIZER
         ASAN_POISON_MEMORY_REGION(
             static_cast<char*>(data->arena),
@@ -1479,6 +1479,7 @@ PYBIND11_MODULE(_functions, pymodule) {
           data->parena = 0;
           data->ncon = 0;
           data->nefc = 0;
+          if (nJ > -1) data->nJ = 0;
           data->contact = static_cast<raw::MjContact*>(data->arena);
 #define X(type, name, nr, nc) data->name = nullptr;
           MJDATA_ARENA_POINTERS_SOLVER
@@ -1486,14 +1487,15 @@ PYBIND11_MODULE(_functions, pymodule) {
 #undef X
         };
 
-        cleanup(data);
+        cleanup(data, nJ);
         data->ncon = ncon;
         data->nefc = nefc;
+        if (nJ > -1) data->nJ = nJ;
         data->contact =
             static_cast<raw::MjContact*>(InterceptMjErrors(::mj_arenaAllocByte)(
                 data, ncon * sizeof(raw::MjContact), alignof(raw::MjContact)));
         if (!data->contact) {
-          cleanup(data);
+          cleanup(data, nJ);
           throw FatalError("insufficient arena memory available");
         }
 
@@ -1505,7 +1507,7 @@ PYBIND11_MODULE(_functions, pymodule) {
   data->name = static_cast<type*>(InterceptMjErrors(::mj_arenaAllocByte)( \
       data, sizeof(type) * (nr) * (nc), alignof(type)));                  \
   if (!data->name) {                                                      \
-    cleanup(data);                                                        \
+    cleanup(data, nJ);                                                        \
     throw FatalError("insufficient arena memory available");              \
   }
 
@@ -1519,7 +1521,7 @@ PYBIND11_MODULE(_functions, pymodule) {
 #undef MJ_M
 #define MJ_M(x) x
       },
-      py::arg("d"), py::arg("ncon"), py::arg("nefc"),
+      py::arg("d"), py::arg("ncon"), py::arg("nefc"), py::arg("nJ") = -1,
       py::call_guard<py::gil_scoped_release>());
 }  // PYBIND11_MODULE NOLINT(readability/fn_size)
 }  // namespace