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

Add tree element support to mjs_firstElement() and mjs_nextElement().

Browse Source 
PiperOrigin-RevId: 671003046
Change-Id: I9e07d7c85c7218e9bded0f50af9db62d0f1e21e1
This commit is contained in:
Alessio Quaglino
2024-09-04 09:45:35 -07:00
committed by Copybara-Service
parent d3dfa6f970
commit ae4987026c
4 changed files with 79 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/user/user_model.cc
+8
View File
@@ -892,6 +892,14 @@ mjsElement* mjCModel::NextObject(mjsElement* object, mjtObj type) {
}
switch (type) {
case mjOBJ_BODY:
case mjOBJ_SITE:
case mjOBJ_GEOM:
case mjOBJ_JOINT:
case mjOBJ_CAMERA:
case mjOBJ_LIGHT:
case mjOBJ_FRAME:
return bodies_[0]->NextChild(object, type, /*recursive=*/true);
case mjOBJ_ACTUATOR:
return GetNext(actuators_, object);
case mjOBJ_SENSOR:
src/user/user_objects.cc
+44 -13
View File
@@ -1253,26 +1253,38 @@ mjCBase* mjCBody::FindObject(mjtObj type, std::string _name, bool recursive) {
template <class T>
static mjsElement* GetNext(std::vector<T*>& list, mjsElement* child) {
mjsElement* mjCBody::GetNext(std::vector<T*>& list, const mjsElement* child, bool recursive) {
if (list.empty()) {
// no children
return nullptr;
}
if (!child) {
if (list.empty()) {
return nullptr;
}
// first child
return list[0]->spec.element;
}
for (unsigned int i = 0; i < list.size()-1; i++) {
// next child is in this body
if (list[i]->spec.element == child) {
return list[i+1]->spec.element;
}
}
if (recursive && list.back()->spec.element == child) {
// next child is in next body
for (int i=0; i<(int)bodies.size(); i++) {
return bodies[i]->NextChild(NULL, child->elemtype, true);
}
}
return nullptr;
}
// get next child of given type
mjsElement* mjCBody::NextChild(mjsElement* child, mjtObj type) {
mjsElement* mjCBody::NextChild(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");
@@ -1283,25 +1295,44 @@ mjsElement* mjCBody::NextChild(mjsElement* child, mjtObj type) {
throw mjCError(this, "child element is not of requested type");
}
mjsElement* candidate = nullptr;
switch (type) {
case mjOBJ_BODY:
case mjOBJ_XBODY:
return GetNext(bodies, child);
candidate = GetNext(bodies, child, recursive);
break;
case mjOBJ_JOINT:
return GetNext(joints, child);
candidate = GetNext(joints, child, recursive);
break;
case mjOBJ_GEOM:
return GetNext(geoms, child);
candidate = GetNext(geoms, child, recursive);
break;
case mjOBJ_SITE:
return GetNext(sites, child);
candidate = GetNext(sites, child, recursive);
break;
case mjOBJ_CAMERA:
return GetNext(cameras, child);
candidate = GetNext(cameras, child, recursive);
break;
case mjOBJ_LIGHT:
return GetNext(lights, child);
candidate = GetNext(lights, child, recursive);
break;
case mjOBJ_FRAME:
return GetNext(frames, child);
candidate = GetNext(frames, child, recursive);
break;
default:
return nullptr;
break;
}
if (!candidate && recursive) {
for (int i=0; i<(int)bodies.size(); i++) {
candidate = bodies[i]->NextChild(child, type, true);
if (candidate) {
return candidate;
}
}
}
return candidate;
}
src/user/user_objects.h
+5 -1
View File
@@ -322,7 +322,7 @@ class mjCBody : public mjCBody_, private mjsBody {
const std::vector<double>& get_userdata() { return userdata_; }
// get next child of given type
mjsElement* NextChild(mjsElement* child, mjtObj type = mjOBJ_UNKNOWN);
mjsElement* NextChild(mjsElement* child, mjtObj type = mjOBJ_UNKNOWN, bool recursive = false);
// reset keyframe references for allowing self-attach
void ForgetKeyframes() const;
@@ -355,6 +355,10 @@ 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(std::vector<T*>& list, const mjsElement* child, bool recursive = false);
};
test/user/user_api_test.cc
+22 -10
View File
@@ -70,6 +70,7 @@ TEST_F(MujocoTest, TreeTraversal) {
mjSpec* spec = mj_makeSpec();
mjsBody* world = mjs_findBody(spec, "world");
mjsBody* body = mjs_addBody(world, 0);
mjsBody* body1 = mjs_addBody(body, 0);
mjsSite* site1 = mjs_addSite(body, 0);
mjsGeom* geom1 = mjs_addGeom(body, 0);
@@ -77,6 +78,7 @@ TEST_F(MujocoTest, TreeTraversal) {
mjsSite* site2 = mjs_addSite(body, 0);
mjsSite* site3 = mjs_addSite(body, 0);
mjsGeom* geom3 = mjs_addGeom(body, 0);
mjsSite* site4 = mjs_addSite(body1, 0);
mjs_setString(site1->name, "site1");
mjs_setString(geom1->name, "geom1");
@@ -84,30 +86,40 @@ TEST_F(MujocoTest, TreeTraversal) {
mjs_setString(site2->name, "site2");
mjs_setString(site3->name, "site3");
mjs_setString(geom3->name, "geom3");
mjs_setString(site4->name, "site4");
mjsElement* a_el1 = mjs_firstElement(spec, mjOBJ_ACTUATOR);
mjsElement* c_el1 = mjs_firstChild(body, mjOBJ_CAMERA);
mjsElement* t_el1 = mjs_firstChild(body, mjOBJ_TENDON);
mjsElement* a_el0 = mjs_firstElement(spec, mjOBJ_ACTUATOR);
mjsElement* l_el0 = mjs_firstElement(spec, mjOBJ_LIGHT);
mjsElement* c_el0 = mjs_firstChild(body, mjOBJ_CAMERA);
mjsElement* t_el0 = mjs_firstChild(body, mjOBJ_TENDON);
mjsElement* s_el1 = mjs_firstChild(body, mjOBJ_SITE);
mjsElement* s_el2 = mjs_nextChild(body, s_el1);
mjsElement* s_el3 = mjs_nextChild(body, s_el2);
mjsElement* s_el4 = mjs_nextChild(body, s_el3);
mjsElement* s_el0 = mjs_nextChild(body, s_el3);
mjsElement* s_el4 = mjs_firstChild(body1, mjOBJ_SITE);
mjsElement* g_el1 = mjs_firstChild(body, mjOBJ_GEOM);
mjsElement* g_el2 = mjs_nextChild(body, g_el1);
mjsElement* g_el3 = mjs_nextChild(body, g_el2);
mjsElement* g_el4 = mjs_nextChild(body, g_el3);
mjsElement* g_el0 = mjs_nextChild(body, g_el3);
EXPECT_EQ(a_el1, nullptr);
EXPECT_EQ(c_el1, nullptr);
EXPECT_EQ(t_el1, nullptr);
EXPECT_EQ(a_el0, nullptr);
EXPECT_EQ(l_el0, nullptr);
EXPECT_EQ(c_el0, nullptr);
EXPECT_EQ(t_el0, nullptr);
EXPECT_EQ(g_el0, nullptr);
EXPECT_EQ(s_el0, nullptr);
EXPECT_EQ(s_el1, site1->element);
EXPECT_EQ(s_el2, site2->element);
EXPECT_EQ(s_el3, site3->element);
EXPECT_EQ(s_el4, site4->element);
EXPECT_EQ(g_el1, geom1->element);
EXPECT_EQ(g_el2, geom2->element);
EXPECT_EQ(g_el3, geom3->element);
EXPECT_EQ(g_el4, nullptr);
EXPECT_EQ(s_el4, nullptr);
EXPECT_EQ(s_el1, mjs_firstElement(spec, mjOBJ_SITE));
EXPECT_EQ(s_el2, mjs_nextElement(spec, s_el1));
EXPECT_EQ(s_el3, mjs_nextElement(spec, s_el2));
EXPECT_EQ(s_el4, mjs_nextElement(spec, s_el3));
EXPECT_EQ(nullptr, mjs_nextElement(spec, s_el4));
EXPECT_EQ(mjs_findElement(spec, mjOBJ_SITE, "site1"), site1->element);
EXPECT_EQ(mjs_findElement(spec, mjOBJ_SITE, "site2"), site2->element);
EXPECT_EQ(mjs_findElement(spec, mjOBJ_SITE, "site3"), site3->element);