diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc index 7c430138..ca0cc9b6 100644 --- a/wasm/codegen/generated/bindings.cc +++ b/wasm/codegen/generated/bindings.cc @@ -12374,18 +12374,6 @@ std::optional mjs_asPlugin_wrapper(MjsElement& element) { EMSCRIPTEN_BINDINGS(mujoco_functions) { function("parseXMLString", &parseXMLString_wrapper, take_ownership()); function("error", &error_wrapper); - function("mj_resetCallbacks", &mj_resetCallbacks); - function("mj_version", &mj_version); - function("mju_bandDiag", &mju_bandDiag); - function("mju_springDamper", &mju_springDamper); - function("mju_min", &mju_min); - function("mju_max", &mju_max); - function("mju_clip", &mju_clip); - function("mju_sign", &mju_sign); - function("mju_round", &mju_round); - function("mju_isBad", &mju_isBad); - function("mju_Halton", &mju_Halton); - function("mju_sigmoid", &mju_sigmoid); function("mj_copyBack", &mj_copyBack_wrapper); function("mj_saveLastXML", &mj_saveLastXML_wrapper); function("mj_step", &mj_step_wrapper); @@ -12403,6 +12391,7 @@ EMSCRIPTEN_BINDINGS(mujoco_functions) { function("mj_resetData", &mj_resetData_wrapper); function("mj_resetDataDebug", &mj_resetDataDebug_wrapper); function("mj_resetDataKeyframe", &mj_resetDataKeyframe_wrapper); + function("mj_resetCallbacks", &mj_resetCallbacks); function("mj_setConst", &mj_setConst_wrapper); function("mj_setLengthRange", &mj_setLengthRange_wrapper); function("mjs_activatePlugin", &mjs_activatePlugin_wrapper); @@ -12494,6 +12483,7 @@ EMSCRIPTEN_BINDINGS(mujoco_functions) { function("mj_local2Global", &mj_local2Global_wrapper); function("mj_getTotalmass", &mj_getTotalmass_wrapper); function("mj_setTotalmass", &mj_setTotalmass_wrapper); + function("mj_version", &mj_version); function("mj_versionString", &mj_versionString_wrapper); function("mj_multiRay", &mj_multiRay_wrapper); function("mj_ray", &mj_ray_wrapper); @@ -12605,6 +12595,7 @@ EMSCRIPTEN_BINDINGS(mujoco_functions) { function("mju_band2Dense", &mju_band2Dense_wrapper); function("mju_dense2Band", &mju_dense2Band_wrapper); function("mju_bandMulMatVec", &mju_bandMulMatVec_wrapper); + function("mju_bandDiag", &mju_bandDiag); function("mju_eig3", &mju_eig3_wrapper); function("mju_boxQP", &mju_boxQP_wrapper); function("mju_muscleGain", &mju_muscleGain_wrapper); @@ -12612,10 +12603,17 @@ EMSCRIPTEN_BINDINGS(mujoco_functions) { function("mju_muscleDynamics", &mju_muscleDynamics_wrapper); function("mju_encodePyramid", &mju_encodePyramid_wrapper); function("mju_decodePyramid", &mju_decodePyramid_wrapper); + function("mju_springDamper", &mju_springDamper); + function("mju_min", &mju_min); + function("mju_max", &mju_max); + function("mju_clip", &mju_clip); + function("mju_sign", &mju_sign); + function("mju_round", &mju_round); function("mju_type2Str", &mju_type2Str_wrapper); function("mju_str2Type", &mju_str2Type_wrapper); function("mju_writeNumBytes", &mju_writeNumBytes_wrapper); function("mju_warningText", &mju_warningText_wrapper); + function("mju_isBad", &mju_isBad); function("mju_isZero", &mju_isZero_wrapper); function("mju_standardNormal", &mju_standardNormal_wrapper); function("mju_f2n", &mju_f2n_wrapper); @@ -12624,6 +12622,8 @@ EMSCRIPTEN_BINDINGS(mujoco_functions) { function("mju_n2d", &mju_n2d_wrapper); function("mju_insertionSort", &mju_insertionSort_wrapper); function("mju_insertionSortInt", &mju_insertionSortInt_wrapper); + function("mju_Halton", &mju_Halton); + function("mju_sigmoid", &mju_sigmoid); function("mjd_transitionFD", &mjd_transitionFD_wrapper); function("mjd_inverseFD", &mjd_inverseFD_wrapper); function("mjd_subQuat", &mjd_subQuat_wrapper); diff --git a/wasm/codegen/generators/functions.py b/wasm/codegen/generators/functions.py index 76f757f7..0bc69ee3 100644 --- a/wasm/codegen/generators/functions.py +++ b/wasm/codegen/generators/functions.py @@ -60,9 +60,12 @@ def get_const_qualifier(func: ast_nodes.FunctionDecl) -> str: def should_be_wrapped(func: ast_nodes.FunctionDecl) -> bool: """Checks if a MuJoCo function needs a wrapper function.""" - return get_pointer_return_inner_value_type(func) is not None or any( - get_inner_value_type(param) for param in func.parameters - ) + 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: @@ -88,11 +91,8 @@ def generate_function_wrapper(func: ast_nodes.FunctionDecl) -> str: if bound_check_code: builder.line(bound_check_code) - c_params_list = get_params_string_maybe_with_conversion(func.parameters) - c_params_str = ", ".join(c_params_list) - c_call = f"{func.name}({c_params_str})" - c_statement = get_compatible_return_call(func, c_call) - builder.line(f"{c_statement};") + for line in get_compatible_return_code(func).splitlines(): + builder.line(line) return builder.to_string() @@ -194,7 +194,7 @@ def get_param_string(p: ast_nodes.FunctionParameterDecl) -> str: def get_params_string_maybe_with_conversion( ast_params: Tuple[ast_nodes.FunctionParameterDecl, ...], -) -> List[str]: +) -> str: """Generates C++ expressions for passing compatible params from JS to MuJoCo C-API functions.""" native_params = [] @@ -217,27 +217,27 @@ def get_params_string_maybe_with_conversion( f"Unhandled parameter type for conversion: {p.type} for param" f" '{p.name}'" ) - return native_params + return ", ".join(native_params) -def get_compatible_return_call( - func: ast_nodes.FunctionDecl, invoker: str -) -> str: +def get_compatible_return_code(func: ast_nodes.FunctionDecl) -> str: """Generates embind compatible return value conversion.""" + c_params = get_params_string_maybe_with_conversion(func.parameters) + c_call = f"{func.name}({c_params})" if isinstance(func.return_type, ast_nodes.ValueType): if func.return_type.name == "void": - return invoker + return f"{c_call};" if func.return_type.name in constants.PRIMITIVE_TYPES: - return f"return {invoker}" + return f"return {c_call};" if inner_type := get_pointer_return_inner_value_type(func): if inner_type.name == "char": - return f"return std::string({invoker})" + return f"return std::string({c_call});" elif inner_type.name == "mjString": - return f"return *{invoker}" + return f"return *{c_call};" elif inner_type.name not in constants.PRIMITIVE_TYPES: - return get_converted_struct_to_class(func, invoker) + return get_optional_return_code(func, c_call) raise RuntimeError( "Failed to calculate return value conversion for function" @@ -262,21 +262,22 @@ def get_compatible_return_type(func: ast_nodes.FunctionDecl) -> str: return "val" -def get_converted_struct_to_class( - func: ast_nodes.FunctionDecl, invoker: str +def get_optional_return_code( + func: ast_nodes.FunctionDecl, c_call: str ) -> str: - """Generates a C++ function invocation for a struct return-type function.""" + """Generates code to return std::optional of the wrapped struct.""" const_qualifier = get_const_qualifier(func) return_type = cast(ast_nodes.PointerType, func.return_type) struct_name = cast(ast_nodes.ValueType, return_type.inner_type).name - class_constructor = common.capitalize(struct_name) - return_str = f"{class_constructor}(result)" - return f"""{const_qualifier}{struct_name}* result = {invoker}; - if (result == nullptr) {{ - return std::nullopt; - }} - return {return_str}""" + + builder = code_builder.CodeBuilder() + 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);") + + return builder.to_string() def is_excluded_function_name(func_name: str) -> bool: @@ -291,63 +292,26 @@ class Generator: """Generates Embind bindings for MuJoCo functions.""" def __init__(self, functions: Mapping[str, ast_nodes.FunctionDecl]): - self.direct_bind_functions: List[ast_nodes.FunctionDecl] = [] - self.wrapper_bind_functions: List[ast_nodes.FunctionDecl] = [] - - for func in functions.values(): - if should_be_wrapped(func): - self.wrapper_bind_functions.append(func) - else: - self.direct_bind_functions.append(func) - - def _generate_wrappers(self) -> str: - """Generates Embind bindings for all functions that need wrappers.""" - - code = [] - for func in self.wrapper_bind_functions: - if func.name not in constants.MANUAL_WRAPPER_FUNCTIONS: - wrapper_code = generate_function_wrapper(func) - code.append(wrapper_code) - - return "\n\n".join(code) - - def _generate_direct_bindable_functions(self) -> list[str]: - """Generates Embind bindings for all directly bindable functions.""" - - result = [] - for func in self.direct_bind_functions: - result.append(self._generate_function_binding(func)) - - return result - - def _generate_function_binding( - self, func: ast_nodes.FunctionDecl, is_wrapper=False - ) -> str: - """Generates the Embind code for a single function.""" - - js_name, cpp_func = func.name, func.name - if is_wrapper: - cpp_func += "_wrapper" - - return f'function("{js_name}", &{cpp_func});' - - def _generate_wrapper_bindable_functions(self) -> list[str]: - """Generates Embind bindings for all functions that need wrappers.""" - - result = [] - for func in self.wrapper_bind_functions: - result.append(self._generate_function_binding(func, True)) - - return result + self.functions = functions def generate(self) -> list[tuple[str, list[str]]]: """Generates the bindings file for all functions.""" - wrapper_functions = self._generate_wrappers() - function_bindings = self._generate_direct_bindable_functions() - function_bindings += self._generate_wrapper_bindable_functions() + function_wrappers = [] + for func in self.functions.values(): + if should_be_wrapped(func): + if func.name not in constants.MANUAL_WRAPPER_FUNCTIONS: + function_wrappers.append(generate_function_wrapper(func)) + + function_bindings = [] + for func in self.functions.values(): + js_name = func.name + cpp_func = func.name + if should_be_wrapped(func): + cpp_func += "_wrapper" + function_bindings.append(f'function("{js_name}", &{cpp_func});') return [ - ("// {{ WRAPPER_FUNCTIONS }}", [wrapper_functions]), + ("// {{ WRAPPER_FUNCTIONS }}", ["\n\n".join(function_wrappers)]), ("// {{ FUNCTION_BINDINGS }}", ["\n".join(function_bindings)]), ] diff --git a/wasm/codegen/tests/generators_test.py b/wasm/codegen/tests/generators_test.py index b7e32683..b68c9987 100644 --- a/wasm/codegen/tests/generators_test.py +++ b/wasm/codegen/tests/generators_test.py @@ -157,17 +157,17 @@ class FunctionUtilsTest(absltest.TestCase): type=ast_nodes.PointerType(ast_nodes.ValueType("customstruct")), ) result = functions.get_params_string_maybe_with_conversion((param,)) - self.assertEqual(result, ["s.get()"]) + self.assertEqual(result, "s.get()") - def test_get_compatible_return_call(self): + def test_get_compatible_return_code(self): func = ast_nodes.FunctionDecl( name="noop", return_type=ast_nodes.ValueType("void"), parameters=tuple(), doc="does nothing", ) - result = functions.get_compatible_return_call(func, "noop()") - self.assertEqual(result, "noop()") + result = functions.get_compatible_return_code(func) + self.assertEqual(result, "noop();") def test_get_compatible_return_type(self): func = ast_nodes.FunctionDecl( @@ -179,14 +179,14 @@ class FunctionUtilsTest(absltest.TestCase): result = functions.get_compatible_return_type(func) self.assertEqual(result.strip(), "std::string") - def test_get_converted_struct_to_class(self): + def test_get_optional_return_code(self): func = ast_nodes.FunctionDecl( name="get_struct", return_type=ast_nodes.PointerType(ast_nodes.ValueType("mystruct")), parameters=tuple(), doc="returns struct", ) - result = functions.get_converted_struct_to_class(func, "get_struct()") + result = functions.get_optional_return_code(func, "get_struct()") self.assertIn("mystruct* result = get_struct();", result) self.assertIn("return Mystruct(result)", result) @@ -836,49 +836,5 @@ class EnumsGeneratorTest(absltest.TestCase): self.assertEqual(actual_code, expected_code) -class FunctionsGeneratorTest(absltest.TestCase): - - def setUp(self): - super().setUp() - self.generator = functions.Generator({}) - self.int_type = ast_nodes.ValueType(name="int") - - def test_generate_function_binding_simple_case(self): - func_simple_void = ast_nodes.FunctionDecl( - name="do_nothing", - return_type=ast_nodes.ValueType(name="void"), - parameters=tuple(), - doc="doc", - ) - self.assertEqual( - self.generator._generate_function_binding(func_simple_void), - 'function("do_nothing", &do_nothing);', - ) - - def test_generate_direct_bindable_functions_simple_filter(self): - direct_bind = ast_nodes.FunctionDecl( - name="direct_bind", - return_type=self.int_type, - parameters=( - ast_nodes.FunctionParameterDecl(name="val", type=self.int_type), - ), - doc="doc", - ) - needs_wrap = ast_nodes.FunctionDecl( - name="needs_wrap", - return_type=ast_nodes.PointerType(inner_type=self.int_type), - parameters=tuple(), - doc="doc", - ) - self.generator = functions.Generator({ - "direct1": direct_bind, - "wrapped1": needs_wrap, - }) - - generated_code = self.generator._generate_direct_bindable_functions() - self.assertIn('function("direct_bind", &direct_bind);', generated_code) - self.assertNotIn('function("needs_wrap", &needs_wrap);', generated_code) - - if __name__ == "__main__": absltest.main()