Dynamically allocate contact and efc_ arrays on a new memory arena.

- Add private function `mj_arenaAlloc`. This is used internally to allocate memory from the arena.

- Add private function `mj_nefc` to count constraints. This function returns a tight upper bound on `d->nefc`. The number of counted constraints can be slightly bigger than exact `d->nefc` in the case of constraints with empty Jacobian, as when placing a frictional tendon between two world sites.

- Add new `memory` attribute to the `size` XML element for specification of arena memory size. This attribute is mutually exclusive with `nstack` and `njmax` specifications, which are now deprecated (but left around for the time being for legacy compatibility).

- Move `d->stack` to the end of the new arena space. The stack now grows in reverse from the end.

PiperOrigin-RevId: 479341539
Change-Id: Ie019c202e0908577ffc6f833a37920858116f667

python/mujoco/structs.h

@@ -92,38 +92,59 @@ class MjWrapper {};
 template <typename T>
 class StructListBase {
  public:
-  StructListBase(T* ptr, int num, pybind11::handle owner) : ptr_(ptr) {
-    for (int i = 0; i < num; ++i) {
-      wrappers_.push_back(std::make_shared<MjWrapper<T>>(&ptr[i], owner));
+  StructListBase(T* ptr, int num, pybind11::handle owner, bool lazy = false)
+      : ptr_(ptr), num_(num), owner_(owner) {
+    if (!lazy) {
+      PopulateUpTo(size());
     }
   }
 
   StructListBase(const StructListBase& other) = delete;
   StructListBase(StructListBase&& other) = default;
 
+  virtual ~StructListBase() = default;
+
   MjWrapper<T>& operator[](int i) {
-    if (i < 0 || i >= wrappers_.size()) {
+    if (i < 0 || i >= size()) {
       throw pybind11::index_error();
     }
+    PopulateUpTo(i);
     return *wrappers_[i];
   }
 
-  int size() const {
-    return wrappers_.size();
+  virtual int size() const {
+    return num_;
   }
 
+  T* get() const { return ptr_; }
+  pybind11::handle owner() const { return owner_; }
+
  protected:
+  void PopulateUpTo(int n) {
+    while (wrappers_.size() <= n) {
+      wrappers_.push_back(
+          std::make_shared<MjWrapper<T>>(&ptr_[wrappers_.size()], owner_));
+    }
+  }
+
   // Slicing
-  StructListBase(StructListBase& other, pybind11::slice slice) {
+  StructListBase(StructListBase& other, pybind11::slice slice)
+      : owner_(other.owner_) {
     pybind11::size_t start, stop, step, slicelength;
-    slice.compute(other.size(), &start, &stop, &step, &slicelength);
+    if (!slice.compute(other.size(), &start, &stop, &step, &slicelength)) {
+      throw pybind11::index_error();
+    }
+    other.PopulateUpTo(stop);
     ptr_ = &other.ptr_[start];
     for (int i = start; i < stop; i += step) {
       wrappers_.push_back(other.wrappers_[i]);
     }
+    num_ = wrappers_.size();
   }
 
   T* ptr_;
+  int num_;
+  pybind11::handle owner_;
 
   // Using shared_ptr here so that we get identical Python objects when slicing.
   std::vector<std::shared_ptr<MjWrapper<T>>> wrappers_;
@@ -277,6 +298,8 @@ class MjStructList<raw::MjWarningStat>
     : public StructListBase<raw::MjWarningStat> {
  public:
   MjStructList(raw::MjWarningStat* ptr, int num, pybind11::handle owner);
+  MjStructList(MjStructList&&) = default;
+  ~MjStructList() override = default;
 
   using StructListBase::operator[];
   using StructListBase::size;
@@ -325,6 +348,8 @@ template <>
 class MjStructList<raw::MjTimerStat> : public StructListBase<raw::MjTimerStat> {
  public:
   MjStructList(raw::MjTimerStat* ptr, int num, pybind11::handle owner);
+  MjStructList(MjStructList&&) = default;
+  ~MjStructList() override = default;
 
   using StructListBase::operator[];
   using StructListBase::size;
@@ -374,6 +399,8 @@ class MjStructList<raw::MjSolverStat>
     : public StructListBase<raw::MjSolverStat> {
  public:
   MjStructList(raw::MjSolverStat* ptr, int num, pybind11::handle owner);
+  MjStructList(MjStructList&&) = default;
+  ~MjStructList() override = default;
 
   using StructListBase::operator[];
   using StructListBase::size;
@@ -489,33 +516,28 @@ struct enable_if_mj_struct<raw::MjContact> { using type = void; };
 template <>
 class MjStructList<raw::MjContact> : public StructListBase<raw::MjContact> {
  public:
-  MjStructList(raw::MjContact* ptr, int num, pybind11::handle owner);
+  MjStructList(raw::MjContact* ptr, int nconmax,
+               int* ncon, pybind11::handle owner);
+  MjStructList(MjStructList&&) = default;
+  ~MjStructList() override = default;
 
   using StructListBase::operator[];
-  using StructListBase::size;
+
+  int size() const override {
+    if (ncon_) {
+      return *ncon_;
+    } else {
+      return StructListBase::size();
+    }
+  }
+
   MjStructList Slice(pybind11::slice slice) {
     return MjStructList(*this, slice);
   }
 
-#define X(type, var) pybind11::array_t<type> var
-  X(mjtNum, dist);
-  X(mjtNum, pos);
-  X(mjtNum, frame);
-  X(mjtNum, includemargin);
-  X(mjtNum, friction);
-  X(mjtNum, solref);
-  X(mjtNum, solimp);
-  X(mjtNum, mu);
-  X(mjtNum, H);
-  X(int, dim);
-  X(int, geom1);
-  X(int, geom2);
-  X(int, exclude);
-  X(int, efc_address);
-#undef X
-
  protected:
   MjStructList(MjStructList& other, pybind11::slice slice);
+  int* ncon_ = nullptr;
 };
 
 using MjContactList = MjStructList<raw::MjContact>;
@@ -539,6 +561,7 @@ class MjWrapper<raw::MjData>: public WrapperBase<raw::MjData> {
   MjWrapper(MjWrapper&&);
   ~MjWrapper();
 
+  const MjDataMetadata& metadata() const { return metadata_; }
   MjDataIndexer& indexer() { return indexer_; }
 
   void Serialize(std::ostream& output) const;
@@ -548,10 +571,13 @@ class MjWrapper<raw::MjData>: public WrapperBase<raw::MjData> {
       "__MUJOCO_STRUCTS_MJDATAWRAPPER_LOOKUP";
   static MjWrapper* FromRawPointer(raw::MjData* m) noexcept;
 
+
 #define X(dtype, var, dim0, dim1) py_array_or_tuple_t<dtype> var;
   MJDATA_POINTERS
 #undef X
 
+  py_array_or_tuple_t<mjContact> contact;
+
   py_array_or_tuple_t<raw::MjWarningStat> warning;
   py_array_or_tuple_t<raw::MjTimerStat> timer;
   py_array_or_tuple_t<raw::MjSolverStat> solver;