Copy only referenced plugins.

Fixes #2111.

PiperOrigin-RevId: 681486400
Change-Id: I384ac63fa76b27d6e84d91fa857127c42ff52e34

src/user/user_model.cc

@@ -297,24 +297,37 @@ void mjCModel::SaveDofOffsets() {
   }
 }
 
-
-
 template <class T>
-static void mapplugin(
-    const std::unordered_map<mjCPlugin*, mjCPlugin*>& plugin_map, std::vector<T*>& list) {
+void mjCModel::CopyPlugin(std::vector<mjCPlugin*>& dest,
+                          const std::vector<mjCPlugin*>& source,
+                          const std::vector<T*>& list) {
+  // store elements that reference a plugin instance
+  std::unordered_map<std::string, T*> instances;
   for (const auto& element : list) {
-      if (element->spec.plugin.element) {
-        mjCPlugin* plugin = static_cast<mjCPlugin*>(element->spec.plugin.element);
-        // the referenced plugin might already exist in the source model
-        if (plugin_map.find(plugin) != plugin_map.end()) {
-          element->spec.plugin.element = plugin_map.at(plugin);
-        }
-      }
+    if (!element->plugin_instance_name.empty()) {
+      instances[element->plugin_instance_name] = element;
     }
+  }
+
+  // only copy plugins that are referenced
+  for (const auto& plugin : source) {
+    if (plugin->instance_name.empty() && plugin->model == this) {
+      continue;
+    }
+    mjCPlugin* candidate = new mjCPlugin(*plugin);
+    candidate->NameSpace(plugin->model);
+    bool referenced = instances.find(candidate->name) != instances.end();
+    auto same_name = [candidate](const mjCPlugin* dest) { return dest->name == candidate->name; };
+    bool instance_exists = std::find_if(dest.begin(), dest.end(), same_name) != dest.end();
+    if (referenced && !instance_exists) {
+      dest.push_back(candidate);
+      instances.at(candidate->name)->spec.plugin.element = candidate;
+    } else {
+      delete candidate;
+    }
+  }
 }
 
-
-
 mjCModel& mjCModel::operator+=(const mjCModel& other) {
   // create global lists
   mjCBody *world = bodies_[0];
@@ -336,22 +349,6 @@ mjCModel& mjCModel::operator+=(const mjCModel& other) {
     for (const auto& key : other.key_pending_) {
       key_pending_.push_back(key);
     }
-
-    // create new plugins and map them
-    std::unordered_map<mjCPlugin*, mjCPlugin*> plugin_map;
-    for (const auto& plugin : other.plugins_) {
-      plugins_.push_back(new mjCPlugin(*plugin));
-      plugins_.back()->NameSpace(&other);
-      plugin_map[plugin] = plugins_.back();
-    }
-    mapplugin(plugin_map, bodies_);
-    mapplugin(plugin_map, geoms_);
-    mapplugin(plugin_map, meshes_);
-    mapplugin(plugin_map, actuators_);
-    mapplugin(plugin_map, sensors_);
-    for (const auto& active_plugin : other.active_plugins_) {
-      active_plugins_.emplace_back(active_plugin);
-    }
   }
   CopyList(flexes_, other.flexes_);
   CopyList(pairs_, other.pairs_);
@@ -364,6 +361,16 @@ mjCModel& mjCModel::operator+=(const mjCModel& other) {
   CopyList(texts_, other.texts_);
   CopyList(tuples_, other.tuples_);
 
+  // create new plugins and map them
+  CopyPlugin(plugins_, other.plugins_, bodies_);
+  CopyPlugin(plugins_, other.plugins_, geoms_);
+  CopyPlugin(plugins_, other.plugins_, meshes_);
+  CopyPlugin(plugins_, other.plugins_, actuators_);
+  CopyPlugin(plugins_, other.plugins_, sensors_);
+  for (const auto& active_plugin : other.active_plugins_) {
+    active_plugins_.emplace_back(active_plugin);
+  }
+
   // restore to the original state
   if (!compiled) {
     ResetTreeLists();

src/user/user_model.h

@@ -364,6 +364,11 @@ class mjCModel : public mjCModel_, private mjSpec {
   template <class T> void CopyList(std::vector<T*>& dest,
                                    const std::vector<T*>& sources);
 
+  // copy vector of plugins to this model
+  template <class T> void CopyPlugin(std::vector<mjCPlugin*>& dest,
+                                     const std::vector<mjCPlugin*>& sources,
+                                     const std::vector<T*>& list);
+
   // delete from list the elements that cause an error
   template <class T> void RemoveFromList(std::vector<T*>& list, const mjCModel& other);
 

src/user/user_objects.h

@@ -824,6 +824,8 @@ class mjCMesh_ : public mjCBase {
 };
 
 class mjCMesh: public mjCMesh_, private mjsMesh {
+  friend class mjCModel;
+
  public:
   mjCMesh(mjCModel* = nullptr, mjCDef* = nullptr);
   mjCMesh(const mjCMesh& other);

test/user/user_api_test.cc

@@ -220,6 +220,7 @@ TEST_F(PluginTest, AttachPlugin) {
         </plugin>
       </extension>
       <worldbody>
+        <body name="empty"/>
         <body name="body">
           <joint name="joint"/>
           <geom size="0.1"/>
@@ -251,8 +252,14 @@ TEST_F(PluginTest, AttachPlugin) {
   mjModel* model_2 = mj_compile(spec_1, nullptr);
   EXPECT_THAT(model_2, NotNull());
 
+  // attach a body not referencing the plugin
+  mjs_attachBody(attachment_frame, mjs_findBody(spec_2, "empty"), "empty-", "");
+  mjModel* model_3 = mj_compile(spec_1, nullptr);
+  EXPECT_THAT(model_3, NotNull());
+
   mj_deleteModel(model_1);
   mj_deleteModel(model_2);
+  mj_deleteModel(model_3);
   mj_deleteSpec(spec_1);
   mj_deleteSpec(spec_2);
 }
@@ -296,7 +303,8 @@ TEST_F(PluginTest, RecompileCompare) {
         std::string xml = p.path().string();
 
         // if file is meant to fail, skip it
-        if (absl::StrContains(p.path().string(), "malformed_") ||
+        if (!absl::StrContains(p.path().string(), "sdf/torus.") ||
+            absl::StrContains(p.path().string(), "malformed_") ||
             absl::StrContains(p.path().string(), "touch_grid") ||
             absl::StrContains(p.path().string(), "cow") ||
             absl::StrContains(p.path().string(), "discardvisual")) {