Add mjm_attachBody to the C API.

PiperOrigin-RevId: 617207000
Change-Id: I1fe59ba0d79938a4a383b115997c31dd6c71b8e5

src/user/user_api.cc

@@ -19,6 +19,7 @@
 #include <map>
 #include <new>
 #include <string>
+#include <string_view>
 #include <utility>
 #include <vector>
 
@@ -28,6 +29,25 @@
 #include "user/user_cache.h"
 #include "xml/xml_util.h"
 
+
+// prepend prefix
+template <typename T>
+static T& operator+(std::string_view prefix, T& base) {
+  base.prefix = std::string(prefix);
+  return base;
+}
+
+
+
+// append suffix
+template <typename T>
+static T& operator+(T& base, std::string_view suffix) {
+  base.suffix = std::string(suffix);
+  return base;
+}
+
+
+
 // create model
 mjSpec* mjm_createSpec() {
   mjCModel* modelC = new mjCModel;
@@ -60,6 +80,17 @@ mjModel* mjm_compile(mjSpec* s, const mjVFS* vfs) {
 
 
 
+// attach body to a frame of the parent
+int mjm_attachBody(mjmFrame* parent, const mjmBody* child,
+                   const char* prefix, const char* suffix) {
+  mjCFrame* frame_parent = reinterpret_cast<mjCFrame*>(parent->element);
+  mjCBody* child_body = reinterpret_cast<mjCBody*>(child->element);
+  *frame_parent += std::string(prefix) + *child_body + std::string(suffix);
+  return 0;
+}
+
+
+
 // get error message from model
 const char* mjm_getError(mjSpec* s) {
   mjCModel* modelC = reinterpret_cast<mjCModel*>(s->element);
@@ -169,6 +200,7 @@ mjmFrame* mjm_addFrame(mjmBody* bodyspec, mjmFrame* parentframe) {
   }
   mjCBody* body = reinterpret_cast<mjCBody*>(bodyspec->element);
   mjCFrame* frameC = body->AddFrame(parentframeC);
+  frameC->SetParent(body);
   return &frameC->spec;
 }
 
@@ -416,11 +448,13 @@ mjmDefault* mjm_getSpecDefault(mjSpec* s) {
 // find body in model by name
 mjmBody* mjm_findBody(mjSpec* s, const char* name) {
   mjCModel* model = reinterpret_cast<mjCModel*>(s->element);
-  mjCBase* body = model->FindObject(mjOBJ_BODY, std::string(name));
-  if (!body) {
-    return 0;
+  mjCBase* body = 0;
+  if (model->IsCompiled()) {
+    body = model->FindObject(mjOBJ_BODY, std::string(name));  // fast lookup
+  } else {
+    body = model->FindBody(model->GetWorld(), std::string(name));  // recursive search
   }
-  return &(static_cast<mjCBody*>(body)->spec);
+  return body ? &(static_cast<mjCBody*>(body)->spec) : nullptr;
 }
 
 
@@ -429,10 +463,7 @@ mjmBody* mjm_findBody(mjSpec* s, const char* name) {
 mjmBody* mjm_findChild(mjmBody* bodyspec, const char* name) {
   mjCBody* body = reinterpret_cast<mjCBody*>(bodyspec->element);
   mjCBase* child = body->FindObject(mjOBJ_BODY, std::string(name));
-  if (!child) {
-    return 0;
-  }
-  return &(static_cast<mjCBody*>(child)->spec);
+  return child ? &(static_cast<mjCBody*>(child)->spec) : nullptr;
 }
 
 
@@ -441,10 +472,16 @@ mjmBody* mjm_findChild(mjmBody* bodyspec, const char* name) {
 mjmMesh* mjm_findMesh(mjSpec* s, const char* name) {
   mjCModel* model = reinterpret_cast<mjCModel*>(s->element);
   mjCMesh* mesh = (mjCMesh*)model->FindObject(mjOBJ_MESH, std::string(name));
-  if (!mesh) {
-    return nullptr;
-  }
-  return &(static_cast<mjCMesh*>(mesh)->spec);
+  return mesh ? &(static_cast<mjCMesh*>(mesh)->spec) : nullptr;
+}
+
+
+
+// find frame by name
+mjmFrame* mjm_findFrame(mjSpec* s, const char* name) {
+  mjCModel* model = reinterpret_cast<mjCModel*>(s->element);
+  mjCFrame* frame = (mjCFrame*)model->FindFrame(model->GetWorld(), std::string(name));
+  return frame ? &(static_cast<mjCFrame*>(frame)->spec) : nullptr;
 }
 
 

src/user/user_api.h

@@ -740,6 +740,10 @@ MJAPI void mjm_copyBack(mjSpec* s, const mjModel* m);
 // Compile model.
 MJAPI mjModel* mjm_compile(mjSpec* s, const mjVFS* vfs);
 
+// Attach child body to a frame of the parent, return 0 if success
+MJAPI int mjm_attachBody(mjmFrame* parent, const mjmBody* child,
+                         const char* prefix, const char* suffix);
+
 // Get error message from model.
 MJAPI const char* mjm_getError(mjSpec* s);
 
@@ -860,6 +864,9 @@ MJAPI mjmBody* mjm_findChild(mjmBody* body, const char* name);
 // Find mesh by name.
 MJAPI mjmMesh* mjm_findMesh(mjSpec* s, const char* name);
 
+// Find frame by name.
+MJAPI mjmFrame* mjm_findFrame(mjSpec* s, const char* name);
+
 // Get element id.
 MJAPI int mjm_getId(mjElement element);
 

src/user/user_model.cc

@@ -162,6 +162,26 @@ mjCModel& mjCModel::operator=(const mjCModel& other) {
     this->spec = other.spec;
     *static_cast<mjCModel_*>(this) = static_cast<const mjCModel_&>(other);
     *static_cast<mjSpec*>(this) = static_cast<const mjSpec&>(other);
+
+    // the world copy constructor takes care of copying the tree
+    mjCBody* world = new mjCBody(*other.bodies[0], this);
+    bodies.push_back(world);
+
+    // add everything else
+    *this += other;
+
+    // copy name maps
+    for (int i=0; i<mjNOBJECT; i++) {
+      ids[i] = other.ids[i];
+    }
+  }
+  return *this;
+}
+
+
+
+mjCModel& mjCModel::operator+=(const mjCModel& other) {
+  if (this != &other) {
     std::map<mjCDef*, int> def_map;
     for (int i = 0; i < other.defaults.size(); i++) {
       defaults.push_back(new mjCDef(*other.defaults[i]));
@@ -252,11 +272,7 @@ mjCModel& mjCModel::operator=(const mjCModel& other) {
     plugins = other.plugins;
     active_plugins = other.active_plugins;
 
-    // the world copy constructor takes care of copying the tree
-    mjCBody* world = new mjCBody(*other.bodies[0], this);
-
     // create global lists
-    bodies.push_back(world);
     MakeLists(bodies[0]);
 
     // update defaults for the copied objects
@@ -279,16 +295,12 @@ mjCModel& mjCModel::operator=(const mjCModel& other) {
       lights[i]->def = defaults[def_map[other.lights[i]->def]];
     }
 
-    // copy name maps
-    for (int i=0; i<mjNOBJECT; i++) {
-      ids[i] = other.ids[i];
-    }
-
     // cast children to mjCBase
     CreateObjectLists();
 
     // restore to the same state as other
     if (!compiled) {
+      mjCBody* world = bodies[0];
       bodies.clear();
       frames.clear();
       joints.clear();
@@ -751,6 +763,44 @@ mjCBase* mjCModel::FindObject(mjtObj type, string name) {
 
 
 
+// find body by name
+mjCBody* mjCModel::FindBody(mjCBody* body, std::string name) {
+  if (body->name == name) {
+    return body;
+  }
+
+  for (auto child : body->bodies) {
+    auto candidate = FindBody(child, name);
+    if (candidate) {
+      return candidate;
+    }
+  }
+
+  return nullptr;
+}
+
+
+
+// find frame by name
+mjCFrame* mjCModel::FindFrame(mjCBody* body, std::string name) {
+  for (auto frame : body->frames) {
+    if (frame->name == name) {
+      return frame;
+    }
+  }
+
+  for (auto body : body->bodies) {
+    auto candidate = FindFrame(body, name);
+    if (candidate) {
+      return candidate;
+    }
+  }
+
+  return nullptr;
+}
+
+
+
 // detect null pose
 bool mjCModel::IsNullPose(const mjtNum* pos, const mjtNum* quat) {
   bool result = true;

src/user/user_model.h

@@ -154,6 +154,7 @@ class mjCModel : public mjCModel_, private mjSpec {
   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();
@@ -202,6 +203,8 @@ class mjCModel : public mjCModel_, private mjSpec {
   mjCDef* FindDef(std::string name);                       // find default class name
   mjCDef* AddDef(std::string name, int parentid);          // add default class to array
   mjCBase* FindObject(mjtObj type, std::string name);      // find object given type and name
+  mjCBody* FindBody(mjCBody* body, std::string name);      // find body given name
+  mjCFrame* FindFrame(mjCBody* body, std::string name);    // find frame given name
   bool IsNullPose(const mjtNum* pos, const mjtNum* quat);  // detect null pose
 
   // accessors

src/user/user_objects.cc

@@ -683,6 +683,7 @@ mjCBody& mjCBody::operator=(const mjCBody& other) {
     for (int i=0; i<other.frames.size(); i++) {
       frames.push_back(new mjCFrame(*other.frames[i]));
       frames.back()->model = model;
+      frames.back()->body = this;
       fmap[other.frames[i]] = i;
     }
 
@@ -731,6 +732,13 @@ mjCBody& mjCBody::operator=(const mjCBody& other) {
 
 
 
+mjCBody& mjCBody::operator+=(mjCBody& other) {
+  bodies.push_back(&other);
+  return *this;
+}
+
+
+
 void mjCBody::PointToLocal() {
   spec.element = (mjElement)this;
   spec.name = (mjString)&name;
@@ -784,6 +792,21 @@ mjCBody::~mjCBody() {
 
 
 
+// apply prefix and suffix, propagate to children
+void mjCBody::SetNameSpace() {
+  if (!name.empty()) {
+    name = prefix + name + suffix;
+  }
+
+  for (auto& body : bodies) {
+    body->prefix = prefix;
+    body->suffix = suffix;
+    body->SetNameSpace();
+  }
+}
+
+
+
 // create child body and add it to body
 mjCBody* mjCBody::AddBody(mjCDef* _def) {
   // create body
@@ -1317,6 +1340,7 @@ mjCFrame::mjCFrame(mjCModel* _model, mjCFrame* _frame) {
   mjm_defaultFrame(spec);
   compiled = false;
   model = _model;
+  body = NULL;
   frame = _frame ? _frame : NULL;
   PointToLocal();
   CopyFromSpec();
@@ -1341,6 +1365,27 @@ mjCFrame& mjCFrame::operator=(const mjCFrame& other) {
 }
 
 
+// attach body to frame
+mjCFrame& mjCFrame::operator+=(const mjCBody& other) {
+  mjCBody* subtree = new mjCBody(other, model);
+  subtree->SetFrame(this);
+  subtree->SetNameSpace();
+
+  // add to tree
+  *body += *subtree;
+
+  // TODO: needs to attach only referencing elements
+  *model += *other.model;
+  return *this;
+}
+
+
+
+void mjCFrame::SetParent(mjCBody* _body) {
+  body = _body;
+}
+
+
 
 void mjCFrame::PointToLocal() {
   spec.element = (mjElement)this;

src/user/user_objects.h

@@ -173,10 +173,12 @@ class mjCBoundingVolumeHierarchy : public mjCBoundingVolumeHierarchy_ {
 
 class mjCBase_ {
  public:
-  std::string name;               // object name
-  std::string classname;          // defaults class name
-  int id;                         // object id
-  std::string info;               // error message info set by the user
+  int id;                 // object id
+  std::string name;       // object name
+  std::string classname;  // defaults class name
+  std::string info;       // error message info set by the user
+  std::string prefix;     // prefix for model operations
+  std::string suffix;     // suffix for model operations
 };
 
 class mjCBase : public mjCBase_ {
@@ -251,6 +253,7 @@ class mjCBody : public mjCBody_, private mjmBody {
   friend class mjCSite;
   friend class mjCCamera;
   friend class mjCComposite;
+  friend class mjCFrame;
   friend class mjCLight;
   friend class mjCFlex;
   friend class mjCFlexcomp;
@@ -272,11 +275,17 @@ class mjCBody : public mjCBody_, private mjmBody {
   mjCCamera*  AddCamera(mjCDef* = 0);
   mjCLight*   AddLight(mjCDef* = 0);
 
+  // API for adding existing objects to body
+  mjCBody& operator+=(mjCBody& other);
+
   // API for accessing objects
   int NumObjects(mjtObj type);
   mjCBase* GetObject(mjtObj type, int id);
   mjCBase* FindObject(mjtObj type, std::string name, bool recursive = true);
 
+  // Propagate suffix and prefix to the whole tree
+  void SetNameSpace();
+
   // set explicitinertial to true
   void MakeInertialExplicit();
 
@@ -346,12 +355,18 @@ class mjCFrame : public mjCFrame_, private mjmFrame {
 
   void CopyFromSpec(void);
   void PointToLocal(void);
+  void SetParent(mjCBody* _body);
+
+  mjCFrame& operator+=(const mjCBody& other);
+  mjCFrame& operator+=(const mjCFrame& other);
 
  private:
   mjCFrame(mjCModel* = 0, mjCFrame* = 0);      // constructor
   mjCFrame(const mjCFrame& other);             // copy constructor
   mjCFrame& operator=(const mjCFrame& other);  // copy assignment
   void Compile(void);                          // compiler
+
+  mjCBody* body;  // body that owns the frame
 };
 
 

src/xml/xml.cc

@@ -15,7 +15,7 @@
 #include "xml/xml.h"
 
 #include <locale.h>
-#include "user/user_api.h"
+#include <cstring>
 
 #if defined(__APPLE__) || defined(__FreeBSD__)
 #include <xlocale.h>
@@ -29,6 +29,7 @@
 
 #include "tinyxml2.h"
 
+#include <mujoco/mujoco.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjplugin.h>
 #include "cc/array_safety.h"
@@ -399,3 +400,33 @@ mjSpec* mjParseXML(const char* filename, const mjVFS* vfs,
 
   return model;
 }
+
+
+static void RegisterResourceProvider() {
+  // register string resource provider if not registered before
+  if (mjp_getResourceProvider("LoadModelFromString:") == nullptr) {
+    mjpResourceProvider resourceProvider;
+    mjp_defaultResourceProvider(&resourceProvider);
+    resourceProvider.prefix = "LoadModelFromString";
+    resourceProvider.open = +[](mjResource* resource) {
+      resource->data = &(resource->name[strlen("LoadModelFromString:")]);
+      return 1;
+    };
+    resourceProvider.read =
+        +[](mjResource* resource, const void** buffer) {
+          *buffer = resource->data;
+          return (int) strlen((const char*) resource->data);
+        };
+    resourceProvider.close = +[](mjResource* resource) {};
+    mjp_registerResourceProvider(&resourceProvider);
+  }
+}
+
+
+mjSpec* ParseSpecFromString(std::string_view xml, char* error,
+                            int error_size, mjVFS* vfs) {
+  RegisterResourceProvider();
+  std::string xml2 = {xml.begin(), xml.end()};
+  std::string str = "LoadModelFromString:" +  xml2;
+  return mjParseXML(str.c_str(), vfs, error, error_size);
+}

src/xml/xml.h

@@ -29,5 +29,10 @@ std::string mjWriteXML(mjSpec* model, char* error, int error_sz);
 // Main parser function
 MJAPI mjSpec* mjParseXML(const char* filename, const mjVFS* vfs, char* error, int error_sz);
 
+// Returns a newly-allocated mjSpec, loaded from the contents of xml.
+// On failure returns nullptr and populates the error array if present.
+MJAPI mjSpec* ParseSpecFromString(std::string_view xml, char* error = nullptr,
+                                  int error_size = 0, mjVFS* vfs = nullptr);
+
 
 #endif  // MUJOCO_SRC_XML_XML_H_

test/user/user_api_test.cc

@@ -195,5 +195,96 @@ TEST_F(PluginTest, RecompileCompareCache) {
   mj_deleteVFS(vfs.get());
 }
 
+// -------------------------------- test attach -------------------------------
+TEST_F(MujocoTest, Attach) {
+  std::array<char, 1000> er;
+  mjtNum tol = 0;
+  std::string field = "";
+
+  static constexpr char xml_parent[] = R"(
+  <mujoco>
+    <worldbody>
+      <body name="sphere">
+        <freejoint/>
+        <geom size=".1"/>
+        <frame name="frame" pos=".1 0 0" euler="0 90 0"/>
+      </body>
+    </worldbody>
+  </mujoco>)";
+
+  static constexpr char xml_child[] = R"(
+  <mujoco>
+    <worldbody>
+      <body name="cylinder">
+        <joint type="hinge"/>
+        <geom type="cylinder" size=".1 1 0"/>
+          <body name="named"/>
+          <body/>
+      </body>
+    </worldbody>
+  </mujoco>)";
+
+  static constexpr char xml_result[] = R"(
+  <mujoco>
+    <worldbody>
+      <body name="sphere">
+        <freejoint/>
+        <geom size=".1"/>
+        <frame name="frame" pos=".1 0 0" euler="0 90 0">
+          <body name="attached-cylinder-1">
+            <joint type="hinge"/>
+            <geom type="cylinder" size=".1 1 0"/>
+              <body name="attached-named-1"/>
+              <body/>
+          </body>
+        </frame>
+      </body>
+    </worldbody>
+  </mujoco>)";
+
+  // model with one free sphere and a frame
+  mjSpec* parent = ParseSpecFromString(xml_parent, er.data(), er.size());
+  EXPECT_THAT(parent, NotNull()) << er.data();
+
+  // get frame
+  mjmFrame* frame = mjm_findFrame(parent, "frame");
+  EXPECT_THAT(frame, NotNull());
+
+  // model with one cylinder and a hinge
+  mjSpec* child = ParseSpecFromString(xml_child, er.data(), er.size());
+  EXPECT_THAT(child, NotNull()) << er.data();
+
+  // get subtree
+  mjmBody* body = mjm_findBody(child, "cylinder");
+  EXPECT_THAT(body, NotNull());
+
+  // attach child to parent frame
+  EXPECT_THAT(
+      mjm_attachBody(frame, body, /*prefix=*/"attached-", /*suffix=*/"-1"), 0);
+
+  // compile new model
+  mjModel* m_attached = mjm_compile(parent, 0);
+  EXPECT_THAT(m_attached, NotNull());
+
+  // check full name stored in mjModel
+  EXPECT_STREQ(mj_id2name(m_attached, mjOBJ_BODY, 2), "attached-cylinder-1");
+
+  // check body 2 is attached to body 1
+  EXPECT_THAT(m_attached->body_parentid[2], 1);
+
+  // 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_attached, m_expected, field), tol)
+            << "Expected and attached models are different!\n"
+            << "Different field: " << field << '\n';;
+
+  // destroy everything
+  mjm_deleteSpec(parent);
+  mjm_deleteSpec(child);
+  mj_deleteModel(m_attached);
+  mj_deleteModel(m_expected);
+}
+
 }  // namespace
 }  // namespace mujoco