Simplified struct code generation for WASM bindings

PiperOrigin-RevId: 827878780
Change-Id: I97ab732eb11c63a6aaf7fa95add0b51bc0a56525

wasm/codegen/generators/structs.py

@@ -18,23 +18,16 @@ from typing import Optional
 
 from wasm.codegen.helpers import constants
 from wasm.codegen.helpers import structs_parser
-from wasm.codegen.helpers import structs_wrappers_data
 
 
 class Generator:
   """Generates C++ code for binding and wrapping MuJoCo structs."""
 
   def __init__(self):
-    # Set up the correct input dict based on the structs we want to bind
-    # and already have a wrapper manually created in the template/bindings.cc
-    wrapped_structs = structs_wrappers_data.create_wrapped_structs_set_up_data(
-        constants.STRUCTS_TO_BIND
-    )
-
     # Traverse the introspect dictionary to get the field
     # wrapper/bindings statements set up for each struct
     self.structs_to_bind_data = structs_parser.generate_wasm_bindings(
-        wrapped_structs
+        constants.STRUCTS_TO_BIND
     )
 
   def generate_header(

wasm/codegen/helpers/struct_constructor_code_builder.py

@@ -25,9 +25,7 @@ from wasm.codegen.helpers import constants
 from wasm.codegen.helpers import structs_wrappers_data
 
 
-def _has_nested_wrapper_members(
+def _has_nested_wrapper_members(struct_info: ast_nodes.StructDecl) -> bool:
-    struct_info: ast_nodes.StructDecl
-) -> bool:
   """Checks if the struct contains other wrapped structs as direct members."""
   for field in struct_info.fields:
     struct_field = cast(ast_nodes.StructFieldDecl, field)
@@ -49,10 +47,12 @@ def _build_struct_header_internal(
     struct_name: str,
     wrapped_fields: List[structs_wrappers_data.WrappedFieldData],
     fields_with_init: List[structs_wrappers_data.WrappedFieldData],
-    use_shallow_copy: bool = False,
     is_mjs: bool = False,
 ):
   """Builds the C++ header file code for a struct."""
+
+  shallow_copy = use_shallow_copy(wrapped_fields)
+
   wrapper_name = common.uppercase_first_letter(struct_name)
   builder = code_builder.CodeBuilder()
   with builder.block(f"struct {wrapper_name}"):
@@ -64,7 +64,7 @@ def _build_struct_header_internal(
     builder.line(f"explicit {wrapper_name}({struct_name} *ptr);")
     builder.line(f"~{wrapper_name}();")
 
-    if use_shallow_copy:
+    if shallow_copy:
       builder.line(f"std::unique_ptr<{wrapper_name}> copy();")
 
     for field in wrapped_fields:
@@ -87,24 +87,62 @@ def _build_struct_header_internal(
       for field in fields_with_init:
         if field.definition:
           builder.line(f"{field.definition}")
-  return builder.to_string()+";"
+  return builder.to_string() + ";"
+
+
+def _get_default_func_name(struct_name: str) -> str:
+  """Returns the default function name for the given struct."""
+  if (
+      struct_name in constants.ANONYMOUS_STRUCTS.keys()
+      or struct_name in constants.NO_DEFAULT_CONSTRUCTORS
+      or (
+          common.uppercase_first_letter(struct_name)
+          in constants.MANUALLY_ADDED_FIELDS_FROM_TEMPLATE.keys()
+      )
+  ):
+    return ""
+  elif struct_name.startswith("mjs"):
+    return f"mjs_default{struct_name.removeprefix('mjs')}"
+  elif struct_name.startswith("mjv"):
+    return f"mjv_default{struct_name.removeprefix('mjv')}"
+  else:
+    return f"mj_default{struct_name.removeprefix('mj')}"
+
+
+def _find_fields_with_init(
+    wrapped_fields: List[structs_wrappers_data.WrappedFieldData],
+) -> List[structs_wrappers_data.WrappedFieldData]:
+  """Finds the fields with initialization in the wrapped fields list."""
+  fields_with_init = []
+  for field in wrapped_fields:
+    if field.initialization:
+      fields_with_init.append(field)
+  return fields_with_init
+
+
+def use_shallow_copy(
+    wrapped_fields: List[structs_wrappers_data.WrappedFieldData],
+) -> bool:
+  """Returns true if the struct fields can be shallow copied."""
+  for field in wrapped_fields:
+    if not field.is_primitive_or_fixed_size:
+      return False
+  return True
 
 
 def build_struct_header(
     struct_name: str,
-    use_shallow_copy: bool = False,
+    wrapped_fields: List[structs_wrappers_data.WrappedFieldData],
-    fields_with_init: List[structs_wrappers_data.WrappedFieldData] = [],
-    wrapped_fields: List[structs_wrappers_data.WrappedFieldData] = [],
 ):
   """Builds the C++ header file code for a struct."""
   struct_info = introspect_structs.STRUCTS.get(struct_name)
 
   if struct_name.startswith("mjs"):
+    fields_with_init = _find_fields_with_init(wrapped_fields)
     return _build_struct_header_internal(
         struct_name,
         wrapped_fields,
         fields_with_init,
-        use_shallow_copy,
         is_mjs=True,
     )
 
@@ -117,27 +155,28 @@ def build_struct_header(
       and not _has_nested_wrapper_members(struct_info)
   ):
     return _build_struct_header_internal(
-        struct_name, wrapped_fields, [], use_shallow_copy, is_mjs=False
+        struct_name, wrapped_fields, [], is_mjs=False
     )
   return ""
 
 
 def build_struct_source(
     struct_name: str,
-    mj_default_func: str | None = None,
+    wrapped_fields: List[structs_wrappers_data.WrappedFieldData],
-    fields_with_init: List[structs_wrappers_data.WrappedFieldData] = [],
-    use_shallow_copy: bool = False,
 ):
   """Builds the C++ .cc file code for a struct."""
   wrapper_name = common.uppercase_first_letter(struct_name)
   is_mjs_struct = "Mjs" in wrapper_name
   builder = code_builder.CodeBuilder()
 
+  fields_with_init = _find_fields_with_init(wrapped_fields)
+  shallow_copy = use_shallow_copy(wrapped_fields)
+  mj_default_func = _get_default_func_name(struct_name)
+
   fields_init = ""
   if fields_with_init:
     fields_init = "".join(
-        field_with_init.initialization
+        field_with_init.initialization for field_with_init in fields_with_init
-        for field_with_init in fields_with_init
     )
   # constructor passing native ptr
   builder.line(
@@ -154,10 +193,9 @@ def build_struct_source(
       if mj_default_func:
         builder.line(f"{mj_default_func}(ptr_);")
   # copy constructor
-  if use_shallow_copy and not is_mjs_struct:
+  if shallow_copy and not is_mjs_struct:
     with builder.block(
-        f"{wrapper_name}::{wrapper_name}(const"
+        f"{wrapper_name}::{wrapper_name}(const {wrapper_name} &other)"
-        f" {wrapper_name} &other)"
         + (f" : {wrapper_name}()" if not is_mjs_struct else "")
     ):
       builder.line("*ptr_ = *other.get();")
@@ -186,12 +224,9 @@ def build_struct_source(
     with builder.block(f"{wrapper_name}::~{wrapper_name}()"):
       builder.line("if (owned_ && ptr_) delete ptr_;")
   # copy function
-  if use_shallow_copy:
+  if shallow_copy:
     with builder.block(
-        f"std::unique_ptr<{wrapper_name}>"
+        f"std::unique_ptr<{wrapper_name}> {wrapper_name}::copy()"
-        f" {wrapper_name}::copy()"
     ):
-      builder.line(
+      builder.line(f"return std::make_unique<{wrapper_name}>(*this);")
-          f"return std::make_unique<{wrapper_name}>(*this);"
-      )
   return builder.to_string()

wasm/codegen/helpers/struct_constructor_code_builder_test.py

@@ -16,37 +16,49 @@ from absl.testing import absltest
 from introspect import ast_nodes
 from wasm.codegen.helpers import struct_constructor_code_builder
 from wasm.codegen.helpers import struct_field_handler
+from wasm.codegen.helpers import structs_parser
 
 
 class StructConstructorCodeBuilderTest(absltest.TestCase):
 
   def test_constructor_code_with_default_function(self):
+    wrapped_structs = structs_parser.generate_wasm_bindings(["mjLROpt"])
     self.assertEqual(
-        struct_constructor_code_builder.build_struct_source(
+        wrapped_structs["mjLROpt"].wrapped_source,
-            "mjLROpt", "mj_defaultLROpt"
-        ),
         """
 MjLROpt::MjLROpt(mjLROpt *ptr) : ptr_(ptr) {}
 MjLROpt::MjLROpt() : ptr_(new mjLROpt) {
   owned_ = true;
   mj_defaultLROpt(ptr_);
 }
+MjLROpt::MjLROpt(const MjLROpt &other) : MjLROpt() {
+  *ptr_ = *other.get();
+}
+MjLROpt& MjLROpt::operator=(const MjLROpt &other) {
+  if (this == &other) {
+    return *this;
+  }
+  *ptr_ = *other.get();
+  return *this;
+}
 MjLROpt::~MjLROpt() {
   if (owned_ && ptr_) delete ptr_;
 }
+std::unique_ptr<MjLROpt> MjLROpt::copy() {
+  return std::make_unique<MjLROpt>(*this);
+}
 """.strip(),
     )
 
   def test_constructor_code_without_default_function(self):
+    wrapped_structs = structs_parser.generate_wasm_bindings(["mjsElement"])
     self.assertEqual(
-        struct_constructor_code_builder.build_struct_source("mjLROpt"),
+        wrapped_structs["mjsElement"].wrapped_source,
         """
-MjLROpt::MjLROpt(mjLROpt *ptr) : ptr_(ptr) {}
+MjsElement::MjsElement(mjsElement *ptr) : ptr_(ptr) {}
-MjLROpt::MjLROpt() : ptr_(new mjLROpt) {
+MjsElement::~MjsElement() {}
-  owned_ = true;
+std::unique_ptr<MjsElement> MjsElement::copy() {
-}
+  return std::make_unique<MjsElement>(*this);
-MjLROpt::~MjLROpt() {
-  if (owned_ && ptr_) delete ptr_;
 }
 """.strip(),
     )
@@ -65,7 +77,6 @@ MjLROpt::~MjLROpt() {
     self.assertEqual(
         struct_constructor_code_builder.build_struct_source(
             "mjsTexture",
-            "mjs_defaultTexture",
             [wrapped_field_data],
         ),
         """
@@ -76,9 +87,7 @@ MjsTexture::~MjsTexture() {}
 
   def test_constructor_code_with_shallow_copy(self):
     self.assertEqual(
-        struct_constructor_code_builder.build_struct_source(
+        struct_constructor_code_builder.build_struct_source("mjvLight", []),
-            "mjvLight", use_shallow_copy=True
-        ),
         """MjvLight::MjvLight(mjvLight *ptr) : ptr_(ptr) {}
 MjvLight::MjvLight() : ptr_(new mjvLight) {
   owned_ = true;
@@ -104,14 +113,14 @@ std::unique_ptr<MjvLight> MjvLight::copy() {
 
 def test_build_struct_header_with_nested_wrappers(self):
   self.assertEqual(
-      struct_constructor_code_builder.build_struct_header("mjData"),
+      struct_constructor_code_builder.build_struct_header("mjData", []),
       "",
   )
 
 
 def test_build_struct_header_basic_struct(self):
   self.assertEqual(
-      struct_constructor_code_builder.build_struct_header("mjLROpt"),
+      struct_constructor_code_builder.build_struct_header("mjLROpt", []),
       """
 struct MjLROpt {
   MjLROpt();

wasm/codegen/helpers/structs_parser.py

@@ -33,11 +33,14 @@ introspect_structs = structs.STRUCTS
 
 
 def generate_wasm_bindings(
-    wrapped_structs: Dict[str, structs_wrappers_data.WrappedStructData],
+    structs_to_bind: List[str],
 ) -> Dict[str, structs_wrappers_data.WrappedStructData]:
   """Generates WASM bindings for MuJoCo structs."""
 
-  for struct_name, wrap_data in wrapped_structs.items():
+  wrapped_structs: Dict[str, structs_wrappers_data.WrappedStructData] = {}
+  for struct_name in structs_to_bind:
+    wrapped_name = common.uppercase_first_letter(struct_name)
+
     if struct_name in introspect_structs:
       struct_fields = introspect_structs[struct_name].fields
     elif struct_name in constants.ANONYMOUS_STRUCTS:
@@ -52,35 +55,33 @@ def generate_wasm_bindings(
 
     debug_print(f"Wrapping struct: {struct_name}")
 
-    fields_with_init: List[structs_wrappers_data.WrappedFieldData] = []
+    wrapped_fields: List[structs_wrappers_data.WrappedFieldData] = []
     for field in struct_fields:
-      field_gen = struct_field_handler.StructFieldHandler(
+      wrapped_field = struct_field_handler.StructFieldHandler(
-          field, wrap_data.wrap_name
+          field, wrapped_name
       ).generate()
-      # If the struct has at least one non-primitive or fixed size field
+      wrapped_fields.append(wrapped_field)
-      # we avoid shallow copy to avoid uninitialized memory.
-      if not field_gen.is_primitive_or_fixed_size:
-        wrap_data.use_shallow_copy = False
-      if field_gen.initialization:
-        fields_with_init.append(field_gen)
-      wrap_data.wrapped_fields.append(field_gen)
 
-    wrap_data.wrapped_header = (
+    wrapped_header = struct_constructor_code_builder.build_struct_header(
-        struct_constructor_code_builder.build_struct_header(
+        struct_name,
-            struct_name,
+        wrapped_fields,
-            wrap_data.use_shallow_copy,
-            fields_with_init,
-            wrap_data.wrapped_fields,
-        )
     )
-    wrap_data.wrapped_source = (
+    wrapped_source = struct_constructor_code_builder.build_struct_source(
-        struct_constructor_code_builder.build_struct_source(
+        struct_name,
-            struct_name,
+        wrapped_fields,
-            get_default_func_name(struct_name),
-            fields_with_init,
-            wrap_data.use_shallow_copy,
-        )
     )
+    wrap_data = structs_wrappers_data.WrappedStructData(
+        wrap_name=wrapped_name,
+        wrapped_fields=wrapped_fields,
+        wrapped_header=wrapped_header,
+        wrapped_source=wrapped_source,
+        use_shallow_copy=struct_constructor_code_builder.use_shallow_copy(
+            wrapped_fields
+        ),
+    )
+
+    wrapped_structs[struct_name] = wrap_data
+
   return wrapped_structs
 
 
@@ -104,25 +105,6 @@ def _get_anonymous_struct_field(
   return target_field
 
 
-def get_default_func_name(struct_name: str) -> str:
-  """Returns the default function name for the given struct."""
-  if (
-      struct_name in constants.ANONYMOUS_STRUCTS.keys()
-      or struct_name in constants.NO_DEFAULT_CONSTRUCTORS
-      or (
-          common.uppercase_first_letter(struct_name)
-          in constants.MANUALLY_ADDED_FIELDS_FROM_TEMPLATE.keys()
-      )
-  ):
-    return ""
-  elif struct_name.startswith("mjs"):
-    return f"mjs_default{struct_name.removeprefix('mjs')}"
-  elif struct_name.startswith("mjv"):
-    return f"mjv_default{struct_name.removeprefix('mjv')}"
-  else:
-    return f"mj_default{struct_name.removeprefix('mj')}"
-
-
 def _get_field_struct_type(field_type):
   """Extracts the base struct name if the field type is a struct or pointer to a struct."""
   if isinstance(field_type, ast_nodes.ValueType):

wasm/codegen/helpers/structs_parser_test.py

@@ -25,15 +25,13 @@ class StructsParserTest(absltest.TestCase):
 
   def setUp(self):
     super().setUp()
-    self.wrapped_structs = structs_parser.generate_wasm_bindings(
+    self.wrapped_structs = structs_parser.generate_wasm_bindings([
-        structs_wrappers_data.create_wrapped_structs_set_up_data([
+        "mjModel",
-            "mjModel",
+        "mjData",
-            "mjData",
+        "mjVisualGlobal",
-            "mjVisualGlobal",
+        "mjVisualQuality",
-            "mjVisualQuality",
+        "mjVisual",
-            "mjVisual",
+    ])
-        ])
-    )
 
   def test_sort_structs_by_dependency(self):
     mock_introspect_structs = {
@@ -83,27 +81,9 @@ class StructsParserTest(absltest.TestCase):
 
   def test_generate_wasm_bindings_with_error(self):
     with self.assertRaises(RuntimeError):
-      structs_parser.generate_wasm_bindings(
+      structs_parser.generate_wasm_bindings(["mjFakeStruct"])
-          structs_wrappers_data.create_wrapped_structs_set_up_data(
-              ["mjFakeStruct"]
-          )
-      )
     with self.assertRaises(RuntimeError):
-      structs_parser.generate_wasm_bindings(
+      structs_parser.generate_wasm_bindings(["mjFakeAnonymousStruct"])
-          structs_wrappers_data.create_wrapped_structs_set_up_data(
-              ["mjFakeAnonymousStruct"]
-          )
-      )
-
-  def test_get_default_func_name_mjv(self):
-    self.assertEqual(
-        structs_parser.get_default_func_name("mjvPerturb"), "mjv_defaultPerturb"
-    )
-
-  def test_get_default_func_name(self):
-    self.assertEqual(
-        structs_parser.get_default_func_name("mjOption"), "mj_defaultOption"
-    )
 
 
 if __name__ == "__main__":

wasm/codegen/helpers/structs_wrappers_data.py

@@ -59,17 +59,3 @@ class WrappedStructData:
   # Whether to use shallow copy for this struct
   use_shallow_copy: bool = True
 
-
-def create_wrapped_structs_set_up_data(
-    struct_names: List[str],
-) -> Dict[str, WrappedStructData]:
-  """Creates a dictionary of WrappedStructData for the given struct names."""
-  return {
-      name: WrappedStructData(
-          wrap_name=common.uppercase_first_letter(name),
-          wrapped_fields=[],
-          wrapped_header="",
-          wrapped_source="",
-      )
-      for name in struct_names
-  }