chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Change mjs_nextChild to be DFS to match XML parsing.

Browse Source 
PiperOrigin-RevId: 778915232
Change-Id: I5e290f47b7d2530597bb6c85b66ff2ae807e5510
This commit is contained in:
Alessio Quaglino
2025-07-03 10:05:46 -07:00
committed by Copybara-Service
parent 67f9ffddbd
commit af7d20c035
4 changed files with 49 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/specs_test.py
+7 -7
View File
@@ -643,7 +643,7 @@ class SpecsTest(absltest.TestCase):
<worldbody>
<body name="body1">
<site name="site1"/>
<body name="body3">
<body name="body2">
<site name="site4"/>
</body>
<geom name="geom1" size="1"/>
@@ -652,7 +652,7 @@ class SpecsTest(absltest.TestCase):
<site name="site3"/>
<geom name="geom3" size="1"/>
</body>
<body name="body2">
<body name="body3">
<site name="site5"/>
</body>
</worldbody>
@@ -669,12 +669,12 @@ class SpecsTest(absltest.TestCase):
self.assertEqual(spec.bodies[2].name, 'body2')
self.assertEqual(spec.bodies[3].name, 'body3')
self.assertEqual(spec.bodies[1].parent, spec.bodies[0])
self.assertEqual(spec.bodies[2].parent, spec.bodies[0])
self.assertEqual(spec.bodies[3].parent, spec.bodies[1])
self.assertEqual(spec.bodies[2].parent, spec.bodies[1])
self.assertEqual(spec.bodies[3].parent, spec.bodies[0])
self.assertLen(spec.worldbody.find_all(bodytype), 3)
self.assertLen(spec.bodies[1].find_all(bodytype), 1)
self.assertEmpty(spec.bodies[3].find_all(bodytype))
self.assertEqual(spec.bodies[1].find_all('body')[0].name, 'body3')
self.assertEqual(spec.bodies[1].find_all('body')[0].name, 'body2')
self.assertEmpty(spec.bodies[3].find_all('body'))
self.assertEmpty(spec.bodies[2].find_all('body'))
for i, body in enumerate(spec.worldbody.find_all('body')):
@@ -687,8 +687,8 @@ class SpecsTest(absltest.TestCase):
self.assertEqual(spec.bodies[1].sites[0].name, 'site1')
self.assertEqual(spec.bodies[1].sites[1].name, 'site2')
self.assertEqual(spec.bodies[1].sites[2].name, 'site3')
self.assertEqual(spec.bodies[3].sites[0].name, 'site4')
self.assertEqual(spec.bodies[2].sites[0].name, 'site5')
self.assertEqual(spec.bodies[2].sites[0].name, 'site4')
self.assertEqual(spec.bodies[3].sites[0].name, 'site5')
for body in spec.bodies:
for site in body.sites:
self.assertEqual(site.parent, body)
src/user/user_objects.cc
+33 -34
View File
@@ -1656,38 +1656,50 @@ mjCBase* mjCBody::FindObject(mjtObj type, std::string _name, bool recursive) {
// return true if child is descendant of this body
bool mjCBody::IsAncestor(const mjCBody* child) const {
if (!child) {
return false;
}
if (child == this) {
return true;
}
return IsAncestor(child->parent);
}
template <class T>
mjsElement* mjCBody::GetNext(const std::vector<T*>& list, const mjsElement* child, bool* found) {
mjsElement* mjCBody::GetNext(const std::vector<T*>& list, const mjsElement* child) {
if (list.empty()) {
// no children
return nullptr;
}
if (!child) {
// first child
return list[0]->spec.element;
}
for (unsigned int i = 0; i < list.size() - (child ? 1 : 0); i++) {
if (!IsAncestor(list[i]->GetParent())) {
continue; // TODO: this recursion is wasteful
}
// first child in this body
if (!child) {
return list[i]->spec.element;
for (unsigned int i = 0; i < list.size()-1; i++) {
// next child is in this body
if (list[i]->spec.element == child) {
*found = true;
} else if (list[i]->spec.element == child && IsAncestor(list[i+1]->GetParent())) {
return list[i+1]->spec.element;
}
}
if (list.back()->spec.element == child) {
// next child is in another body
*found = true;
}
return nullptr;
}
// get next child of given type
mjsElement* mjCBody::NextChild(const mjsElement* child, mjtObj type, bool recursive, bool* found) {
mjsElement* mjCBody::NextChild(const mjsElement* child, mjtObj type, bool recursive) {
if (type == mjOBJ_UNKNOWN) {
if (!child) {
throw mjCError(this, "child type must be specified if no child element is given");
@@ -1698,34 +1710,30 @@ mjsElement* mjCBody::NextChild(const mjsElement* child, mjtObj type, bool recurs
throw mjCError(this, "child element is not of requested type");
}
bool found_ = false;
if (!found) {
found = &found_;
}
mjsElement* candidate = nullptr;
switch (type) {
case mjOBJ_BODY:
case mjOBJ_XBODY:
candidate = GetNext(bodies, child, found);
candidate = GetNext(recursive ? model->bodies_ : bodies, child);
break;
case mjOBJ_JOINT:
candidate = GetNext(joints, child, found);
candidate = GetNext(recursive ? model->joints_ : joints, child);
break;
case mjOBJ_GEOM:
candidate = GetNext(geoms, child, found);
candidate = GetNext(recursive ? model->geoms_ : geoms, child);
break;
case mjOBJ_SITE:
candidate = GetNext(sites, child, found);
candidate = GetNext(recursive ? model->sites_ : sites, child);
break;
case mjOBJ_CAMERA:
candidate = GetNext(cameras, child, found);
candidate = GetNext(recursive ? model->cameras_ : cameras, child);
break;
case mjOBJ_LIGHT:
candidate = GetNext(lights, child, found);
candidate = GetNext(recursive ? model->lights_ : lights, child);
break;
case mjOBJ_FRAME:
candidate = GetNext(frames, child, found);
candidate = GetNext(recursive ? model->frames_ : frames, child);
break;
default:
throw mjCError(this,
@@ -1734,15 +1742,6 @@ mjsElement* mjCBody::NextChild(const mjsElement* child, mjtObj type, bool recurs
break;
}
if (!candidate && recursive) {
for (int i=0; i < (int)bodies.size(); i++) {
candidate = bodies[i]->NextChild(*found ? nullptr : child, type, recursive, found);
if (candidate) {
return candidate;
}
}
}
return candidate;
}
src/user/user_objects.h
+4 -2
View File
@@ -475,7 +475,7 @@ class mjCBody : public mjCBody_, private mjsBody {
// returns nullptr if the next child is not found or if `child` is the last element, returns
// the next child after the input `child` otherwise
mjsElement* NextChild(const mjsElement* child, mjtObj type = mjOBJ_UNKNOWN,
bool recursive = false, bool* found = nullptr);
bool recursive = false);
// reset keyframe references for allowing self-attach
void ForgetKeyframes() const;
@@ -534,7 +534,9 @@ class mjCBody : public mjCBody_, private mjsBody {
// gets next child of the same type in this body
template <class T>
mjsElement* GetNext(const std::vector<T*>& list, const mjsElement* child, bool* found);
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
+5
View File
@@ -98,6 +98,9 @@ TEST_F(MujocoTest, TreeTraversal) {
mjsBody* body1 = mjs_findBody(spec, "body1");
mjsBody* body2 = mjs_findBody(spec, "body2");
mjsBody* body3 = mjs_findBody(spec, "body3");
mjsElement* b1 = body1->element;
mjsElement* b2 = body2->element;
mjsElement* b3 = body3->element;
mjsElement* site1 = mjs_findElement(spec, mjOBJ_SITE, "site1");
mjsElement* site2 = mjs_findElement(spec, mjOBJ_SITE, "site2");
mjsElement* site3 = mjs_findElement(spec, mjOBJ_SITE, "site3");
@@ -144,6 +147,8 @@ TEST_F(MujocoTest, TreeTraversal) {
EXPECT_EQ(mjs_nextElement(spec, geom3), nullptr);
// text 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));
EXPECT_EQ(site3, mjs_nextChild(body1, site2, /*recursive=*/true));
EXPECT_EQ(site4, mjs_nextChild(body1, site3, /*recursive=*/true));