diff --git a/wasm/codegen/generators/enums_test.py b/wasm/codegen/generators/enums_test.py deleted file mode 100644 index 1d09de6f..00000000 --- a/wasm/codegen/generators/enums_test.py +++ /dev/null @@ -1,62 +0,0 @@ -# Copyright 2025 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from absl.testing import absltest - -from introspect import ast_nodes - -from wasm.codegen.generators import enums - - -class EnumsGeneratorTest(absltest.TestCase): - - def test_generate_enum_bindings(self): - - generator = enums.Generator({ - "TestEnum": ast_nodes.EnumDecl( - name="TestEnum", - declname="enum TestEnum_", - values={"FIRST_VAL": 0, "SECOND_VAL": 1, "THIRD_VAL": 2}, - ), - "AnotherEnum": ast_nodes.EnumDecl( - name="AnotherEnum", - declname="enum AnotherEnum_", - values={"ALPHA": 100, "BETA": 200}, - ), - "EmptyEnum": ast_nodes.EnumDecl( - name="EmptyEnum", - declname="enum EmptyEnum_", - values={}, - ), - }) - - expected_code = """ enum_("TestEnum") - .value("FIRST_VAL", FIRST_VAL) - .value("SECOND_VAL", SECOND_VAL) - .value("THIRD_VAL", THIRD_VAL); - - enum_("AnotherEnum") - .value("ALPHA", ALPHA) - .value("BETA", BETA); - - enum_("EmptyEnum"); -""" - - actual_code = generator.generate() - - self.assertEqual(actual_code, expected_code) - - -if __name__ == "__main__": - absltest.main() diff --git a/wasm/codegen/generators/functions_test.py b/wasm/codegen/generators/generators_test.py similarity index 67% rename from wasm/codegen/generators/functions_test.py rename to wasm/codegen/generators/generators_test.py index e1ac87a3..8ef097a0 100644 --- a/wasm/codegen/generators/functions_test.py +++ b/wasm/codegen/generators/generators_test.py @@ -16,9 +16,48 @@ from absl.testing import absltest from introspect import ast_nodes +from wasm.codegen.generators import enums from wasm.codegen.generators import functions +class EnumsGeneratorTest(absltest.TestCase): + + def test_generate_enum_bindings(self): + + generator = enums.Generator({ + "TestEnum": ast_nodes.EnumDecl( + name="TestEnum", + declname="enum TestEnum_", + values={"FIRST_VAL": 0, "SECOND_VAL": 1, "THIRD_VAL": 2}, + ), + "AnotherEnum": ast_nodes.EnumDecl( + name="AnotherEnum", + declname="enum AnotherEnum_", + values={"ALPHA": 100, "BETA": 200}, + ), + "EmptyEnum": ast_nodes.EnumDecl( + name="EmptyEnum", + declname="enum EmptyEnum_", + values={}, + ), + }) + + expected_code = """ enum_("TestEnum") + .value("FIRST_VAL", FIRST_VAL) + .value("SECOND_VAL", SECOND_VAL) + .value("THIRD_VAL", THIRD_VAL); + + enum_("AnotherEnum") + .value("ALPHA", ALPHA) + .value("BETA", BETA); + + enum_("EmptyEnum"); +""" + + actual_code = generator.generate() + + self.assertEqual(actual_code, expected_code) + class FunctionsGeneratorTest(absltest.TestCase): def setUp(self): diff --git a/wasm/codegen/helpers/code_builder_test.py b/wasm/codegen/helpers/code_builder_test.py deleted file mode 100644 index a1b0783a..00000000 --- a/wasm/codegen/helpers/code_builder_test.py +++ /dev/null @@ -1,51 +0,0 @@ -# Copyright 2025 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Tests for the code_builder module.""" - -from absl.testing import absltest -from wasm.codegen.helpers import code_builder - - -class CodeBuilderTest(absltest.TestCase): - - def test_code_builder_functionality(self): - """Test nested indentation blocks.""" - builder = code_builder.CodeBuilder(indent_str=" ") - builder.line("let a = 1") - with builder.block("function myFunc()"): - builder.line("let flag = true") - with builder.block("while (flag)"): - builder.line("a++") - builder.line("flag = a < 10") - builder.line("return a") - builder.line("print('Done')") - - expected_lines = [ - "let a = 1", - "function myFunc() {", - " let flag = true", - " while (flag) {", - " a++", - " flag = a < 10", - " }", - " return a", - "}", - "print('Done')", - ] - self.assertEqual(builder.to_string(), "\n".join(expected_lines)) - - -if __name__ == "__main__": - absltest.main() diff --git a/wasm/codegen/helpers/common_test.py b/wasm/codegen/helpers/common_test.py deleted file mode 100644 index cdb3d4f4..00000000 --- a/wasm/codegen/helpers/common_test.py +++ /dev/null @@ -1,37 +0,0 @@ -# Copyright 2025 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from absl.testing import absltest -from wasm.codegen.helpers import common - - -class CommonUtilsTest(absltest.TestCase): - - def test_uppercase_first_letter(self): - self.assertEqual(common.uppercase_first_letter(""), "") - self.assertEqual(common.uppercase_first_letter("hello"), "Hello") - self.assertEqual(common.uppercase_first_letter("1st place"), "1st place") - self.assertEqual(common.uppercase_first_letter("!wow"), "!wow") - self.assertEqual( - common.uppercase_first_letter(" leading space"), " leading space" - ) - - def test_try_cast_to_scalar_type(self): - self.assertEqual(common.try_cast_to_scalar_type("123"), 123) - self.assertEqual(common.try_cast_to_scalar_type("123.456"), 123.456) - self.assertEqual(common.try_cast_to_scalar_type("abc"), "abc") - - -if __name__ == "__main__": - absltest.main() diff --git a/wasm/codegen/helpers/function_utils_test.py b/wasm/codegen/helpers/function_utils_test.py deleted file mode 100644 index 7023760a..00000000 --- a/wasm/codegen/helpers/function_utils_test.py +++ /dev/null @@ -1,241 +0,0 @@ -# Copyright 2025 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from absl.testing import absltest -from introspect import ast_nodes - -from wasm.codegen.helpers import constants -from wasm.codegen.helpers import function_utils - - -class FunctionUtilsTest(absltest.TestCase): - - def setUp(self): - super().setUp() - self.struct_type = ast_nodes.ValueType("MyStruct") - self.ptr_to_int = ast_nodes.PointerType(ast_nodes.ValueType("int")) - self.func_ret_ptr_int = ast_nodes.FunctionDecl( - "func_pi", ast_nodes.PointerType(ast_nodes.ValueType("int")), [], "doc" - ) - self.func_ret_ptr_struct = ast_nodes.FunctionDecl( - "func_ps", - ast_nodes.PointerType(ast_nodes.ValueType("MyStruct")), - [], - "doc", - ) - - def test_return_is_value_of_type(self): - self.assertTrue( - function_utils.return_is_value_of_type( - ast_nodes.FunctionDecl( - "func_i", ast_nodes.ValueType("int"), [], "doc" - ), - constants.PRIMITIVE_TYPES, - ) - ) - self.assertFalse( - function_utils.return_is_value_of_type( - ast_nodes.FunctionDecl( - "func_s", ast_nodes.ValueType("MyStruct"), [], "doc" - ), - constants.PRIMITIVE_TYPES, - ) - ) - - def test_return_is_pointer_to_struct(self): - self.assertTrue( - function_utils.return_is_pointer_to_struct(self.func_ret_ptr_struct) - ) - self.assertFalse( - function_utils.return_is_pointer_to_struct(self.func_ret_ptr_int) - ) - - def test_return_is_pointer_to_primitive(self): - self.assertTrue( - function_utils.return_is_pointer_to_primitive(self.func_ret_ptr_int) - ) - self.assertFalse( - function_utils.return_is_pointer_to_primitive(self.func_ret_ptr_struct) - ) - - def test_param_is_primitive_value(self): - param_prim_val = ast_nodes.FunctionParameterDecl( - "prim_v", ast_nodes.ValueType("int") - ) - param_arr = ast_nodes.FunctionParameterDecl( - "arr_v", ast_nodes.ArrayType(ast_nodes.ValueType("int"), extents=(10,)) - ) - - self.assertTrue(function_utils.param_is_primitive_value(param_prim_val)) - self.assertFalse(function_utils.param_is_primitive_value(param_arr)) - - def test_param_is_pointer_to_primitive_value(self): - param_ptr_to_prim = ast_nodes.FunctionParameterDecl( - "p_prim", self.ptr_to_int - ) - param_arr_of_prim = ast_nodes.FunctionParameterDecl( - "a_prim", ast_nodes.ArrayType(ast_nodes.ValueType("int"), extents=(10,)) - ) - param_ptr_to_struct = ast_nodes.FunctionParameterDecl( - name="p_struct", type=ast_nodes.PointerType(inner_type=self.struct_type) - ) - self.assertTrue( - function_utils.param_is_pointer_to_primitive_value(param_ptr_to_prim) - ) - self.assertTrue( - function_utils.param_is_pointer_to_primitive_value(param_arr_of_prim) - ) - self.assertFalse( - function_utils.param_is_pointer_to_primitive_value(param_ptr_to_struct) - ) - - def test_param_is_pointer_to_struct(self): - param_arr_of_struct = ast_nodes.FunctionParameterDecl( - "a_struct", ast_nodes.ArrayType(self.struct_type, extents=(5,)) - ) - param_ptr_to_struct = ast_nodes.FunctionParameterDecl( - "p_struct", ast_nodes.PointerType(self.struct_type) - ) - param_ptr_to_ptr = ast_nodes.FunctionParameterDecl( - "p_ptr", ast_nodes.PointerType(self.ptr_to_int) - ) - self.assertTrue( - function_utils.param_is_pointer_to_struct(param_arr_of_struct) - ) - self.assertTrue( - function_utils.param_is_pointer_to_struct(param_ptr_to_struct) - ) - self.assertFalse( - function_utils.param_is_pointer_to_struct(param_ptr_to_ptr) - ) - - def test_should_be_wrapped_with_primitive_ptr_return(self): - func = ast_nodes.FunctionDecl( - name="get_data", - return_type=ast_nodes.PointerType(ast_nodes.ValueType("int")), - parameters=tuple(), - doc="Returns int pointer", - ) - self.assertTrue(function_utils.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", - ) - result = function_utils.generate_function_wrapper(func) - self.assertEqual( - result, - """int get_id_wrapper() -{ - return get_id(); -}""", - ) - - 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), - ), - ), - 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", - ) - result = function_utils.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); -}""", - ) - - def test_get_params_string_with_struct_ptr(self): - param = ast_nodes.FunctionParameterDecl( - name="my_struct", - type=ast_nodes.PointerType(ast_nodes.ValueType("mystruct")), - ) - result = function_utils.get_params_string((param,)) - self.assertEqual(result, ["Mystruct& my_struct"]) - - def test_get_params_string_maybe_with_conversion_struct_ptr(self): - param = ast_nodes.FunctionParameterDecl( - name="s", - type=ast_nodes.PointerType(ast_nodes.ValueType("customstruct")), - ) - result = function_utils.get_params_string_maybe_with_conversion((param,)) - self.assertEqual(result, ["s.get()"]) - - def test_get_compatible_return_call(self): - func = ast_nodes.FunctionDecl( - name="noop", - return_type=ast_nodes.ValueType("void"), - parameters=tuple(), - doc="does nothing", - ) - result = function_utils.get_compatible_return_call(func, "noop()") - self.assertEqual(result, "noop()") - - def test_get_compatible_return_type(self): - func = ast_nodes.FunctionDecl( - name="get_name", - return_type=ast_nodes.PointerType(ast_nodes.ValueType("char")), - parameters=tuple(), - doc="returns name", - ) - result = function_utils.get_compatible_return_type(func) - self.assertEqual(result.strip(), "std::string") - - def test_get_converted_struct_to_class(self): - func = ast_nodes.FunctionDecl( - name="get_struct", - return_type=ast_nodes.PointerType(ast_nodes.ValueType("mystruct")), - parameters=tuple(), - doc="returns struct", - ) - result = function_utils.get_converted_struct_to_class(func, "get_struct()") - self.assertIn("mystruct* result = get_struct();", result) - self.assertIn("return Mystruct(result)", result) - - def test_is_excluded_function_name(self): - self.assertTrue(function_utils.is_excluded_function_name("mjr_function")) - self.assertTrue(function_utils.is_excluded_function_name("mjui_function")) - self.assertTrue(function_utils.is_excluded_function_name("mju_malloc")) - self.assertTrue(function_utils.is_excluded_function_name("mj_makeData")) - self.assertFalse( - function_utils.is_excluded_function_name("mjv_updateScene") - ) - self.assertFalse( - function_utils.is_excluded_function_name("mj_normalFunction") - ) - self.assertFalse( - function_utils.is_excluded_function_name("mju_someOtherFunction") - ) - - -if __name__ == "__main__": - absltest.main() diff --git a/wasm/codegen/helpers/helpers_test.py b/wasm/codegen/helpers/helpers_test.py new file mode 100644 index 00000000..8b02c997 --- /dev/null +++ b/wasm/codegen/helpers/helpers_test.py @@ -0,0 +1,907 @@ +# Copyright 2025 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for helper modules.""" + +from absl.testing import absltest +from introspect import ast_nodes + +from wasm.codegen.helpers import code_builder +from wasm.codegen.helpers import common +from wasm.codegen.helpers import constants +from wasm.codegen.helpers import function_utils +from wasm.codegen.helpers import struct_constructor_code_builder +from wasm.codegen.helpers import struct_field_code_builder +from wasm.codegen.helpers import struct_field_handler +from wasm.codegen.helpers import structs_parser + + +class CodeBuilderTest(absltest.TestCase): + + def test_code_builder_functionality(self): + """Test nested indentation blocks.""" + builder = code_builder.CodeBuilder(indent_str=" ") + builder.line("let a = 1") + with builder.block("function myFunc()"): + builder.line("let flag = true") + with builder.block("while (flag)"): + builder.line("a++") + builder.line("flag = a < 10") + builder.line("return a") + builder.line("print('Done')") + + expected_lines = [ + "let a = 1", + "function myFunc() {", + " let flag = true", + " while (flag) {", + " a++", + " flag = a < 10", + " }", + " return a", + "}", + "print('Done')", + ] + self.assertEqual(builder.to_string(), "\n".join(expected_lines)) + + +class CommonUtilsTest(absltest.TestCase): + + def test_uppercase_first_letter(self): + self.assertEqual(common.uppercase_first_letter(""), "") + self.assertEqual(common.uppercase_first_letter("hello"), "Hello") + self.assertEqual(common.uppercase_first_letter("1st place"), "1st place") + self.assertEqual(common.uppercase_first_letter("!wow"), "!wow") + self.assertEqual( + common.uppercase_first_letter(" leading space"), " leading space" + ) + + def test_try_cast_to_scalar_type(self): + self.assertEqual(common.try_cast_to_scalar_type("123"), 123) + self.assertEqual(common.try_cast_to_scalar_type("123.456"), 123.456) + self.assertEqual(common.try_cast_to_scalar_type("abc"), "abc") + + +class FunctionUtilsTest(absltest.TestCase): + + def setUp(self): + super().setUp() + self.struct_type = ast_nodes.ValueType("MyStruct") + self.ptr_to_int = ast_nodes.PointerType(ast_nodes.ValueType("int")) + self.func_ret_ptr_int = ast_nodes.FunctionDecl( + "func_pi", ast_nodes.PointerType(ast_nodes.ValueType("int")), [], "doc" + ) + self.func_ret_ptr_struct = ast_nodes.FunctionDecl( + "func_ps", + ast_nodes.PointerType(ast_nodes.ValueType("MyStruct")), + [], + "doc", + ) + + def test_return_is_value_of_type(self): + self.assertTrue( + function_utils.return_is_value_of_type( + ast_nodes.FunctionDecl( + "func_i", ast_nodes.ValueType("int"), [], "doc" + ), + constants.PRIMITIVE_TYPES, + ) + ) + self.assertFalse( + function_utils.return_is_value_of_type( + ast_nodes.FunctionDecl( + "func_s", ast_nodes.ValueType("MyStruct"), [], "doc" + ), + constants.PRIMITIVE_TYPES, + ) + ) + + def test_return_is_pointer_to_struct(self): + self.assertTrue( + function_utils.return_is_pointer_to_struct(self.func_ret_ptr_struct) + ) + self.assertFalse( + function_utils.return_is_pointer_to_struct(self.func_ret_ptr_int) + ) + + def test_return_is_pointer_to_primitive(self): + self.assertTrue( + function_utils.return_is_pointer_to_primitive(self.func_ret_ptr_int) + ) + self.assertFalse( + function_utils.return_is_pointer_to_primitive(self.func_ret_ptr_struct) + ) + + def test_param_is_primitive_value(self): + param_prim_val = ast_nodes.FunctionParameterDecl( + "prim_v", ast_nodes.ValueType("int") + ) + param_arr = ast_nodes.FunctionParameterDecl( + "arr_v", ast_nodes.ArrayType(ast_nodes.ValueType("int"), extents=(10,)) + ) + + self.assertTrue(function_utils.param_is_primitive_value(param_prim_val)) + self.assertFalse(function_utils.param_is_primitive_value(param_arr)) + + def test_param_is_pointer_to_primitive_value(self): + param_ptr_to_prim = ast_nodes.FunctionParameterDecl( + "p_prim", self.ptr_to_int + ) + param_arr_of_prim = ast_nodes.FunctionParameterDecl( + "a_prim", ast_nodes.ArrayType(ast_nodes.ValueType("int"), extents=(10,)) + ) + param_ptr_to_struct = ast_nodes.FunctionParameterDecl( + name="p_struct", type=ast_nodes.PointerType(inner_type=self.struct_type) + ) + self.assertTrue( + function_utils.param_is_pointer_to_primitive_value(param_ptr_to_prim) + ) + self.assertTrue( + function_utils.param_is_pointer_to_primitive_value(param_arr_of_prim) + ) + self.assertFalse( + function_utils.param_is_pointer_to_primitive_value(param_ptr_to_struct) + ) + + def test_param_is_pointer_to_struct(self): + param_arr_of_struct = ast_nodes.FunctionParameterDecl( + "a_struct", ast_nodes.ArrayType(self.struct_type, extents=(5,)) + ) + param_ptr_to_struct = ast_nodes.FunctionParameterDecl( + "p_struct", ast_nodes.PointerType(self.struct_type) + ) + param_ptr_to_ptr = ast_nodes.FunctionParameterDecl( + "p_ptr", ast_nodes.PointerType(self.ptr_to_int) + ) + self.assertTrue( + function_utils.param_is_pointer_to_struct(param_arr_of_struct) + ) + self.assertTrue( + function_utils.param_is_pointer_to_struct(param_ptr_to_struct) + ) + self.assertFalse( + function_utils.param_is_pointer_to_struct(param_ptr_to_ptr) + ) + + def test_should_be_wrapped_with_primitive_ptr_return(self): + func = ast_nodes.FunctionDecl( + name="get_data", + return_type=ast_nodes.PointerType(ast_nodes.ValueType("int")), + parameters=tuple(), + doc="Returns int pointer", + ) + self.assertTrue(function_utils.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", + ) + result = function_utils.generate_function_wrapper(func) + self.assertEqual( + result, + """int get_id_wrapper() +{ + return get_id(); +}""", + ) + + 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), + ), + ), + 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", + ) + result = function_utils.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); +}""", + ) + + def test_get_params_string_with_struct_ptr(self): + param = ast_nodes.FunctionParameterDecl( + name="my_struct", + type=ast_nodes.PointerType(ast_nodes.ValueType("mystruct")), + ) + result = function_utils.get_params_string((param,)) + self.assertEqual(result, ["Mystruct& my_struct"]) + + def test_get_params_string_maybe_with_conversion_struct_ptr(self): + param = ast_nodes.FunctionParameterDecl( + name="s", + type=ast_nodes.PointerType(ast_nodes.ValueType("customstruct")), + ) + result = function_utils.get_params_string_maybe_with_conversion((param,)) + self.assertEqual(result, ["s.get()"]) + + def test_get_compatible_return_call(self): + func = ast_nodes.FunctionDecl( + name="noop", + return_type=ast_nodes.ValueType("void"), + parameters=tuple(), + doc="does nothing", + ) + result = function_utils.get_compatible_return_call(func, "noop()") + self.assertEqual(result, "noop()") + + def test_get_compatible_return_type(self): + func = ast_nodes.FunctionDecl( + name="get_name", + return_type=ast_nodes.PointerType(ast_nodes.ValueType("char")), + parameters=tuple(), + doc="returns name", + ) + result = function_utils.get_compatible_return_type(func) + self.assertEqual(result.strip(), "std::string") + + def test_get_converted_struct_to_class(self): + func = ast_nodes.FunctionDecl( + name="get_struct", + return_type=ast_nodes.PointerType(ast_nodes.ValueType("mystruct")), + parameters=tuple(), + doc="returns struct", + ) + result = function_utils.get_converted_struct_to_class(func, "get_struct()") + self.assertIn("mystruct* result = get_struct();", result) + self.assertIn("return Mystruct(result)", result) + + def test_is_excluded_function_name(self): + self.assertTrue(function_utils.is_excluded_function_name("mjr_function")) + self.assertTrue(function_utils.is_excluded_function_name("mjui_function")) + self.assertTrue(function_utils.is_excluded_function_name("mju_malloc")) + self.assertTrue(function_utils.is_excluded_function_name("mj_makeData")) + self.assertFalse( + function_utils.is_excluded_function_name("mjv_updateScene") + ) + self.assertFalse( + function_utils.is_excluded_function_name("mj_normalFunction") + ) + self.assertFalse( + function_utils.is_excluded_function_name("mju_someOtherFunction") + ) + + +class StructConstructorCodeBuilderTest(absltest.TestCase): + + def test_constructor_code_with_default_function(self): + wrapped_structs = structs_parser.generate_wasm_bindings(["mjLROpt"]) + self.assertEqual( + 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( + wrapped_structs["mjsElement"].wrapped_source, + """ +MjsElement::MjsElement(mjsElement *ptr) : ptr_(ptr) {} +MjsElement::~MjsElement() {} +std::unique_ptr MjsElement::copy() { + return std::make_unique(*this); +} +""".strip(), + ) + + def test_constructor_code_with_fields_with_init(self): + field_with_init = ast_nodes.StructFieldDecl( + name="element", + type=ast_nodes.PointerType( + inner_type=ast_nodes.ValueType(name="mjsElement"), + ), + doc="", + ) + wrapped_field_data = struct_field_handler.StructFieldHandler( + field_with_init, "MjsTexture" + ).generate() + self.assertEqual( + struct_constructor_code_builder.build_struct_source( + "mjsTexture", + [wrapped_field_data], + ), + """ +MjsTexture::MjsTexture(mjsTexture *ptr) : ptr_(ptr), element(ptr_->element) {} +MjsTexture::~MjsTexture() {} +""".strip(), + ) + + def test_constructor_code_with_shallow_copy(self): + self.assertEqual( + struct_constructor_code_builder.build_struct_source("mjvLight", []), + """MjvLight::MjvLight(mjvLight *ptr) : ptr_(ptr) {} +MjvLight::MjvLight() : ptr_(new mjvLight) { + owned_ = true; +} +MjvLight::MjvLight(const MjvLight &other) : MjvLight() { + *ptr_ = *other.get(); +} +MjvLight& MjvLight::operator=(const MjvLight &other) { + if (this == &other) { + return *this; + } + *ptr_ = *other.get(); + return *this; +} +MjvLight::~MjvLight() { + if (owned_ && ptr_) delete ptr_; +} +std::unique_ptr MjvLight::copy() { + return std::make_unique(*this); +}""".strip(), + ) + + +def test_build_struct_header_with_nested_wrappers(self): + self.assertEqual( + 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 MjLROpt { + MjLROpt(); + MjLROpt(const MjLROpt &); + MjLROpt &operator=(const MjLROpt &); + explicit MjLROpt(mjLROpt *ptr); + ~MjLROpt(); + mjLROpt* get() const { return ptr_; } + void set(mjLROpt* ptr) { ptr_ = ptr; } + + private: + mjLROpt* ptr_; + bool owned_ = false; +}; +""".strip(), + ) + + +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( + struct_field_code_builder.build_primitive_type_definition(field), + """ +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( + struct_field_code_builder.build_memory_view_definition( + 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)); +} +""".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( + struct_field_code_builder.build_string_field_definition(field), + """ +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( + struct_field_code_builder.build_mjvec_pointer_definition( + field, "mjDoubleVec" + ), + """ +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( + struct_field_code_builder.build_mjvec_pointer_definition( + field, "mjByteVec" + ), + """ +std::vector &vector_field() const { + return *(reinterpret_cast*>(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( + struct_field_code_builder.build_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( + struct_field_code_builder.build_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( + struct_field_code_builder.build_simple_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( + struct_field_code_builder.build_primitive_type_definition(field), + """ +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( + struct_field_code_builder.build_memory_view_definition( + 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)); +} +""".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( + struct_field_code_builder.build_string_field_definition(field), + """ +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( + struct_field_code_builder.build_mjvec_pointer_definition( + field, "mjDoubleVec" + ), + """ +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( + struct_field_code_builder.build_mjvec_pointer_definition( + field, "mjByteVec" + ), + """ +std::vector &vector_field() const { + return *(reinterpret_cast*>(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( + struct_field_code_builder.build_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( + struct_field_code_builder.build_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( + struct_field_code_builder.build_simple_property_binding( + field, + "MjModel", + add_setter=True, + add_return_value_policy_as_ref=True, + ), + '.property("ngeom", &MjModel::ngeom, &MjModel::set_ngeom, reference())', + ) + + +class StructFieldHandlerTest(absltest.TestCase): + + def test_scalar_field(self): + """Test that a scalar type field is handled correctly.""" + field_scalar = ast_nodes.StructFieldDecl( + name="ngeom", + type=ast_nodes.ValueType(name="int"), + doc="number of geoms", + ) + + field_handler_scalar = struct_field_handler.StructFieldHandler( + field_scalar, "MjModel" + ) + wrapped_field_data = field_handler_scalar.generate() + self.assertEqual( + wrapped_field_data.definition, + """ +int ngeom() const { + return ptr_->ngeom; +} +void set_ngeom(int value) { + ptr_->ngeom = value; +} +""".strip(), + ) + self.assertEqual( + wrapped_field_data.binding, + '.property("ngeom", &MjModel::ngeom, &MjModel::set_ngeom, reference())', + ) + + def test_pointer_type_field(self): + """Test that a pointer type field is handled correctly.""" + 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), + ) + wrapped_field_data = struct_field_handler.StructFieldHandler( + field, "MjModel" + ).generate() + + self.assertEqual( + wrapped_field_data.definition, + (""" +emscripten::val geom_rgba() const { + return emscripten::val(emscripten::typed_memory_view(ptr_->ngeom * 4, ptr_->geom_rgba)); +} +""".strip()), + ) + + self.assertEqual( + wrapped_field_data.binding, + '.property("geom_rgba", &MjModel::geom_rgba)', + ) + + def test_pointer_type_field_for_byte_type(self): + """Test that a pointer type field for a byte type is handled correctly.""" + field = ast_nodes.StructFieldDecl( + name="buffer", + type=ast_nodes.PointerType( + inner_type=ast_nodes.ValueType(name="void"), + ), + doc="main buffer; all pointers point in it (nbuffer bytes)", + ) + wrapped_field_data = struct_field_handler.StructFieldHandler( + field, "MjData" + ).generate() + + self.assertEqual( + wrapped_field_data.definition, + (""" +emscripten::val buffer() const { + return emscripten::val(emscripten::typed_memory_view(model->nbuffer, static_cast(ptr_->buffer))); +} +""".strip()), + ) + assert wrapped_field_data.binding == '.property("buffer", &MjData::buffer)' + + def test_pointer_type_field_for_mj_struct(self): + """Test that a pointer type field for a mj struct is handled correctly.""" + field = ast_nodes.StructFieldDecl( + name="element", + type=ast_nodes.PointerType( + inner_type=ast_nodes.ValueType(name="mjsElement"), + ), + doc="", + ) + wrapped_field_data = struct_field_handler.StructFieldHandler( + field, "MjsTexture" + ).generate() + self.assertEqual(wrapped_field_data.definition, "MjsElement element;") + + self.assertEqual( + wrapped_field_data.binding, + '.property("element", &MjsTexture::element, reference())', + ) + self.assertEqual( + wrapped_field_data.initialization, + ", element(ptr_->element)", + ) + + def test_array_type_field(self): + """Test that an array type field is handled correctly.""" + field = ast_nodes.StructFieldDecl( + name="gravity", + type=ast_nodes.ArrayType( + inner_type=ast_nodes.ValueType(name="mjtNum"), + extents=(3,), + ), + doc="gravitational acceleration", + ) + wrapped_field_data = struct_field_handler.StructFieldHandler( + field, "MjOption" + ).generate() + + self.assertEqual( + wrapped_field_data.definition, + (""" +emscripten::val gravity() const { + return emscripten::val(emscripten::typed_memory_view(3, ptr_->gravity)); +} +""".strip()), + ) + self.assertEqual( + wrapped_field_data.binding, + '.property("gravity", &MjOption::gravity)', + ) + + def test_array_field_with_multi_dimensional_array(self): + """Test that multi-dimensional arrays are handled correctly.""" + field = ast_nodes.StructFieldDecl( + name="multi_dim_array", + type=ast_nodes.ArrayType( + inner_type=ast_nodes.ValueType(name="float"), + extents=(3, 4), + ), + doc="description", + ) + wrapped_field_data = struct_field_handler.StructFieldHandler( + field, "MjModel" + ).generate() + self.assertEqual( + wrapped_field_data.definition, + """ +emscripten::val multi_dim_array() const { + return emscripten::val(emscripten::typed_memory_view(12, reinterpret_cast(ptr_->multi_dim_array))); +} +""".strip(), + ) + self.assertEqual( + wrapped_field_data.binding, + '.property("multi_dim_array", &MjModel::multi_dim_array)', + ) + + def test_parse_array_extent(self): + """Test that parse_array_extent handles various cases correctly.""" + self.assertEqual( + struct_field_handler.parse_array_extent((1, 2), "MjModel", "geom_rgba"), + "1 * 2", + ) + self.assertEqual( + struct_field_handler.parse_array_extent( + (1, "ngeom"), "MjModel", "geom_rgba" + ), + "1 * ptr_->ngeom", + ) + self.assertEqual( + struct_field_handler.parse_array_extent( + (1, "mjConstant"), "MjModel", "geom_rgba" + ), + "1 * mjConstant", + ) + + def test_resolve_extent(self): + """Test that resolve_extent handles various cases correctly.""" + # for integer just return the number + self.assertEqual( + struct_field_handler.resolve_extent(1, "MjModel", "geom_rgba"), "1" + ) + + # for string that does not start with mj, it's a member of the struct + self.assertEqual( + struct_field_handler.resolve_extent("ngeom", "MjModel", "geom_rgba"), + "ptr_->ngeom", + ) + + # when it's MjData and the field is in MJDATA_SIZES, it should use ptr_-> + self.assertEqual( + struct_field_handler.resolve_extent("size_value", "MjData", "efc_AR"), + "ptr_->size_value", + ) + # when it's MjData and the field is not in MJDATA_SIZES, + # it should use model-> + self.assertEqual( + struct_field_handler.resolve_extent( + "size_value", "MjData", "data_field" + ), + "model->size_value", + ) + self.assertEqual( + struct_field_handler.resolve_extent( + "mjConstant", "MjModel", "model_field" + ), + "mjConstant", + ) + + +if __name__ == "__main__": + absltest.main() diff --git a/wasm/codegen/helpers/struct_constructor_code_builder_test.py b/wasm/codegen/helpers/struct_constructor_code_builder_test.py deleted file mode 100644 index 27e5313b..00000000 --- a/wasm/codegen/helpers/struct_constructor_code_builder_test.py +++ /dev/null @@ -1,143 +0,0 @@ -# Copyright 2025 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -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( - 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( - wrapped_structs["mjsElement"].wrapped_source, - """ -MjsElement::MjsElement(mjsElement *ptr) : ptr_(ptr) {} -MjsElement::~MjsElement() {} -std::unique_ptr MjsElement::copy() { - return std::make_unique(*this); -} -""".strip(), - ) - - def test_constructor_code_with_fields_with_init(self): - field_with_init = ast_nodes.StructFieldDecl( - name="element", - type=ast_nodes.PointerType( - inner_type=ast_nodes.ValueType(name="mjsElement"), - ), - doc="", - ) - wrapped_field_data = struct_field_handler.StructFieldHandler( - field_with_init, "MjsTexture" - ).generate() - self.assertEqual( - struct_constructor_code_builder.build_struct_source( - "mjsTexture", - [wrapped_field_data], - ), - """ -MjsTexture::MjsTexture(mjsTexture *ptr) : ptr_(ptr), element(ptr_->element) {} -MjsTexture::~MjsTexture() {} -""".strip(), - ) - - def test_constructor_code_with_shallow_copy(self): - self.assertEqual( - struct_constructor_code_builder.build_struct_source("mjvLight", []), - """MjvLight::MjvLight(mjvLight *ptr) : ptr_(ptr) {} -MjvLight::MjvLight() : ptr_(new mjvLight) { - owned_ = true; -} -MjvLight::MjvLight(const MjvLight &other) : MjvLight() { - *ptr_ = *other.get(); -} -MjvLight& MjvLight::operator=(const MjvLight &other) { - if (this == &other) { - return *this; - } - *ptr_ = *other.get(); - return *this; -} -MjvLight::~MjvLight() { - if (owned_ && ptr_) delete ptr_; -} -std::unique_ptr MjvLight::copy() { - return std::make_unique(*this); -}""".strip(), - ) - - -def test_build_struct_header_with_nested_wrappers(self): - self.assertEqual( - 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 MjLROpt { - MjLROpt(); - MjLROpt(const MjLROpt &); - MjLROpt &operator=(const MjLROpt &); - explicit MjLROpt(mjLROpt *ptr); - ~MjLROpt(); - mjLROpt* get() const { return ptr_; } - void set(mjLROpt* ptr) { ptr_ = ptr; } - - private: - mjLROpt* ptr_; - bool owned_ = false; -}; -""".strip(), - ) - - -if __name__ == "__main__": - absltest.main() diff --git a/wasm/codegen/helpers/struct_field_code_builder_test.py b/wasm/codegen/helpers/struct_field_code_builder_test.py deleted file mode 100644 index 36ee7247..00000000 --- a/wasm/codegen/helpers/struct_field_code_builder_test.py +++ /dev/null @@ -1,166 +0,0 @@ -# Copyright 2025 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from absl.testing import absltest -from introspect import ast_nodes - -from wasm.codegen.helpers import struct_field_code_builder - - -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( - struct_field_code_builder.build_primitive_type_definition(field), - """ -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( - struct_field_code_builder.build_memory_view_definition( - 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)); -} -""".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( - struct_field_code_builder.build_string_field_definition(field), - """ -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( - struct_field_code_builder.build_mjvec_pointer_definition( - field, "mjDoubleVec" - ), - """ -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( - struct_field_code_builder.build_mjvec_pointer_definition( - field, "mjByteVec" - ), - """ -std::vector &vector_field() const { - return *(reinterpret_cast*>(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( - struct_field_code_builder.build_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( - struct_field_code_builder.build_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( - struct_field_code_builder.build_simple_property_binding( - field, - "MjModel", - add_setter=True, - add_return_value_policy_as_ref=True, - ), - '.property("ngeom", &MjModel::ngeom,' - " &MjModel::set_ngeom," - " reference())", - ) - - -if __name__ == "__main__": - absltest.main() diff --git a/wasm/codegen/helpers/struct_field_handler_test.py b/wasm/codegen/helpers/struct_field_handler_test.py deleted file mode 100644 index 6e0ca244..00000000 --- a/wasm/codegen/helpers/struct_field_handler_test.py +++ /dev/null @@ -1,235 +0,0 @@ -# Copyright 2025 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from absl.testing import absltest -from introspect import ast_nodes - -from wasm.codegen.helpers import struct_field_handler - - -class StructFieldHandlerTest(absltest.TestCase): - - def test_scalar_field(self): - """Test that a scalar type field is handled correctly.""" - field_scalar = ast_nodes.StructFieldDecl( - name='ngeom', - type=ast_nodes.ValueType(name='int'), - doc='number of geoms', - ) - - field_handler_scalar = struct_field_handler.StructFieldHandler( - field_scalar, 'MjModel' - ) - wrapped_field_data = field_handler_scalar.generate() - self.assertEqual( - wrapped_field_data.definition, - """ -int ngeom() const { - return ptr_->ngeom; -} -void set_ngeom(int value) { - ptr_->ngeom = value; -} -""".strip(), - ) - self.assertEqual( - wrapped_field_data.binding, - '.property("ngeom", &MjModel::ngeom, &MjModel::set_ngeom, reference())', - ) - - def test_pointer_type_field(self): - """Test that a pointer type field is handled correctly.""" - 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), - ) - wrapped_field_data = struct_field_handler.StructFieldHandler( - field, 'MjModel' - ).generate() - - self.assertEqual( - wrapped_field_data.definition, - (""" -emscripten::val geom_rgba() const { - return emscripten::val(emscripten::typed_memory_view(ptr_->ngeom * 4, ptr_->geom_rgba)); -} -""".strip()), - ) - - self.assertEqual( - wrapped_field_data.binding, - '.property("geom_rgba", &MjModel::geom_rgba)', - ) - - def test_pointer_type_field_for_byte_type(self): - """Test that a pointer type field for a byte type is handled correctly.""" - field = ast_nodes.StructFieldDecl( - name='buffer', - type=ast_nodes.PointerType( - inner_type=ast_nodes.ValueType(name='void'), - ), - doc='main buffer; all pointers point in it (nbuffer bytes)', - ) - wrapped_field_data = struct_field_handler.StructFieldHandler( - field, 'MjData' - ).generate() - - self.assertEqual( - wrapped_field_data.definition, - (""" -emscripten::val buffer() const { - return emscripten::val(emscripten::typed_memory_view(model->nbuffer, static_cast(ptr_->buffer))); -} -""".strip()), - ) - assert ( - wrapped_field_data.binding - == '.property("buffer", &MjData::buffer)' - ) - - def test_pointer_type_field_for_mj_struct(self): - """Test that a pointer type field for a mj struct is handled correctly.""" - field = ast_nodes.StructFieldDecl( - name='element', - type=ast_nodes.PointerType( - inner_type=ast_nodes.ValueType(name='mjsElement'), - ), - doc='', - ) - wrapped_field_data = struct_field_handler.StructFieldHandler( - field, 'MjsTexture' - ).generate() - self.assertEqual(wrapped_field_data.definition, "MjsElement element;") - - self.assertEqual( - wrapped_field_data.binding, - '.property("element", &MjsTexture::element, reference())', - ) - self.assertEqual( - wrapped_field_data.initialization, - ', element(ptr_->element)', - ) - - def test_array_type_field(self): - """Test that an array type field is handled correctly.""" - field = ast_nodes.StructFieldDecl( - name='gravity', - type=ast_nodes.ArrayType( - inner_type=ast_nodes.ValueType(name='mjtNum'), - extents=(3,), - ), - doc='gravitational acceleration', - ) - wrapped_field_data = struct_field_handler.StructFieldHandler( - field, 'MjOption' - ).generate() - - self.assertEqual( - wrapped_field_data.definition, - (""" -emscripten::val gravity() const { - return emscripten::val(emscripten::typed_memory_view(3, ptr_->gravity)); -} -""".strip()), - ) - self.assertEqual( - wrapped_field_data.binding, - '.property("gravity", &MjOption::gravity)', - ) - - def test_array_field_with_multi_dimensional_array(self): - """Test that multi-dimensional arrays are handled correctly.""" - field = ast_nodes.StructFieldDecl( - name='multi_dim_array', - type=ast_nodes.ArrayType( - inner_type=ast_nodes.ValueType(name='float'), - extents=(3, 4), - ), - doc='description', - ) - wrapped_field_data = struct_field_handler.StructFieldHandler( - field, 'MjModel' - ).generate() - self.assertEqual( - wrapped_field_data.definition, - """ -emscripten::val multi_dim_array() const { - return emscripten::val(emscripten::typed_memory_view(12, reinterpret_cast(ptr_->multi_dim_array))); -} -""".strip(), - ) - self.assertEqual( - wrapped_field_data.binding, - '.property("multi_dim_array", &MjModel::multi_dim_array)', - ) - - def test_parse_array_extent(self): - """Test that parse_array_extent handles various cases correctly.""" - self.assertEqual( - struct_field_handler.parse_array_extent((1, 2), 'MjModel', 'geom_rgba'), - '1 * 2', - ) - self.assertEqual( - struct_field_handler.parse_array_extent( - (1, 'ngeom'), 'MjModel', 'geom_rgba' - ), - '1 * ptr_->ngeom', - ) - self.assertEqual( - struct_field_handler.parse_array_extent( - (1, 'mjConstant'), 'MjModel', 'geom_rgba' - ), - '1 * mjConstant', - ) - - def test_resolve_extent(self): - """Test that resolve_extent handles various cases correctly.""" - # for integer just return the number - self.assertEqual( - struct_field_handler.resolve_extent(1, 'MjModel', 'geom_rgba'), '1' - ) - - # for string that does not start with mj, it's a member of the struct - self.assertEqual( - struct_field_handler.resolve_extent('ngeom', 'MjModel', 'geom_rgba'), - 'ptr_->ngeom', - ) - - # when it's MjData and the field is in MJDATA_SIZES, it should use ptr_-> - self.assertEqual( - struct_field_handler.resolve_extent('size_value', 'MjData', 'efc_AR'), - 'ptr_->size_value', - ) - # when it's MjData and the field is not in MJDATA_SIZES, - # it should use model-> - self.assertEqual( - struct_field_handler.resolve_extent( - 'size_value', 'MjData', 'data_field' - ), - 'model->size_value', - ) - self.assertEqual( - struct_field_handler.resolve_extent( - 'mjConstant', 'MjModel', 'model_field' - ), - 'mjConstant', - ) - - -if __name__ == '__main__': - absltest.main() diff --git a/wasm/codegen/helpers/structs_parser_test.py b/wasm/codegen/helpers/structs_parser_test.py deleted file mode 100644 index 864928bc..00000000 --- a/wasm/codegen/helpers/structs_parser_test.py +++ /dev/null @@ -1,90 +0,0 @@ -# Copyright 2025 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Tests for structs_parser.""" - -from absl.testing import absltest -from introspect import ast_nodes -from introspect import structs as introspect_structs -from wasm.codegen.helpers import structs_parser -from wasm.codegen.helpers import structs_wrappers_data - - -class StructsParserTest(absltest.TestCase): - - def setUp(self): - super().setUp() - self.wrapped_structs = structs_parser.generate_wasm_bindings([ - "mjModel", - "mjData", - "mjVisualGlobal", - "mjVisualQuality", - "mjVisual", - ]) - - def test_sort_structs_by_dependency(self): - mock_introspect_structs = { - "mjA": ast_nodes.StructDecl( - name="mjA", - declname="mjA", - fields=[ - ast_nodes.StructFieldDecl( - doc="", name="b_field", type=ast_nodes.ValueType(name="mjB") - ) - ], - ), - "mjB": ast_nodes.StructDecl( - name="mjB", - declname="mjB", - fields=[ - ast_nodes.StructFieldDecl( - doc="", name="c_field", type=ast_nodes.ValueType(name="mjC") - ) - ], - ), - "mjC": ast_nodes.StructDecl(name="mjC", declname="mjC", fields=[]), - "mjD": ast_nodes.StructDecl(name="mjD", declname="mjD", fields=[]), - } - with absltest.mock.patch.dict( - introspect_structs.STRUCTS, mock_introspect_structs - ): - struct_names = ["mjA", "mjB", "mjC", "mjD"] - sorted_names = structs_parser.sort_structs_by_dependency(struct_names) - self.assertEqual(sorted_names, ["mjC", "mjD", "mjB", "mjA"]) - - def test_generate_wasm_bindings(self): - self.assertEqual(self.wrapped_structs["mjModel"].wrap_name, "MjModel") - self.assertEqual(self.wrapped_structs["mjData"].wrap_name, "MjData") - self.assertEqual( - self.wrapped_structs["mjVisualGlobal"].wrap_name, "MjVisualGlobal" - ) - self.assertEqual( - self.wrapped_structs["mjVisualQuality"].wrap_name, "MjVisualQuality" - ) - self.assertEqual(self.wrapped_structs["mjVisual"].wrap_name, "MjVisual") - self.assertNotEmpty(self.wrapped_structs["mjModel"].wrapped_fields) - self.assertNotEmpty(self.wrapped_structs["mjData"].wrapped_fields) - self.assertNotEmpty(self.wrapped_structs["mjVisualGlobal"].wrapped_fields) - self.assertNotEmpty(self.wrapped_structs["mjVisualQuality"].wrapped_fields) - self.assertNotEmpty(self.wrapped_structs["mjVisual"].wrapped_fields) - - def test_generate_wasm_bindings_with_error(self): - with self.assertRaises(RuntimeError): - structs_parser.generate_wasm_bindings(["mjFakeStruct"]) - with self.assertRaises(RuntimeError): - structs_parser.generate_wasm_bindings(["mjFakeAnonymousStruct"]) - - -if __name__ == "__main__": - absltest.main()