From e79efa60d6afd191610f5462dc9a9b47aac98278 Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Mon, 22 Jul 2024 08:29:31 -0700
Subject: [PATCH] Reset tree indices before recompiling.

Fixes #1822.

PiperOrigin-RevId: 654762203
Change-Id: I7f656f7dad376d4899229a8741315c8f821bfa56
---
 src/user/user_model.cc     |  1 +
 test/user/user_api_test.cc | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/src/user/user_model.cc b/src/user/user_model.cc
index f365a994..377ed8e6 100644
--- a/src/user/user_model.cc
+++ b/src/user/user_model.cc
@@ -3206,6 +3206,7 @@ mjModel* mjCModel::Compile(const mjVFS* vfs, mjModel** m) {
       bodies_[i]->subtreedofs = 0;
     }
     mjCBody* world = bodies_[0];
+    ResetTreeLists();
     Clear();
     bodies_.push_back(world);
   }
diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc
index 093c07c1..39f224f6 100644
--- a/test/user/user_api_test.cc
+++ b/test/user/user_api_test.cc
@@ -186,6 +186,38 @@ TEST_F(PluginTest, RecompileCompare) {
   }
 }
 
+TEST_F(PluginTest, RecompileEdit) {
+  static constexpr char xml[] = R"(
+  <mujoco>
+    <worldbody>
+      <body>
+        <freejoint/>
+        <geom size=".1"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+
+  std::array<char, 1000> er;
+  mjSpec *spec = mj_parseXMLString(xml, 0, er.data(), er.size());
+  EXPECT_THAT(spec, NotNull()) << er.data();
+  mjModel *m1 = mj_compile(spec, nullptr);
+  EXPECT_THAT(m1, NotNull());
+
+  // add a geom
+  mjsBody *world = mjs_findBody(spec, "world");
+  mjsGeom *geom = mjs_addGeom(world, nullptr);
+  geom->size[0] = 1;
+
+  // compile again
+  mjModel *m2 = mj_compile(spec, nullptr);
+  EXPECT_THAT(m2, NotNull());
+
+  mj_deleteModel(m1);
+  mj_deleteModel(m2);
+  mj_deleteSpec(spec);
+}
+
 // ------------------- test cache with modified assets -------------------------
 
 TEST_F(PluginTest, RecompileCompareObjCache) {