From 488041d362b0314551ca4128032ecfafce39b743 Mon Sep 17 00:00:00 2001
From: Saran Tunyasuvunakool <stunya@google.com>
Date: Fri, 18 Aug 2023 13:52:15 -0700
Subject: [PATCH] Add a private function for resizing contact and efc-sized
 arrays in Python.

This function is currently used for internal experimentation.

PiperOrigin-RevId: 558238153
Change-Id: I32eec8f13ac2d995e90c5421f2a157412e2148f2
---
 python/mujoco/bindings_test.py | 22 +++++++++++++
 python/mujoco/functions.cc     | 58 ++++++++++++++++++++++++++++++++++
 2 files changed, 80 insertions(+)

diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py
index 629face1..6d278773 100644
--- a/python/mujoco/bindings_test.py
+++ b/python/mujoco/bindings_test.py
@@ -540,6 +540,27 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     self.assertLen(expected_H, expected_ncon)
     np.testing.assert_array_equal(self.data.contact.H, expected_H)
 
+  def test_realloc_con_efc(self):
+    self.assertEmpty(self.data.contact)
+
+    ncon = 9
+    nefc = 11
+    mujoco._functions._realloc_con_efc(self.data, ncon, nefc)
+
+    ncon = 13
+    nefc = 17
+    mujoco._functions._realloc_con_efc(self.data, ncon=ncon, nefc=nefc)
+
+    self.assertLen(self.data.contact, ncon)
+    self.assertEqual(self.data.efc_id.shape, (nefc,))
+    self.assertEqual(self.data.efc_KBIP.shape, (nefc, 4))
+
+    expected_error = 'insufficient arena memory available'
+    with self.assertRaisesWithLiteralMatch(mujoco.FatalError, expected_error):
+      mujoco._functions._realloc_con_efc(self.data, 100000000, 100000000)
+    self.assertEmpty(self.data.contact)
+    self.assertEmpty(self.data.efc_id)
+
   def test_mj_struct_list_equality(self):
     model2 = mujoco.MjModel.from_xml_string(TEST_XML)
     data2 = mujoco.MjData(model2)
@@ -1281,5 +1302,6 @@ Euler integrator, semi-implicit in velocity.
   def test_load_plugin(self):
     mujoco.MjModel.from_xml_string(TEST_XML_PLUGIN)
 
+
 if __name__ == '__main__':
   absltest.main()
diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc
index edd4d4d8..580595df 100644
--- a/python/mujoco/functions.cc
+++ b/python/mujoco/functions.cc
@@ -20,8 +20,10 @@
 #include <optional>
 
 #include <Eigen/Core>
+#include <mujoco/mjxmacro.h>
 #include "function_traits.h"
 #include "functions.h"
+#include "private.h"
 #include "raw.h"
 #include <pybind11/eigen.h>
 #include <pybind11/pybind11.h>
@@ -1327,6 +1329,62 @@ PYBIND11_MODULE(_functions, pymodule) {
       });
   Def<traits::mjd_subQuat>(pymodule);
   Def<traits::mjd_quatIntegrate>(pymodule);
+
+  pymodule.def(
+      "_realloc_con_efc",
+      [](MjDataWrapper& d, int ncon, int nefc) {
+        raw::MjData* data = d.get();
+
+        auto cleanup = [](raw::MjData* data) {
+#ifdef ADDRESS_SANITIZER
+        ASAN_POISON_MEMORY_REGION(
+            static_cast<char*>(data->arena),
+            (data->nstack - data->pstack) * sizeof(mjtNum));
+#endif
+          data->ncon = 0;
+          data->nefc = 0;
+          data->contact = static_cast<raw::MjContact*>(data->arena);
+#define X(type, name, nr, nc) data->name = nullptr;
+          MJDATA_ARENA_POINTERS_PRIMAL
+          MJDATA_ARENA_POINTERS_DUAL
+#undef X
+        };
+
+        cleanup(data);
+        data->ncon = ncon;
+        data->nefc = nefc;
+        data->contact =
+            static_cast<raw::MjContact*>(InterceptMjErrors(::mj_arenaAlloc)(
+                data, ncon * sizeof(raw::MjContact), alignof(raw::MjContact)));
+        if (!data->contact) {
+          cleanup(data);
+          throw FatalError("insufficient arena memory available");
+        }
+
+#undef MJ_M
+#define MJ_M(x) d.metadata().x
+#undef MJ_D
+#define MJ_D(x) data->x
+#define X(type, name, nr, nc)                                         \
+  data->name = static_cast<type*>(InterceptMjErrors(::mj_arenaAlloc)( \
+      data, sizeof(type) * (nr) * (nc), _Alignof(type)));             \
+  if (!data->name) {                                                  \
+    cleanup(data);                                                    \
+    throw FatalError("insufficient arena memory available");          \
+  }
+
+        MJDATA_ARENA_POINTERS_PRIMAL
+        if (d.metadata().is_dual) {
+          MJDATA_ARENA_POINTERS_DUAL
+        }
+#undef X
+#undef MJ_D
+#define MJ_D(x) x
+#undef MJ_M
+#define MJ_M(x) x
+      },
+      py::arg("d"), py::arg("ncon"), py::arg("nefc"),
+      py::call_guard<py::gil_scoped_release>());
 }  // PYBIND11_MODULE NOLINT(readability/fn_size)
 }  // namespace
 }  // namespace mujoco::python