Change from parent ID to parent pointer in mjCBody.

This is required to have parent body info during attach and during body to frame tranformation for uncompiled models.

PiperOrigin-RevId: 703038426
Change-Id: Ib482059ba419a020860c649d6ea1b1088045f2e1

src/user/user_model.cc

@@ -153,7 +153,7 @@ mjCModel::mjCModel() {
   world->mass = 0;
   mjuu_zerovec(world->inertia, 3);
   world->id = 0;
-  world->parentid = 0;
+  world->parent = nullptr;
   world->weldid = 0;
   world->name = "world";
   world->classname = "main";
@@ -2067,10 +2067,10 @@ void mjCModel::CopyTree(mjModel* m) {
   for (int i=0; i<nbody; i++) {
     // get body and parent pointers
     mjCBody* pb = bodies_[i];
-    mjCBody* par = bodies_[pb->parentid];
+    mjCBody* par = pb->parent;
 
     // set body fields
-    m->body_parentid[i] = pb->parentid;
+    m->body_parentid[i] = pb->parent ? pb->parent->id : 0;
     m->body_weldid[i] = pb->weldid;
     m->body_mocapid[i] = pb->mocapid;
     m->body_jntnum[i] = (int)pb->joints.size();
@@ -2115,19 +2115,19 @@ void mjCModel::CopyTree(mjModel* m) {
     if (cntfree>1 || (cntfree==1 && pb->joints.size()>1)) {
       throw mjCError(pb, "free joint can only appear by itself");
     }
-    if (cntfree && pb->parentid) {
+    if (cntfree && par && par->name != "world") {
       throw mjCError(pb, "free joint can only be used on top level");
     }
 
     // rootid: self if world or child of world, otherwise parent's rootid
-    if (i==0 || pb->parentid==0) {
+    if (i==0 || (par && par->name == "world")) {
       m->body_rootid[i] = i;
     } else {
-      m->body_rootid[i] = m->body_rootid[pb->parentid];
+      m->body_rootid[i] = m->body_rootid[par->id];
     }
 
     // init lastdof from parent
-    pb->lastdof = par->lastdof;
+    pb->lastdof = par ? par->lastdof : -1;
 
     // set sameframe
     mjtSameFrame sameframe;
@@ -2447,7 +2447,9 @@ void mjCModel::CopyTree(mjModel* m) {
     bodies_[i]->subtreedofs += bodies_[i]->dofnum;
 
     // add to parent count
-    bodies_[bodies_[i]->parentid]->subtreedofs += bodies_[i]->subtreedofs;
+    if (bodies_[i]->parent) {
+      bodies_[i]->parent->subtreedofs += bodies_[i]->subtreedofs;
+    }
   }
 
   // make sure all dofs are in world "subtree", SHOULD NOT OCCUR
@@ -2462,10 +2464,10 @@ void mjCModel::CopyTree(mjModel* m) {
     nB += bodies_[i]->subtreedofs;
 
     // add dofs in ancestor bodies
-    int j = bodies_[i]->parentid;
+    int j = bodies_[i]->parent ? bodies_[i]->parent->id : 0;
     while (j > 0) {
       nB += bodies_[j]->dofnum;
-      j = bodies_[j]->parentid;
+      j = bodies_[j]->parent ? bodies_[j]->parent->id : 0;
     }
   }
   m->nB = nB;
@@ -3413,7 +3415,7 @@ static void changeframe(double childpos[3], double childquat[4],
 void mjCModel::FuseReindex(mjCBody* body) {
   // set parentid and weldid of children
   for (int i=0; i<body->bodies.size(); i++) {
-    body->bodies[i]->parentid = body->id;
+    body->bodies[i]->parent = body;
     body->bodies[i]->weldid = (!body->bodies[i]->joints.empty() ?
                                body->bodies[i]->id : body->weldid);
   }
@@ -3450,7 +3452,7 @@ void mjCModel::FuseStatic(void) {
   for (int i=1; i<bodies_.size(); i++) {
     // get body and parent
     mjCBody* body = bodies_[i];
-    mjCBody* par = bodies_[body->parentid];
+    mjCBody* par = body->parent;
 
     // skip if body has joints or mocap
     if (!body->joints.empty() || body->mocap) {
@@ -3459,7 +3461,7 @@ void mjCModel::FuseStatic(void) {
 
     //------------- add mass and inertia (if parent not world)
 
-    if (body->parentid>0 && body->mass>=mjMINVAL) {
+    if (body->parent && body->parent->name != "world" && body->mass>=mjMINVAL) {
       // body_ipose = body_pose * body_ipose
       changeframe(body->ipos, body->iquat, body->pos, body->quat);
 

src/user/user_objects.cc

@@ -765,7 +765,7 @@ mjCBody::mjCBody(mjCModel* _model) {
 
   mjs_defaultBody(&spec);
   elemtype = mjOBJ_BODY;
-  parentid = -1;
+  parent = nullptr;
   weldid = -1;
   dofnum = 0;
   lastdof = -1;
@@ -850,6 +850,7 @@ mjCBody& mjCBody::operator+=(const mjCBody& other) {
 
   for (int i=0; i<other.bodies.size(); i++) {
     bodies.push_back(new mjCBody(*other.bodies[i], model));  // triggers recursive call
+    bodies.back()->parent = this;
     bodies.back()->frame =
         other.bodies[i]->frame ? frames[fmap[other.bodies[i]->frame]] : nullptr;
   }
@@ -918,6 +919,7 @@ mjCBody& mjCBody::operator+=(const mjCFrame& other) {
       continue;
     }
     bodies.push_back(new mjCBody(*subtree->bodies[i], model));  // triggers recursive call
+    bodies.back()->parent = this;
     bodies.back()->frame =
         subtree->bodies[i]->frame ? frames[fmap[subtree->bodies[i]->frame]] : nullptr;
     bodies.back()->NameSpace_(other.model, /*propagate=*/ false);
@@ -1098,6 +1100,7 @@ mjCBody* mjCBody::AddBody(mjCDef* _def) {
   obj->classname = _def ? _def->name : classname;
 
   bodies.push_back(obj);
+  obj->parent = this;
   return obj;
 }
 
@@ -1199,11 +1202,6 @@ mjCLight* mjCBody::AddLight(mjCDef* _def) {
 
 // create a frame in the parent body and move all contents of this body into it
 mjCFrame* mjCBody::ToFrame() {
-  if (parentid < 0) {
-    // TODO: store the parent pointer instead of using the id
-    throw mjCError(this, "parent body is not defined, please compile the model first");
-  }
-  mjCBody* parent = model->Bodies()[parentid];
   mjCFrame* newframe = parent->AddFrame(frame);
   mjuu_copyvec(newframe->spec.pos, spec.pos, 3);
   mjuu_copyvec(newframe->spec.quat, spec.quat, 4);
@@ -1218,6 +1216,11 @@ mjCFrame* mjCBody::ToFrame() {
       std::remove_if(parent->bodies.begin(), parent->bodies.end(),
                      [this](mjCBody* body) { return body == this; }),
       parent->bodies.end());
+  if (model->IsCompiled()) {
+    mjCBody *world = model->bodies_[0];
+    model->ResetTreeLists();
+    model->MakeLists(world);
+  }
   return newframe;
 }
 
@@ -1573,7 +1576,6 @@ void mjCBody::Compile(void) {
 
   // set parentid and weldid of children
   for (int i=0; i<bodies.size(); i++) {
-    bodies[i]->parentid = id;
     bodies[i]->weldid = (!bodies[i]->joints.empty() ? bodies[i]->id : weldid);
   }
 
@@ -1721,13 +1723,12 @@ void mjCBody::Compile(void) {
   }
 
   // make sure mocap body is fixed child of world
-  if (mocap && (dofnum || parentid)) {
+  if (mocap && (dofnum || (parent && parent->name != "world"))) {
     throw mjCError(this, "mocap body '%s' is not a fixed child of world", name.c_str());
   }
 
   // compute body global pose (no joint transformations in qpos0)
   if (id>0) {
-    mjCBody* parent = model->Bodies()[parentid];
     mjuu_rotVecQuat(xpos0, pos, parent->xquat0);
     mjuu_addtovec(xpos0, parent->xpos0, 3);
     mjuu_mulquat(xquat0, parent->xquat0, quat);
@@ -1852,6 +1853,7 @@ mjCFrame& mjCFrame::operator+=(const mjCBody& other) {
   other.ForgetKeyframes();
   other.model->prefix = subtree->prefix;
   other.model->suffix = subtree->suffix;
+  subtree->SetParent(body);
   subtree->SetFrame(this);
   subtree->NameSpace(other.model);
 

src/user/user_objects.h

@@ -236,8 +236,9 @@ class mjCBase : public mjCBase_ {
 
 class mjCBody_ : public mjCBase {
  protected:
+  mjCBody* parent;
+
   // variables computed by 'Compile' and 'AddXXX'
-  int parentid;                   // parent index in global array
   int weldid;                     // top index of body we are welded to
   int dofnum;                     // number of motion dofs for body
   int mocapid;                    // mocap id, -1: not mocap
@@ -345,7 +346,7 @@ class mjCBody : public mjCBody_, private mjsBody {
   mjsFrame* last_attached;  // last attached frame to this body
 
   // set parent of this body
-  void SetParent(const mjCBody* _body) { parentid = _body->id; }
+  void SetParent(mjCBody* _body) { parent = _body; }
 
  private:
   mjCBody(const mjCBody& other, mjCModel* _model);  // copy constructor

test/user/user_api_test.cc

@@ -1316,7 +1316,9 @@ TEST_F(MujocoTest, AttachWorld) {
   static constexpr char xml_parent[] = R"(
   <mujoco>
     <worldbody>
-      <frame name="frame" pos="1 2 3"/>
+      <body name="body">
+        <frame name="frame" pos="1 2 3"/>
+      </body>
     </worldbody>
   </mujoco>)";
 
@@ -1334,15 +1336,21 @@ TEST_F(MujocoTest, AttachWorld) {
   static constexpr char xml_result[] = R"(
   <mujoco>
     <worldbody>
-      <frame name="frame" pos="1 2 3">
+      <body name="body">
-        <frame name="attached-world-1">
+        <frame name="frame" pos="1 2 3">
           <body name="attached-sphere-1">
             <joint type="slide"/>
             <geom size=".1"/>
           </body>
-          <camera pos="0 0 0" quat="1 0 0 0"/>
+          <frame name="attached-world-2">
+            <body name="attached-sphere-2">
+              <joint type="slide"/>
+              <geom size=".1"/>
+            </body>
+            <camera pos="0 0 0" quat="1 0 0 0"/>
+          </frame>
         </frame>
-      </frame>
+      </body>
     </worldbody>
   </mujoco>)";
 
@@ -1351,27 +1359,38 @@ TEST_F(MujocoTest, AttachWorld) {
   mjSpec* child = mj_parseXMLString(xml_child, 0, er.data(), er.size());
   EXPECT_THAT(child, NotNull()) << er.data();
 
+  // attach a body to the frame
   mjsFrame* frame = mjs_findFrame(parent, "frame");
   EXPECT_THAT(frame, NotNull());
+  mjsBody* body = mjs_findBody(child, "sphere");
+  EXPECT_THAT(body, NotNull());
+  mjsBody* attached = mjs_attachBody(frame, body, "attached-", "-1");
+  EXPECT_THAT(attached, NotNull());
+  mjModel* model1 = mj_compile(parent, 0);
+  EXPECT_THAT(model1, NotNull());
+
+  // attach the world to the same frame and convert it to a frame
   mjsBody* world = mjs_findBody(child, "world");
   EXPECT_THAT(world, NotNull());
-  mjsBody* child_world = mjs_attachBody(frame, world, "attached-", "-1");
+  mjsBody* child_world = mjs_attachBody(frame, world, "attached-", "-2");
   EXPECT_THAT(child_world, NotNull());
   mjsFrame* frame_world = mjs_bodyToFrame(&child_world);
   EXPECT_THAT(frame_world, NotNull());
   EXPECT_THAT(child_world, IsNull());
 
-  mjModel* model = mj_compile(parent, 0);
+  // compile and compare
-  EXPECT_THAT(model, NotNull());
+  mjModel* model2 = mj_compile(parent, 0);
+  EXPECT_THAT(model2, NotNull());
   mjModel* expected = LoadModelFromString(xml_result, er.data(), er.size());
   EXPECT_THAT(expected, NotNull()) << er.data();
-  EXPECT_LE(CompareModel(model, expected, field), tol)
+  EXPECT_LE(CompareModel(model2, expected, field), tol)
             << "Expected and attached models are different!\n"
             << "Different field: " << field << '\n';
 
   mj_deleteSpec(parent);
   mj_deleteSpec(child);
-  mj_deleteModel(model);
+  mj_deleteModel(model1);
+  mj_deleteModel(model2);
   mj_deleteModel(expected);
 }