Attach the tree of defaults recursively and with namespacing.

PiperOrigin-RevId: 650664018
Change-Id: I17751a02a406dcb2a2ca218c316b5eb0370071b8

src/user/user_composite.cc

@@ -491,7 +491,7 @@ bool mjCComposite::MakeParticle(mjCModel* model, mjsBody* body, char* error, int
 
     // add user-specified joints
     else {
-      for (auto defjnt : defjoint[mjCOMPKIND_PARTICLE]) {
+      for (auto& defjnt : defjoint[mjCOMPKIND_PARTICLE]) {
         mjsJoint* jnt = mjs_addJoint(b, &defjnt.spec);
         mjs_setDefault(jnt->element, mjs_getDefault(body->element));
       }

src/user/user_model.cc

@@ -147,7 +147,9 @@ mjCModel& mjCModel::operator=(const mjCModel& other) {
     // add everything else
     *this += other;
 
-    // update the default map
+    // create new default tree
+    mjCDef* subtree = new mjCDef(*other.defaults_[0]);
+    *this += *subtree;
     def_map["main"] = Default();
 
     // copy name maps
@@ -163,9 +165,7 @@ mjCModel& mjCModel::operator=(const mjCModel& other) {
 // copy vector of elements from another model to this model
 template <class T>
 void mjCModel::CopyList(std::vector<T*>& dest,
-                        const std::vector<T*>& source,
-                        std::map<mjCDef*, int>& def_map,
-                        const std::vector<mjCDef*>& defaults) {
+                        const std::vector<T*>& source) {
   // loop over the elements from the other model
   int nsource = (int)source.size();
   for (int i = 0; i < nsource; i++) {
@@ -220,34 +220,26 @@ mjCModel& mjCModel::operator+=(const mjCModel& other) {
   ProcessLists(/*checkrepeat=*/false);
 
   // copy all elements not in the tree
-  std::map<mjCDef*, int> def_map;
-  int ndefaults = (int)other.defaults_.size();
-  for (int i = 0; i < ndefaults; i++) {
-    if (this != &other) {
-      defaults_.push_back(new mjCDef(*other.defaults_[i]));
-    }
-    def_map[other.defaults_[i]] = i;
-  }
-
   if (this != &other) {
     // do not copy assets for self-attach
-    CopyList(meshes_, other.meshes_, def_map, defaults_);
-    CopyList(skins_, other.skins_, def_map, defaults_);
-    CopyList(hfields_, other.hfields_, def_map, defaults_);
-    CopyList(textures_, other.textures_, def_map, defaults_);
-    CopyList(materials_, other.materials_, def_map, defaults_);
-    CopyList(keys_, other.keys_, def_map, defaults_);
+    // TODO: asset should be copied only when referenced
+    CopyList(meshes_, other.meshes_);
+    CopyList(skins_, other.skins_);
+    CopyList(hfields_, other.hfields_);
+    CopyList(textures_, other.textures_);
+    CopyList(materials_, other.materials_);
+    CopyList(keys_, other.keys_);
   }
-  CopyList(flexes_, other.flexes_, def_map, defaults_);
-  CopyList(pairs_, other.pairs_, def_map, defaults_);
-  CopyList(excludes_, other.excludes_, def_map, defaults_);
-  CopyList(tendons_, other.tendons_, def_map, defaults_);
-  CopyList(equalities_, other.equalities_, def_map, defaults_);
-  CopyList(actuators_, other.actuators_, def_map, defaults_);
-  CopyList(sensors_, other.sensors_, def_map, defaults_);
-  CopyList(numerics_, other.numerics_, def_map, defaults_);
-  CopyList(texts_, other.texts_, def_map, defaults_);
-  CopyList(tuples_, other.tuples_, def_map, defaults_);
+  CopyList(flexes_, other.flexes_);
+  CopyList(pairs_, other.pairs_);
+  CopyList(excludes_, other.excludes_);
+  CopyList(tendons_, other.tendons_);
+  CopyList(equalities_, other.equalities_);
+  CopyList(actuators_, other.actuators_);
+  CopyList(sensors_, other.sensors_);
+  CopyList(numerics_, other.numerics_);
+  CopyList(texts_, other.texts_);
+  CopyList(tuples_, other.tuples_);
 
   // plugins are global
   plugins_ = other.plugins_;
@@ -358,6 +350,24 @@ mjCModel& mjCModel::operator-=(const mjCBody& subtree) {
 
 
 
+// add default tree to this model
+mjCModel_& mjCModel::operator+=(mjCDef& subtree) {
+  defaults_.push_back(&subtree);
+
+  // set parent to the main default if this is not the only default in the model
+  if (!subtree.parent && &subtree != defaults_[0]) {
+    subtree.parent = defaults_[0];
+    defaults_[0]->child.push_back(&subtree);
+  }
+
+  for (auto def : subtree.child) {
+    *this += *def;  // triggers recursive call
+  }
+  return *this;
+}
+
+
+
 // TODO: we should not use C-type casting with multiple C++ inheritance
 void mjCModel::CreateObjectLists() {
   for (int i = 0; i < mjNOBJECT; ++i) {
@@ -806,7 +816,7 @@ mjCDef* mjCModel::AddDefault(string name, mjCDef* parent) {
   // initialize contents
   if (parent && parent->id<thisid) {
     parent->CopyFromSpec();
-    *def = *parent;
+    def->CopyWithoutChildren(*parent);
     parent->child.push_back(def);
   }
   def->parent = parent;

src/user/user_model.h

@@ -165,6 +165,7 @@ class mjCModel : public mjCModel_, private mjSpec {
   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
+  mjCModel_& operator+=(mjCDef& subtree);  // add default tree to this model
 
   mjSpec spec;
 
@@ -333,9 +334,7 @@ class mjCModel : public mjCModel_, private mjSpec {
 
   // copy vector of elements to this model
   template <class T> void CopyList(std::vector<T*>& dest,
-                                   const std::vector<T*>& sources,
-                                   std::map<mjCDef*, int>& def_map,
-                                   const std::vector<mjCDef*>& defaults);
+                                   const std::vector<T*>& sources);
 
   // 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);

src/user/user_objects.cc

@@ -564,28 +564,58 @@ void mjCDef::Compile(const mjCModel* model) {
 // assignment operator
 mjCDef& mjCDef::operator=(const mjCDef& other) {
   if (this != &other) {
-    name = other.name;
-    parent = other.parent;
-    child = other.child;
-    joint_ = other.joint_;
-    geom_ = other.geom_;
-    site_ = other.site_;
-    camera_ = other.camera_;
-    light_ = other.light_;
-    flex_ = other.flex_;
-    mesh_ = other.mesh_;
-    material_ = other.material_;
-    pair_ = other.pair_;
-    equality_ = other.equality_;
-    tendon_ = other.tendon_;
-    actuator_ = other.actuator_;
+    CopyWithoutChildren(other);
+
+    // copy the rest of the default tree
+    *this += other;
   }
-  PointToLocal();
   return *this;
 }
 
 
 
+mjCDef& mjCDef::operator+=(const mjCDef& other) {
+  for (unsigned int i=0; i<other.child.size(); i++) {
+    child.push_back(new mjCDef(*other.child[i]));  // triggers recursive call
+    child.back()->parent = this;
+  }
+  return *this;
+}
+
+
+
+void mjCDef::NameSpace(const mjCModel* m) {
+  if (!name.empty()) {
+    name = m->prefix + name + m->suffix;
+  }
+  for (auto c : child) {
+    c->NameSpace(m);
+  }
+}
+
+
+
+void mjCDef::CopyWithoutChildren(const mjCDef& other) {
+  name = other.name;
+  parent = nullptr;
+  child.clear();
+  joint_ = other.joint_;
+  geom_ = other.geom_;
+  site_ = other.site_;
+  camera_ = other.camera_;
+  light_ = other.light_;
+  flex_ = other.flex_;
+  mesh_ = other.mesh_;
+  material_ = other.material_;
+  pair_ = other.pair_;
+  equality_ = other.equality_;
+  tendon_ = other.tendon_;
+  actuator_ = other.actuator_;
+  PointToLocal();
+}
+
+
+
 void mjCDef::PointToLocal() {
   joint_.PointToLocal();
   geom_.PointToLocal();
@@ -817,6 +847,13 @@ mjCBody& mjCBody::operator+=(const mjCFrame& other) {
   other.model->prefix = other.prefix;
   other.model->suffix = other.suffix;
 
+  // attach defaults
+  if (other.model != model) {
+    mjCDef* subdef = new mjCDef(*other.model->Default());
+    subdef->NameSpace(other.model);
+    *model += *subdef;
+  }
+
   // copy input frame
   frames.push_back(new mjCFrame(other));
   frames.back()->body = this;
@@ -1591,6 +1628,13 @@ mjCFrame& mjCFrame::operator+=(const mjCBody& other) {
   subtree->SetFrame(this);
   subtree->NameSpace(other.model);
 
+  // attach defaults
+  if (other.model != model) {
+    mjCDef* subdef = new mjCDef(*other.model->Default());
+    subdef->NameSpace(other.model);
+    *model += *subdef;
+  }
+
   // add to body children
   body->bodies.push_back(subtree);
 

src/user/user_objects.h

@@ -1611,9 +1611,12 @@ class mjCDef : public mjsElement {
   mjCDef();
   mjCDef(const mjCDef& other);
   mjCDef& operator=(const mjCDef& other);
+  mjCDef& operator+=(const mjCDef& other);
 
+  void CopyWithoutChildren(const mjCDef& other);
   void PointToLocal(void);
   void CopyFromSpec(void);
+  void NameSpace(const mjCModel* m);
 
   void Compile(const mjCModel* model);
 

test/user/user_api_test.cc

@@ -559,6 +559,11 @@ TEST_F(MujocoTest, AttachDifferent) {
   // check body 2 is attached to body 1
   EXPECT_THAT(m_attached->body_parentid[2], 1);
 
+  // check that the correct defaults are present
+  EXPECT_THAT(mjs_findDefault(parent, "main"), NotNull());
+  EXPECT_THAT(mjs_findDefault(parent, "geom_size"), NotNull());
+  EXPECT_THAT(mjs_findDefault(parent, "attached-cylinder-1"), NotNull());
+
   // compare with expected XML
   mjModel* m_expected = LoadModelFromString(xml_result, er.data(), er.size());
   EXPECT_THAT(m_expected, NotNull()) << er.data();