diff --git a/src/user/user_api.cc b/src/user/user_api.cc index 0bf072f1..34cfa704 100644 --- a/src/user/user_api.cc +++ b/src/user/user_api.cc @@ -554,6 +554,22 @@ void mjs_setDefault(mjElement* element, mjsDefault* defspec) { +// return first child of selected type +mjElement* mjs_firstChild(mjsBody* body, mjtObj type) { + mjCBody* bodyC = static_cast(body->element); + return bodyC->NextChild(NULL, type); +} + + + +// return body's next child; return NULL if child is last +mjElement* mjs_nextChild(mjsBody* body, mjElement* child) { + mjCBody* bodyC = static_cast(body->element); + return bodyC->NextChild(child); +} + + + // set string void mjs_setString(mjString dest, const char* text) { std::string* str = reinterpret_cast(dest); diff --git a/src/user/user_api.h b/src/user/user_api.h index 400325c4..3f6ae15e 100644 --- a/src/user/user_api.h +++ b/src/user/user_api.h @@ -916,6 +916,15 @@ MJAPI mjsDefault* mjs_getSpecDefault(mjSpec* s); MJAPI int mjs_getId(mjElement* element); +//---------------------------------- Tree traversal ------------------------------------------------ + +// Return body's first child of given type. +MJAPI mjElement* mjs_firstChild(mjsBody* body, mjtObj type); + +// Return body's next child of the same type; return NULL if child is last. +MJAPI mjElement* mjs_nextChild(mjsBody* body, mjElement* child); + + //---------------------------------- Attribute setters --------------------------------------------- // Copy text to string. diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index 6691143f..034abae6 100644 --- a/src/user/user_objects.cc +++ b/src/user/user_objects.cc @@ -1255,6 +1255,53 @@ mjCBase* mjCBody::FindObject(mjtObj type, string _name, bool recursive) { +template +static mjElement* GetNext(std::vector& list, mjElement* child) { + for (unsigned int i = 0; i < list.size()-1; i++) { + if (list[i]->spec.element == child) { + return list[i+1]->spec.element; + } + } + return nullptr; +} + + + +// get next child of given type +mjElement* mjCBody::NextChild(mjElement* child, mjtObj type) { + if (type == mjOBJ_UNKNOWN) { + if (!child) { + throw mjCError(this, "child type must be specified if no child element is given"); + } else { + type = child->elemtype; + } + } else if (child && child->elemtype != type) { + throw mjCError(this, "child element is not of requested type"); + } + + switch (type) { + case mjOBJ_BODY: + case mjOBJ_XBODY: + return child ? GetNext(bodies, child) : bodies[0]; + case mjOBJ_JOINT: + return child ? GetNext(joints, child) : joints[0]; + case mjOBJ_GEOM: + return child ? GetNext(geoms, child) : geoms[0]; + case mjOBJ_SITE: + return child ? GetNext(sites, child) : sites[0]; + case mjOBJ_CAMERA: + return child ? GetNext(cameras, child) : cameras[0]; + case mjOBJ_LIGHT: + return child ? GetNext(lights, child) : lights[0]; + case mjOBJ_FRAME: + return child ? GetNext(frames, child) : frames[0]; + default: + return nullptr; + } +} + + + // compute geom inertial frame: ipos, iquat, mass, inertia void mjCBody::GeomFrame(void) { int sz; diff --git a/src/user/user_objects.h b/src/user/user_objects.h index 0ed56b4e..d5773515 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -308,6 +308,9 @@ class mjCBody : public mjCBody_, private mjsBody { // used by mjXWriter and mjCModel const std::vector& get_userdata() { return userdata_; } + // get next child of given type + mjElement* NextChild(mjElement* child, mjtObj type = mjOBJ_UNKNOWN); + private: mjCBody(const mjCBody& other, mjCModel* _model); // copy constructor mjCBody& operator=(const mjCBody& other); // copy assignment diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc index 6db3b558..440f0c96 100644 --- a/test/user/user_api_test.cc +++ b/test/user/user_api_test.cc @@ -36,9 +36,9 @@ using ::testing::HasSubstr; using ::testing::NotNull; -// ----------------------------- test set/get --------------------------------- +// -------------------------- test model manipulation ------------------------- -TEST_F(MujocoTest, ReadWriteData) { +TEST_F(MujocoTest, GetSetData) { mjSpec* spec = mjs_createSpec(); mjsBody* world = mjs_findBody(spec, "world"); mjsBody* body = mjs_addBody(world, 0); @@ -63,6 +63,41 @@ TEST_F(MujocoTest, ReadWriteData) { mjs_deleteSpec(spec); } +TEST_F(MujocoTest, TreeTraversal) { + mjSpec* spec = mjs_createSpec(); + mjsBody* world = mjs_findBody(spec, "world"); + mjsBody* body = mjs_addBody(world, 0); + + mjsSite* site1 = mjs_addSite(body, 0); + mjsGeom* geom1 = mjs_addGeom(body, 0); + mjsGeom* geom2 = mjs_addGeom(body, 0); + mjsSite* site2 = mjs_addSite(body, 0); + mjsSite* site3 = mjs_addSite(body, 0); + mjsGeom* geom3 = mjs_addGeom(body, 0); + + mjElement* t_el1 = mjs_firstChild(body, mjOBJ_TENDON); + mjElement* s_el1 = mjs_firstChild(body, mjOBJ_SITE); + mjElement* s_el2 = mjs_nextChild(body, s_el1); + mjElement* s_el3 = mjs_nextChild(body, s_el2); + mjElement* s_el4 = mjs_nextChild(body, s_el3); + mjElement* g_el1 = mjs_firstChild(body, mjOBJ_GEOM); + mjElement* g_el2 = mjs_nextChild(body, g_el1); + mjElement* g_el3 = mjs_nextChild(body, g_el2); + mjElement* g_el4 = mjs_nextChild(body, g_el3); + + EXPECT_EQ(t_el1, nullptr); + EXPECT_EQ(s_el1, site1->element); + EXPECT_EQ(s_el2, site2->element); + EXPECT_EQ(s_el3, site3->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); + + mjs_deleteSpec(spec); +} + // ------------------- test recompilation multiple files ----------------------- TEST_F(PluginTest, RecompileCompare) { mjtNum tol = 0;