Add tree traversal utilities to C API.

PiperOrigin-RevId: 639285787
Change-Id: Ieb52da87fb5e1c674985f03c9ff2b18783d1d657

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<mjCBody*>(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<mjCBody*>(body->element);
+  return bodyC->NextChild(child);
+}
+
+
+
 // set string
 void mjs_setString(mjString dest, const char* text) {
   std::string* str = reinterpret_cast<std::string*>(dest);

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.

src/user/user_objects.cc

@@ -1255,6 +1255,53 @@ mjCBase* mjCBody::FindObject(mjtObj type, string _name, bool recursive) {
 
 
 
+template <class T>
+static mjElement* GetNext(std::vector<T*>& 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;

src/user/user_objects.h

@@ -308,6 +308,9 @@ class mjCBody : public mjCBody_, private mjsBody {
   // used by mjXWriter and mjCModel
   const std::vector<double>& 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

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;