Set parent frame of new frames after copying all frames.

PiperOrigin-RevId: 737620918
Change-Id: I8f49570cab39bd84b2da908d3683dddc3b29cb3c

Author: quagla

src/user/user_objects.cc

@@ -847,7 +847,7 @@ mjCBody& mjCBody::operator+=(const mjCBody& other) {
   // map other frames to indices
   std::map<mjCFrame*, int> fmap;
   for (int i=0; i<other.frames.size(); i++) {
-    fmap[other.frames[i]] = i;
+    fmap[other.frames[i]] = i + frames.size();
   }
 
   // copy frames, needs to happen first
@@ -977,6 +977,7 @@ template <typename T>
 void mjCBody::CopyList(std::vector<T*>& dst, const std::vector<T*>& src,
                        std::map<mjCFrame*, int>& fmap, const mjCFrame* pframe) {
   int nsrc = (int)src.size();
+  int ndst = (int)dst.size();
   for (int i=0; i<nsrc; i++) {
     if (pframe && !pframe->IsAncestor(src[i]->frame)) {
       continue;  // skip if the element is not inside pframe
@@ -996,12 +997,19 @@ void mjCBody::CopyList(std::vector<T*>& dst, const std::vector<T*>& src,
       dst.back()->AddRef();
     }
 
-    // assign dst frame to src frame
-    dst.back()->frame = src[i]->frame ? frames[fmap[src[i]->frame]] : nullptr;
-
     // set namespace
     dst.back()->NameSpace(src[i]->model);
   }
+
+  // assign dst frame to src frame
+  // needs to be done after the copy in case T is an mjCFrame
+  int j = 0;
+  for (int i = 0; i < src.size(); i++) {
+    if (pframe && !pframe->IsAncestor(src[i]->frame)) {
+      continue;  // skip if the element is not inside pframe
+    }
+    dst[ndst + j++]->frame = src[i]->frame ? frames[fmap[src[i]->frame]] : nullptr;
+  }
 }
 
 

test/user/user_api_test.cc

@@ -2569,5 +2569,21 @@ TEST_F(MujocoTest, ErrorWhenCompilingOrphanedSpec) {
   mj_deleteSpec(child);
 }
 
+TEST_F(MujocoTest, SetFrameReverseOrder) {
+  mjSpec* spec = mj_makeSpec();
+  mjsBody* world = mjs_findBody(spec, "world");
+  mjsFrame* child = mjs_addFrame(world, nullptr);
+  mjsFrame* parent = mjs_addFrame(world, nullptr);
+  mjs_setString(child->name, "child");
+  mjs_setString(parent->name, "parent");
+  mjs_setFrame(child->element, parent);
+  mjSpec* copy = mj_copySpec(spec);
+  EXPECT_THAT(copy, NotNull());
+  EXPECT_THAT(mjs_findFrame(copy, "child"), NotNull());
+  EXPECT_THAT(mjs_findFrame(copy, "parent"), NotNull());
+  mj_deleteSpec(spec);
+  mj_deleteSpec(copy);
+}
+
 }  // namespace
 }  // namespace mujoco