Associate children of deleted body to the newly created frame in mjs_bodyToFrame.

Also associate all children that need a parent body to the parent of the deleted body.

This is a bug that was causing the former children of the body that gets transformed to a frame to not be children of the new frame.

PiperOrigin-RevId: 702702853
Change-Id: I36d48f0446bc87c665b18d34e0b0609d82c70f51

src/user/user_objects.cc

@@ -134,6 +134,18 @@ PNGImage PNGImage::Load(const mjCBase* obj, mjResource* resource,
   return image;
 }
 
+// associate all child list elements with a frame and copy them to parent list, clear child list
+template <typename T>
+void MapFrame(std::vector<T*>& parent, std::vector<T*>& child,
+              mjCFrame* frame, mjCBody* parent_body) {
+  std::for_each(child.begin(), child.end(), [frame, parent_body](T* element) {
+    element->SetFrame(frame);
+    element->SetParent(parent_body);
+  });
+  parent.insert(parent.end(), child.begin(), child.end());
+  child.clear();
+}
+
 }  // namespace
 
 
@@ -1195,20 +1207,13 @@ mjCFrame* mjCBody::ToFrame() {
   mjCFrame* newframe = parent->AddFrame(frame);
   mjuu_copyvec(newframe->spec.pos, spec.pos, 3);
   mjuu_copyvec(newframe->spec.quat, spec.quat, 4);
-  parent->bodies.insert(parent->bodies.end(), bodies.begin(), bodies.end());
+  MapFrame(parent->bodies, bodies, newframe, parent);
-  parent->geoms.insert(parent->geoms.end(), geoms.begin(), geoms.end());
+  MapFrame(parent->geoms, geoms, newframe, parent);
-  parent->joints.insert(parent->joints.end(), joints.begin(), joints.end());
+  MapFrame(parent->joints, joints, newframe, parent);
-  parent->sites.insert(parent->sites.end(), sites.begin(), sites.end());
+  MapFrame(parent->sites, sites, newframe, parent);
-  parent->cameras.insert(parent->cameras.end(), cameras.begin(), cameras.end());
+  MapFrame(parent->cameras, cameras, newframe, parent);
-  parent->lights.insert(parent->lights.end(), lights.begin(), lights.end());
+  MapFrame(parent->lights, lights, newframe, parent);
-  parent->frames.insert(parent->frames.end(), frames.begin(), frames.end());
+  MapFrame(parent->frames, frames, newframe, parent);
-  bodies.clear();
-  geoms.clear();
-  joints.clear();
-  sites.clear();
-  cameras.clear();
-  lights.clear();
-  frames.clear();
   parent->bodies.erase(
       std::remove_if(parent->bodies.begin(), parent->bodies.end(),
                      [this](mjCBody* body) { return body == this; }),
@@ -1892,12 +1897,6 @@ bool mjCFrame::IsAncestor(const mjCFrame* child) const {
 
 
 
-void mjCFrame::SetParent(mjCBody* _body) {
-  body = _body;
-}
-
-
-
 void mjCFrame::PointToLocal() {
   spec.element = static_cast<mjsElement*>(this);
   spec.name = &name;

src/user/user_objects.h

@@ -344,6 +344,9 @@ 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; }
+
  private:
   mjCBody(const mjCBody& other, mjCModel* _model);  // copy constructor
   mjCBody& operator=(const mjCBody& other);         // copy assignment
@@ -406,7 +409,7 @@ class mjCFrame : public mjCFrame_, private mjsFrame {
 
   void CopyFromSpec(void);
   void PointToLocal(void);
-  void SetParent(mjCBody* _body);
+  void SetParent(mjCBody* _body) { body = _body; }
 
   mjCFrame& operator+=(const mjCBody& other);
 
@@ -457,6 +460,7 @@ class mjCJoint : public mjCJoint_, private mjsJoint {
   using mjCBase::info;
 
   void CopyFromSpec(void);
+  void SetParent(mjCBody* _body) { body = _body; }
 
   // used by mjXWriter and mjCModel
   const std::vector<double>& get_userdata() const { return userdata_; }
@@ -537,6 +541,7 @@ class mjCGeom : public mjCGeom_, private mjsGeom {
   void SetInertia(void);              // compute and set geom inertia
   bool IsVisual(void) const { return visual_; }
   void SetNotVisual(void) { visual_ = false; }
+  void SetParent(mjCBody* _body) { body = _body; }
   mjtGeom Type() const { return type; }
 
   // Compute all coefs modeling the interaction with the surrounding fluid.
@@ -601,6 +606,7 @@ class mjCSite : public mjCSite_, private mjsSite {
 
   // site's body
   mjCBody* Body() const { return body; }
+  void SetParent(mjCBody* _body) { body = _body; }
 
   // use strings from mjCBase rather than mjStrings from mjsSite
   using mjCBase::name;
@@ -653,6 +659,8 @@ class mjCCamera : public mjCCamera_, private mjsCamera {
   const std::string& get_targetbody() const { return targetbody_; }
   const std::vector<double>& get_userdata() const { return userdata_; }
 
+  void SetParent(mjCBody* _body) { body = _body; }
+
  private:
   void Compile(void);                     // compiler
   void CopyFromSpec(void);
@@ -691,6 +699,8 @@ class mjCLight : public mjCLight_, private mjsLight {
   // used by mjXWriter and mjCModel
   const std::string& get_targetbody() const { return targetbody_; }
 
+  void SetParent(mjCBody* _body) { body = _body; }
+
  private:
   void Compile(void);                     // compiler
   void CopyFromSpec(void);

test/user/user_api_test.cc

@@ -1316,7 +1316,7 @@ TEST_F(MujocoTest, AttachWorld) {
   static constexpr char xml_parent[] = R"(
   <mujoco>
     <worldbody>
-      <frame name="frame"/>
+      <frame name="frame" pos="1 2 3"/>
     </worldbody>
   </mujoco>)";
 
@@ -1327,18 +1327,20 @@ TEST_F(MujocoTest, AttachWorld) {
         <joint type="slide"/>
         <geom size=".1"/>
       </body>
+      <camera pos="0 0 0" quat="1 0 0 0"/>
     </worldbody>
   </mujoco>)";
 
   static constexpr char xml_result[] = R"(
   <mujoco>
     <worldbody>
-      <frame name="frame">
+      <frame name="frame" pos="1 2 3">
         <frame name="attached-world-1">
           <body name="attached-sphere-1">
             <joint type="slide"/>
             <geom size=".1"/>
           </body>
+          <camera pos="0 0 0" quat="1 0 0 0"/>
         </frame>
       </frame>
     </worldbody>