Keep tree lists updated at all times.

This simplifies the compiler logic since an updated tree list was necessary at many stages (e.g. attach and keyframes resizing) and it will be anyway required when computing the mjSpec signature.

PiperOrigin-RevId: 740356447
Change-Id: I7f2ec25b27b8d4ca4364801c9a401c40c6d84569

src/user/user_model.cc

@@ -449,9 +449,8 @@ static bool IsPluginActive(
 
 mjCModel& mjCModel::operator+=(const mjCModel& other) {
   // create global lists
-  mjCBody *world = bodies_[0];
   ResetTreeLists();
-  MakeLists(world);
+  MakeTreeLists();
   ProcessLists(/*checkrepeat=*/false);
 
   // copy all elements not in the tree
@@ -500,11 +499,6 @@ mjCModel& mjCModel::operator+=(const mjCModel& other) {
     nq = nv = na = nu = nmocap = 0;
   }
 
-  // restore to the original state
-  if (!compiled) {
-    ResetTreeLists();
-  }
-
   PointToLocal();
   return *this;
 }
@@ -609,13 +603,11 @@ mjCModel& mjCModel::operator-=(const mjCBody& subtree) {
 
   // create global lists in the old model if not compiled
   if (!oldmodel.IsCompiled()) {
-    oldmodel.MakeLists(oldmodel.bodies_[0]);
     oldmodel.ProcessLists(/*checkrepeat=*/false);
   }
 
   // create global lists in this model if not compiled
   if (!IsCompiled()) {
-    MakeLists(bodies_[0]);
     ProcessLists(/*checkrepeat=*/false);
   }
 
@@ -629,7 +621,7 @@ mjCModel& mjCModel::operator-=(const mjCBody& subtree) {
 
   // update global lists
   ResetTreeLists();
-  MakeLists(world);
+  MakeTreeLists();
   ProcessLists(/*checkrepeat=*/false);
 
   // check if we have to remove anything else
@@ -641,11 +633,6 @@ mjCModel& mjCModel::operator-=(const mjCBody& subtree) {
   RemoveFromList(sensors_, oldmodel);
   RemovePlugins();
 
-  // restore to the original state
-  if (!compiled) {
-    ResetTreeLists();
-  }
-
   return *this;
 }
 
@@ -742,18 +729,16 @@ void deletefromlist(std::vector<T*>* list, mjsElement* element) {
 
 // discard all invalid elements from all lists
 void mjCModel::DeleteElement(mjsElement* el) {
-  mjCBody *world = nullptr;
-  if (compiled) {
-    world = bodies_[0];
-    ResetTreeLists();
-  }
+  ResetTreeLists();
 
   switch (el->elemtype) {
     case mjOBJ_BODY:
+      MakeTreeLists();  // rebuild lists that were reset at the beginning of the function
       throw mjCError(nullptr, "bodies cannot be deleted, use detach instead");
       break;
 
     case mjOBJ_DEFAULT:
+      MakeTreeLists();  // rebuild lists that were reset at the beginning of the function
       throw mjCError(nullptr, "defaults cannot be deleted, use detach instead");
       break;
 
@@ -818,11 +803,9 @@ void mjCModel::DeleteElement(mjsElement* el) {
       break;
   }
 
-  if (compiled) {
-    ResetTreeLists();  // in case of a nested delete
-    MakeLists(world);
-    ProcessLists(/*checkrepeat=*/false);
-  }
+  ResetTreeLists();  // in case of a nested delete
+  MakeTreeLists();
+  ProcessLists(/*checkrepeat=*/false);
 }
 
 
@@ -1020,15 +1003,6 @@ void mjCModel::Clear() {
   nconmax = -1;
   nmocap = 0;
 
-  // pointer lists created by Compile
-  bodies_.clear();
-  joints_.clear();
-  geoms_.clear();
-  sites_.clear();
-  cameras_.clear();
-  lights_.clear();
-  frames_.clear();
-
   // internal variables
   hasImplicitPluginElem = false;
   compiled = false;
@@ -1486,7 +1460,11 @@ mjSpec* mjCModel::GetSourceSpec() const {
 //------------------------------- COMPILER PHASES --------------------------------------------------
 
 // make lists of objects in tree: bodies, geoms, joints, sites, cameras, lights
-void mjCModel::MakeLists(mjCBody* body) {
+void mjCModel::MakeTreeLists(mjCBody* body) {
+  if (body == nullptr) {
+    body = bodies_[0];
+  }
+
   // add this body if not world
   if (body != bodies_[0]) {
     bodies_.push_back(body);
@@ -1501,7 +1479,7 @@ void mjCModel::MakeLists(mjCBody* body) {
   for (mjCFrame *frame : body->frames) frames_.push_back(frame);
 
   // recursive call to all child bodies
-  for (mjCBody* body : body->bodies) MakeLists(body);
+  for (mjCBody* body : body->bodies) MakeTreeLists(body);
 }
 
 
@@ -3657,21 +3635,12 @@ template void mjCModel::RestoreState<mjtNum>(
 
 // resolve keyframe references
 void mjCModel::StoreKeyframes(mjCModel* dest) {
-  bool resetlists = false;
-
   if (this != dest && !key_pending_.empty()) {
     mju_warning(
       "Child model has pending keyframes. They will not be namespaced correctly. "
       "To prevent this, compile the child model before attaching it again.");
   }
 
-  // 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;
-  }
-
   // do not change compilation quantities in case the user wants to recompile preserving the state
   if (!compiled) {
     SaveDofOffsets(/*computesize=*/true);
@@ -3719,10 +3688,6 @@ void mjCModel::StoreKeyframes(mjCModel* dest) {
               key->spec_mpos_.data(), key->spec_mquat_.data());
   }
 
-  if (resetlists) {
-    ResetTreeLists();
-  }
-
   if (!compiled) {
     nq = nv = na = nu = nmocap = 0;
   }
@@ -4054,11 +4019,7 @@ static void warninghandler(const char* msg) {
 // compiler
 mjModel* mjCModel::Compile(const mjVFS* vfs, mjModel** m) {
   if (compiled) {
-    // clear kinematic tree
-    mjCBody* world = bodies_[0];
-    ResetTreeLists();
     Clear();
-    bodies_.push_back(world);
   }
 
   CopyFromSpec();
@@ -4105,9 +4066,7 @@ mjModel* mjCModel::Compile(const mjVFS* vfs, mjModel** m) {
     // deallocate everything allocated in Compile
     mj_deleteModel(model);
     mj_deleteData(data);
-    mjCBody* world = bodies_[0];
     Clear();
-    bodies_.push_back(world);
 
     // save error info
     errInfo = err;
@@ -4354,9 +4313,6 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
     AddKey();
   }
 
-  // make lists of objects created in kinematic tree
-  MakeLists(bodies_[0]);
-
   // clear subtreedofs
   for (int i=0; i < bodies_.size(); i++) {
     bodies_[i]->subtreedofs = 0;

src/user/user_model.h

@@ -334,9 +334,11 @@ class mjCModel : public mjCModel_, private mjSpec {
   // list of active plugins
   std::vector<std::pair<const mjpPlugin*, int>> active_plugins_;
 
+  // make lists of bodies and children
+  void MakeTreeLists(mjCBody* body = nullptr);
+
   // compile phases
   void TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs);
-  void MakeLists(mjCBody* body);        // make lists of bodies, geoms, joints, sites
   void SetNuser();                      // set nuser fields
   void IndexAssets(bool discard);       // convert asset names into indices
   void CheckEmptyNames();               // check empty names

src/user/user_objects.cc

@@ -1231,6 +1231,11 @@ mjCBody* mjCBody::AddBody(mjCDef* _def) {
   obj->classname = _def ? _def->name : classname;
 
   bodies.push_back(obj);
+
+  // recompute lists
+  model->ResetTreeLists();
+  model->MakeTreeLists();
+
   obj->parent = this;
   return obj;
 }
@@ -1241,6 +1246,8 @@ mjCBody* mjCBody::AddBody(mjCDef* _def) {
 mjCFrame* mjCBody::AddFrame(mjCFrame* _frame) {
   mjCFrame* obj = new mjCFrame(model, _frame ? _frame : NULL);
   frames.push_back(obj);
+  model->ResetTreeLists();
+  model->MakeTreeLists();
   return obj;
 }
 
@@ -1256,6 +1263,11 @@ mjCJoint* mjCBody::AddFreeJoint() {
   obj->body = this;
 
   joints.push_back(obj);
+
+  // recompute lists
+  model->ResetTreeLists();
+  model->MakeTreeLists();
+
   return obj;
 }
 
@@ -1270,6 +1282,11 @@ mjCJoint* mjCBody::AddJoint(mjCDef* _def) {
   obj->body = this;
 
   joints.push_back(obj);
+
+  // recompute lists
+  model->ResetTreeLists();
+  model->MakeTreeLists();
+
   return obj;
 }
 
@@ -1284,6 +1301,11 @@ mjCGeom* mjCBody::AddGeom(mjCDef* _def) {
   obj->body = this;
 
   geoms.push_back(obj);
+
+  // recompute lists
+  model->ResetTreeLists();
+  model->MakeTreeLists();
+
   return obj;
 }
 
@@ -1298,6 +1320,11 @@ mjCSite* mjCBody::AddSite(mjCDef* _def) {
   obj->body = this;
 
   sites.push_back(obj);
+
+  // recompute lists
+  model->ResetTreeLists();
+  model->MakeTreeLists();
+
   return obj;
 }
 
@@ -1312,6 +1339,11 @@ mjCCamera* mjCBody::AddCamera(mjCDef* _def) {
   obj->body = this;
 
   cameras.push_back(obj);
+
+  // recompute lists
+  model->ResetTreeLists();
+  model->MakeTreeLists();
+
   return obj;
 }
 
@@ -1326,6 +1358,11 @@ mjCLight* mjCBody::AddLight(mjCDef* _def) {
   obj->body = this;
 
   lights.push_back(obj);
+
+  // recompute lists
+  model->ResetTreeLists();
+  model->MakeTreeLists();
+
   return obj;
 }
 
@@ -1347,11 +1384,8 @@ mjCFrame* mjCBody::ToFrame() {
       std::remove_if(parent->bodies.begin(), parent->bodies.end(),
                      [this](mjCBody* body) { return body == this; }),
       parent->bodies.end());
-  if (model->IsCompiled()) {
-    mjCBody *world = model->bodies_[0];
-    model->ResetTreeLists();
-    model->MakeLists(world);
-  }
+  model->ResetTreeLists();
+  model->MakeTreeLists();
   return newframe;
 }