From 231d9632cd8ca54e834758f42efd8d7ed4b7ff2a Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Thu, 18 Jul 2024 05:07:13 -0700
Subject: [PATCH] Avoid duplicated code in mjCModel by adding ResetTreeLists().

PiperOrigin-RevId: 653575505
Change-Id: I61ee1462deae2a371cf6f36ca1cf3dbbd88d174c
---
 src/user/user_model.cc | 55 +++++++++++++++---------------------------
 src/user/user_model.h  |  3 +++
 2 files changed, 22 insertions(+), 36 deletions(-)

diff --git a/src/user/user_model.cc b/src/user/user_model.cc
index 4ada3658..a810f6c3 100644
--- a/src/user/user_model.cc
+++ b/src/user/user_model.cc
@@ -188,19 +188,26 @@ static void resetlist(std::vector<T*>& list) {
 
 
 
+void mjCModel::ResetTreeLists() {
+  mjCBody *world = bodies_[0];
+  resetlist(bodies_);
+  resetlist(joints_);
+  resetlist(geoms_);
+  resetlist(sites_);
+  resetlist(cameras_);
+  resetlist(lights_);
+  resetlist(frames_);
+  world->id = 0;
+  bodies_.push_back(world);
+}
+
+
+
 mjCModel& mjCModel::operator+=(const mjCModel& other) {
   // create global lists
   mjCBody *world = bodies_[0];
   if (compiled) {
-    resetlist(bodies_);
-    resetlist(joints_);
-    resetlist(geoms_);
-    resetlist(sites_);
-    resetlist(cameras_);
-    resetlist(lights_);
-    resetlist(frames_);
-    world->id = 0;
-    bodies_.push_back(world);
+    ResetTreeLists();
   }
   MakeLists(world);
   ProcessLists(/*checkrepeat=*/false);
@@ -233,15 +240,7 @@ mjCModel& mjCModel::operator+=(const mjCModel& other) {
 
   // restore to the original state
   if (!compiled) {
-    resetlist(bodies_);
-    resetlist(joints_);
-    resetlist(geoms_);
-    resetlist(sites_);
-    resetlist(cameras_);
-    resetlist(lights_);
-    resetlist(frames_);
-    world->id = 0;
-    bodies_.push_back(world);
+    ResetTreeLists();
   }
 
   PointToLocal();
@@ -291,15 +290,7 @@ mjCModel& mjCModel::operator-=(const mjCBody& subtree) {
 
   // create global lists
   if (compiled) {
-    resetlist(bodies_);
-    resetlist(joints_);
-    resetlist(geoms_);
-    resetlist(sites_);
-    resetlist(cameras_);
-    resetlist(lights_);
-    resetlist(frames_);
-    world->id = 0;
-    bodies_.push_back(world);
+    ResetTreeLists();
   }
   MakeLists(world);
   ProcessLists(/*checkrepeat=*/false);
@@ -314,15 +305,7 @@ mjCModel& mjCModel::operator-=(const mjCBody& subtree) {
 
   // restore to the original state
   if (!compiled) {
-    resetlist(bodies_);
-    resetlist(joints_);
-    resetlist(geoms_);
-    resetlist(sites_);
-    resetlist(cameras_);
-    resetlist(lights_);
-    resetlist(frames_);
-    world->id = 0;
-    bodies_.push_back(world);
+    ResetTreeLists();
   }
 
   PointToLocal();
diff --git a/src/user/user_model.h b/src/user/user_model.h
index 7e36d6ec..f6c1c782 100644
--- a/src/user/user_model.h
+++ b/src/user/user_model.h
@@ -353,6 +353,9 @@ class mjCModel : public mjCModel_, private mjSpec {
   // populate objects ids
   void ProcessLists(bool checkrepeat = true);
 
+  // reset lists of kinematic tree
+  void ResetTreeLists();
+
   mjListKeyMap ids;   // map from object names to ids
   mjCError errInfo;   // last error info
   bool plugin_owner;  // this class allocated the plugins