Allow to fusestatic a body if it doesn't generate a referencing error.

PiperOrigin-RevId: 769189929
Change-Id: I62f512813ea330da088b28d39373df39be10a602

src/user/user_model.cc

@@ -3824,24 +3824,77 @@ void mjCModel::FuseReindex(mjCBody* body) {
 
 
 
+template <class T>
+void mjCModel::ReassignChild(std::vector<T*>& dest, std::vector<T*>& list,
+                             mjCBody* parent, mjCBody* body) {
+  for (int j=0; j < list.size(); j++) {
+    // assign
+    list[j]->body = parent;
+    dest.push_back(list[j]);
+
+    // change frame
+    changeframe(list[j]->pos, list[j]->quat, body->pos, body->quat);
+  }
+  list.clear();
+}
+
+
+
+template <class T>
+void mjCModel::ResolveReferences(std::vector<T*>& list, mjCBody* body) {
+  for (auto& item : list) {
+    item->CopyFromSpec();
+    item->ResolveReferences(this);
+  }
+}
+
+
+
+template <>
+void mjCModel::ResolveReferences(std::vector<mjCSensor*>& list, mjCBody* body) {
+  for (auto& item : list) {
+    item->CopyFromSpec();
+    item->ResolveReferences(this);
+  }
+  for (mjCSensor* sensor : list) {
+    if (sensor->objtype == mjOBJ_SITE &&
+        (sensor->type == mjSENS_FORCE || sensor->type == mjSENS_TORQUE) &&
+        static_cast<mjCSite*>(sensor->obj)->body == body) {
+      throw mjCError(sensor, "cannot fuse a body used by a force/torque sensor");
+    }
+  }
+}
+
+
+
 // fuse static bodies with their parent
 void mjCModel::FuseStatic(void) {
-  // skip if model has potential to reference elements with changed ids
-  if (!skins_.empty()        ||
-      !pairs_.empty()        ||
-      !excludes_.empty()     ||
-      !equalities_.empty()   ||
-      !tendons_.empty()      ||
-      !actuators_.empty()    ||
-      !sensors_.empty()      ||
-      !tuples_.empty()       ||
-      !cameras_.empty()      ||
-      !lights_.empty()) {
-    return;
-  }
-
-  // process fusable bodies
   for (int i=1; i < bodies_.size(); i++) {
+    // check if the body can be fused
+    if (!bodies_[i]->name.empty()) {
+      ids[mjOBJ_BODY].erase(bodies_[i]->name);
+
+      // try to resolve references without the name of this body, if it fails, skip
+      try {
+        ResolveReferences(cameras_);
+        ResolveReferences(lights_);
+        ResolveReferences(skins_);
+        ResolveReferences(pairs_);
+        ResolveReferences(excludes_);
+        ResolveReferences(equalities_);
+        ResolveReferences(tendons_);
+        ResolveReferences(actuators_);
+        ResolveReferences(sensors_, bodies_[i]);
+        ResolveReferences(tuples_);
+      } catch (mjCError err) {
+        ids[mjOBJ_BODY].insert({bodies_[i]->name, i});
+        continue;
+      }
+
+      // put body back the body name in the map
+      ids[mjOBJ_BODY].insert({bodies_[i]->name, i});
+    }
+
     // get body and parent
     mjCBody* body = bodies_[i];
     mjCBody* par = body->parent;
@@ -3891,25 +3944,8 @@ void mjCModel::FuseStatic(void) {
 
     //------------- assign geoms and sites to parent, change frames
 
-    // geoms
-    for (int j=0; j < body->geoms.size(); j++) {
-      // assign
-      body->geoms[j]->body = par;
-      par->geoms.push_back(body->geoms[j]);
-
-      // change frame
-      changeframe(body->geoms[j]->pos, body->geoms[j]->quat, body->pos, body->quat);
-    }
-
-    // sites
-    for (int j=0; j < body->sites.size(); j++) {
-      // assign
-      body->sites[j]->body = par;
-      par->sites.push_back(body->sites[j]);
-
-      // change frame
-      changeframe(body->sites[j]->pos, body->sites[j]->quat, body->pos, body->quat);
-    }
+    ReassignChild(par->geoms, body->geoms, par, body);
+    ReassignChild(par->sites, body->sites, par, body);
 
     //------------- remove from global body list, reduce global counts
 
@@ -3960,15 +3996,21 @@ void mjCModel::FuseStatic(void) {
 
     //------------- delete body (without deleting children)
 
+    // remove body name from map
+    if (!body->name.empty()) {
+      ids[mjOBJ_BODY].erase(body->name);
+    }
+
     // delete allocation
     body->bodies.clear();
-    body->geoms.clear();
-    body->sites.clear();
     delete body;
 
     // check index i again (we have a new body at this index)
     i--;
   }
+
+  // remove empty names
+  processlist(ids, bodies_, mjOBJ_BODY, true);
 }
 
 
@@ -4392,6 +4434,17 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
     bodies_[i]->Compile();  // also compiles joints, geoms, sites, cameras, lights, frames
   }
 
+  // fuse static if enabled
+  if (compiler.fusestatic) {
+    FuseStatic();
+    for (int i=0; i < lights_.size(); i++) {
+      lights_[i]->Compile();
+    }
+    for (int i=0; i < cameras_.size(); i++) {
+      cameras_[i]->Compile();
+    }
+  }
+
   // compile all other objects except for keyframes
   for (auto flex : flexes_) flex->Compile(vfs);
   for (auto skin : skins_) skin->Compile(vfs);
@@ -4423,10 +4476,6 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
   // resolve asset references, compute sizes
   IndexAssets(compiler.discardvisual);
   SetSizes();
-  // fuse static if enabled
-  if (compiler.fusestatic) {
-    FuseStatic();
-  }
 
   // set nmocap and body.mocapid
   for (mjCBody* body : bodies_) {

src/user/user_model.h

@@ -444,6 +444,14 @@ class mjCModel : public mjCModel_, private mjSpec {
   // generate a signature for the model
   uint64_t Signature();
 
+  // reassign children of a body to a new parent
+  template <class T>
+  void ReassignChild(std::vector<T*>& dest, std::vector<T*>& list, mjCBody* parent, mjCBody* body);
+
+  // resolve references in a list of objects
+  template <class T>
+  void ResolveReferences(std::vector<T*>& list, mjCBody* body = nullptr);
+
   mjListKeyMap ids;   // map from object names to ids
   mjCError errInfo;   // last error info
   std::vector<mjKeyInfo> key_pending_;  // attached keyframes

src/user/user_objects.cc

@@ -3562,6 +3562,19 @@ void mjCCamera::CopyFromSpec() {
 
 
 
+void mjCCamera::ResolveReferences(const mjCModel* m) {
+  if (!targetbody_.empty()) {
+    mjCBody* tb = (mjCBody*)m->FindObject(mjOBJ_BODY, targetbody_);
+    if (tb) {
+      targetbodyid = tb->id;
+    } else {
+      throw mjCError(this, "unknown target body in camera");
+    }
+  }
+}
+
+
+
 // compiler
 void mjCCamera::Compile(void) {
   CopyFromSpec();
@@ -3587,14 +3600,7 @@ void mjCCamera::Compile(void) {
   mjuu_normvec(quat, 4);
 
   // get targetbodyid
-  if (!targetbody_.empty()) {
-    mjCBody* tb = (mjCBody*)model->FindObject(mjOBJ_BODY, targetbody_);
-    if (tb) {
-      targetbodyid = tb->id;
-    } else {
-      throw mjCError(this, "unknown target body in camera");
-    }
-  }
+  ResolveReferences(model);
 
   // make sure the image size is finite
   if (fovy >= 180) {
@@ -3716,6 +3722,27 @@ void mjCLight::CopyFromSpec() {
 
 
 
+void mjCLight::ResolveReferences(const mjCModel* m) {
+  if (!targetbody_.empty()) {
+    mjCBody* tb = (mjCBody*)m->FindObject(mjOBJ_BODY, targetbody_);
+    if (tb) {
+      targetbodyid = tb->id;
+    } else {
+      throw mjCError(this, "unknown target body in light");
+    }
+  }
+  if (!texture_.empty()) {
+    mjCTexture* tex = (mjCTexture*)m->FindObject(mjOBJ_TEXTURE, texture_);
+    if (tex) {
+      texid = tex->id;
+    } else {
+      throw mjCError(this, "unknown texture in light");
+    }
+  }
+}
+
+
+
 // compiler
 void mjCLight::Compile(void) {
   CopyFromSpec();
@@ -3735,25 +3762,8 @@ void mjCLight::Compile(void) {
     throw mjCError(this, "zero direction in light");
   }
 
-  // get targetbodyid
-  if (!targetbody_.empty()) {
-    mjCBody* tb = (mjCBody*)model->FindObject(mjOBJ_BODY, targetbody_);
-    if (tb) {
-      targetbodyid = tb->id;
-    } else {
-      throw mjCError(this, "unknown target body in light");
-    }
-  }
-
-  // get texture
-  if (!texture_.empty()) {
-    mjCTexture* tex = (mjCTexture*)model->FindObject(mjOBJ_TEXTURE, texture_);
-    if (tex) {
-      texid = tex->id;
-    } else {
-      throw mjCError(this, "unknown target body in light");
-    }
-  }
+  // get targetbodyid and texid
+  ResolveReferences(model);
 }
 
 
@@ -6493,6 +6503,8 @@ void mjCSensor::CopyPlugin() {
 
 
 void mjCSensor::ResolveReferences(const mjCModel* m) {
+  obj = nullptr;
+  ref = nullptr;
   objname_ = prefix + objname_ + suffix;
   refname_ = prefix + refname_ + suffix;
 

src/user/user_objects.h

@@ -786,6 +786,7 @@ class mjCCamera : public mjCCamera_, private mjsCamera {
   void CopyFromSpec(void);
   void PointToLocal(void);
   void NameSpace(const mjCModel* m);
+  void ResolveReferences(const mjCModel* m);
 };
 
 
@@ -831,6 +832,7 @@ class mjCLight : public mjCLight_, private mjsLight {
   void CopyFromSpec(void);
   void PointToLocal(void);
   void NameSpace(const mjCModel* m);
+  void ResolveReferences(const mjCModel* m);
 };