diff --git a/wasm/codegen/generators/structs.py b/wasm/codegen/generators/structs.py index 062f13ed..e5baa770 100644 --- a/wasm/codegen/generators/structs.py +++ b/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( diff --git a/wasm/codegen/helpers/struct_constructor_code_builder.py b/wasm/codegen/helpers/struct_constructor_code_builder.py index 38506bb8..03bd16aa 100644 --- a/wasm/codegen/helpers/struct_constructor_code_builder.py +++ b/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( - struct_info: ast_nodes.StructDecl -) -> bool: +def _has_nested_wrapper_members(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, - fields_with_init: List[structs_wrappers_data.WrappedFieldData] = [], - wrapped_fields: 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, - fields_with_init: List[structs_wrappers_data.WrappedFieldData] = [], - use_shallow_copy: bool = False, + wrapped_fields: List[structs_wrappers_data.WrappedFieldData], ): """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 - for field_with_init in fields_with_init + field_with_init.initialization 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} &other)" + f"{wrapper_name}::{wrapper_name}(const {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" {wrapper_name}::copy()" + f"std::unique_ptr<{wrapper_name}> {wrapper_name}::copy()" ): - builder.line( - f"return std::make_unique<{wrapper_name}>(*this);" - ) + builder.line(f"return std::make_unique<{wrapper_name}>(*this);") return builder.to_string() diff --git a/wasm/codegen/helpers/struct_constructor_code_builder_test.py b/wasm/codegen/helpers/struct_constructor_code_builder_test.py index f7255e8d..27e5313b 100644 --- a/wasm/codegen/helpers/struct_constructor_code_builder_test.py +++ b/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( - "mjLROpt", "mj_defaultLROpt" - ), + wrapped_structs["mjLROpt"].wrapped_source, """ 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::copy() { + return std::make_unique(*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) {} -MjLROpt::MjLROpt() : ptr_(new mjLROpt) { - owned_ = true; -} -MjLROpt::~MjLROpt() { - if (owned_ && ptr_) delete ptr_; +MjsElement::MjsElement(mjsElement *ptr) : ptr_(ptr) {} +MjsElement::~MjsElement() {} +std::unique_ptr MjsElement::copy() { + return std::make_unique(*this); } """.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( - "mjvLight", use_shallow_copy=True - ), + struct_constructor_code_builder.build_struct_source("mjvLight", []), """MjvLight::MjvLight(mjvLight *ptr) : ptr_(ptr) {} MjvLight::MjvLight() : ptr_(new mjvLight) { owned_ = true; @@ -104,14 +113,14 @@ std::unique_ptr 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(); diff --git a/wasm/codegen/helpers/structs_parser.py b/wasm/codegen/helpers/structs_parser.py index 46cd8df6..1303dcd2 100644 --- a/wasm/codegen/helpers/structs_parser.py +++ b/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( - field, wrap_data.wrap_name + wrapped_field = struct_field_handler.StructFieldHandler( + field, wrapped_name ).generate() - # If the struct has at least one non-primitive or fixed size 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) + wrapped_fields.append(wrapped_field) - wrap_data.wrapped_header = ( - struct_constructor_code_builder.build_struct_header( - struct_name, - wrap_data.use_shallow_copy, - fields_with_init, - wrap_data.wrapped_fields, - ) + wrapped_header = struct_constructor_code_builder.build_struct_header( + struct_name, + wrapped_fields, ) - wrap_data.wrapped_source = ( - struct_constructor_code_builder.build_struct_source( - struct_name, - get_default_func_name(struct_name), - fields_with_init, - wrap_data.use_shallow_copy, - ) + wrapped_source = struct_constructor_code_builder.build_struct_source( + struct_name, + wrapped_fields, ) + 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): diff --git a/wasm/codegen/helpers/structs_parser_test.py b/wasm/codegen/helpers/structs_parser_test.py index efacd730..864928bc 100644 --- a/wasm/codegen/helpers/structs_parser_test.py +++ b/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( - structs_wrappers_data.create_wrapped_structs_set_up_data([ - "mjModel", - "mjData", - "mjVisualGlobal", - "mjVisualQuality", - "mjVisual", - ]) - ) + self.wrapped_structs = structs_parser.generate_wasm_bindings([ + "mjModel", + "mjData", + "mjVisualGlobal", + "mjVisualQuality", + "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_wrappers_data.create_wrapped_structs_set_up_data( - ["mjFakeStruct"] - ) - ) + structs_parser.generate_wasm_bindings(["mjFakeStruct"]) with self.assertRaises(RuntimeError): - structs_parser.generate_wasm_bindings( - 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" - ) + structs_parser.generate_wasm_bindings(["mjFakeAnonymousStruct"]) if __name__ == "__main__": diff --git a/wasm/codegen/helpers/structs_wrappers_data.py b/wasm/codegen/helpers/structs_wrappers_data.py index f6f6d308..6d4ae717 100644 --- a/wasm/codegen/helpers/structs_wrappers_data.py +++ b/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 - }