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