Generate MjVisual, MjModel, MjData and MjSpec structs declaration

PiperOrigin-RevId: 834264723
Change-Id: Id33d7568a1f9f0cfb1271cf9dd2d8713c00eee34

wasm/codegen/generators/constants.py

@@ -235,11 +235,7 @@ SKIPPED_STRUCTS: List[str] = [
 # or some of their fields need to be handled manually for now;
 # making their wrapper constructors/destructors non-trivial.
 MANUAL_STRUCTS: List[str] = [
-    "MjData",
-    "MjModel",
     "MjvScene",
-    "MjSpec",
-    "MjVisual",
 ]
 
 # Dictionary that maps anonymous structs to their parent struct and field name.
@@ -364,9 +360,7 @@ SKIPPED_FIELDS: Dict[str, List[str]] = {}
 # Fields handled manually in template file struct declaration.
 MANUAL_FIELDS: Dict[str, List[str]] = {
     # go/keep-sorted start
-    "MjData": ["solver", "timer", "warning", "contact"],
+    "MjData": ["contact"],
-    "MjModel": ["opt", "vis", "stat"],
-    "MjSpec": ["option", "visual", "stat", "element", "compiler"],
     "MjvScene": [
         # go/keep-sorted start
         "camera",

wasm/codegen/generators/structs.py

@@ -129,7 +129,6 @@ def _generate_field_data(
     )
 
   elif isinstance(f.type, ast_nodes.ValueType) and f.type.name.startswith("mj"):
-
     return WrappedFieldData(
         definition=f"{common.uppercase_first_letter(f.type.name)} {f.name};",
         typename=_get_field_struct_type(f.type),
@@ -156,7 +155,10 @@ def _generate_field_data(
     if anonymous_struct_name in constants.STRUCTS_TO_BIND:
       return WrappedFieldData(
           binding=_simple_property_binding(f, w, setter=False, reference=True),
-          typename=_get_field_struct_type(f.type),
+          typename=anonymous_struct_name,
+          definition=(
+              f"{common.uppercase_first_letter(anonymous_struct_name)} {f.name};"
+          ),
           ptr_initialization=f"{f.name}(&ptr_->{f.name})",
           ptr_copy_reset=f"{f.name}.set(&ptr_->{f.name});",
           is_primitive_or_fixed_size=True,
@@ -166,7 +168,11 @@ def _generate_field_data(
 
     inner_type = f.type.inner_type
     size = math.prod(f.type.extents)
-
+    inner_type_name = (
+        f.type.inner_type.name
+        if isinstance(f.type.inner_type, ast_nodes.ValueType)
+        else ""
+    )
     if (
         isinstance(inner_type, ast_nodes.ValueType)
         and inner_type.name in constants.PRIMITIVE_TYPES
@@ -192,9 +198,18 @@ def _generate_field_data(
           binding=_simple_property_binding(f, w),
           is_primitive_or_fixed_size=True,
       )
+    elif inner_type_name.startswith("mj"):
+      return WrappedFieldData(
+          definition=(
+              f"std::vector<{common.uppercase_first_letter(inner_type_name)}>"
+              f" {f.name};"
+          ),
+          ptr_initialization=f"{f.name}(&ptr_->{f.name})",
+          typename=_get_field_struct_type(f.type),
+          binding=_simple_property_binding(f, w, reference=True),
+      )
 
   elif isinstance(f.type, ast_nodes.PointerType):
-
     inner_type_name = (
         f.type.inner_type.name
         if isinstance(f.type.inner_type, ast_nodes.ValueType)
@@ -279,13 +294,7 @@ def _generate_field_data(
         binding=_simple_property_binding(f, w),
     )
 
-  # SHOULD NOT OCCUR
+  raise RuntimeError(f"Field {f.name} from struct {w} not properly handled")
-  print("Error: field {f.name} not properly handled")
-  return WrappedFieldData(
-      definition=f"// Error: field {f.name} not properly handled.",
-      typename=_get_field_struct_type(f.type),
-      binding=f"// Error: field {f.name} not properly handled.",
-  )
 
 
 def _default_function_statement(struct_name: str) -> str:
@@ -353,42 +362,67 @@ def build_struct_header(
   ):
     raise RuntimeError(f"Struct {s} not found in introspect structs")
 
-  is_mjs = w.startswith("Mjs")
+  not_mjs = not w.startswith("Mjs")
   member_inits = _find_member_inits(wrapped_fields)
-  shallow_copy = use_shallow_copy(wrapped_fields)
+  shallow_copy = not_mjs and use_shallow_copy(wrapped_fields)
   builder = code_builder.CodeBuilder()
 
   with builder.struct(f"{w}"):
-    builder.line(f"explicit {w}({s} *ptr);")
+    # destructor
-
+    if not_mjs:
-    if not is_mjs:
       builder.line(f"~{w}();")
+
+    # default constructor
+    if not_mjs and w not in ["MjData", "MjModel"]:
       builder.line(f"{w}();")
 
-    if shallow_copy and not is_mjs:
+    # constructor passing native ptr
+    if w != "MjData":
+      builder.line(f"explicit {w}({s} *ptr);")
+    else:
+      # special handling for MjData
+      builder.line("MjData(MjModel *m);")
+      builder.line("explicit MjData(const MjModel &, const MjData &);")
+      builder.line("std::vector<MjContact> contact() const;")
+
+    # copy constructor
+    if shallow_copy or w in ["MjSpec", "MjModel"]:
       builder.line(f"{w}(const {w} &);")
+
+    # assignment operator
+    if shallow_copy or w == "MjSpec":
       builder.line(f"{w} &operator=(const {w} &);")
+
+    # explicit copy function
+    if shallow_copy or w in ["MjSpec", "MjData", "MjModel"]:
       builder.line(f"std::unique_ptr<{w}> copy();")
 
+    # C struct getter/setter
     builder.line(f"{s}* get() const;")
     builder.line(f"void set({s}* ptr);")
 
+    # field declarations
     for field in wrapped_fields:
       if field.definition and field not in member_inits:
         for line in field.definition.splitlines():
           builder.line(line)
 
+    # define private struct members
     builder.private()
     builder.line(f"{s}* ptr_;")
-    if not is_mjs:
+    if not_mjs and w not in ["MjData", "MjModel"]:
       builder.line("bool owned_ = false;")
 
+    # define public struct members
     if member_inits:
       builder.public()
       for f in member_inits:
         if f.definition:
           builder.line(f"{f.definition}")
 
+      if w == "MjData":
+        builder.line("mjModel *model;")
+
   return builder.to_string() + ";"
 
 
@@ -649,6 +683,9 @@ def sort_structs_by_dependency(
   sorted_struct_names = sorted(struct_names)
 
   for struct_name in sorted_struct_names:
+    if struct_name == "mjData":
+      adj["mjModel"].append("mjData")
+      in_degree["mjData"] += 1
     for field in struct_wrappers[struct_name].wrapped_fields:
 
       field_type_name = field.typename

wasm/codegen/templates/bindings.cc

@@ -109,72 +109,9 @@ std::vector<WrapperType> InitWrapperArray(ArrayType* array, SizeType size) {
 
 // {{ AUTOGENNED_STRUCTS_HEADER }}
 
-struct MjVisual {
-  MjVisual();
-  explicit MjVisual(mjVisual *ptr_);
-  MjVisual(const MjVisual &);
-  MjVisual &operator=(const MjVisual &);
-  ~MjVisual();
-  std::unique_ptr<MjVisual> copy();
-  mjVisual* get() const;
-  void set(mjVisual* ptr);
-  // INSERT-GENERATED-MjVisual-DEFINITIONS
-
- private:
-  mjVisual* ptr_;
-  bool owned_ = false;
-
- public:
-  MjVisualGlobal global;
-  MjVisualQuality quality;
-  MjVisualHeadlight headlight;
-  MjVisualMap map;
-  MjVisualScale scale;
-  MjVisualRgba rgba;
-};
-
-struct MjModel {
-  explicit MjModel(mjModel *m);
-  explicit MjModel(const MjModel &other);
-  ~MjModel();
-  std::unique_ptr<MjModel> copy();
-  mjModel* get() const;
-  void set(mjModel* ptr);
-  // INSERT-GENERATED-MjModel-DEFINITIONS
-
- private:
-  mjModel* ptr_;
-
- public:
-  MjOption opt;
-  MjStatistic stat;
-  MjVisual vis;
-};
-
-struct MjData {
-  MjData(MjModel *m);
-  explicit MjData(const MjModel &, const MjData &);
-  ~MjData();
-  std::vector<MjContact> contact() const;
-  std::unique_ptr<MjData> copy();
-  mjData* get() const;
-  void set(mjData* ptr);
-  // INSERT-GENERATED-MjData-DEFINITIONS
-
- private:
-  mjData* ptr_;
-
- public:
-  mjModel *model;
-  std::vector<MjSolverStat> solver;
-  std::vector<MjTimerStat> timer;
-  std::vector<MjWarningStat> warning;
-};
-
 struct MjvScene {
   MjvScene();
   MjvScene(MjModel *m, int maxgeom);
-  // MjvScene(const MjvScene &);
   ~MjvScene();
   std::unique_ptr<MjvScene> copy();
   int GetSumFlexFaces() const;
@@ -268,29 +205,6 @@ struct MjvScene {
   std::vector<MjvGLCamera> camera;
 };
 
-struct MjSpec {
-  MjSpec();
-  explicit MjSpec(mjSpec *ptr);
-  MjSpec(const MjSpec &);
-  MjSpec &operator=(const MjSpec &);
-  ~MjSpec();
-  std::unique_ptr<MjSpec> copy();
-  mjSpec* get() const;
-  void set(mjSpec* ptr);
-  // INSERT-GENERATED-MjSpec-DEFINITIONS
-
- private:
-  mjSpec* ptr_;
-  bool owned_ = false;
-
- public:
-  MjOption option;
-  MjVisual visual;
-  MjStatistic stat;
-  MjsCompiler compiler;
-  MjsElement element;
-};
-
 val get_mjDISABLESTRING() { return MakeValArray(mjDISABLESTRING); }
 val get_mjENABLESTRING() { return MakeValArray(mjENABLESTRING); }
 val get_mjTIMERSTRING() { return MakeValArray(mjTIMERSTRING); }
@@ -348,13 +262,13 @@ EMSCRIPTEN_BINDINGS(mujoco_enums) {
 // {{ AUTOGENNED_STRUCTS_SOURCE }}
 
 // =============== MjModel =============== //
-MjModel::MjModel(mjModel *m)
+MjModel::MjModel(mjModel *ptr)
-    : ptr_(m), opt(&m->opt), stat(&m->stat), vis(&m->vis) {}
+    : ptr_(ptr), opt(&ptr->opt), vis(&ptr->vis), stat(&ptr->stat) {}
 MjModel::MjModel(const MjModel &other)
     : ptr_(mj_copyModel(nullptr, other.get())),
       opt(&ptr_->opt),
-      stat(&ptr_->stat),
+      vis(&ptr_->vis),
-      vis(&ptr_->vis) {}
+      stat(&ptr_->stat) {}
 MjModel::~MjModel() {
   if (ptr_) {
     mj_deleteModel(ptr_);
@@ -462,30 +376,30 @@ std::vector<MjvGeom> MjvScene::geoms() const {
 
 MjSpec::MjSpec()
     : ptr_(mj_makeSpec()),
+      element(ptr_->element),
+      compiler(&ptr_->compiler),
       option(&ptr_->option),
       visual(&ptr_->visual),
-      stat(&ptr_->stat),
+      stat(&ptr_->stat) {
-      compiler(&ptr_->compiler),
-      element(ptr_->element) {
   owned_ = true;
   mjs_defaultSpec(ptr_);
 };
 
 MjSpec::MjSpec(mjSpec *ptr)
     : ptr_(ptr),
+      element(ptr_->element),
+      compiler(&ptr_->compiler),
       option(&ptr_->option),
       visual(&ptr_->visual),
-      stat(&ptr_->stat),
+      stat(&ptr_->stat) {}
-      compiler(&ptr_->compiler),
-      element(ptr_->element) {}
 
 MjSpec::MjSpec(const MjSpec &other)
     : ptr_(mj_copySpec(other.get())),
+      element(ptr_->element),
+      compiler(&ptr_->compiler),
       option(&ptr_->option),
       visual(&ptr_->visual),
-      stat(&ptr_->stat),
+      stat(&ptr_->stat) {
-      compiler(&ptr_->compiler),
-      element(ptr_->element) {
   owned_ = true;
 }