chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Various minor cleanups to WASM code generation

Browse Source 
* Add newline function to code_builder
* Remove type aliases which made jump to definition less ergonomic
* Add types to constants

PiperOrigin-RevId: 827863199
Change-Id: Ib4390d398926466ecce491a3d0ebe308ed5e0cbe
This commit is contained in:
Matija Kecman
2025-11-04 02:06:05 -08:00
committed by Copybara-Service
parent 10130297c0
commit c22d94e470
12 changed files with 190 additions and 218 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wasm/codegen/generators/functions.py
+5 -11
View File
@@ -15,26 +15,20 @@
"""Generates Embind bindings for MuJoCo functions."""
import pathlib
from typing import List, Mapping, TypeAlias
from typing import List, Mapping
from introspect import ast_nodes
from wasm.codegen.helpers import code_builder
from wasm.codegen.helpers import function_utils
FunctionDecl: TypeAlias = ast_nodes.FunctionDecl
FunctionParameterDecl: TypeAlias = ast_nodes.FunctionParameterDecl
PointerType: TypeAlias = ast_nodes.PointerType
ValueType: TypeAlias = ast_nodes.ValueType
Path: TypeAlias = pathlib.Path
class Generator:
"""Generates Embind bindings for MuJoCo functions."""
def __init__(self, functions: Mapping[str, FunctionDecl]):
self.direct_bind_functions: List[FunctionDecl] = []
self.wrapper_bind_functions: List[FunctionDecl] = []
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 function_utils.should_be_wrapped(func):
@@ -63,7 +57,7 @@ class Generator:
return result
def _generate_function_binding(
self, func: FunctionDecl, is_wrapper=False
self, func: ast_nodes.FunctionDecl, is_wrapper=False
) -> str:
"""Generates the Embind code for a single function."""
wasm/codegen/helpers/code_builder.py
+4
View File
@@ -36,6 +36,10 @@ class CodeBuilder:
else:
self._lines.append("")
def newline(self) -> None:
"""Adds a newline."""
self.line("")
def to_string(self) -> str:
"""Returns the complete code string."""
return "\n".join(self._lines)
wasm/codegen/helpers/common.py
+1 -3
View File
@@ -19,13 +19,11 @@ import pathlib
from wasm.codegen.helpers import constants
Path = pathlib.Path
def get_default_output_dir() -> str:
"""Gets the default output directory (sibling of 'generated' folder)."""
# Get the directory of the current file (generator/base.py)
current_dir = Path(__file__).parent
current_dir = pathlib.Path(__file__).parent
# Go up one level to the project root and then down to 'generated'
default_output_dir = str(current_dir.parent / "generated")
return default_output_dir
wasm/codegen/helpers/constants.py
+23 -22
View File
@@ -14,9 +14,10 @@
"""Constants used in the code generation process."""
from typing import List, Set, Dict
from introspect import structs as introspect_structs
PRIMITIVE_TYPES = {
PRIMITIVE_TYPES: Set[str] = {
# go/keep-sorted start
"char",
"double",
@@ -36,7 +37,7 @@ PRIMITIVE_TYPES = {
# go/keep-sorted end
}
_PLUGIN_FUNCTIONS = [
_PLUGIN_FUNCTIONS: List[str] = [
# go/keep-sorted start
"mj_getPluginConfig",
"mj_loadAllPluginLibraries",
@@ -61,7 +62,7 @@ _PLUGIN_FUNCTIONS = [
]
# Functions that are bound as class methods
_CLASS_METHODS = [
_CLASS_METHODS: List[str] = [
# go/keep-sorted start
"mj_compile",
"mj_copyData",
@@ -83,12 +84,12 @@ _CLASS_METHODS = [
]
# Omitted because not very useful
_WRITABLE_ERROR = [
_WRITABLE_ERROR: List[str] = [
"mj_printSchema",
]
# Omitted thread management functions
_THREAD_FUNCTIONS = [
_THREAD_FUNCTIONS: List[str] = [
# go/keep-sorted start
"mju_bindThreadPool",
"mju_defaultTask",
@@ -100,7 +101,7 @@ _THREAD_FUNCTIONS = [
]
# Omitted asset cache functions
_ASSET_CACHE_FUNCTIONS = [
_ASSET_CACHE_FUNCTIONS: List[str] = [
# go/keep-sorted start
"mj_clearCache",
"mj_getCache",
@@ -111,7 +112,7 @@ _ASSET_CACHE_FUNCTIONS = [
]
# Omitted Virtual Filesystem (VFS) functions
_VFS_FUNCTIONS = [
_VFS_FUNCTIONS: List[str] = [
# go/keep-sorted start
"mj_addBufferVFS",
"mj_addFileVFS",
@@ -122,7 +123,7 @@ _VFS_FUNCTIONS = [
]
# Omitted irrelevant visual functions
_VISUAL_FUNCTIONS = [
_VISUAL_FUNCTIONS: List[str] = [
# go/keep-sorted start
"mjv_averageCamera",
"mjv_copyData",
@@ -133,7 +134,7 @@ _VISUAL_FUNCTIONS = [
# go/keep-sorted end
]
_MEMORY_FUNCTIONS = [
_MEMORY_FUNCTIONS: List[str] = [
# go/keep-sorted start
"mj_freeLastXML",
"mj_freeStack",
@@ -159,7 +160,7 @@ _MEMORY_FUNCTIONS = [
# go/keep-sorted end
]
_GETTERS_AND_SETTERS = [
_GETTERS_AND_SETTERS: List[str] = [
# go/keep-sorted start
"mjs_appendFloatVec",
"mjs_appendIntVec",
@@ -180,14 +181,14 @@ _GETTERS_AND_SETTERS = [
# go/keep-sorted end
]
_UTILITY_FUNCTIONS = [
_UTILITY_FUNCTIONS: List[str] = [
# go/keep-sorted start
"mju_getXMLDependencies",
# go/keep-sorted end
]
# List of functions that should be skipped during the code generation process.
SKIPPED_FUNCTIONS = (
SKIPPED_FUNCTIONS: List[str] = (
_CLASS_METHODS +
_THREAD_FUNCTIONS +
_MEMORY_FUNCTIONS +
@@ -203,7 +204,7 @@ SKIPPED_FUNCTIONS = (
# Functions that require special wrappers to infer sizes and make additional
# validation checks. These functions are not bound automatically but are
# written by hand instead.
BOUNDCHECK_FUNCS = [
BOUNDCHECK_FUNCS: List[str] = [
# go/keep-sorted start
"mj_addM",
"mj_angmomMat",
@@ -287,7 +288,7 @@ BOUNDCHECK_FUNCS = [
]
# List of structs that should be skipped during the code generation process.
SKIPPED_STRUCTS = [
SKIPPED_STRUCTS: List[str] = [
# go/keep-sorted start
"mjCache",
"mjSDF",
@@ -309,7 +310,7 @@ SKIPPED_STRUCTS = [
# Anonymous structs are not defined as independent structs in the MuJoCo
# codebase, but they are part of other structs. This dictionary is used to
# handle them as if they were independent structs.
ANONYMOUS_STRUCTS = {
ANONYMOUS_STRUCTS: Dict[str, Dict[str, str]] = {
# go/keep-sorted start
"mjVisualGlobal": {"parent": "mjVisual", "field_name": "global"},
"mjVisualHeadlight": {"parent": "mjVisual", "field_name": "headlight"},
@@ -322,14 +323,14 @@ ANONYMOUS_STRUCTS = {
# This list is created by subtracting the skipped structs from the list of all
# structs and adding the anonymous structs.
STRUCTS_TO_BIND = list(
STRUCTS_TO_BIND: List[str] = list(
(set(introspect_structs.STRUCTS.keys()) - set(SKIPPED_STRUCTS)).union(
ANONYMOUS_STRUCTS.keys()
)
)
# List of structs that do not have a default constructor.
NO_DEFAULT_CONSTRUCTORS = [
NO_DEFAULT_CONSTRUCTORS: List[str] = [
# go/keep-sorted start
"mjContact",
"mjSolverStat",
@@ -349,7 +350,7 @@ NO_DEFAULT_CONSTRUCTORS = [
# List of `mjData` fields where the array size should be obtained from other
# `mjData` members, instead of from `mjModel` members. This is typically the
# case for fields that are dynamically allocated during the simulation.
MJDATA_SIZES = [
MJDATA_SIZES: List[str] = [
# go/keep-sorted start
"contact",
"efc_AR",
@@ -421,7 +422,7 @@ MJDATA_SIZES = [
# Dictionary where keys are the struct names and the values are lists of the
# fields that are manually specified in the bindings.cc template file.
MANUALLY_ADDED_FIELDS_FROM_TEMPLATE = {
MANUALLY_ADDED_FIELDS_FROM_TEMPLATE: Dict[str, List[str]] = {
# go/keep-sorted start
"MjData": ["solver", "timer", "warning", "contact"],
"MjSpec": ["option", "visual", "stat", "element", "compiler"],
@@ -456,7 +457,7 @@ MANUALLY_ADDED_FIELDS_FROM_TEMPLATE = {
# When generating the code for these fields, a specific cast to `uint8_t*` is
# required for embind. This dictionary is used to register those fields and
# their sizes.
BYTE_FIELDS = {
BYTE_FIELDS: Dict[str, Dict[str, str]] = {
"buffer": {"size": "nbuffer"},
"arena": {"size": "narena"},
}
@@ -464,12 +465,12 @@ BYTE_FIELDS = {
# Boolean flag to enable debug prints during the struct wrapper and binding
# generation process. When set to `True`, it will print additional information
# about the steps being executed.
STRUCT_DEBUG_MODE = False
STRUCT_DEBUG_MODE: bool = False
# These structs require specific function calls for creation and/or deletion,
# or some of their fields need to be handled manually for now;
# making their wrapper constructors/destructors non-trivial.
HARDCODED_WRAPPER_STRUCTS = [
HARDCODED_WRAPPER_STRUCTS: List[str] = [
"MjData",
"MjModel",
"MjvScene",
wasm/codegen/helpers/function_utils_test.py
+61 -56
View File
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TypeAlias
from absl.testing import absltest
from introspect import ast_nodes
@@ -21,37 +19,37 @@ from wasm.codegen.helpers import constants
from wasm.codegen.helpers import function_utils
PrimitiveTypes: TypeAlias = constants.PRIMITIVE_TYPES
ValueType: TypeAlias = ast_nodes.ValueType
PointerType: TypeAlias = ast_nodes.PointerType
ArrayType: TypeAlias = ast_nodes.ArrayType
FunctionParameterDecl: TypeAlias = ast_nodes.FunctionParameterDecl
FunctionDecl: TypeAlias = ast_nodes.FunctionDecl
class FunctionUtilsTest(absltest.TestCase):
def setUp(self):
super().setUp()
self.struct_type = ValueType("MyStruct")
self.ptr_to_int = PointerType(ValueType("int"))
self.func_ret_ptr_int = FunctionDecl(
"func_pi", PointerType(ValueType("int")), [], "doc"
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 = FunctionDecl(
"func_ps", PointerType(ValueType("MyStruct")), [], "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(
FunctionDecl("func_i", ValueType("int"), [], "doc"), PrimitiveTypes
ast_nodes.FunctionDecl(
"func_i", ast_nodes.ValueType("int"), [], "doc"
),
constants.PRIMITIVE_TYPES,
)
)
self.assertFalse(
function_utils.return_is_value_of_type(
FunctionDecl("func_s", ValueType("MyStruct"), [], "doc"),
PrimitiveTypes,
ast_nodes.FunctionDecl(
"func_s", ast_nodes.ValueType("MyStruct"), [], "doc"
),
constants.PRIMITIVE_TYPES,
)
)
@@ -72,21 +70,25 @@ class FunctionUtilsTest(absltest.TestCase):
)
def test_param_is_primitive_value(self):
param_prim_val = FunctionParameterDecl("prim_v", ValueType("int"))
param_arr = FunctionParameterDecl(
"arr_v", ArrayType(ValueType("int"), extents=(10,))
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 = FunctionParameterDecl("p_prim", self.ptr_to_int)
param_arr_of_prim = FunctionParameterDecl(
"a_prim", ArrayType(ValueType("int"), extents=(10,))
param_ptr_to_prim = ast_nodes.FunctionParameterDecl(
"p_prim", self.ptr_to_int
)
param_ptr_to_struct = FunctionParameterDecl(
name="p_struct", type=PointerType(inner_type=self.struct_type)
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)
@@ -99,14 +101,14 @@ class FunctionUtilsTest(absltest.TestCase):
)
def test_param_is_pointer_to_struct(self):
param_arr_of_struct = FunctionParameterDecl(
"a_struct", ArrayType(self.struct_type, extents=(5,))
param_arr_of_struct = ast_nodes.FunctionParameterDecl(
"a_struct", ast_nodes.ArrayType(self.struct_type, extents=(5,))
)
param_ptr_to_struct = FunctionParameterDecl(
"p_struct", PointerType(self.struct_type)
param_ptr_to_struct = ast_nodes.FunctionParameterDecl(
"p_struct", ast_nodes.PointerType(self.struct_type)
)
param_ptr_to_ptr = FunctionParameterDecl(
"p_ptr", PointerType(self.ptr_to_int)
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)
@@ -119,43 +121,46 @@ class FunctionUtilsTest(absltest.TestCase):
)
def test_should_be_wrapped_with_primitive_ptr_return(self):
func = FunctionDecl(
func = ast_nodes.FunctionDecl(
name="get_data",
return_type=PointerType(ValueType("int")),
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 = FunctionDecl(
func = ast_nodes.FunctionDecl(
name="get_id",
return_type=ValueType("int"),
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()
self.assertEqual(
result,
"""int get_id_wrapper()
{
return get_id();
}""")
}""",
)
def test_generate_function_wrapper_checking_param(self):
parameters = (
FunctionParameterDecl(
ast_nodes.FunctionParameterDecl(
name="mat",
type=PointerType(
inner_type=ValueType(name="mjtNum", is_const=True),
type=ast_nodes.PointerType(
inner_type=ast_nodes.ValueType(name="mjtNum", is_const=True),
),
),
FunctionParameterDecl(
ast_nodes.FunctionParameterDecl(
name="nr",
type=ValueType(name="int"),
type=ast_nodes.ValueType(name="int"),
),
)
func = FunctionDecl(
func = ast_nodes.FunctionDecl(
name="get_id",
return_type=ValueType("int"),
return_type=ast_nodes.ValueType("int"),
parameters=parameters,
doc="Returns an integer ID",
)
@@ -170,25 +175,25 @@ class FunctionUtilsTest(absltest.TestCase):
)
def test_get_params_string_with_struct_ptr(self):
param = FunctionParameterDecl(
param = ast_nodes.FunctionParameterDecl(
name="my_struct",
type=PointerType(ValueType("mystruct")),
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 = FunctionParameterDecl(
param = ast_nodes.FunctionParameterDecl(
name="s",
type=PointerType(ValueType("customstruct")),
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 = FunctionDecl(
func = ast_nodes.FunctionDecl(
name="noop",
return_type=ValueType("void"),
return_type=ast_nodes.ValueType("void"),
parameters=tuple(),
doc="does nothing",
)
@@ -196,9 +201,9 @@ class FunctionUtilsTest(absltest.TestCase):
self.assertEqual(result, "noop()")
def test_get_compatible_return_type(self):
func = FunctionDecl(
func = ast_nodes.FunctionDecl(
name="get_name",
return_type=PointerType(ValueType("char")),
return_type=ast_nodes.PointerType(ast_nodes.ValueType("char")),
parameters=tuple(),
doc="returns name",
)
@@ -206,9 +211,9 @@ class FunctionUtilsTest(absltest.TestCase):
self.assertEqual(result.strip(), "std::string")
def test_get_converted_struct_to_class(self):
func = FunctionDecl(
func = ast_nodes.FunctionDecl(
name="get_struct",
return_type=PointerType(ValueType("mystruct")),
return_type=ast_nodes.PointerType(ast_nodes.ValueType("mystruct")),
parameters=tuple(),
doc="returns struct",
)
wasm/codegen/helpers/struct_constructor_code_builder.py
+3 -2
View File
@@ -74,14 +74,15 @@ def _build_struct_header_internal(
builder.line(f"{struct_name}* get() const {{ return ptr_; }}")
builder.line(f"void set({struct_name}* ptr) {{ ptr_ = ptr; }}")
builder.line("")
builder.newline()
builder.line("private:")
builder.line(f"{struct_name}* ptr_;")
if not is_mjs:
builder.line("bool owned_ = false;")
if is_mjs and fields_with_init:
builder.line("")
builder.newline()
builder.line("public:")
for field in fields_with_init:
if field.definition:
wasm/codegen/helpers/struct_constructor_code_builder_test.py
+3 -9
View File
@@ -17,12 +17,6 @@ from introspect import ast_nodes
from wasm.codegen.helpers import struct_constructor_code_builder
from wasm.codegen.helpers import struct_field_handler
StructFieldDecl = ast_nodes.StructFieldDecl
ValueType = ast_nodes.ValueType
PointerType = ast_nodes.PointerType
ArrayType = ast_nodes.ArrayType
StructDecl = ast_nodes.StructDecl
class StructConstructorCodeBuilderTest(absltest.TestCase):
@@ -58,10 +52,10 @@ MjLROpt::~MjLROpt() {
)
def test_constructor_code_with_fields_with_init(self):
field_with_init = StructFieldDecl(
field_with_init = ast_nodes.StructFieldDecl(
name="element",
type=PointerType(
inner_type=ValueType(name="mjsElement"),
type=ast_nodes.PointerType(
inner_type=ast_nodes.ValueType(name="mjsElement"),
),
doc="",
)
wasm/codegen/helpers/struct_field_code_builder.py
+6 -9
View File
@@ -17,13 +17,10 @@
from introspect import ast_nodes
from wasm.codegen.helpers import code_builder
StructFieldDecl = ast_nodes.StructFieldDecl
ValueType = ast_nodes.ValueType
def build_primitive_type_definition(field: StructFieldDecl) -> str:
def build_primitive_type_definition(field: ast_nodes.StructFieldDecl) -> str:
"""Builds the C++ code for a primitive type field wrapper."""
if not isinstance(field.type, ValueType):
if not isinstance(field.type, ast_nodes.ValueType):
raise ValueError(f"{field.type} must be ValueType.")
builder = code_builder.CodeBuilder()
# build getter for primitive type field
@@ -36,7 +33,7 @@ def build_primitive_type_definition(field: StructFieldDecl) -> str:
def build_memory_view_definition(
field: StructFieldDecl, array_size_str: str, ptr_expr: str
field: ast_nodes.StructFieldDecl, array_size_str: str, ptr_expr: str
) -> str:
"""Builds the C++ code for a pointer type field wrapper."""
builder = code_builder.CodeBuilder()
@@ -49,7 +46,7 @@ def build_memory_view_definition(
return builder.to_string()
def build_string_field_definition(field: StructFieldDecl) -> str:
def build_string_field_definition(field: ast_nodes.StructFieldDecl) -> str:
"""Builds the C++ code for a string type field wrapper."""
builder = code_builder.CodeBuilder()
with builder.block(f"mjString {field.name}() const"):
@@ -63,7 +60,7 @@ def build_string_field_definition(field: StructFieldDecl) -> str:
def build_mjvec_pointer_definition(
field: StructFieldDecl, vector_type: str
field: ast_nodes.StructFieldDecl, vector_type: str
) -> str:
"""Builds the C++ code for a mjVec type field wrapper."""
ptr_field_expr = f"*(ptr_->{field.name})"
@@ -79,7 +76,7 @@ def build_mjvec_pointer_definition(
def build_simple_property_binding(
field: StructFieldDecl,
field: ast_nodes.StructFieldDecl,
struct_wrapper_name: str,
add_setter: bool = False,
add_return_value_policy_as_ref: bool = False,
wasm/codegen/helpers/struct_field_code_builder_test.py
+20 -25
View File
@@ -17,18 +17,13 @@ from introspect import ast_nodes
from wasm.codegen.helpers import struct_field_code_builder
StructFieldDecl = ast_nodes.StructFieldDecl
ValueType = ast_nodes.ValueType
PointerType = ast_nodes.PointerType
ArrayType = ast_nodes.ArrayType
class StructFieldCodeBuilderTest(absltest.TestCase):
def test_primitive_type_definition(self):
field = StructFieldDecl(
field = ast_nodes.StructFieldDecl(
name="ngeom",
type=ValueType(name="int"),
type=ast_nodes.ValueType(name="int"),
doc="number of geoms",
)
self.assertEqual(
@@ -44,10 +39,10 @@ void set_ngeom(int value) {
)
def test_memory_view_definition(self):
field = StructFieldDecl(
field = ast_nodes.StructFieldDecl(
name="geom_rgba",
type=PointerType(
inner_type=ValueType(name="float"),
type=ast_nodes.PointerType(
inner_type=ast_nodes.ValueType(name="float"),
),
doc="rgba when material is omitted",
array_extent=("ngeom", 4),
@@ -64,10 +59,10 @@ emscripten::val geom_rgba() const {
)
def test_string_field_definition(self):
field = StructFieldDecl(
field = ast_nodes.StructFieldDecl(
name="string_field",
type=PointerType(
inner_type=ValueType(name="mjString"),
type=ast_nodes.PointerType(
inner_type=ast_nodes.ValueType(name="mjString"),
),
doc="rgba when material is omitted",
)
@@ -86,10 +81,10 @@ void set_string_field(const mjString& value) {
)
def test_mjvec_pointer_definition(self):
field = StructFieldDecl(
field = ast_nodes.StructFieldDecl(
name="vector_field",
type=PointerType(
inner_type=ValueType(name="mjDoubleVec"),
type=ast_nodes.PointerType(
inner_type=ast_nodes.ValueType(name="mjDoubleVec"),
),
doc="",
)
@@ -104,10 +99,10 @@ mjDoubleVec &vector_field() const {
)
def test_mjbyte_vec_pointer_definition(self):
field = StructFieldDecl(
field = ast_nodes.StructFieldDecl(
name="vector_field",
type=PointerType(
inner_type=ValueType(name="mjByteVec"),
type=ast_nodes.PointerType(
inner_type=ast_nodes.ValueType(name="mjByteVec"),
),
doc="",
)
@@ -122,9 +117,9 @@ std::vector<uint8_t> &vector_field() const {
)
def test_simple_property_binding(self):
field = StructFieldDecl(
field = ast_nodes.StructFieldDecl(
name="ngeom",
type=ValueType(name="int"),
type=ast_nodes.ValueType(name="int"),
doc="number of geoms",
)
self.assertEqual(
@@ -135,9 +130,9 @@ std::vector<uint8_t> &vector_field() const {
)
def test_simple_property_binding_with_setter(self):
field = StructFieldDecl(
field = ast_nodes.StructFieldDecl(
name="ngeom",
type=ValueType(name="int"),
type=ast_nodes.ValueType(name="int"),
doc="",
)
self.assertEqual(
@@ -149,9 +144,9 @@ std::vector<uint8_t> &vector_field() const {
)
def test_simple_property_binding_with_return_value_policy_as_ref(self):
field = StructFieldDecl(
field = ast_nodes.StructFieldDecl(
name="ngeom",
type=ValueType(name="int"),
type=ast_nodes.ValueType(name="int"),
doc="",
)
self.assertEqual(
wasm/codegen/helpers/struct_field_handler.py
+38 -45
View File
@@ -22,13 +22,6 @@ from wasm.codegen.helpers import constants
from wasm.codegen.helpers import struct_field_code_builder
from wasm.codegen.helpers import structs_wrappers_data
AnonymousStructDecl = ast_nodes.AnonymousStructDecl
ArrayType = ast_nodes.ArrayType
PointerType = ast_nodes.PointerType
StructFieldDecl = ast_nodes.StructFieldDecl
ValueType = ast_nodes.ValueType
WrappedFieldData = structs_wrappers_data.WrappedFieldData
debug_print = common.debug_print
@@ -37,7 +30,7 @@ class StructFieldHandler:
def __init__(
self,
field: StructFieldDecl,
field: ast_nodes.StructFieldDecl,
struct_wrapper_name: str,
):
self.field = field
@@ -53,27 +46,27 @@ class StructFieldHandler:
)
)
def generate(self) -> WrappedFieldData:
def generate(self) -> structs_wrappers_data.WrappedFieldData:
"""Generates the C++ definition and binding code for the struct field."""
field_type = self.field.type
if isinstance(field_type, ValueType) and (
if isinstance(field_type, ast_nodes.ValueType) and (
field_type.name in constants.PRIMITIVE_TYPES
or field_type.name.startswith("mjt")
):
return self._handle_primitive()
elif isinstance(field_type, PointerType):
elif isinstance(field_type, ast_nodes.PointerType):
return self._handle_pointer()
elif isinstance(field_type, ArrayType):
elif isinstance(field_type, ast_nodes.ArrayType):
return self._handle_array()
elif isinstance(field_type, ValueType) and field_type.name.startswith("mj"):
elif isinstance(field_type, ast_nodes.ValueType) and field_type.name.startswith("mj"):
return self._handle_mj_struct()
elif isinstance(field_type, AnonymousStructDecl):
elif isinstance(field_type, ast_nodes.AnonymousStructDecl):
return self._handle_anonymous_struct()
return self._undefined()
def _handle_primitive(self) -> WrappedFieldData:
def _handle_primitive(self) -> structs_wrappers_data.WrappedFieldData:
"""Handles the generation of C++ definition and binding code for primitive fields."""
return WrappedFieldData(
return structs_wrappers_data.WrappedFieldData(
definition=(
struct_field_code_builder.build_primitive_type_definition(
self.field
@@ -88,17 +81,17 @@ class StructFieldHandler:
is_primitive_or_fixed_size=True,
)
def _handle_pointer(self) -> WrappedFieldData:
def _handle_pointer(self) -> structs_wrappers_data.WrappedFieldData:
"""Handles the generation of C++ definition and binding code for pointer fields."""
if not isinstance(self.field.type, PointerType):
if not isinstance(self.field.type, ast_nodes.PointerType):
raise ValueError(
f"Expected PointerType, got {type(self.field.type)} for field"
f" {self.field.name}"
)
field_type: PointerType = self.field.type
field_type: ast_nodes.PointerType = self.field.type
inner_type_name = (
field_type.inner_type.name
if isinstance(field_type.inner_type, ValueType)
if isinstance(field_type.inner_type, ast_nodes.ValueType)
else ""
)
ptr_field_expr = f"ptr_->{self.field.name}"
@@ -111,9 +104,7 @@ class StructFieldHandler:
elif self.field.name in constants.BYTE_FIELDS.keys():
# for byte fields, we need to cast the pointer to uint8_t*
# so embind can correctly interpret the memory view
ptr_field_expr = (
f"static_cast<uint8_t*>({ptr_field_expr})"
)
ptr_field_expr = f"static_cast<uint8_t*>({ptr_field_expr})"
# for these byte fields, there is no array_extent, so we add the size of
# in the config file based in the documentation
extent = (constants.BYTE_FIELDS[self.field.name]["size"],)
@@ -121,7 +112,7 @@ class StructFieldHandler:
extent, self.struct_wrapper_name, self.field.name
)
elif inner_type_name == "mjString":
return WrappedFieldData(
return structs_wrappers_data.WrappedFieldData(
definition=struct_field_code_builder.build_string_field_definition(
self.field
),
@@ -133,7 +124,7 @@ class StructFieldHandler:
),
)
elif inner_type_name.startswith("mj") and inner_type_name.endswith("Vec"):
return WrappedFieldData(
return structs_wrappers_data.WrappedFieldData(
definition=struct_field_code_builder.build_mjvec_pointer_definition(
self.field, inner_type_name
),
@@ -164,11 +155,11 @@ class StructFieldHandler:
and self.struct_wrapper_name
not in constants.MANUALLY_ADDED_FIELDS_FROM_TEMPLATE.keys()
):
ptr_field = cast(PointerType, self.field.type)
ptr_field = cast(ast_nodes.PointerType, self.field.type)
wrapper_field_name = common.uppercase_first_letter(
cast(ValueType, ptr_field.inner_type).name
cast(ast_nodes.ValueType, ptr_field.inner_type).name
)
return WrappedFieldData(
return structs_wrappers_data.WrappedFieldData(
definition=f"{wrapper_field_name} {self.field.name};",
binding=struct_field_code_builder.build_simple_property_binding(
self.field,
@@ -185,7 +176,7 @@ class StructFieldHandler:
)
return self._get_manual_definition(comment_type="complex pointer field")
return WrappedFieldData(
return structs_wrappers_data.WrappedFieldData(
definition=(
struct_field_code_builder.build_memory_view_definition(
self.field, array_size_str, ptr_field_expr
@@ -194,10 +185,10 @@ class StructFieldHandler:
binding=self.simple_property_binding,
)
def _handle_array(self) -> WrappedFieldData:
def _handle_array(self) -> structs_wrappers_data.WrappedFieldData:
"""Handles the generation of C++ definition and binding code for array fields."""
field_type = self.field.type
if not isinstance(field_type, ArrayType):
if not isinstance(field_type, ast_nodes.ArrayType):
raise ValueError(
f"Expected ArrayType, got {type(field_type)} for field"
f" {self.field.name}"
@@ -205,7 +196,7 @@ class StructFieldHandler:
inner_type = field_type.inner_type
size = math.prod(field_type.extents)
if isinstance(inner_type, ValueType):
if isinstance(inner_type, ast_nodes.ValueType):
if inner_type.name in constants.PRIMITIVE_TYPES:
ptr_expr = f"ptr_->{self.field.name}"
if len(field_type.extents) > 1:
@@ -213,7 +204,7 @@ class StructFieldHandler:
# to a pointer, so embind can correctly interpret the memory
# view
ptr_expr = f"reinterpret_cast<{inner_type.name}*>({ptr_expr})"
return WrappedFieldData(
return structs_wrappers_data.WrappedFieldData(
definition=(
struct_field_code_builder.build_memory_view_definition(
self.field, str(size), ptr_expr
@@ -229,17 +220,17 @@ class StructFieldHandler:
return self._get_manual_definition(comment_type="array field")
debug_print(f"\tNOT IMPLEMENTED ARRAY field: {self.field.name}")
return WrappedFieldData(
return structs_wrappers_data.WrappedFieldData(
definition=(
f"// TODO: NOT IMPLEMENTED ARRAY wrapper for {self.field.name}"
),
binding=f"// TODO: NOT IMPLEMENTED ARRAY binding for {self.field.name}",
)
def _handle_mj_struct(self) -> WrappedFieldData:
def _handle_mj_struct(self) -> structs_wrappers_data.WrappedFieldData:
"""Handles the generation of C++ definition and binding code for mj struct fields."""
if (
isinstance(self.field.type, ValueType)
isinstance(self.field.type, ast_nodes.ValueType)
and self.field.name not in self.manually_added_fields
and self.field.type.name in constants.STRUCTS_TO_BIND
):
@@ -249,7 +240,7 @@ class StructFieldHandler:
if self.struct_wrapper_name not in constants.HARDCODED_WRAPPER_STRUCTS:
wrapper_field_name = common.uppercase_first_letter(self.field.type.name)
definition = f"{wrapper_field_name} {self.field.name};"
return WrappedFieldData(
return structs_wrappers_data.WrappedFieldData(
definition=definition,
binding=struct_field_code_builder.build_simple_property_binding(
self.field,
@@ -263,7 +254,7 @@ class StructFieldHandler:
)
return self._get_manual_definition(comment_type="struct field")
def _handle_anonymous_struct(self) -> WrappedFieldData:
def _handle_anonymous_struct(self) -> structs_wrappers_data.WrappedFieldData:
"""Handles the generation of C++ definition and binding code for anonymous struct fields."""
anonymous_struct_name = ""
@@ -277,11 +268,11 @@ class StructFieldHandler:
break
if (
isinstance(self.field.type, AnonymousStructDecl)
isinstance(self.field.type, ast_nodes.AnonymousStructDecl)
and self.field.name not in self.manually_added_fields
and anonymous_struct_name in constants.STRUCTS_TO_BIND
):
return WrappedFieldData(
return structs_wrappers_data.WrappedFieldData(
binding=struct_field_code_builder.build_simple_property_binding(
self.field,
self.struct_wrapper_name,
@@ -294,24 +285,26 @@ class StructFieldHandler:
)
return self._get_manual_definition(comment_type="anonymous struct field")
def _undefined(self) -> WrappedFieldData:
def _undefined(self) -> structs_wrappers_data.WrappedFieldData:
"""This function adds a TODO comment for fields that are not handled by this class yet."""
return WrappedFieldData(
return structs_wrappers_data.WrappedFieldData(
definition=f"// TODO: UNDEFINED definition for {self.field.name}",
binding=f"// TODO: UNDEFINED binding for {self.field.name}",
)
def _get_manual_definition(self, comment_type: str = "") -> WrappedFieldData:
def _get_manual_definition(
self, comment_type: str = ""
) -> structs_wrappers_data.WrappedFieldData:
"""Helper method to generate a comment as a definition for manually added fields."""
if self.field.name in self.manually_added_fields:
return WrappedFieldData(
return structs_wrappers_data.WrappedFieldData(
definition=(
f"// {comment_type} is defined manually. {self.field.name}"
),
binding=self.simple_property_binding,
)
return WrappedFieldData(
return structs_wrappers_data.WrappedFieldData(
definition=(
f"// TODO: Define {comment_type} manually for {self.field.name}"
),
wasm/codegen/helpers/struct_field_handler_test.py
+17 -23
View File
@@ -18,19 +18,13 @@ from introspect import ast_nodes
from wasm.codegen.helpers import struct_field_handler
StructFieldDecl = ast_nodes.StructFieldDecl
ValueType = ast_nodes.ValueType
PointerType = ast_nodes.PointerType
ArrayType = ast_nodes.ArrayType
class StructFieldHandlerTest(absltest.TestCase):
def test_scalar_field(self):
"""Test that a scalar type field is handled correctly."""
field_scalar = StructFieldDecl(
field_scalar = ast_nodes.StructFieldDecl(
name='ngeom',
type=ValueType(name='int'),
type=ast_nodes.ValueType(name='int'),
doc='number of geoms',
)
@@ -56,10 +50,10 @@ void set_ngeom(int value) {
def test_pointer_type_field(self):
"""Test that a pointer type field is handled correctly."""
field = StructFieldDecl(
field = ast_nodes.StructFieldDecl(
name='geom_rgba',
type=PointerType(
inner_type=ValueType(name='float'),
type=ast_nodes.PointerType(
inner_type=ast_nodes.ValueType(name='float'),
),
doc='rgba when material is omitted',
array_extent=('ngeom', 4),
@@ -84,10 +78,10 @@ emscripten::val geom_rgba() const {
def test_pointer_type_field_for_byte_type(self):
"""Test that a pointer type field for a byte type is handled correctly."""
field = StructFieldDecl(
field = ast_nodes.StructFieldDecl(
name='buffer',
type=PointerType(
inner_type=ValueType(name='void'),
type=ast_nodes.PointerType(
inner_type=ast_nodes.ValueType(name='void'),
),
doc='main buffer; all pointers point in it (nbuffer bytes)',
)
@@ -110,10 +104,10 @@ emscripten::val buffer() const {
def test_pointer_type_field_for_mj_struct(self):
"""Test that a pointer type field for a mj struct is handled correctly."""
field = StructFieldDecl(
field = ast_nodes.StructFieldDecl(
name='element',
type=PointerType(
inner_type=ValueType(name='mjsElement'),
type=ast_nodes.PointerType(
inner_type=ast_nodes.ValueType(name='mjsElement'),
),
doc='',
)
@@ -133,10 +127,10 @@ emscripten::val buffer() const {
def test_array_type_field(self):
"""Test that an array type field is handled correctly."""
field = StructFieldDecl(
field = ast_nodes.StructFieldDecl(
name='gravity',
type=ArrayType(
inner_type=ValueType(name='mjtNum'),
type=ast_nodes.ArrayType(
inner_type=ast_nodes.ValueType(name='mjtNum'),
extents=(3,),
),
doc='gravitational acceleration',
@@ -160,10 +154,10 @@ emscripten::val gravity() const {
def test_array_field_with_multi_dimensional_array(self):
"""Test that multi-dimensional arrays are handled correctly."""
field = StructFieldDecl(
field = ast_nodes.StructFieldDecl(
name='multi_dim_array',
type=ArrayType(
inner_type=ValueType(name='float'),
type=ast_nodes.ArrayType(
inner_type=ast_nodes.ValueType(name='float'),
extents=(3, 4),
),
doc='description',
wasm/codegen/helpers/structs_parser.py
+9 -13
View File
@@ -27,20 +27,14 @@ from wasm.codegen.helpers import struct_field_handler
from wasm.codegen.helpers import structs_wrappers_data
WrappedFieldData = structs_wrappers_data.WrappedFieldData
WrappedStructData = structs_wrappers_data.WrappedStructData
StructFieldHandler = struct_field_handler.StructFieldHandler
AnonymousStructDecl = ast_nodes.AnonymousStructDecl
StructFieldDecl = ast_nodes.StructFieldDecl
debug_print = common.debug_print
introspect_structs = structs.STRUCTS
def generate_wasm_bindings(
wrapped_structs: Dict[str, WrappedStructData],
) -> Dict[str, WrappedStructData]:
wrapped_structs: Dict[str, structs_wrappers_data.WrappedStructData],
) -> Dict[str, structs_wrappers_data.WrappedStructData]:
"""Generates WASM bindings for MuJoCo structs."""
for struct_name, wrap_data in wrapped_structs.items():
@@ -49,7 +43,7 @@ def generate_wasm_bindings(
elif struct_name in constants.ANONYMOUS_STRUCTS:
anonymous_struct = _get_anonymous_struct_field(struct_name)
if not anonymous_struct or not isinstance(
anonymous_struct.type, AnonymousStructDecl
anonymous_struct.type, ast_nodes.AnonymousStructDecl
):
raise RuntimeError(f"Anonymous struct not found: {struct_name}")
struct_fields = anonymous_struct.type.fields
@@ -58,9 +52,11 @@ def generate_wasm_bindings(
debug_print(f"Wrapping struct: {struct_name}")
fields_with_init: List[WrappedFieldData] = []
fields_with_init: List[structs_wrappers_data.WrappedFieldData] = []
for field in struct_fields:
field_gen = StructFieldHandler(field, wrap_data.wrap_name).generate()
field_gen = struct_field_handler.StructFieldHandler(
field, wrap_data.wrap_name
).generate()
# If the struct has at least one non-primitive or fixed size field
# we avoid shallow copy to avoid uninitialized memory.
if not field_gen.is_primitive_or_fixed_size:
@@ -90,7 +86,7 @@ def generate_wasm_bindings(
def _get_anonymous_struct_field(
anonymous_structs_key: str,
) -> StructFieldDecl | None:
) -> ast_nodes.StructFieldDecl | None:
"""Looks up the given key in the anonymous_structs dict and generates bindings for its fields."""
info = constants.ANONYMOUS_STRUCTS[anonymous_structs_key]
parent_decl = introspect_structs[info["parent"]]
@@ -101,7 +97,7 @@ def _get_anonymous_struct_field(
if hasattr(f, "name")
and f.name == info["field_name"]
and hasattr(f, "type")
and isinstance(f.type, AnonymousStructDecl)
and isinstance(f.type, ast_nodes.AnonymousStructDecl)
),
None,
)