Do not rely on global object lists in mjCBody::NextChild().

PiperOrigin-RevId: 782963753
Change-Id: I0afecbab4b7796f077963ba66eb7140d8f83ad89

doc/changelog.rst

@@ -12,6 +12,12 @@ General
 - Removed the SdfLib plugin and the dependency on `SdfLib <https://github.com/UPC-ViRVIG/SdfLib>`__. SDFs are now
   supported natively in mjModel.
 
+Bug fixes
+^^^^^^^^^
+- Fixed a bug that caused object lists in the child to have missing elements after attaching an mjSpec. This was caused
+  by adding to the lists only the objects that belong to the tree of the requested body, but this causes to skip objects
+  that were attached, since they belong to the tree of the parent.
+
 Version 3.3.4 (July 8, 2025)
 ----------------------------
 

python/mujoco/specs_test.py

@@ -1152,6 +1152,38 @@ class SpecsTest(absltest.TestCase):
     with self.assertRaisesRegex(ValueError, 'Frame not found.'):
       parent.attach(child4, frame='invalid_frame', prefix='child3-')
 
+  def test_attach_valid_child_lists(self):
+    xml1 = """
+    <mujoco>
+      <worldbody>
+        <body name="b1">
+          <geom name="g1"/>
+          <joint name="j1" type="hinge"/>
+        </body>
+      </worldbody>
+    </mujoco>
+    """
+
+    xml2 = """
+    <mujoco>
+      <worldbody>
+        <body name="b2">
+          <geom name="g2"/>
+          <joint name="j2" type="hinge"/>
+        </body>
+      </worldbody>
+    </mujoco>
+    """
+
+    parent = mujoco.MjSpec.from_string(xml1)
+    child = mujoco.MjSpec.from_string(xml2)
+    self.assertLen(child.joints, 1)
+    self.assertLen(child.geoms, 1)
+    frame = parent.worldbody.add_frame()
+    parent.attach(child, prefix='', frame=frame)
+    self.assertLen(child.joints, 1)
+    self.assertLen(child.geoms, 1)
+
   def test_bind(self):
     spec = mujoco.MjSpec.from_string("""
     <mujoco>

src/user/user_objects.cc

@@ -1660,40 +1660,68 @@ mjCBase* mjCBody::FindObject(mjtObj type, std::string _name, bool recursive) {
 
 
 
-// return true if child is descendant of this body
+// get list of a given type
-bool mjCBody::IsAncestor(const mjCBody* child) const {
+template<>
-  if (!child) {
+const std::vector<mjCBody*>& mjCBody::GetList<mjCBody>() const {
-    return false;
+  return bodies;
-  }
+}
 
-  if (child == this) {
+template<>
-    return true;
+const std::vector<mjCJoint*>& mjCBody::GetList<mjCJoint>() const {
-  }
+  return joints;
+}
 
-  return IsAncestor(child->parent);
+template<>
+const std::vector<mjCGeom*>& mjCBody::GetList<mjCGeom>() const {
+  return geoms;
+}
+
+template<>
+const std::vector<mjCSite*>& mjCBody::GetList<mjCSite>() const {
+  return sites;
+}
+
+template<>
+const std::vector<mjCCamera*>& mjCBody::GetList<mjCCamera>() const {
+  return cameras;
+}
+
+template<>
+const std::vector<mjCLight*>& mjCBody::GetList<mjCLight>() const {
+  return lights;
+}
+
+template<>
+const std::vector<mjCFrame*>& mjCBody::GetList<mjCFrame>() const {
+  return frames;
 }
 
 
 
+// gets next child of the same type, recursively depth first if requested
 template <class T>
-mjsElement* mjCBody::GetNext(const std::vector<T*>& list, const mjsElement* child) {
+static mjsElement* GetNext(const mjCBody* body, const mjsElement* child,
-  if (list.empty()) {
+                           bool* found, bool recursive) {
-    // no children
+  std::vector<T*> list = body->GetList<T>();
+
+  for (unsigned int i = 0; i < list.size(); i++) {
+    if (*found) {
+      return list[i]->spec.element;
+    }
+
+    if (list[i]->spec.element == child) {
+      *found = true;
+    }
+  }
+
+  if (!recursive) {
     return nullptr;
   }
 
-  for (unsigned int i = 0; i < list.size() - (child ? 1 : 0); i++) {
+  for (auto& other : body->Bodies()) {
-    if (!IsAncestor(list[i]->GetParent())) {
+    mjsElement* candidate = GetNext<T>(other, child, found, true);
-      continue;  // TODO: this recursion is wasteful
+    if (candidate) {
-    }
+      return candidate;
-
-    // first child in this body
-    if (!child) {
-      return list[i]->spec.element;
-
-    // next child is in this body
-    } else if (list[i]->spec.element == child && IsAncestor(list[i+1]->GetParent())) {
-      return list[i+1]->spec.element;
     }
   }
 
@@ -1702,6 +1730,34 @@ mjsElement* mjCBody::GetNext(const std::vector<T*>& list, const mjsElement* chil
 
 
 
+// get next body depth first
+static mjsElement* GetNextBody(const mjCBody* body, const mjsElement* child,
+                              bool* found, bool recursive) {
+  for (auto& other : body->Bodies()) {
+    if (*found) {
+      return other->spec.element;
+    }
+
+    if (other->spec.element == child) {
+      *found = true;
+    }
+
+    if (!recursive) {
+      continue;
+    }
+
+    mjsElement* candidate = GetNextBody(other, child, found, true);
+    if (candidate) {
+      return candidate;
+    }
+  }
+
+  return nullptr;
+}
+
+
+
+
 // get next child of given type
 mjsElement* mjCBody::NextChild(const mjsElement* child, mjtObj type, bool recursive) {
   if (type == mjOBJ_UNKNOWN) {
@@ -1714,30 +1770,30 @@ mjsElement* mjCBody::NextChild(const mjsElement* child, mjtObj type, bool recurs
     throw mjCError(this, "child element is not of requested type");
   }
 
-
   mjsElement* candidate = nullptr;
+  bool found = child == nullptr;
   switch (type) {
     case mjOBJ_BODY:
     case mjOBJ_XBODY:
-      candidate = GetNext(recursive ? model->bodies_ : bodies, child);
+      candidate = GetNextBody(this, child, &found, recursive);
       break;
     case mjOBJ_JOINT:
-      candidate = GetNext(recursive ? model->joints_ : joints, child);
+      candidate = GetNext<mjCJoint>(this, child, &found, recursive);
       break;
     case mjOBJ_GEOM:
-      candidate = GetNext(recursive ? model->geoms_ : geoms, child);
+      candidate = GetNext<mjCGeom>(this, child, &found, recursive);
       break;
     case mjOBJ_SITE:
-      candidate = GetNext(recursive ? model->sites_ : sites, child);
+      candidate = GetNext<mjCSite>(this, child, &found, recursive);
       break;
     case mjOBJ_CAMERA:
-      candidate = GetNext(recursive ? model->cameras_ : cameras, child);
+      candidate = GetNext<mjCCamera>(this, child, &found, recursive);
       break;
     case mjOBJ_LIGHT:
-      candidate = GetNext(recursive ? model->lights_ : lights, child);
+      candidate = GetNext<mjCLight>(this, child, &found, recursive);
       break;
     case mjOBJ_FRAME:
-      candidate = GetNext(recursive ? model->frames_ : frames, child);
+      candidate = GetNext<mjCFrame>(this, child, &found, recursive);
       break;
     default:
       throw mjCError(this,

src/user/user_objects.h

@@ -505,6 +505,10 @@ class mjCBody : public mjCBody_, private mjsBody {
   // getters
   std::vector<mjCBody*> Bodies() const { return bodies; }
 
+  // get list of a given type
+  template <class T>
+  const std::vector<T*>& GetList() const;
+
   // accumulate inertia of another body into this body, if `result` is not nullptr, the accumulated
   // inertia will be stored in `result`, otherwise the body's private spec will be used.
   void AccumulateInertia(const mjsBody* other, mjsBody* result = nullptr);
@@ -534,12 +538,6 @@ class mjCBody : public mjCBody_, private mjsBody {
   template <typename T>
   void CopyList(std::vector<T*>& dst, const std::vector<T*>& src,
                 std::map<mjCFrame*, int>& fmap, const mjCFrame* pframe = nullptr);
-
-  // gets next child of the same type in this body
-  template <class T>
-  mjsElement* GetNext(const std::vector<T*>& list, const mjsElement* child);
-
-  bool IsAncestor(const mjCBody* child) const;  // true if child is a descendant of this body
 };
 
 

test/user/user_api_test.cc

@@ -130,7 +130,7 @@ TEST_F(MujocoTest, TreeTraversal) {
   EXPECT_EQ(site4, mjs_firstChild(body2, mjOBJ_SITE, /*recurse=*/true));
   EXPECT_EQ(site5, mjs_firstChild(body3, mjOBJ_SITE, /*recurse=*/true));
 
-  // text next, nonrecursive
+  // test next, nonrecursive
   EXPECT_EQ(site2, mjs_nextChild(body1, site1, /*recursive=*/false));
   EXPECT_EQ(site3, mjs_nextChild(body1, site2, /*recursive=*/false));
   EXPECT_EQ(nullptr, mjs_nextChild(body1, site3, /*recursive=*/false));
@@ -146,7 +146,7 @@ TEST_F(MujocoTest, TreeTraversal) {
   EXPECT_EQ(mjs_nextElement(spec, geom2), geom3);
   EXPECT_EQ(mjs_nextElement(spec, geom3), nullptr);
 
-  // text next, recursive
+  // test next, recursive
   EXPECT_EQ(b2, mjs_nextChild(world, b1, /*recursive=*/true));
   EXPECT_EQ(b3, mjs_nextChild(world, b2, /*recursive=*/true));
   EXPECT_EQ(site2, mjs_nextChild(body1, site1, /*recursive=*/true));