From 4e206c29c1afe9b0cf0dd879811e2b5fda31bbc3 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Tue, 25 Mar 2025 08:20:14 -0700 Subject: [PATCH] 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 | 72 ++++++++-------------------------------- src/user/user_model.h | 4 ++- src/user/user_objects.cc | 44 +++++++++++++++++++++--- 3 files changed, 56 insertions(+), 64 deletions(-) diff --git a/src/user/user_model.cc b/src/user/user_model.cc index bc12a17e..2a27d430 100644 --- a/src/user/user_model.cc +++ b/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* 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( // 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; diff --git a/src/user/user_model.h b/src/user/user_model.h index 1078080f..c97216e5 100644 --- a/src/user/user_model.h +++ b/src/user/user_model.h @@ -334,9 +334,11 @@ class mjCModel : public mjCModel_, private mjSpec { // list of active plugins std::vector> 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 diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index 18c573d2..a493bf30 100644 --- a/src/user/user_objects.cc +++ b/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; }