From 4b88e9bebf43f92de4b0f7a51116b7f3460b273e Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Thu, 15 Aug 2024 06:09:45 -0700
Subject: [PATCH] Do not delete keyframes during mjs_detachBody.

PiperOrigin-RevId: 663281082
Change-Id: I6eedf8c5a3a422bcd78086f2fe30c2099f0b9d16
---
 src/user/user_model.cc     | 34 +++++++++++++++++++++++++---------
 src/user/user_objects.cc   |  3 ---
 test/user/user_api_test.cc |  5 +++++
 3 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/src/user/user_model.cc b/src/user/user_model.cc
index bfcce1fd..cbd37927 100644
--- a/src/user/user_model.cc
+++ b/src/user/user_model.cc
@@ -325,6 +325,16 @@ void mjCModel::RemoveFromList(std::vector<T*>& list, const mjCModel& other) {
 
 
 
+template <>
+void mjCModel::DeleteAll<mjCKey>(std::vector<mjCKey*>& elements) {
+  for (mjCKey* element : elements) {
+    delete element;
+  }
+  elements.clear();
+}
+
+
+
 mjCModel& mjCModel::operator-=(const mjCBody& subtree) {
   mjCModel oldmodel(*this);
 
@@ -345,6 +355,10 @@ mjCModel& mjCModel::operator-=(const mjCBody& subtree) {
   MakeLists(world);
   ProcessLists(/*checkrepeat=*/false);
 
+  // store keyframes in the old model
+  oldmodel.key_pending_.clear();
+  oldmodel.StoreKeyframes();
+
   // check if we have to remove anything else
   RemoveFromList(pairs_, oldmodel);
   RemoveFromList(excludes_, oldmodel);
@@ -353,6 +367,12 @@ mjCModel& mjCModel::operator-=(const mjCBody& subtree) {
   RemoveFromList(actuators_, oldmodel);
   RemoveFromList(sensors_, oldmodel);
 
+  // move all keyframes to pending so that they will be resized
+  DeleteAll(keys_);
+  for (const auto& key : oldmodel.key_pending_) {
+    key_pending_.push_back(key);
+  }
+
   // restore to the original state
   if (!compiled) {
     ResetTreeLists();
@@ -1163,13 +1183,6 @@ void mjCModel::DeleteAll<mjCTexture>(std::vector<mjCTexture*>& elements) {
   elements.clear();
 }
 
-template <>
-void mjCModel::DeleteAll<mjCKey>(std::vector<mjCKey*>& elements) {
-  for (mjCKey* element : elements) {
-    delete element;
-  }
-  elements.clear();
-}
 
 // set nuser fields
 void mjCModel::SetNuser() {
@@ -2983,14 +2996,17 @@ template void mjCModel::RestoreState<mjtNum>(
 // resolve keyframe references
 void mjCModel::StoreKeyframes() {
   bool resetlists = false;
-  if (joints_.empty()) {
+
+  // create tree lists if they are empty, occurs if an uncompiled model is attached
+  if (bodies_.size() == 1 && geoms_.empty() && sites_.empty() && joints_.empty() &&
+      cameras_.empty() && lights_.empty() && frames_.empty()) {
     MakeLists(bodies_[0]);
     resetlists = true;
   }
 
   SaveDofOffsets();
 
-  for (auto key : keys_) {
+  for (auto& key : keys_) {
     mjKeyInfo info;
     info.name = prefix + key->name + suffix;
     info.time = key->spec.time;
diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc
index 1c177d5a..cc0c7c2f 100644
--- a/src/user/user_objects.cc
+++ b/src/user/user_objects.cc
@@ -937,9 +937,6 @@ mjCBody& mjCBody::operator-=(const mjCBody& subtree) {
     *bodies[i] -= subtree;
   }
 
-  // (b/350784262) delete keyframes
-  model->DeleteAll<mjCKey>(model->keys_);
-
   return *this;
 }
 
diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc
index 8f539ee8..5eb68356 100644
--- a/test/user/user_api_test.cc
+++ b/test/user/user_api_test.cc
@@ -816,6 +816,11 @@ void TestDetachBody(bool compile) {
     <sensor>
       <framepos name="ignore" objtype="body" objname="ignore"/>
     </sensor>
+
+    <keyframe>
+      <key name="two" time="2"/>
+      <key name="three" time="3"/>
+    </keyframe>
   </mujoco>)";
 
   // model with one cylinder and a hinge