From 4a7649653343daacfb061a2597374dd8ffc4d1bb Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Wed, 17 Apr 2024 09:54:57 -0700
Subject: [PATCH] Compile frames during compilation instead of parsing.

PiperOrigin-RevId: 625717564
Change-Id: Ib82230f6cb423bb3091efd04a7fa0b1411e7cdf1
---
 src/user/user_model.cc   |  2 +-
 src/user/user_objects.cc | 21 +++++++++++++++++----
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/src/user/user_model.cc b/src/user/user_model.cc
index 0a539376..20f81f32 100644
--- a/src/user/user_model.cc
+++ b/src/user/user_model.cc
@@ -3118,7 +3118,7 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
 
   // compile objects in kinematic tree
   for (int i=0; i<bodies.size(); i++) {
-    bodies[i]->Compile();  // also compiles joints, geoms, sites, cameras, lights
+    bodies[i]->Compile();  // also compiles joints, geoms, sites, cameras, lights, frames
   }
 
   // compile all other objects except for keyframes
diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc
index 9f9053f6..e102b001 100644
--- a/src/user/user_objects.cc
+++ b/src/user/user_objects.cc
@@ -628,7 +628,6 @@ void mjCBase::SetFrame(mjCFrame* _frame) {
     return;
   }
   frame = _frame;
-  frame->Compile();
 }
 
 
@@ -735,10 +734,19 @@ mjCBody& mjCBody::operator+=(const mjCFrame& other) {
   other.model->prefix = other.prefix;
   other.model->suffix = other.suffix;
 
-  // map old frames to indices
+  // copy input frame
+  frames.push_back(new mjCFrame(other));
+  frames.back()->body = this;
+  frames.back()->model = model;
+  frames.back()->frame = other.frame;
+
+  // map input frames to index in this->frames
   std::map<mjCFrame*, int> fmap;
-  for (int i=0; i<subtree->frames.size(); i++) {
-    fmap[subtree->frames[i]] = i;
+  for (auto frame : subtree->frames) {
+    if (frame == static_cast<const mjCFrame*>(&other)) {
+      fmap[frame] = frames.size() - 1;
+      break;
+    }
   }
 
   // copy children that are inside the input frame
@@ -1212,6 +1220,11 @@ void mjCBody::MakeInertialExplicit() {
 void mjCBody::Compile(void) {
   CopyFromSpec();
 
+  // compile all frames
+  for (int i=0; i<frames.size(); i++) {
+    frames[i]->Compile();
+  }
+
   // resize userdata
   if (userdata_.size() > model->nuser_body) {
     throw mjCError(this, "user has more values than nuser_body in body '%s' (id = %d)",