From 5c955b8fe264e5235c2e2c89a6532259e80cae4c Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Wed, 26 Mar 2025 05:35:33 -0700
Subject: [PATCH] Add frame support to the cable composite.

Also cleanup user_composite.cc includes.

PiperOrigin-RevId: 740730596
Change-Id: I05df59926d8e61de45fd6b2a3a63121b7a390448
---
 src/user/user_composite.cc   | 18 +++++++-----------
 src/user/user_composite.h    |  1 +
 src/xml/xml_native_reader.cc |  5 +++--
 src/xml/xml_native_reader.h  |  3 ++-
 4 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/src/user/user_composite.cc b/src/user/user_composite.cc
index 0c09164b..98c3ed41 100644
--- a/src/user/user_composite.cc
+++ b/src/user/user_composite.cc
@@ -12,25 +12,17 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include <mujoco/mjspec.h>
 #include "user/user_composite.h"
 
-#include <algorithm>
 #include <cmath>
-#include <cstddef>
-#include <cstdio>
 #include <cstring>
-#include <map>
-#include <stdexcept>
 #include <string>
-#include <utility>
 #include <vector>
 
-#include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
+#include <mujoco/mjspec.h>
 #include <mujoco/mjtnum.h>
 #include "cc/array_safety.h"
-#include "engine/engine_io.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"
@@ -40,9 +32,9 @@
 #include "user/user_util.h"
 
 namespace {
+
 namespace mju = ::mujoco::util;
-using mujoco::user::VectorToString;
-using mujoco::user::StringToVector;
+
 }  // namespace
 
 // strncpy with 0, return false
@@ -60,6 +52,7 @@ mjCComposite::mjCComposite(void) {
   type = mjCOMPTYPE_PARTICLE;
   count[0] = count[1] = count[2] = 1;
   mjuu_setvec(offset, 0, 0, 0);
+  frame = nullptr;
 
   // plugin variables
   mjs_defaultPlugin(&plugin);
@@ -381,6 +374,9 @@ mjsBody* mjCComposite::AddCableBody(mjCModel* model, mjsBody* body, int ix,
                            offset[1]+uservert[3*ix+1],
                            offset[2]+uservert[3*ix+2]);
     mjuu_copyvec(body->quat, this_quat, 4);
+    if (frame) {
+      mjs_setFrame(body->element, frame);
+    }
   } else {
     mjuu_setvec(body->pos, length_prev, 0, 0);
     double negquat[4] = {prev_quat[0], -prev_quat[1], -prev_quat[2], -prev_quat[3]};
diff --git a/src/user/user_composite.h b/src/user/user_composite.h
index 383dd3a2..0f826760 100644
--- a/src/user/user_composite.h
+++ b/src/user/user_composite.h
@@ -80,6 +80,7 @@ class mjCComposite {
   std::vector<float> uservert;    // user-specified vertex positions
   double size[3];                 // rope size (meaning depends on the shape)
   mjtCompShape curve[3];          // geometric shape
+  mjsFrame* frame;                // frame where the composite is defined
 
   // body names used in the skin
   std::vector<std::string> username;
diff --git a/src/xml/xml_native_reader.cc b/src/xml/xml_native_reader.cc
index 4f941c15..44efc4ec 100644
--- a/src/xml/xml_native_reader.cc
+++ b/src/xml/xml_native_reader.cc
@@ -2388,7 +2388,7 @@ void mjXReader::OneActuator(XMLElement* elem, mjsActuator* actuator) {
 
 
 // make composite
-void mjXReader::OneComposite(XMLElement* elem, mjsBody* body, const mjsDefault* def) {
+void mjXReader::OneComposite(XMLElement* elem, mjsBody* body, mjsFrame* frame, const mjsDefault* def) {
   string text;
   int n;
 
@@ -2402,6 +2402,7 @@ void mjXReader::OneComposite(XMLElement* elem, mjsBody* body, const mjsDefault*
   }
   ReadAttr(elem, "count", 3, comp.count, text, false, false);
   ReadAttr(elem, "offset", 3, comp.offset, text);
+  comp.frame = frame;
 
   // plugin
   XMLElement* eplugin = FirstChildElement(elem, "plugin");
@@ -3467,7 +3468,7 @@ void mjXReader::Body(XMLElement* section, mjsBody* body, mjsFrame* frame,
     // composite sub-element
     else if (name == "composite") {
       // parse composite
-      OneComposite(elem, body, def);
+      OneComposite(elem, body, frame, def);
     }
 
     // flexcomp sub-element
diff --git a/src/xml/xml_native_reader.h b/src/xml/xml_native_reader.h
index 41aacbbb..0e568f99 100644
--- a/src/xml/xml_native_reader.h
+++ b/src/xml/xml_native_reader.h
@@ -80,7 +80,8 @@ class mjXReader : public mjXBase {
   void OneEquality(tinyxml2::XMLElement* elem, mjsEquality* pequality);
   void OneTendon(tinyxml2::XMLElement* elem, mjsTendon* ptendon);
   void OneActuator(tinyxml2::XMLElement* elem, mjsActuator* pactuator);
-  void OneComposite(tinyxml2::XMLElement* elem, mjsBody* pbody, const mjsDefault* def);
+  void OneComposite(tinyxml2::XMLElement* elem, mjsBody* pbody, mjsFrame* pframe,
+                    const mjsDefault* def);
   void OneFlexcomp(tinyxml2::XMLElement* elem, mjsBody* pbody, const mjVFS* vfs);
   void OnePlugin(tinyxml2::XMLElement* elem, mjsPlugin* plugin);