Add ASAN and MSAN instrumentation for arena and mjData stack.

PiperOrigin-RevId: 535635490
Change-Id: Ia1e58771bae25bf72c861f57452fb8faa2b918fb

python/mujoco/CMakeLists.txt

@@ -203,6 +203,11 @@ add_subdirectory(simulate)
 
 add_subdirectory(util)
 
+add_library(mujoco_private_api_header INTERFACE)
+target_sources(mujoco_private_api_header INTERFACE private.h)
+set_target_properties(mujoco_private_api_header PROPERTIES PUBLIC_HEADER private.h)
+target_link_libraries(mujoco_private_api_header INTERFACE mujoco)
+
 if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/enum_traits.h)
   add_library(enum_traits INTERFACE)
   target_sources(enum_traits INTERFACE enum_traits.h)
@@ -242,7 +247,13 @@ target_link_libraries(function_traits INTERFACE mujoco absl::core_headers)
 add_library(errors_header INTERFACE)
 target_sources(errors_header INTERFACE errors.h)
 set_target_properties(errors_header PROPERTIES PUBLIC_HEADER errors.h)
-target_link_libraries(errors_header INTERFACE crossplatform func_wrap mujoco)
+target_link_libraries(
+  errors_header
+  INTERFACE crossplatform
+            func_wrap
+            mujoco
+            mujoco_private_api_header
+)
 
 add_library(raw INTERFACE)
 target_sources(raw INTERFACE raw.h)
@@ -381,6 +392,7 @@ target_link_libraries(
   PRIVATE absl::flat_hash_map
           crossplatform
           mujoco
+          mujoco_private_api_header
           raw
           errors_header
           func_wrap

python/mujoco/bindings_test.py

@@ -515,8 +515,8 @@ class MuJoCoBindingsTest(parameterized.TestCase):
   def test_mj_contact_list(self):
     self.assertEmpty(self.data.contact)
 
-    expected_ncon = 1234
-    self.data.ncon = expected_ncon
+    expected_ncon = 4
+    mujoco.mj_forward(self.model, self.data)
     self.assertLen(self.data.contact, expected_ncon)
 
     expected_pos = []

python/mujoco/errors.h

@@ -21,16 +21,11 @@
 #include <type_traits>
 
 #include <mujoco/mjexport.h>
+#include "private.h"
 #include "util/crossplatform.h"
 #include "util/func_wrap.h"
 #include <pybind11/pybind11.h>
 
-// DO NOT USE THIS FUNCTION ELSEWHERE.
-// It should be regarded as part of MuJoCo's internal implementation detail.
-extern "C" {
-MJAPI void _mjPRIVATE__set_tls_error_fn(void (*h)(const char*));
-}
-
 // When building for Linux and statically linking against a "hermetic" libc++abi
 // (i.e. where libc++/libc++abi symbols all have "hidden" visibility), exception
 // types do not propagate correctly across shared library boundaries.

python/mujoco/private.h

@@ -0,0 +1,28 @@
+// Copyright 2023 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_PYTHON_PRIVATE_H_
+#define MUJOCO_PYTHON_PRIVATE_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+
+// DO NOT USE THESE FUNCTIONS ELSEWHERE.
+// They should be regarded as part of MuJoCo's internal implementation detail.
+extern "C" {
+MJAPI void _mjPRIVATE__set_tls_error_fn(void (*h)(const char*));
+MJAPI void* mj_arenaAlloc(mjData* d, int bytes, int alignment);
+}
+
+#endif  // MUJOCO_PYTHON_PRIVATE_H_

python/mujoco/structs.cc

@@ -38,6 +38,7 @@
 #include "function_traits.h"
 #include "indexers.h"
 #include "mjdata_meta.h"
+#include "private.h"
 #include "raw.h"
 #include "serialization.h"
 #include <pybind11/numpy.h>
@@ -695,7 +696,6 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
 
   // Write struct and scalar fields
 #define X(var) WriteBytes(output, &ptr_->var, sizeof(ptr_->var))
-  X(parena);
   X(maxuse_stack);
   X(maxuse_arena);
   X(maxuse_con);
@@ -727,7 +727,6 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
 #define MJ_D(x) this->ptr_->x
 #define X(type, name, nr, nc)                                   \
   if ((nr) * (nc)) {                                            \
-    WriteInt(output, PTRDIFF(ptr_->name, ptr_->arena));         \
     WriteBytes(output, ptr_->name, sizeof(type) * (nr) * (nc)); \
   }
 
@@ -786,7 +785,6 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
   ReadBytes(input, (void*) &d->var, sizeof(d->var)); \
   CheckInput(input, "mjData");
 
-  X(parena);
   X(maxuse_stack);
   X(maxuse_arena);
   X(maxuse_con);
@@ -816,11 +814,11 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
 #define MJ_M(x) m.x
 #undef MJ_D
 #define MJ_D(x) d->x
-#define X(type, name, nr, nc)                              \
-  if ((nr) * (nc)) {                                       \
-    d->name = reinterpret_cast<decltype(d->name)>(         \
-        static_cast<char*>(d->arena) + ReadInt(input));    \
-    ReadBytes(input, d->name, sizeof(type) * (nr) * (nc)); \
+#define X(type, name, nr, nc)                                          \
+  if ((nr) * (nc)) {                                                   \
+    d->name = static_cast<decltype(d->name)>(                          \
+        mj_arenaAlloc(d, sizeof(type) * (nr) * (nc), alignof(type)));  \
+    ReadBytes(input, d->name, sizeof(type) * (nr) * (nc));             \
   }
 
     MJDATA_ARENA_POINTERS_CONTACT