Refactor MuJoCo WASM struct codegen.

PiperOrigin-RevId: 837530063
Change-Id: I57518b2fbb536f356a5e15839d8e17f455bccde9

wasm/codegen/generators/constants.py

@@ -37,7 +37,7 @@ PRIMITIVE_TYPES: Set[str] = {
     # go/keep-sorted end
 }
 
-_SKIPPED_PLUGIN_FUNCTIONS: List[str] = [
+_SKIPPED_PLUGIN_FUNCTIONS = (
     # go/keep-sorted start
     "mj_getPluginConfig",
     "mj_loadAllPluginLibraries",
@@ -59,10 +59,10 @@ _SKIPPED_PLUGIN_FUNCTIONS: List[str] = [
     "mjp_registerResourceProvider",
     "mjp_resourceProviderCount",
     # go/keep-sorted end
-]
+)
 
 # Functions that are bound as class methods
-_SKIPPED_CLASS_METHODS: List[str] = [
+_SKIPPED_CLASS_METHODS = (
     # go/keep-sorted start
     "mj_compile",
     "mj_copyData",
@@ -81,15 +81,15 @@ _SKIPPED_CLASS_METHODS: List[str] = [
     "mj_saveXML",  # TODO(manevi): Bind this function.
     "mj_saveXMLString",  # TODO(manevi): Bind this function.
     # go/keep-sorted end
-]
+)
 
 # Omitted because not very useful
-_SKIPPED_WRITABLE_ERROR: List[str] = [
+_SKIPPED_WRITABLE_ERROR = (
     "mj_printSchema",
-]
+)
 
 # Omitted thread management functions
-_SKIPPED_THREAD_FUNCTIONS: List[str] = [
+_SKIPPED_THREAD_FUNCTIONS = (
     # go/keep-sorted start
     "mju_bindThreadPool",
     "mju_defaultTask",
@@ -98,10 +98,10 @@ _SKIPPED_THREAD_FUNCTIONS: List[str] = [
     "mju_threadPoolDestroy",
     "mju_threadPoolEnqueue",
     # go/keep-sorted end
-]
+)
 
 # Omitted asset cache functions
-_SKIPPED_ASSET_CACHE_FUNCTIONS: List[str] = [
+_SKIPPED_ASSET_CACHE_FUNCTIONS = (
     # go/keep-sorted start
     "mj_clearCache",
     "mj_getCache",
@@ -109,10 +109,10 @@ _SKIPPED_ASSET_CACHE_FUNCTIONS: List[str] = [
     "mj_getCacheSize",
     "mj_setCacheCapacity",
     # go/keep-sorted end
-]
+)
 
 # Omitted Virtual Filesystem (VFS) functions
-_SKIPPED_VFS_FUNCTIONS: List[str] = [
+_SKIPPED_VFS_FUNCTIONS = (
     # go/keep-sorted start
     "mj_addBufferVFS",
     "mj_addFileVFS",
@@ -120,10 +120,10 @@ _SKIPPED_VFS_FUNCTIONS: List[str] = [
     "mj_deleteFileVFS",
     "mj_deleteVFS",
     # go/keep-sorted end
-]
+)
 
 # Omitted irrelevant visual functions
-_SKIPPED_VISUAL_FUNCTIONS: List[str] = [
+_SKIPPED_VISUAL_FUNCTIONS = (
     # go/keep-sorted start
     "mjv_averageCamera",
     "mjv_copyData",
@@ -132,9 +132,9 @@ _SKIPPED_VISUAL_FUNCTIONS: List[str] = [
     "mjv_freeScene",
     "mjv_makeScene",
     # go/keep-sorted end
-]
+)
 
-_SKIPPED_MEMORY_FUNCTIONS: List[str] = [
+_SKIPPED_MEMORY_FUNCTIONS = (
     # go/keep-sorted start
     "mj_freeLastXML",
     "mj_freeStack",
@@ -158,9 +158,9 @@ _SKIPPED_MEMORY_FUNCTIONS: List[str] = [
     "mju_warning_i",
     "mju_warning_s",
     # go/keep-sorted end
-]
+)
 
-_SKIPPED_GETTERS_AND_SETTERS: List[str] = [
+_SKIPPED_GETTERS_AND_SETTERS = (
     # go/keep-sorted start
     "mjs_appendFloatVec",
     "mjs_appendIntVec",
@@ -179,26 +179,26 @@ _SKIPPED_GETTERS_AND_SETTERS: List[str] = [
     "mjs_setStringVec",
     "mjs_setUserValue",
     # go/keep-sorted end
-]
+)
 
-_SKIPPED_UTILITY_FUNCTIONS: List[str] = [
+_SKIPPED_UTILITY_FUNCTIONS = (
     # go/keep-sorted start
     "mju_getXMLDependencies",
     # go/keep-sorted end
-]
+)
 
 # Functions that require special wrappers.
 # These functions are not bound automatically but are written by hand instead.
-MANUAL_WRAPPER_FUNCTIONS: List[str] = [
+MANUAL_WRAPPER_FUNCTIONS = (
     # go/keep-sorted start
     "mj_saveLastXML",
     "mj_setLengthRange",
     "mju_error",
     # go/keep-sorted end
-]
+)
 
 # List of functions that should be skipped during the code generation process.
-SKIPPED_FUNCTIONS: List[str] = (
+SKIPPED_FUNCTIONS = (
     _SKIPPED_CLASS_METHODS
     + _SKIPPED_THREAD_FUNCTIONS
     + _SKIPPED_MEMORY_FUNCTIONS
@@ -212,7 +212,7 @@ SKIPPED_FUNCTIONS: List[str] = (
 )
 
 # List of structs that should be skipped during the code generation process.
-SKIPPED_STRUCTS: List[str] = [
+SKIPPED_STRUCTS = (
     # go/keep-sorted start
     "mjCache",
     "mjSDF",
@@ -229,14 +229,22 @@ SKIPPED_STRUCTS: List[str] = [
     "mjuiThemeColor",
     "mjuiThemeSpacing"
     # go/keep-sorted end
-]
+)
 
-# These structs require specific function calls for creation and/or deletion,
+# Structs for which header generation is done manually.
-# or some of their fields need to be handled manually for now;
+# mjvScene is included here because buffer sizes need to be calculated based on
-# making their wrapper constructors/destructors non-trivial.
+# introspect doc strings, which was considered a brittle unreliable solution in
-MANUAL_STRUCTS: List[str] = [
+# the past.
-    "MjvScene",
+MANUAL_STRUCTS_HEADERS = (
-]
+    "mjvScene",
+)
+# Structs for which source code generation is done manually.
+MANUAL_STRUCTS_SOURCES = (
+    "mjData",
+    "mjModel",
+    "mjvScene",
+    "mjSpec",
+)
 
 # Dictionary that maps anonymous structs to their parent struct and field name.
 # Anonymous structs are not defined as independent structs in the MuJoCo
@@ -262,7 +270,7 @@ STRUCTS_TO_BIND: List[str] = list(
 )
 
 # List of structs that do not have a default constructor.
-NO_DEFAULT_CONSTRUCTORS: List[str] = [
+NO_DEFAULT_CONSTRUCTORS = (
     # go/keep-sorted start
     "mjContact",
     "mjSolverStat",
@@ -277,12 +285,12 @@ NO_DEFAULT_CONSTRUCTORS: List[str] = [
     "mjvGLCamera",
     "mjvLight",
     # go/keep-sorted end
-]
+)
 
 # List of `mjData` fields where the array size should be obtained from other
 # `mjData` members, instead of from `mjModel` members. This is typically the
 # case for fields that are dynamically allocated during the simulation.
-MJDATA_SIZES: List[str] = [
+MJDATA_SIZES = (
     # go/keep-sorted start
     "contact",
     "efc_AR",
@@ -352,7 +360,7 @@ MJDATA_SIZES: List[str] = [
     "map_efc2iefc",
     "map_iefc2efc",
     # go/keep-sorted end
-]
+)
 
 # Fields that should be entirely omitted from the bindings.
 SKIPPED_FIELDS: Dict[str, List[str]] = {}

wasm/codegen/generators/structs.py

@@ -17,7 +17,7 @@
 import collections
 import dataclasses
 import math
-from typing import Dict, List, Tuple, Union, cast
+from typing import Tuple, Union, cast
 
 from introspect import ast_nodes
 from introspect import structs as introspect_structs
@@ -34,8 +34,8 @@ class WrappedFieldData:
   # Line for struct field binding
   binding: str = ""
 
-  # Line for struct field definition
+  # Line for struct field declaration
-  definition: str = ""
+  declaration: str = ""
 
   # Used for constructor fields that are class types and need to be initialized
   # with a pointer to the corresponding member in the native struct
@@ -53,27 +53,7 @@ class WrappedFieldData:
   typename: str = ""
 
 
-@dataclasses.dataclass
+def _get_property_binding(
-class WrappedStructData:
-  """Data class for struct wrapper definition and binding."""
-
-  # Name of wrapper struct
-  wrap_name: str
-
-  # List of WrappedFieldData for this struct
-  wrapped_fields: List[WrappedFieldData]
-
-  # Struct header code
-  wrapped_header: str
-
-  # Struct source code
-  wrapped_source: str
-
-  # Struct bindings code
-  bindings: str = ""
-
-
-def _simple_property_binding(
     field: ast_nodes.StructFieldDecl,
     struct_wrapper_name: str,
     setter: bool = False,
@@ -99,15 +79,15 @@ def _generate_field_data(
     # Note: Manually handled MjModel fields are special cased so that a
     # by-reference embind return value policy is used.
     return WrappedFieldData(
-        typename=_get_field_struct_type(f.type),
+        typename=_get_field_struct_type(f, s),
-        definition=f"// {f.name} field is handled manually in template file struct declaration",  # pylint: disable=line-too-long
+        declaration=f"// {f.name} field is handled manually in template file struct declaration",  # pylint: disable=line-too-long
-        binding=_simple_property_binding(f, w, reference=(w == "MjModel")),
+        binding=_get_property_binding(f, w, reference=(w == "MjModel")),
     )
 
   if f.name in constants.SKIPPED_FIELDS.get(w, []):
     return WrappedFieldData(
         typename="",
-        definition=f"// {f.name} field is skipped.",
+        declaration=f"// {f.name} field is skipped.",
         binding=f"// {f.name} field is skipped.",
     )
 
@@ -122,45 +102,32 @@ def _generate_field_data(
       builder.line(f"ptr_->{f.name} = value;")
 
     return WrappedFieldData(
-        definition=builder.to_string(),
+        declaration=builder.to_string(),
-        typename=_get_field_struct_type(f.type),
+        typename=_get_field_struct_type(f, s),
-        binding=_simple_property_binding(f, w, setter=True, reference=True),
+        binding=_get_property_binding(f, w, setter=True, reference=True),
         is_primitive_or_fixed_size=True,
     )
 
   elif isinstance(f.type, ast_nodes.ValueType) and f.type.name.startswith("mj"):
     return WrappedFieldData(
-        definition=f"{common.capitalize(f.type.name)} {f.name};",
+        declaration=f"{common.capitalize(f.type.name)} {f.name};",
-        typename=_get_field_struct_type(f.type),
+        typename=_get_field_struct_type(f, s),
-        binding=_simple_property_binding(f, w, setter=False, reference=True),
+        binding=_get_property_binding(f, w, setter=False, reference=True),
         ptr_initialization=f"{f.name}(&ptr_->{f.name})",
         ptr_copy_reset=f"{f.name}.set(&ptr_->{f.name});",
         is_primitive_or_fixed_size=True,
     )
 
   elif isinstance(f.type, ast_nodes.AnonymousStructDecl):
-
+    anonymous_struct_name = _get_field_struct_type(f, s)
-    anonymous_struct_name = ""
+    return WrappedFieldData(
-    for name, value in constants.ANONYMOUS_STRUCTS.items():
+        binding=_get_property_binding(f, w, setter=False, reference=True),
-      if value["parent"] == s and value["field_name"] == f.name:
+        typename=anonymous_struct_name,
-        anonymous_struct_name = name
+        declaration=f"{common.capitalize(anonymous_struct_name)} {f.name};",
-        break
+        ptr_initialization=f"{f.name}(&ptr_->{f.name})",
-
+        ptr_copy_reset=f"{f.name}.set(&ptr_->{f.name});",
-    if not anonymous_struct_name:
+        is_primitive_or_fixed_size=True,
-      raise RuntimeError(
+    )
-          f"Anonymous struct for field {f.name} in {w} not found in"
-          " ANONYMOUS_STRUCTS."
-      )
-
-    if anonymous_struct_name in constants.STRUCTS_TO_BIND:
-      return WrappedFieldData(
-          binding=_simple_property_binding(f, w, setter=False, reference=True),
-          typename=anonymous_struct_name,
-          definition=f"{common.capitalize(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,
-      )
 
   elif isinstance(f.type, ast_nodes.ArrayType):
 
@@ -191,19 +158,19 @@ def _generate_field_data(
         )
 
       return WrappedFieldData(
-          definition=builder.to_string(),
+          declaration=builder.to_string(),
-          typename=_get_field_struct_type(f.type),
+          typename=_get_field_struct_type(f, s),
-          binding=_simple_property_binding(f, w),
+          binding=_get_property_binding(f, w),
           is_primitive_or_fixed_size=True,
       )
     elif inner_type_name.startswith("mj"):
       return WrappedFieldData(
-          definition=(
+          declaration=(
               f"std::vector<{common.capitalize(inner_type_name)}> {f.name};"
           ),
           ptr_initialization=f"{f.name}(&ptr_->{f.name})",
-          typename=_get_field_struct_type(f.type),
+          typename=_get_field_struct_type(f, s),
-          binding=_simple_property_binding(f, w, reference=True),
+          binding=_get_property_binding(f, w, reference=True),
       )
 
   elif isinstance(f.type, ast_nodes.PointerType):
@@ -212,21 +179,10 @@ def _generate_field_data(
         if isinstance(f.type.inner_type, ast_nodes.ValueType)
         else ""
     )
-    ptr_field_expr = f"ptr_->{f.name}"
+    is_dynamically_sized = bool(f.array_extent)
-    array_size_str = ""
-
-    if f.array_extent:
-      array_size_str = parse_array_extent(f.array_extent, w, f.name)
-    elif f.name in constants.BYTE_FIELDS.keys():
-      # for byte fields, we need to cast the pointer to uint8_t*
-      # so embind can correctly interpret the memory view
-      ptr_field_expr = f"static_cast<uint8_t*>({ptr_field_expr})"
-      # for these byte fields, there is no array_extent, so we add the size of
-      # in the config file based in the documentation
-      extent = (constants.BYTE_FIELDS[f.name]["size"],)
-      array_size_str = parse_array_extent(extent, w, f.name)
-    elif inner_type_name == "mjString":
 
+    # Case 1: mjString
+    if inner_type_name == "mjString":
       builder = code_builder.CodeBuilder()
       with builder.function(f"mjString {f.name}() const"):
         builder.line(
@@ -235,12 +191,13 @@ def _generate_field_data(
       with builder.function(f"void set_{f.name}(const mjString& value)"):
         with builder.block(f"if (ptr_ && ptr_->{f.name})"):
           builder.line(f"*(ptr_->{f.name}) = value;")
-
       return WrappedFieldData(
-          definition=builder.to_string(),
+          declaration=builder.to_string(),
-          typename=_get_field_struct_type(f.type),
+          typename=_get_field_struct_type(f, s),
-          binding=_simple_property_binding(f, w, setter=True, reference=True),
+          binding=_get_property_binding(f, w, setter=True, reference=True),
       )
+
+    # Case 2: mj*Vec
     elif inner_type_name.startswith("mj") and inner_type_name.endswith("Vec"):
       ptr_field_expr_vec = f"*(ptr_->{f.name})"
       vector_type = inner_type_name
@@ -249,47 +206,62 @@ def _generate_field_data(
         ptr_field_expr_vec = (
             f"*(reinterpret_cast<std::vector<uint8_t>*>(ptr_->{f.name}))"
         )
-
       builder = code_builder.CodeBuilder()
       with builder.function(f"{vector_type} &{f.name}() const"):
         builder.line(f"return {ptr_field_expr_vec};")
-
       return WrappedFieldData(
-          definition=builder.to_string(),
+          declaration=builder.to_string(),
-          typename=_get_field_struct_type(f.type),
+          typename=_get_field_struct_type(f, s),
-          binding=_simple_property_binding(f, w, setter=False, reference=True),
+          binding=_get_property_binding(f, w, setter=False, reference=True),
       )
 
-    if (
+    # Case 3: Non-dynamically sized pointer fields to other structs.
-        inner_type_name.startswith("mj")
+    elif (
+        not is_dynamically_sized
         and inner_type_name not in constants.PRIMITIVE_TYPES
-        and not f.array_extent
         and w not in constants.MANUAL_FIELDS.keys()
     ):
+      # These are wrapped as direct members of the wrapper class, initialized
+      # with a pointer to the corresponding member in the native struct.
       ptr_field = cast(ast_nodes.PointerType, f.type)
       wrapper_field_name = common.capitalize(
           cast(ast_nodes.ValueType, ptr_field.inner_type).name
       )
       return WrappedFieldData(
-          definition=f"{wrapper_field_name} {f.name};",
+          declaration=f"{wrapper_field_name} {f.name};",
-          typename=_get_field_struct_type(f.type),
+          typename=_get_field_struct_type(f, s),
-          binding=_simple_property_binding(f, w, setter=False, reference=True),
+          binding=_get_property_binding(f, w, setter=False, reference=True),
           ptr_initialization=f"{f.name}(ptr_->{f.name})",
       )
 
-    builder = code_builder.CodeBuilder()
+    # Case 4: Dynamically sized pointer fields and other pointer types that are
-    with builder.function(f"emscripten::val {f.name}() const"):
+    # exposed as emscripten::typed_memory_view.
-      builder.line(
+    else:
-          "return"
+      ptr_field_expr = f"ptr_->{f.name}"
-          f" emscripten::val(emscripten::typed_memory_view({array_size_str},"
+      array_size_str = ""
-          f" {ptr_field_expr}));"
+      if is_dynamically_sized:
-      )
+        array_size_str = parse_array_extent(f.array_extent, w, f.name)
+      elif f.name in constants.BYTE_FIELDS.keys():
+        # For byte fields, we need to cast the pointer to uint8_t*
+        # so embind can correctly interpret the memory view
+        ptr_field_expr = f"static_cast<uint8_t*>({ptr_field_expr})"
+        # Byte fields lack `array_extent`, so their size is defined in
+        # `constants.BYTE_FIELDS` based on the MuJoCo documentation.
+        extent = (constants.BYTE_FIELDS[f.name]["size"],)
+        array_size_str = parse_array_extent(extent, w, f.name)
 
-    return WrappedFieldData(
+      builder = code_builder.CodeBuilder()
-        definition=builder.to_string(),
+      with builder.function(f"emscripten::val {f.name}() const"):
-        typename=_get_field_struct_type(f.type),
+        builder.line(
-        binding=_simple_property_binding(f, w),
+            "return"
-    )
+            f" emscripten::val(emscripten::typed_memory_view({array_size_str},"
+            f" {ptr_field_expr}));"
+        )
+      return WrappedFieldData(
+          declaration=builder.to_string(),
+          typename=_get_field_struct_type(f, s),
+          binding=_get_property_binding(f, w),
+      )
 
   raise RuntimeError(f"Field {f.name} from struct {w} not properly handled")
 
@@ -322,8 +294,8 @@ def _delete_ptr_statement(struct_name: str) -> str:
 
 
 def _find_member_inits(
-    wrapped_fields: List[WrappedFieldData],
+    wrapped_fields: list[WrappedFieldData],
-) -> List[WrappedFieldData]:
+) -> list[WrappedFieldData]:
   """Finds the fields with ptr_initialization in the wrapped fields list."""
   member_inits = []
   for field in wrapped_fields:
@@ -333,7 +305,7 @@ def _find_member_inits(
 
 
 def use_shallow_copy(
-    wrapped_fields: List[WrappedFieldData],
+    wrapped_fields: list[WrappedFieldData],
 ) -> bool:
   """Returns true if the struct fields can be shallow copied."""
   for field in wrapped_fields:
@@ -344,15 +316,12 @@ def use_shallow_copy(
 
 def build_struct_header(
     struct_name: str,
-    wrapped_fields: List[WrappedFieldData],
+    wrapped_fields: list[WrappedFieldData],
 ):
   """Builds the C++ header file code for a struct."""
   s = struct_name
   w = common.capitalize(s)
 
-  if w in constants.MANUAL_STRUCTS:
-    return ""
-
   if (
       s not in constants.ANONYMOUS_STRUCTS
       and s not in introspect_structs.STRUCTS
@@ -400,8 +369,8 @@ def build_struct_header(
 
     # field declarations
     for field in wrapped_fields:
-      if field.definition and field not in member_inits:
+      if field.declaration and field not in member_inits:
-        for line in field.definition.splitlines():
+        for line in field.declaration.splitlines():
           builder.line(line)
 
     # define private struct members
@@ -414,8 +383,8 @@ def build_struct_header(
     if member_inits:
       builder.public()
       for f in member_inits:
-        if f.definition:
+        if f.declaration:
-          builder.line(f"{f.definition}")
+          builder.line(f"{f.declaration}")
 
       if w == "MjData":
         builder.line("mjModel *model;")
@@ -425,19 +394,9 @@ def build_struct_header(
 
 def build_struct_source(
     struct_name: str,
-    wrapped_fields: List[WrappedFieldData],
+    wrapped_fields: list[WrappedFieldData],
 ):
   """Builds the C++ .cc file code for a struct."""
-  # These structs require specific function calls for creation and/or deletion
-  # which, for now, are hardcoded in the template file.
-  if struct_name in [
-      "mjData",
-      "mjModel",
-      "mjvScene",
-      "mjSpec",
-  ]:
-    return ""
-
   s = struct_name
   w = common.capitalize(s)
   is_mjs = w.startswith("Mjs")
@@ -505,7 +464,7 @@ def build_struct_source(
 
 def _build_struct_bindings(
     struct_name: str,
-    wrapped_fields: List[WrappedFieldData],
+    wrapped_fields: list[WrappedFieldData],
 ):
   """Builds the C++ bindings for a struct."""
   w = common.capitalize(struct_name)
@@ -536,12 +495,15 @@ def _build_struct_bindings(
     if shallow_copy and not is_mjs:
       builder.line(f'.function("copy", &{w}::copy, take_ownership())')
 
-    wrapped_fields.sort(key=lambda field: field.binding)
+    bindings = sorted(
-    for field in wrapped_fields[:-1]:
+        [f.binding for f in wrapped_fields if f.binding.startswith(".")]
-      if field.binding:
+    )
-        builder.line(field.binding)
+    for binding in bindings[:-1]:
-    if wrapped_fields:
+      builder.line(binding)
-      builder.line(f"{wrapped_fields[-1].binding};")
+    if bindings:
+      builder.line(f"{bindings[-1]};")
+    elif builder._lines:
+      builder._lines[-1] += ";"
 
   return builder.to_string()
 
@@ -585,45 +547,6 @@ def resolve_extent(
   return f"{var_name}->{extent}"
 
 
-def generate_wasm_bindings(
-    structs_to_bind: List[str],
-) -> Dict[str, WrappedStructData]:
-  """Generates WASM bindings for MuJoCo structs."""
-
-  wrapped_structs: Dict[str, WrappedStructData] = {}
-  for struct_name in structs_to_bind:
-    s = struct_name
-    w = common.capitalize(s)
-
-    if s in introspect_structs.STRUCTS:
-      struct_fields = introspect_structs.STRUCTS[s].fields
-    elif s in constants.ANONYMOUS_STRUCTS:
-      anonymous_struct = _get_anonymous_struct_field(s)
-      if not anonymous_struct or not isinstance(
-          anonymous_struct.type, ast_nodes.AnonymousStructDecl
-      ):
-        raise RuntimeError(f"Anonymous struct not found: {s}")
-      struct_fields = anonymous_struct.type.fields
-    else:
-      raise RuntimeError(f"Struct not found: {s}")
-
-    wrapped_fields: List[WrappedFieldData] = []
-    for field in struct_fields:
-      wrapped_fields.append(_generate_field_data(field, w))
-
-    wrap_data = WrappedStructData(
-        wrap_name=w,
-        wrapped_fields=wrapped_fields,
-        wrapped_header=build_struct_header(s, wrapped_fields),
-        wrapped_source=build_struct_source(s, wrapped_fields),
-        bindings=_build_struct_bindings(s, wrapped_fields),
-    )
-
-    wrapped_structs[s] = wrap_data
-
-  return wrapped_structs
-
-
 def _get_anonymous_struct_field(
     anonymous_structs_key: str,
 ) -> ast_nodes.StructFieldDecl | None:
@@ -644,30 +567,64 @@ def _get_anonymous_struct_field(
   return target_field
 
 
-def _get_field_struct_type(field_type):
+def _get_field_struct_type(
+    field: ast_nodes.StructFieldDecl, struct_name: str
+) -> str | None:
   """Extracts the base struct name if the field type is a struct or pointer to a struct."""
-  if isinstance(field_type, ast_nodes.ValueType):
+  s = struct_name
-    return field_type.name
+  w = common.capitalize(s)
-  if isinstance(field_type, ast_nodes.PointerType):
+  if isinstance(field.type, ast_nodes.AnonymousStructDecl):
-    if isinstance(field_type.inner_type, ast_nodes.ValueType):
+    anonymous_struct_name = ""
-      return field_type.inner_type.name
+    for name, value in constants.ANONYMOUS_STRUCTS.items():
+      if value["parent"] == s and value["field_name"] == field.name:
+        anonymous_struct_name = name
+        break
+
+    if not anonymous_struct_name:
+      raise RuntimeError(
+          f"Anonymous struct for field {field.name} in {w} not found in"
+          " ANONYMOUS_STRUCTS."
+      )
+    return anonymous_struct_name
+  elif isinstance(field.type, ast_nodes.ValueType):
+    return field.type.name
+  if isinstance(field.type, ast_nodes.PointerType):
+    if isinstance(field.type.inner_type, ast_nodes.ValueType):
+      return field.type.inner_type.name
   return None
 
 
+def get_introspect_struct_fields(struct_name: str):
+  """Retrieves the fields of a struct from the introspect data."""
+  s = struct_name
+
+  if s in introspect_structs.STRUCTS:
+    return introspect_structs.STRUCTS[s].fields
+  elif s in constants.ANONYMOUS_STRUCTS:
+    anonymous_struct = _get_anonymous_struct_field(s)
+    if not anonymous_struct or not isinstance(
+        anonymous_struct.type, ast_nodes.AnonymousStructDecl
+    ):
+      raise RuntimeError(f"Anonymous struct not found: {s}")
+    return anonymous_struct.type.fields
+  else:
+    raise RuntimeError(f"Struct not found: {s}")
+
+
 def sort_structs_by_dependency(
-    struct_wrappers: dict[str, WrappedStructData],
+    struct_with_fields: dict[str, list[WrappedFieldData]],
-) -> List[str]:
+) -> list[str]:
   """Sorts structs based on their field dependencies using topological sort.
 
-  Structs with no dependencies on other structs in the list come first.
+    Structs with no dependencies on other structs in the list come first.
   Struct A has a dependency on struct B if struct A has a field where the
   underlying_type is B. Note that this definition is stricter than the C++
   struct dependency criterion where forward declarations can be used to
   eliminate dependencies A and B if A only has a pointer to B.
 
   Args:
-    struct_wrappers: A dictionary mapping struct names to their
+    struct_with_fields: A dictionary mapping struct names to their
-      WrappedStructData.
+      WrappedFieldData.
 
   Returns:
     A new list of struct names sorted by dependency.
@@ -677,25 +634,26 @@ def sort_structs_by_dependency(
   """
   adj = collections.defaultdict(list)
   in_degree = collections.defaultdict(int)
-  struct_names = struct_wrappers.keys()
+  struct_names = struct_with_fields.keys()
   struct_set = set(struct_names)
   sorted_struct_names = sorted(struct_names)
 
-  for struct_name in sorted_struct_names:
+  for s in sorted_struct_names:
-    if struct_name == "mjData":
+    if s == "mjData":
       adj["mjModel"].append("mjData")
       in_degree["mjData"] += 1
-    for field in struct_wrappers[struct_name].wrapped_fields:
 
+    fields = struct_with_fields[s]
+    for field in fields:
       field_type_name = field.typename
       if (
           field_type_name
-          and field_type_name != struct_name
+          and field_type_name != s
           and field_type_name in struct_set
       ):
-        if struct_name not in adj[field_type_name]:
+        if s not in adj[field_type_name]:
-          adj[field_type_name].append(struct_name)
+          adj[field_type_name].append(s)
-          in_degree[struct_name] += 1
+          in_degree[s] += 1
 
   queue = collections.deque(
       [name for name in sorted_struct_names if in_degree[name] == 0]
@@ -720,14 +678,8 @@ def sort_structs_by_dependency(
     )
 
 
-def generate(struct_to_bind: List[str]) -> list[tuple[str, list[str]]]:
+def generate(struct_to_bind: list[str]) -> list[tuple[str, list[str]]]:
   """Generates C++ header file for binding and wrapping MuJoCo structs."""
-
-  # Traverse the introspect dictionary to get the field
-  # wrapper/bindings statements set up for each struct
-  structs_to_bind_data = generate_wasm_bindings(struct_to_bind)
-
-  markers_and_content = []
   typedefs = []
   for type_name in sorted(constants.ANONYMOUS_STRUCTS):
     s = constants.ANONYMOUS_STRUCTS[type_name]
@@ -735,48 +687,63 @@ def generate(struct_to_bind: List[str]) -> list[tuple[str, list[str]]]:
         f"using {type_name} = decltype(::{s['parent']}::{s['field_name']});"
     )
 
-  # Sort by struct name by dependency to ensure deterministic output order
+  wrapped_structs_with_fields: dict[str, list[WrappedFieldData]] = {}
+  for s in struct_to_bind:
+    fields: list[WrappedFieldData] = []
+    introspect_fields = get_introspect_struct_fields(s)
+    for field in introspect_fields:
+      fields.append(_generate_field_data(field, common.capitalize(s)))
+    wrapped_structs_with_fields[s] = fields
+
   dependency_sorted_struct_names = sort_structs_by_dependency(
-      structs_to_bind_data
+      wrapped_structs_with_fields
   )
 
-  structs_header = []
+  headers = []
   for s in dependency_sorted_struct_names:
-    struct_data = structs_to_bind_data[s]
+    fields = wrapped_structs_with_fields[s]
-    if struct_data.wrapped_header:
-      structs_header.append(struct_data.wrapped_header + "\n")
-    else:
-      definitions: list[str] = []
-      for f in sorted(struct_data.wrapped_fields, key=lambda f: f.definition):
-        if f.definition:
-          definitions.append(f.definition)
-      markers_and_content.append((
-          f"// INSERT-GENERATED-{struct_data.wrap_name}-DEFINITIONS",
-          definitions,
-      ))
 
-  structs_source = []
+    if s not in constants.MANUAL_STRUCTS_HEADERS:
+      header = build_struct_header(s, fields)
+      headers.append(header + "\n")
+
+  sources = []
   for s in dependency_sorted_struct_names:
-    struct_data = structs_to_bind_data[s]
+    fields = wrapped_structs_with_fields[s]
-    if struct_data.wrapped_source:
+    if s not in constants.MANUAL_STRUCTS_SOURCES:
-      structs_source.append(struct_data.wrapped_source + "\n")
+      source = build_struct_source(s, fields)
+      sources.append(source + "\n")
+
+  bindings = []
+  alphabetically_sorted_struct_names = sorted(
+      wrapped_structs_with_fields.keys()
+  )
+  for s in alphabetically_sorted_struct_names:
+    fields = wrapped_structs_with_fields[s]
+    bindings.append(_build_struct_bindings(s, fields))
 
-  struct_bindings = []
-  alphabetically_sorted_struct_names = sorted(structs_to_bind_data.keys())
   for s in alphabetically_sorted_struct_names:
-    struct_data = structs_to_bind_data[s]
-    struct_bindings.append(struct_data.bindings)
-  for s in sorted(struct_to_bind):
     w = common.capitalize(s)
     if w.startswith("Mjs") or w == "MjSpec":
-      struct_bindings.append(f"emscripten::register_optional<{w}>();")
+      bindings.append(f"emscripten::register_optional<{w}>();")
 
-  # Combine all the markers and content into a single list
+  manual_struct_field_declarations = []
-  markers_and_content += [
+  for s in dependency_sorted_struct_names:
+    w = common.capitalize(s)
+    fields = wrapped_structs_with_fields[s]
+    if s in constants.MANUAL_STRUCTS_HEADERS:
+      decls: list[str] = []
+      for f in sorted(fields, key=lambda f: f.declaration):
+        if f.declaration:
+          decls.append(f.declaration)
+      manual_struct_field_declarations.append(
+          (f"// INSERT-GENERATED-{w}-DECLARATION", decls)
+      )
+
+  return [
+      *manual_struct_field_declarations,
       ("// {{ ANONYMOUS_STRUCT_TYPEDEFS }}", typedefs),
-      ("// {{ STRUCTS_HEADER }}", structs_header),
+      ("// {{ STRUCTS_HEADER }}", headers),
-      ("// {{ STRUCTS_SOURCE }}", structs_source),
+      ("// {{ STRUCTS_SOURCE }}", sources),
-      ("// {{ STRUCTS_BINDINGS }}", struct_bindings),
+      ("// {{ STRUCTS_BINDINGS }}", bindings),
   ]
-
-  return markers_and_content

wasm/codegen/templates/bindings.cc

@@ -199,7 +199,7 @@ struct MjvScene {
     return emscripten::val(emscripten::typed_memory_view(
         6 * MjvScene::GetSumFlexFaces(), ptr_->flextexcoord));
   }
-  // INSERT-GENERATED-MjvScene-DEFINITIONS
+  // INSERT-GENERATED-MjvScene-DECLARATION
 
  private:
   mjvScene* ptr_;

wasm/codegen/tests/generators_test.py

@@ -203,58 +203,6 @@ class FunctionUtilsTest(absltest.TestCase):
 
 class StructConstructorCodeBuilderTest(absltest.TestCase):
 
-  def test_constructor_code_with_default_function(self):
-    wrapped_structs = structs.generate_wasm_bindings(["mjLROpt"])
-    self.assertEqual(
-        wrapped_structs["mjLROpt"].wrapped_source,
-        """
-MjLROpt::MjLROpt(mjLROpt *ptr) : ptr_(ptr) {}
-MjLROpt::~MjLROpt() {
-  if (owned_ && ptr_) {
-    delete 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;
-}
-std::unique_ptr<MjLROpt> MjLROpt::copy() {
-  return std::make_unique<MjLROpt>(*this);
-}
-mjLROpt* MjLROpt::get() const {
-  return ptr_;
-}
-void MjLROpt::set(mjLROpt* ptr) {
-  ptr_ = ptr;
-}
-""".strip(),
-    )
-
-  def test_constructor_code_without_default_function(self):
-    wrapped_structs = structs.generate_wasm_bindings(["mjsElement"])
-    self.assertEqual(
-        wrapped_structs["mjsElement"].wrapped_source,
-        """
-MjsElement::MjsElement(mjsElement *ptr) : ptr_(ptr) {}
-mjsElement* MjsElement::get() const {
-  return ptr_;
-}
-void MjsElement::set(mjsElement* ptr) {
-  ptr_ = ptr;
-}
-""".strip(),
-    )
-
   def test_constructor_code_with_fields_with_init(self):
     field_with_init = ast_nodes.StructFieldDecl(
         name="element",
@@ -354,7 +302,7 @@ class StructFieldCodeBuilderTest(absltest.TestCase):
         doc="number of geoms",
     )
     self.assertEqual(
-        structs.build_primitive_type_definition(field),
+        structs._generate_field_data(field, "ngeom").declaration,
         """
 int ngeom() const {
   return ptr_->ngeom;
@@ -375,9 +323,7 @@ void set_ngeom(int value) {
         array_extent=("ngeom", 4),
     )
     self.assertEqual(
-        structs.build_memory_view_definition(
+        structs._generate_field_data(field, "geom_rgba").declaration,
-            field, "ptr_->ngeom * 4", "ptr_->geom_rgba"
-        ),
         """
 emscripten::val geom_rgba() const {
   return emscripten::val(emscripten::typed_memory_view(ptr_->ngeom * 4, ptr_->geom_rgba));
@@ -394,7 +340,7 @@ emscripten::val geom_rgba() const {
         doc="rgba when material is omitted",
     )
     self.assertEqual(
-        structs.build_string_field_definition(field),
+        structs._generate_field_data(field, "MjString").declaration,
         """
 mjString string_field() const {
   return (ptr_ && ptr_->string_field) ? *(ptr_->string_field) : "";
@@ -416,7 +362,7 @@ void set_string_field(const mjString& value) {
         doc="",
     )
     self.assertEqual(
-        structs.build_mjvec_pointer_definition(field, "mjDoubleVec"),
+        structs._generate_field_data(field, "mjDoubleVec").declaration,
         """
 mjDoubleVec &vector_field() const {
   return *(ptr_->vector_field);
@@ -432,174 +378,43 @@ mjDoubleVec &vector_field() const {
         doc="",
     )
     self.assertEqual(
-        structs.build_mjvec_pointer_definition(field, "mjByteVec"),
+        structs._generate_field_data(field, "mjByteVec").declaration,
         """
 std::vector<uint8_t> &vector_field() const {
   return *(reinterpret_cast<std::vector<uint8_t>*>(ptr_->vector_field));
 }""".strip(),
     )
 
-  def test_simple_property_binding(self):
+  def test_get_property_binding(self):
     field = ast_nodes.StructFieldDecl(
         name="ngeom",
         type=ast_nodes.ValueType(name="int"),
         doc="number of geoms",
     )
     self.assertEqual(
-        structs.build_simple_property_binding(field, "MjModel"),
+        structs._get_property_binding(field, "MjModel"),
         '.property("ngeom", &MjModel::ngeom)',
     )
 
-  def test_simple_property_binding_with_setter(self):
+  def test_get_property_binding_with_setter(self):
     field = ast_nodes.StructFieldDecl(
         name="ngeom",
         type=ast_nodes.ValueType(name="int"),
         doc="",
     )
     self.assertEqual(
-        structs.build_simple_property_binding(field, "MjModel", True),
+        structs._get_property_binding(field, "MjModel", True),
         '.property("ngeom", &MjModel::ngeom, &MjModel::set_ngeom)',
     )
 
-  def test_simple_property_binding_with_return_value_policy_as_ref(self):
+  def test_get_property_binding_with_return_value_policy_as_ref(self):
     field = ast_nodes.StructFieldDecl(
         name="ngeom",
         type=ast_nodes.ValueType(name="int"),
         doc="",
     )
     self.assertEqual(
-        structs.build_simple_property_binding(
+        structs._get_property_binding(
-            field,
-            "MjModel",
-            add_setter=True,
-            add_return_value_policy_as_ref=True,
-        ),
-        '.property("ngeom", &MjModel::ngeom, &MjModel::set_ngeom, reference())',
-    )
-
-
-class StructFieldCodeBuilderTest(absltest.TestCase):
-
-  def test_primitive_type_definition(self):
-    field = ast_nodes.StructFieldDecl(
-        name="ngeom",
-        type=ast_nodes.ValueType(name="int"),
-        doc="number of geoms",
-    )
-    self.assertEqual(
-        structs._generate_field_data(field, "ngeom").definition,
-        """
-int ngeom() const {
-  return ptr_->ngeom;
-}
-void set_ngeom(int value) {
-  ptr_->ngeom = value;
-}
-""".strip(),
-    )
-
-  def test_memory_view_definition(self):
-    field = ast_nodes.StructFieldDecl(
-        name="geom_rgba",
-        type=ast_nodes.PointerType(
-            inner_type=ast_nodes.ValueType(name="float"),
-        ),
-        doc="rgba when material is omitted",
-        array_extent=("ngeom", 4),
-    )
-    self.assertEqual(
-        structs._generate_field_data(field, "geom_rgba").definition,
-        """
-emscripten::val geom_rgba() const {
-  return emscripten::val(emscripten::typed_memory_view(ptr_->ngeom * 4, ptr_->geom_rgba));
-}
-""".strip(),
-    )
-
-  def test_string_field_definition(self):
-    field = ast_nodes.StructFieldDecl(
-        name="string_field",
-        type=ast_nodes.PointerType(
-            inner_type=ast_nodes.ValueType(name="mjString"),
-        ),
-        doc="rgba when material is omitted",
-    )
-    self.assertEqual(
-        structs._generate_field_data(field, "MjString").definition,
-        """
-mjString string_field() const {
-  return (ptr_ && ptr_->string_field) ? *(ptr_->string_field) : "";
-}
-void set_string_field(const mjString& value) {
-  if (ptr_ && ptr_->string_field) {
-    *(ptr_->string_field) = value;
-  }
-}
-""".strip(),
-    )
-
-  def test_mjvec_pointer_definition(self):
-    field = ast_nodes.StructFieldDecl(
-        name="vector_field",
-        type=ast_nodes.PointerType(
-            inner_type=ast_nodes.ValueType(name="mjDoubleVec"),
-        ),
-        doc="",
-    )
-    self.assertEqual(
-        structs._generate_field_data(field, "mjDoubleVec").definition,
-        """
-mjDoubleVec &vector_field() const {
-  return *(ptr_->vector_field);
-}""".strip(),
-    )
-
-  def test_mjbyte_vec_pointer_definition(self):
-    field = ast_nodes.StructFieldDecl(
-        name="vector_field",
-        type=ast_nodes.PointerType(
-            inner_type=ast_nodes.ValueType(name="mjByteVec"),
-        ),
-        doc="",
-    )
-    self.assertEqual(
-        structs._generate_field_data(field, "mjByteVec").definition,
-        """
-std::vector<uint8_t> &vector_field() const {
-  return *(reinterpret_cast<std::vector<uint8_t>*>(ptr_->vector_field));
-}""".strip(),
-    )
-
-  def test_simple_property_binding(self):
-    field = ast_nodes.StructFieldDecl(
-        name="ngeom",
-        type=ast_nodes.ValueType(name="int"),
-        doc="number of geoms",
-    )
-    self.assertEqual(
-        structs._simple_property_binding(field, "MjModel"),
-        '.property("ngeom", &MjModel::ngeom)',
-    )
-
-  def test_simple_property_binding_with_setter(self):
-    field = ast_nodes.StructFieldDecl(
-        name="ngeom",
-        type=ast_nodes.ValueType(name="int"),
-        doc="",
-    )
-    self.assertEqual(
-        structs._simple_property_binding(field, "MjModel", True),
-        '.property("ngeom", &MjModel::ngeom, &MjModel::set_ngeom)',
-    )
-
-  def test_simple_property_binding_with_return_value_policy_as_ref(self):
-    field = ast_nodes.StructFieldDecl(
-        name="ngeom",
-        type=ast_nodes.ValueType(name="int"),
-        doc="",
-    )
-    self.assertEqual(
-        structs._simple_property_binding(
             field,
             "MjModel",
             setter=True,
@@ -621,7 +436,7 @@ class StructFieldHandlerTest(absltest.TestCase):
 
     wrapped_field_data = structs._generate_field_data(field_scalar, "MjModel")
     self.assertEqual(
-        wrapped_field_data.definition,
+        wrapped_field_data.declaration,
         """
 int ngeom() const {
   return ptr_->ngeom;
@@ -649,7 +464,7 @@ void set_ngeom(int value) {
     wrapped_field_data = structs._generate_field_data(field, "MjModel")
 
     self.assertEqual(
-        wrapped_field_data.definition,
+        wrapped_field_data.declaration,
         ("""
 emscripten::val geom_rgba() const {
   return emscripten::val(emscripten::typed_memory_view(ptr_->ngeom * 4, ptr_->geom_rgba));
@@ -674,7 +489,7 @@ emscripten::val geom_rgba() const {
     wrapped_field_data = structs._generate_field_data(field, "MjData")
 
     self.assertEqual(
-        wrapped_field_data.definition,
+        wrapped_field_data.declaration,
         ("""
 emscripten::val buffer() const {
   return emscripten::val(emscripten::typed_memory_view(model->nbuffer, static_cast<uint8_t*>(ptr_->buffer)));
@@ -693,7 +508,7 @@ emscripten::val buffer() const {
         doc="",
     )
     wrapped_field_data = structs._generate_field_data(field, "MjsTexture")
-    self.assertEqual(wrapped_field_data.definition, "MjsElement element;")
+    self.assertEqual(wrapped_field_data.declaration, "MjsElement element;")
 
     self.assertEqual(
         wrapped_field_data.binding,
@@ -717,7 +532,7 @@ emscripten::val buffer() const {
     wrapped_field_data = structs._generate_field_data(field, "MjOption")
 
     self.assertEqual(
-        wrapped_field_data.definition,
+        wrapped_field_data.declaration,
         ("""
 emscripten::val gravity() const {
   return emscripten::val(emscripten::typed_memory_view(3, ptr_->gravity));
@@ -741,7 +556,7 @@ emscripten::val gravity() const {
     )
     wrapped_field_data = structs._generate_field_data(field, "MjModel")
     self.assertEqual(
-        wrapped_field_data.definition,
+        wrapped_field_data.declaration,
         """
 emscripten::val multi_dim_array() const {
   return emscripten::val(emscripten::typed_memory_view(12, reinterpret_cast<float*>(ptr_->multi_dim_array)));