Add mjs_detachBody to C API.

PiperOrigin-RevId: 626947850
Change-Id: I7c897048098ecd4f2247725ef2d378dbd00ef4a0

src/user/user_api.cc

@@ -110,6 +110,16 @@ const char* mjs_getError(mjSpec* s) {
 
 
 
+// Detach body from mjSpec, return 0 if success.
+int mjs_detachBody(mjSpec* s, const mjsBody* b) {
+  mjCModel* model = static_cast<mjCModel*>(s->element);
+  mjCBody* body = static_cast<mjCBody*>(b->element);
+  *model -= *body;
+  return 0;
+}
+
+
+
 // check if model has warnings
 int mjs_isWarning(mjSpec* s) {
   mjCModel* modelC = static_cast<mjCModel*>(s->element);
@@ -126,6 +136,14 @@ void mjs_deleteSpec(mjSpec* s) {
 
 
 
+// delete body
+void mjs_deleteBody(mjsBody* b) {
+  mjCBody* body = static_cast<mjCBody*>(b->element);
+  delete body;
+}
+
+
+
 // add child body to body, return child spec
 mjsBody* mjs_addBody(mjsBody* bodyspec, mjsDefault* defspec) {
   mjCDef* def = defspec ? static_cast<mjCDef*>(defspec->element) : 0;

src/user/user_api.h

@@ -762,10 +762,13 @@ MJAPI void mjs_deleteSpec(mjSpec* s);
 MJAPI int mjs_attachBody(mjsFrame* parent, const mjsBody* child,
                          const char* prefix, const char* suffix);
 
-// Attach child frame to a parent body, return 0 if success.
+// Attach child frame to a parent body, return 0 on success.
 MJAPI int mjs_attachFrame(mjsBody* parent, const mjsFrame* child,
                           const char* prefix, const char* suffix);
 
+// Detach body from mjSpec, remove all references, return 0 on success.
+MJAPI int mjs_detachBody(mjSpec* s, const mjsBody* b);
+
 
 //---------------------------------- Add tree elements ---------------------------------------------
 
@@ -793,6 +796,9 @@ MJAPI mjsLight* mjs_addLight(mjsBody* body, mjsDefault* def);
 // Add frame to body.
 MJAPI mjsFrame* mjs_addFrame(mjsBody* body, mjsFrame* parentframe);
 
+// Delete body. TODO: make this a general mjs_deleteElement function
+MJAPI void mjs_deleteBody(mjsBody* b);
+
 
 //---------------------------------- Add non-tree elements -----------------------------------------
 

src/user/user_model.cc

@@ -216,6 +216,16 @@ void mjCModel::CopyList(std::vector<T*>& dest,
 
 
 
+template <class T>
+static void resetlist(std::vector<T*>& list) {
+  for (auto element : list) {
+    element->id = -1;
+  }
+  list.clear();
+}
+
+
+
 mjCModel& mjCModel::operator+=(const mjCModel& other) {
   // create global lists
   MakeLists(bodies[0]);
@@ -273,13 +283,85 @@ mjCModel& mjCModel::operator+=(const mjCModel& other) {
   // restore to the same state as other
   if (!compiled) {
     mjCBody* world = bodies[0];
-    bodies.clear();
+    resetlist(bodies);
-    frames.clear();
+    resetlist(joints);
-    joints.clear();
+    resetlist(geoms);
-    geoms.clear();
+    resetlist(sites);
-    sites.clear();
+    resetlist(cameras);
-    cameras.clear();
+    resetlist(lights);
-    lights.clear();
+    resetlist(frames);
+    world->id = 0;
+    bodies.push_back(world);
+  }
+
+  PointToLocal();
+  return *this;
+}
+
+
+
+template <class T>
+void mjCModel::RemoveFromList(std::vector<T*>& list, const mjCModel& other) {
+  int nlist = (int)list.size();
+  int removed = 0;
+  for (int i = 0; i < nlist; i++) {
+    T* element = list[i];
+    element->id -= removed;
+    try {
+      // check if the element contains an error
+      element->CopyFromSpec();
+      element->ResolveReferences(&other);
+    } catch (mjCError err) {
+      continue;
+    }
+    try {
+      // check if the element references something that was removed
+      element->ResolveReferences(this);
+    } catch (mjCError err) {
+      delete element;
+      list.erase(list.begin() + i);
+      nlist--;
+      i--;
+      removed++;
+    }
+  }
+}
+
+
+
+mjCModel& mjCModel::operator-=(const mjCBody& subtree) {
+  mjCModel oldmodel(*this);
+  oldmodel.MakeLists(oldmodel.bodies[0]);
+  oldmodel.CreateObjectLists();
+  oldmodel.ProcessLists();
+
+  // remove body from tree
+  *bodies[0] -= subtree;
+
+  // create global lists
+  MakeLists(bodies[0]);
+  CreateObjectLists();
+  ProcessLists();
+
+  // check if we have to remove anything else
+  RemoveFromList(pairs, oldmodel);
+  RemoveFromList(excludes, oldmodel);
+  RemoveFromList(tendons, oldmodel);
+  RemoveFromList(equalities, oldmodel);
+  RemoveFromList(actuators, oldmodel);
+  RemoveFromList(sensors, oldmodel);
+
+  // restore to the same state as before call
+  if (!compiled) {
+    mjCBody* world = bodies[0];
+    resetlist(bodies);
+    resetlist(joints);
+    resetlist(geoms);
+    resetlist(sites);
+    resetlist(cameras);
+    resetlist(lights);
+    resetlist(frames);
+    world->id = 0;
     bodies.push_back(world);
   }
 
@@ -722,6 +804,9 @@ static T* findobject(std::string_view name, const vector<T*>& list, const mjKeyM
   if (id == ids.end()) {
     return nullptr;
   }
+  if (id->second > (int)list.size() - 1) {
+    throw mjCError(0, "object not found");
+  }
   return list[id->second];
 }
 

src/user/user_model.h

@@ -158,12 +158,14 @@ class mjCModel : public mjCModel_, private mjSpec {
  public:
   mjCModel();
   mjCModel(const mjCModel& other);
-  mjCModel& operator=(const mjCModel& other);
-  mjCModel& operator+=(const mjCModel& other);
   ~mjCModel();
   void CopyFromSpec();  // copy spec to private attributes
   void PointToLocal();
 
+  mjCModel& operator=(const mjCModel& other);  // copy other into this, if they are not the same
+  mjCModel& operator+=(const mjCModel& other);  // add other into this, even if they are the same
+  mjCModel& operator-=(const mjCBody& subtree);  // remove subtree and all references from model
+
   mjSpec spec;
 
   mjModel* Compile(const mjVFS* vfs = nullptr);  // construct mjModel
@@ -196,6 +198,9 @@ class mjCModel : public mjCModel_, private mjSpec {
                                    std::map<mjCDef*, int>& def_map,
                                    const std::vector<mjCDef*>& defaults);
 
+  // delete from list the elements that are compatible with other but not this model
+  template <class T> void RemoveFromList(std::vector<T*>& list, const mjCModel& other);
+
   // delete elements marked as discard=true
   template <class T> void Delete(std::vector<T*>& elements,
                                  const std::vector<bool>& discard);

src/user/user_objects.cc

@@ -689,6 +689,7 @@ mjCBody& mjCBody::operator=(const mjCBody& other) {
     sites.clear();
     cameras.clear();
     lights.clear();
+    id = other.id;
 
     // add elements to lists
     *this += other;
@@ -807,6 +808,20 @@ void mjCBody::CopyList(std::vector<T*>& dst, const std::vector<T*>& src,
 
 
 
+// find and remove subtree
+mjCBody& mjCBody::operator-=(const mjCBody& subtree) {
+  for (int i=0; i<bodies.size(); i++) {
+    if (bodies[i] == &subtree) {
+      bodies.erase(bodies.begin() + i);
+      break;
+    }
+    *bodies[i] -= subtree;
+  }
+  return *this;
+}
+
+
+
 void mjCBody::PointToLocal() {
   spec.element = static_cast<mjElement*>(this);
   spec.name = (mjString)&name;

src/user/user_objects.h

@@ -263,6 +263,9 @@ class mjCBody : public mjCBody_, private mjsBody {
   friend class mjXURDF;
 
  public:
+  mjCBody(mjCModel*);  // constructor
+  ~mjCBody();          // destructor
+
   // API for adding objects to body
   mjCBody*    AddBody(mjCDef* = 0);
   mjCFrame*   AddFrame(mjCFrame* = 0);
@@ -273,9 +276,10 @@ class mjCBody : public mjCBody_, private mjsBody {
   mjCCamera*  AddCamera(mjCDef* = 0);
   mjCLight*   AddLight(mjCDef* = 0);
 
-  // API for adding existing objects to body
+  // API for adding/removing objects to body
   mjCBody& operator+=(const mjCBody& other);
   mjCBody& operator+=(const mjCFrame& other);
+  mjCBody& operator-=(const mjCBody& subtree);
 
   // API for accessing objects
   int NumObjects(mjtObj type);
@@ -304,10 +308,8 @@ class mjCBody : public mjCBody_, private mjsBody {
   const std::vector<double>& get_userdata() { return userdata_; }
 
  private:
-  mjCBody(mjCModel*);                               // constructor
   mjCBody(const mjCBody& other, mjCModel* _model);  // copy constructor
   mjCBody& operator=(const mjCBody& other);         // copy assignment
-  ~mjCBody();                                       // destructor
 
   void Compile(void);             // compiler
   void GeomFrame(void);           // get inertial info from geoms

test/user/user_api_test.cc

@@ -512,5 +512,54 @@ TEST_F(MujocoTest, AttachFrame) {
   mj_deleteModel(m_expected);
 }
 
+TEST_F(MujocoTest, DetachBody) {
+  std::array<char, 1000> er;
+  mjtNum tol = 0;
+  std::string field = "";
+
+  static constexpr char xml_result[] = R"(
+  <mujoco>
+    <worldbody>
+      <frame name="pframe">
+        <frame name="cframe">
+        </frame>
+      </frame>
+      <body name="ignore"/>
+      <frame name="frame" pos=".1 0 0" euler="0 90 0"/>
+    </worldbody>
+    <sensor>
+      <framepos name="ignore" objtype="body" objname="ignore"/>
+    </sensor>
+  </mujoco>)";
+
+  // model with one cylinder and a hinge
+  mjSpec* child = ParseSpecFromString(xml_child, er.data(), er.size());
+  EXPECT_THAT(child, NotNull()) << er.data();
+
+  // get subtree
+  mjsBody* body = mjs_findBody(child, "body");
+  EXPECT_THAT(body, NotNull());
+
+  // detach subtree
+  EXPECT_THAT(mjs_detachBody(child, body), 0);
+
+  // compile new model
+  mjModel* m_detached = mjs_compile(child, 0);
+  EXPECT_THAT(m_detached, NotNull());
+
+  // compare with expected XML
+  mjModel* m_expected = LoadModelFromString(xml_result, er.data(), er.size());
+  EXPECT_THAT(m_expected, NotNull()) << er.data();
+  EXPECT_LE(CompareModel(m_detached, m_expected, field), tol)
+            << "Expected and attached models are different!\n"
+            << "Different field: " << field << '\n';
+
+  // destroy everything
+  mjs_deleteSpec(child);
+  mjs_deleteBody(body);
+  mj_deleteModel(m_detached);
+  mj_deleteModel(m_expected);
+}
+
 }  // namespace
 }  // namespace mujoco