From 203602b01f6c81903ac4183ae2c4305c26311f6e Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Tue, 26 Nov 2024 10:16:16 -0800
Subject: [PATCH] Move computation of flex_vert0 to user_mesh.cc.

This removes the dependency of flex_vert0 from mj_flex and mjData.

PiperOrigin-RevId: 700382297
Change-Id: I0b6d023c2e1ba14fa92417b65bfc288f37c223ae
---
 src/engine/engine_setconst.c | 12 ------------
 src/user/user_mesh.cc        | 10 ++++++++++
 src/user/user_model.cc       |  3 +++
 src/user/user_objects.h      |  2 ++
 4 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/src/engine/engine_setconst.c b/src/engine/engine_setconst.c
index c1926673..04fcec99 100644
--- a/src/engine/engine_setconst.c
+++ b/src/engine/engine_setconst.c
@@ -120,18 +120,6 @@ static void set0(mjModel* m, mjData* d) {
     m->light_mode[i] = lightmode[i];
   }
 
-  // compute bounding box coordinates
-  for (int i=0; i < m->nflex; i++) {
-    int bvhadr = m->flex_bvhadr[i];
-    const mjtNum* bvh = d->bvh_aabb_dyn + 6*(bvhadr - m->nbvhstatic);
-    for (int j=0; j < m->nflexvert; j++) {
-      for (int k=0; k < 3; k++) {
-        mjtNum size = 2*(bvh[3+k] - m->flex_radius[i]);
-        m->flex_vert0[3*j+k] = (d->flexvert_xpos[3*j+k] - bvh[k]) / size + 0.5;
-      }
-    }
-  }
-
   // copy fields
   mju_copy(m->flexedge_length0, d->flexedge_length, m->nflexedge);
   mju_copy(m->tendon_length0, d->ten_length, m->ntendon);
diff --git a/src/user/user_mesh.cc b/src/user/user_mesh.cc
index 766e4976..f92f8c66 100644
--- a/src/user/user_mesh.cc
+++ b/src/user/user_mesh.cc
@@ -2944,6 +2944,16 @@ void mjCFlex::Compile(const mjVFS* vfs) {
 
   // create bounding volume hierarchy
   CreateBVH();
+
+  // compute bounding box coordinates
+  vert0_.assign(3*nvert, 0);
+  const mjtNum* bvh = tree.Bvh().data();
+  for (int j=0; j < nvert; j++) {
+    for (int k=0; k < 3; k++) {
+      double size = 2*(bvh[k+3] - radius);
+      vert0_[3*j+k] = (vertxpos[3*j+k] - bvh[k]) / size + 0.5;
+    }
+  }
 }
 
 
diff --git a/src/user/user_model.cc b/src/user/user_model.cc
index fb21b1b1..f3223a0f 100644
--- a/src/user/user_model.cc
+++ b/src/user/user_model.cc
@@ -2825,6 +2825,9 @@ void mjCModel::CopyObjects(mjModel* m) {
       mjuu_copyvec(m->flex_vert + 3*vert_adr, pfl->vert_.data(), 3*pfl->nvert);
     }
 
+    // copy vert0
+    mjuu_copyvec(m->flex_vert0 + 3*vert_adr, pfl->vert0_.data(), 3*pfl->nvert);
+
     // copy or set vertbodyid
     if (pfl->rigid) {
       for (int k=0; k<pfl->nvert; k++) {
diff --git a/src/user/user_objects.h b/src/user/user_objects.h
index 3238f1ff..d5dd09b9 100644
--- a/src/user/user_objects.h
+++ b/src/user/user_objects.h
@@ -774,6 +774,8 @@ class mjCFlex: public mjCFlex_, private mjsFlex {
   void Compile(const mjVFS* vfs);         // compiler
   void CreateBVH(void);                   // create flex BVH
   void CreateShellPair(void);             // create shells and evpairs
+
+  std::vector<double> vert0_;             // vertex positions in [0, 1]^d in the bounding box
 };