diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc index 198a156d..d0b7b2ee 100644 --- a/wasm/codegen/generated/bindings.cc +++ b/wasm/codegen/generated/bindings.cc @@ -7803,7 +7803,7 @@ MjSpec::~MjSpec() { mjSpec *MjSpec::get() const { return ptr_; } void MjSpec::set(mjSpec *ptr) { ptr_ = ptr; } -std::unique_ptr loadFromXML_wrapper(std::string filename) { +std::unique_ptr mj_loadXML_wrapper(std::string filename) { char error[1000]; mjModel *model = mj_loadXML(filename.c_str(), nullptr, error, sizeof(error)); if (!model) { @@ -10940,7 +10940,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { .property("useexisting", &MjLROpt::useexisting, &MjLROpt::set_useexisting, reference()) .property("uselimit", &MjLROpt::uselimit, &MjLROpt::set_uselimit, reference()); emscripten::class_("MjModel") - .class_function("loadFromXML", &loadFromXML_wrapper, take_ownership()) + .class_function("mj_loadXML", &mj_loadXML_wrapper, take_ownership()) .constructor() .property("B_colind", &MjModel::B_colind) .property("B_rowadr", &MjModel::B_rowadr) diff --git a/wasm/codegen/generators/binding_builder.py b/wasm/codegen/generators/binding_builder.py index 0d052621..b0f229e4 100644 --- a/wasm/codegen/generators/binding_builder.py +++ b/wasm/codegen/generators/binding_builder.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from introspect import ast_nodes from introspect import enums as introspect_enums from introspect import functions as introspect_functions diff --git a/wasm/codegen/generators/common.py b/wasm/codegen/generators/common.py index c809d265..69d450c0 100644 --- a/wasm/codegen/generators/common.py +++ b/wasm/codegen/generators/common.py @@ -15,6 +15,7 @@ """Utility functions for code generation.""" import os +from introspect import ast_nodes def write_to_file(filepath: str, content: str) -> None: @@ -64,3 +65,52 @@ def replace_lines_containing_marker( ) return lines[:i] + replacement_lines + lines[i + 1 :] return lines + + +def get_pointer_return_inner_value_type( + func: ast_nodes.FunctionDecl, +) -> ast_nodes.ValueType | None: + if not isinstance(func.return_type, ast_nodes.PointerType): + return None + if not isinstance(func.return_type.inner_type, ast_nodes.ValueType): + return None + return func.return_type.inner_type + + +def get_inner_value_type( + param: ast_nodes.FunctionParameterDecl, +) -> ast_nodes.ValueType | None: + if not isinstance(param.type, (ast_nodes.PointerType, ast_nodes.ArrayType)): + return None + if not isinstance(param.type.inner_type, ast_nodes.ValueType): + return None + return param.type.inner_type + + +def should_be_wrapped(func: ast_nodes.FunctionDecl) -> bool: + """Checks if a MuJoCo function needs a wrapper function.""" + if get_pointer_return_inner_value_type(func): + return True + for param in func.parameters: + if get_inner_value_type(param): + return True + return False + + +def wrapped_struct_name(c_struct_name: str) -> str: + """Returns the name of the struct wrapping the given C struct.""" + return capitalize(c_struct_name) + + +def wrapped_function_name(func: ast_nodes.FunctionDecl) -> str: + """Returns the name of the function wrapping the given C function. + + Hard-coded wrappers in the template need changing if the implementation of the + wrapped name is changed + + Args: + func: The FunctionDecl of function to wrap. + """ + if should_be_wrapped(func): + return f"{func.name}_wrapper" + return func.name diff --git a/wasm/codegen/generators/constants.py b/wasm/codegen/generators/constants.py index b60d9422..d5b73dcc 100644 --- a/wasm/codegen/generators/constants.py +++ b/wasm/codegen/generators/constants.py @@ -14,7 +14,7 @@ """Constants used in the code generation process.""" -from typing import Dict, List, Set +from typing import Dict, Set from introspect import structs as introspect_structs PRIMITIVE_TYPES: Set[str] = { @@ -263,7 +263,7 @@ ANONYMOUS_STRUCTS: Dict[str, Dict[str, str]] = { # This list is created by subtracting the skipped structs from the list of all # structs and adding the anonymous structs. -STRUCTS_TO_BIND: List[str] = list( +STRUCTS_TO_BIND: list[str] = list( set(introspect_structs.STRUCTS.keys()) .union(ANONYMOUS_STRUCTS.keys()) .difference(set(SKIPPED_STRUCTS)) @@ -363,10 +363,10 @@ MJDATA_SIZES = ( ) # Fields that should be entirely omitted from the bindings. -SKIPPED_FIELDS: Dict[str, List[str]] = {} +SKIPPED_FIELDS: Dict[str, list[str]] = {} # Fields handled manually in template file struct declaration. -MANUAL_FIELDS: Dict[str, List[str]] = { +MANUAL_FIELDS: Dict[str, list[str]] = { # go/keep-sorted start "MjData": ["contact"], "MjvScene": [ diff --git a/wasm/codegen/generators/enums.py b/wasm/codegen/generators/enums.py index 85ff9194..e427b8b4 100644 --- a/wasm/codegen/generators/enums.py +++ b/wasm/codegen/generators/enums.py @@ -14,15 +14,13 @@ """Generates Embind bindings for MuJoCo enums.""" -from typing import List - from introspect import ast_nodes from wasm.codegen.generators import code_builder def generate( - enums: List[ast_nodes.EnumDecl], + enums: list[ast_nodes.EnumDecl], ) -> list[tuple[str, list[str]]]: """Generates all Embind code for the provided enums.""" diff --git a/wasm/codegen/generators/functions.py b/wasm/codegen/generators/functions.py index ead8705c..49104568 100644 --- a/wasm/codegen/generators/functions.py +++ b/wasm/codegen/generators/functions.py @@ -14,7 +14,7 @@ """Helper functions for processing and generating bindings for MuJoCo functions.""" -from typing import List, Tuple, cast +from typing import Tuple, cast from introspect import ast_nodes @@ -23,26 +23,6 @@ from wasm.codegen.generators import common from wasm.codegen.generators import constants -def get_inner_value_type( - param: ast_nodes.FunctionParameterDecl, -) -> ast_nodes.ValueType | None: - if not isinstance(param.type, (ast_nodes.PointerType, ast_nodes.ArrayType)): - return None - if not isinstance(param.type.inner_type, ast_nodes.ValueType): - return None - return param.type.inner_type - - -def get_pointer_return_inner_value_type( - func: ast_nodes.FunctionDecl, -) -> ast_nodes.ValueType | None: - if not isinstance(func.return_type, ast_nodes.PointerType): - return None - if not isinstance(func.return_type.inner_type, ast_nodes.ValueType): - return None - return func.return_type.inner_type - - def param_is_primitive_value(param: ast_nodes.FunctionParameterDecl) -> bool: """Checks if param is a primitive value type.""" if isinstance(param.type, ast_nodes.ValueType): @@ -52,22 +32,12 @@ def param_is_primitive_value(param: ast_nodes.FunctionParameterDecl) -> bool: def get_const_qualifier(func: ast_nodes.FunctionDecl) -> str: """Returns the const qualifier of func's return type.""" - inner_type = get_pointer_return_inner_value_type(func) + inner_type = common.get_pointer_return_inner_value_type(func) if inner_type and inner_type.is_const: return "const " return "" -def should_be_wrapped(func: ast_nodes.FunctionDecl) -> bool: - """Checks if a MuJoCo function needs a wrapper function.""" - if get_pointer_return_inner_value_type(func): - return True - for param in func.parameters: - if get_inner_value_type(param): - return True - return False - - def generate_function_wrapper(func: ast_nodes.FunctionDecl) -> str: """Generates C++ code for a wrapper function.""" @@ -78,7 +48,8 @@ def generate_function_wrapper(func: ast_nodes.FunctionDecl) -> str: wrapper_params = ", ".join([get_param_string(p) for p in wrapper_parameters]) ret_type = get_compatible_return_type(func) builder = code_builder.CodeBuilder() - with builder.function(f"{ret_type} {func.name}_wrapper({wrapper_params})"): + w = common.wrapped_function_name(func) + with builder.function(f"{ret_type} {w}({wrapper_params})"): for p in wrapper_parameters: if c_notnullable := get_param_notnullable(p): @@ -119,7 +90,7 @@ def get_param_unpack_statement( ) -> str: """Generates C++ statements to unpack JS values for pointer/array parameters.""" - inner_type = get_inner_value_type(p) + inner_type = common.get_inner_value_type(p) if not inner_type: return "" @@ -199,7 +170,7 @@ def get_params_string_maybe_with_conversion( native_params = [] for p in ast_params: - if inner_type := get_inner_value_type(p): + if inner_type := common.get_inner_value_type(p): if inner_type.name in constants.PRIMITIVE_TYPES: if inner_type.name == "char": const_qualifier = "const " if inner_type.is_const else "" @@ -231,7 +202,7 @@ def get_compatible_return_code(func: ast_nodes.FunctionDecl) -> str: if func.return_type.name in constants.PRIMITIVE_TYPES: return f"return {c_call};" - if inner_type := get_pointer_return_inner_value_type(func): + if inner_type := common.get_pointer_return_inner_value_type(func): if inner_type.name == "char": return f"return std::string({c_call});" elif inner_type.name == "mjString": @@ -248,7 +219,7 @@ def get_compatible_return_code(func: ast_nodes.FunctionDecl) -> str: def get_compatible_return_type(func: ast_nodes.FunctionDecl) -> str: """Creates embind compatible return type.""" - if inner_type := get_pointer_return_inner_value_type(func): + if inner_type := common.get_pointer_return_inner_value_type(func): if inner_type.name in ["char", "mjString"]: return "std::string" if inner_type.name not in constants.PRIMITIVE_TYPES: @@ -275,7 +246,7 @@ def get_optional_return_code( builder.line(f"{const_qualifier}{struct_name}* result = {c_call};") with builder.block("if (result == nullptr)"): builder.line("return std::nullopt;") - builder.line(f"return {common.capitalize(struct_name)}(result);") + builder.line(f"return {common.wrapped_struct_name(struct_name)}(result);") return builder.to_string() @@ -289,20 +260,20 @@ def is_excluded_function_name(func_name: str) -> bool: def generate( - functions: List[ast_nodes.FunctionDecl], + functions: list[ast_nodes.FunctionDecl], ) -> list[tuple[str, list[str]]]: """Generates Embind bindings for MuJoCo functions.""" wrapper_functions = [] for func in sorted(functions, key=lambda f: f.name): - if should_be_wrapped(func): + if common.should_be_wrapped(func): if func.name not in constants.MANUAL_WRAPPER_FUNCTIONS: wrapper_functions.append(generate_function_wrapper(func)) wrapper_content = "\n\n".join(wrapper_functions) function_bindings = [] for func in sorted(functions, key=lambda f: f.name): - suffix = "_wrapper" if should_be_wrapped(func) else "" - function_bindings.append(f'function("{func.name}", &{func.name}{suffix});') + w = common.wrapped_function_name(func) + function_bindings.append(f'function("{func.name}", &{w});') bindings_content = "\n".join(function_bindings) return [ diff --git a/wasm/codegen/generators/structs.py b/wasm/codegen/generators/structs.py index 9e36f1f2..acae9cec 100644 --- a/wasm/codegen/generators/structs.py +++ b/wasm/codegen/generators/structs.py @@ -21,6 +21,7 @@ from typing import Tuple, Union, cast from introspect import ast_nodes from introspect import structs as introspect_structs +from introspect import functions as introspect_functions from wasm.codegen.generators import code_builder from wasm.codegen.generators import common @@ -123,7 +124,9 @@ def _generate_field_data( return WrappedFieldData( binding=_get_property_binding(f, w, setter=False, reference=True), typename=anonymous_struct_name, - declaration=f"{common.capitalize(anonymous_struct_name)} {f.name};", + declaration=( + f"{common.wrapped_struct_name(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, @@ -320,7 +323,7 @@ def build_struct_header( ): """Builds the C++ header file code for a struct.""" s = struct_name - w = common.capitalize(s) + w = common.wrapped_struct_name(s) if ( s not in constants.ANONYMOUS_STRUCTS @@ -398,7 +401,7 @@ def build_struct_source( ): """Builds the C++ .cc file code for a struct.""" s = struct_name - w = common.capitalize(s) + w = common.wrapped_struct_name(s) is_mjs = w.startswith("Mjs") member_inits = _find_member_inits(wrapped_fields) @@ -467,7 +470,7 @@ def _build_struct_bindings( wrapped_fields: list[WrappedFieldData], ): """Builds the C++ bindings for a struct.""" - w = common.capitalize(struct_name) + w = common.wrapped_struct_name(struct_name) is_mjs = w.startswith("Mjs") builder = code_builder.CodeBuilder() @@ -478,10 +481,10 @@ def _build_struct_bindings( builder.line(".constructor()") builder.line(".constructor()") elif w == "MjModel": - builder.line( - '.class_function("loadFromXML", &loadFromXML_wrapper,' - " take_ownership())" + w = common.wrapped_function_name( + introspect_functions.FUNCTIONS["mj_loadXML"] ) + builder.line(f'.class_function("mj_loadXML", &{w}, take_ownership())') builder.line(".constructor()") elif w == "MjSpec": builder.line(".constructor()") @@ -572,7 +575,7 @@ def _get_field_struct_type( ) -> str | None: """Extracts the base struct name if the field type is a struct or pointer to a struct.""" s = struct_name - w = common.capitalize(s) + w = common.wrapped_struct_name(s) if isinstance(field.type, ast_nodes.AnonymousStructDecl): anonymous_struct_name = "" for name, value in constants.ANONYMOUS_STRUCTS.items(): @@ -692,7 +695,7 @@ def generate(struct_to_bind: list[str]) -> list[tuple[str, list[str]]]: fields: list[WrappedFieldData] = [] introspect_fields = get_introspect_struct_fields(s) for field in introspect_fields: - fields.append(_generate_field_data(field, common.capitalize(s))) + fields.append(_generate_field_data(field, common.wrapped_struct_name(s))) wrapped_structs_with_fields[s] = fields dependency_sorted_struct_names = sort_structs_by_dependency( @@ -723,13 +726,13 @@ def generate(struct_to_bind: list[str]) -> list[tuple[str, list[str]]]: bindings.append(_build_struct_bindings(s, fields)) for s in alphabetically_sorted_struct_names: - w = common.capitalize(s) + w = common.wrapped_struct_name(s) if w.startswith("Mjs") or w == "MjSpec": bindings.append(f"emscripten::register_optional<{w}>();") manual_struct_field_declarations = [] for s in dependency_sorted_struct_names: - w = common.capitalize(s) + w = common.wrapped_struct_name(s) fields = wrapped_structs_with_fields[s] if s in constants.MANUAL_STRUCTS_HEADERS: decls: list[str] = [] diff --git a/wasm/codegen/templates/bindings.cc b/wasm/codegen/templates/bindings.cc index 301fee3b..e74ab501 100644 --- a/wasm/codegen/templates/bindings.cc +++ b/wasm/codegen/templates/bindings.cc @@ -386,7 +386,7 @@ MjSpec::~MjSpec() { mjSpec *MjSpec::get() const { return ptr_; } void MjSpec::set(mjSpec *ptr) { ptr_ = ptr; } -std::unique_ptr loadFromXML_wrapper(std::string filename) { +std::unique_ptr mj_loadXML_wrapper(std::string filename) { char error[1000]; mjModel *model = mj_loadXML(filename.c_str(), nullptr, error, sizeof(error)); if (!model) { diff --git a/wasm/codegen/tests/generators_test.py b/wasm/codegen/tests/generators_test.py index 9a365379..1d21bc89 100644 --- a/wasm/codegen/tests/generators_test.py +++ b/wasm/codegen/tests/generators_test.py @@ -97,48 +97,75 @@ class FunctionUtilsTest(absltest.TestCase): parameters=tuple(), doc="Returns int pointer", ) - self.assertTrue(functions.should_be_wrapped(func)) + self.assertTrue(common.should_be_wrapped(func)) def test_generate_function_wrapper_for_simple_func(self): func = ast_nodes.FunctionDecl( - name="get_id", - return_type=ast_nodes.ValueType("int"), - parameters=tuple(), - doc="Returns an integer ID", + name="mj_defaultLROpt", + return_type=ast_nodes.ValueType("void"), + parameters=( + ast_nodes.FunctionParameterDecl( + name="opt", + type=ast_nodes.PointerType( + inner_type=ast_nodes.ValueType("mjLROpt"), + ), + ), + ), + doc="Set default options for length range computation.", ) result = functions.generate_function_wrapper(func) self.assertEqual( result, - """int get_id_wrapper() { - return get_id(); + """void mj_defaultLROpt_wrapper(MjLROpt& opt) { + mj_defaultLROpt(opt.get()); }""", ) def test_generate_function_wrapper_checking_param(self): - parameters = ( - ast_nodes.FunctionParameterDecl( - name="mat", - type=ast_nodes.PointerType( - inner_type=ast_nodes.ValueType(name="mjtNum", is_const=True), + func = ast_nodes.FunctionDecl( + name="mj_extractState", + return_type=ast_nodes.ValueType(name="void"), + parameters=( + ast_nodes.FunctionParameterDecl( + name="m", + type=ast_nodes.PointerType( + inner_type=ast_nodes.ValueType( + name="mjModel", is_const=True + ), + ), + ), + ast_nodes.FunctionParameterDecl( + name="src", + type=ast_nodes.PointerType( + inner_type=ast_nodes.ValueType( + name="mjtNum", is_const=True + ), + ), + ), + ast_nodes.FunctionParameterDecl( + name="srcsig", + type=ast_nodes.ValueType(name="unsigned int"), + ), + ast_nodes.FunctionParameterDecl( + name="dst", + type=ast_nodes.PointerType( + inner_type=ast_nodes.ValueType(name="mjtNum"), + ), + ), + ast_nodes.FunctionParameterDecl( + name="dstsig", + type=ast_nodes.ValueType(name="unsigned int"), ), ), - ast_nodes.FunctionParameterDecl( - name="nr", - type=ast_nodes.ValueType(name="int"), - ), - ) - func = ast_nodes.FunctionDecl( - name="get_id", - return_type=ast_nodes.ValueType("int"), - parameters=parameters, - doc="Returns an integer ID", + doc="Extract a subset of components from a state previously obtained via mj_getState.", # pylint: disable=line-too-long ) result = functions.generate_function_wrapper(func) self.assertEqual( result, - """int get_id_wrapper(const NumberArray& mat, int nr) { - UNPACK_ARRAY(mjtNum, mat); - return get_id(mat_.data(), nr); + """void mj_extractState_wrapper(const MjModel& m, const NumberArray& src, unsigned int srcsig, const val& dst, unsigned int dstsig) { + UNPACK_ARRAY(mjtNum, src); + UNPACK_VALUE(mjtNum, dst); + mj_extractState(m.get(), src_.data(), srcsig, dst_.data(), dstsig); }""", ) diff --git a/wasm/tests/bindings_test.ts b/wasm/tests/bindings_test.ts index 84163e80..cc14f9dd 100644 --- a/wasm/tests/bindings_test.ts +++ b/wasm/tests/bindings_test.ts @@ -119,7 +119,7 @@ describe('MuJoCo WASM Bindings', () => { writeXMLFile(tempXmlFilename, TEST_XML); - model = mujoco.MjModel!.loadFromXML(tempXmlFilename); + model = mujoco.MjModel!.mj_loadXML(tempXmlFilename); if (!model) { unlinkXMLFile(tempXmlFilename); throw new Error('Failed to load model from XML'); @@ -182,7 +182,7 @@ describe('MuJoCo WASM Bindings', () => { beforeEach(() => { writeXMLFile(tempXmlFilename, simpleXmlContent); - simpleModel = mujoco.MjModel!.loadFromXML(tempXmlFilename); + simpleModel = mujoco.MjModel!.mj_loadXML(tempXmlFilename); assertExists(simpleModel); simpleData = new mujoco.MjData(simpleModel); assertExists(simpleData); @@ -345,7 +345,7 @@ describe('MuJoCo WASM Bindings', () => { `; const tempXmlFilename = '/tmp/model_c.xml'; writeXMLFile(tempXmlFilename, xmlString); - const model = mujoco.MjModel!.loadFromXML(tempXmlFilename); + const model = mujoco.MjModel!.mj_loadXML(tempXmlFilename); expect(model).not.toBeNull(); const data = new mujoco.MjData(model!); expect(data).not.toBeNull(); @@ -879,7 +879,7 @@ describe('MuJoCo WASM Bindings', () => { writeXMLFile(model1XmlFilename, xml2); writeXMLFile(model2XmlFilename, xml3); - const model = mujoco.MjModel!.loadFromXML(modelXmlFilename); + const model = mujoco.MjModel!.mj_loadXML(modelXmlFilename); try { expect(model).toBeDefined(); @@ -1078,7 +1078,7 @@ describe('MuJoCo WASM Bindings', () => { it('should check MjContactVec equality', () => { const tempXmlFilename2 = '/tmp/model2.xml'; writeXMLFile(tempXmlFilename2, TEST_XML); - const model2 = mujoco.MjModel!.loadFromXML(tempXmlFilename2); + const model2 = mujoco.MjModel!.mj_loadXML(tempXmlFilename2); const data2 = new mujoco.MjData(model2); try { mujoco.mj_forward(model!, data!); @@ -1635,7 +1635,7 @@ describe('MuJoCo WASM Bindings', () => { `; writeXMLFile(tempXmlFilename, TEST_XML_TEXTURE); - const model = mujoco.MjModel!.loadFromXML(tempXmlFilename); + const model = mujoco.MjModel!.mj_loadXML(tempXmlFilename); try { expect(model).toBeDefined(); expect(model!.tex_height).toEqual(new Int32Array([512])); @@ -1754,7 +1754,7 @@ describe('MuJoCo WASM Bindings', () => { `; writeXMLFile(tempXmlFilename, xml); - const model = mujoco.MjModel!.loadFromXML(tempXmlFilename); + const model = mujoco.MjModel!.mj_loadXML(tempXmlFilename); try { mujoco.mj_saveLastXML(tempXmlFilename, model!); const savedXmlContent = @@ -1781,7 +1781,7 @@ describe('MuJoCo WASM Bindings', () => { `; writeXMLFile(tempXmlFilename, actuatorXml); - const model = mujoco.MjModel.loadFromXML(tempXmlFilename); + const model = mujoco.MjModel.mj_loadXML(tempXmlFilename); assertExists(model); const data = new mujoco.MjData(model); assertExists(data); diff --git a/wasm/tests/sandbox/main.ts b/wasm/tests/sandbox/main.ts index 9ec14845..79702e5f 100644 --- a/wasm/tests/sandbox/main.ts +++ b/wasm/tests/sandbox/main.ts @@ -43,7 +43,7 @@ async function main() { try { console.log('Hello world!: Loading model'); - model = mujoco.MjModel.loadFromXML('/working/hello.xml'); + model = mujoco.MjModel.mj_loadXML('/working/hello.xml'); if (!model) { throw new Error('Failed to load model'); }