Add JavaScript bindings and WASM support

Co-authored-by: Matija Kecman <matijak@google.com>
Co-authored-by: Sebastian Noreña Rendón <sebas.norena@creativa77.com.ar>
Co-authored-by: Kyle Bayes <kylebayes@google.com>
PiperOrigin-RevId: 826094938
Change-Id: Id2658a187bc078056f4410c71202696753319e64
This commit is contained in:
Matias Manevi
2025-10-30 10:37:06 -07:00
committed by Copybara-Service
parent 10f4b36fb7
commit 76e64d96bb
56 changed files with 24362 additions and 22 deletions
+73
View File
@@ -0,0 +1,73 @@
# 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.
"""Helper class to build code string line by line with indentation."""
INDENT = " "
class CodeBuilder:
"""Helper class to build code string line by line with indentation."""
def __init__(self, indent_str: str = INDENT):
self._lines = []
self._indent_level = 0
self._indent_str = indent_str
def line(self, line_content: str) -> None:
"""Adds a line with indentation, special-casing "private:" and "public:"."""
indent = self._indent_str * self._indent_level
content = line_content.strip()
if content == "private:" or content == "public:":
self._lines.append(indent[:-1] + line_content)
elif content:
self._lines.append(indent + line_content)
else:
self._lines.append("")
def to_string(self) -> str:
"""Returns the complete code string."""
return "\n".join(self._lines)
class IndentBlock:
"""Helper class to manage indentation within a `with` statement."""
def __init__(self, builder: "CodeBuilder", header_line=""):
self._builder = builder
self._header_line = header_line
def __enter__(self):
line = self._header_line
line += " {" if line else "{"
self._builder.line(line)
self._builder._indent_level += 1
return self._builder
def __exit__(self, exc_type, exc_val, exc_tb):
if self._builder._indent_level > 0:
self._builder._indent_level -= 1
self._builder.line("}")
def block(self, header_line="") -> IndentBlock:
"""Creates a block including braces and an optional header before the opening brace.
Use via a `with` statement.
Args:
header_line: Optional header line to add before the opening brace.
Returns:
An IndentBlock instance that manages the indentation.
"""
return self.IndentBlock(self, header_line)
+51
View File
@@ -0,0 +1,51 @@
# 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()
+128
View File
@@ -0,0 +1,128 @@
# 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.
"""Utility functions for code generation."""
import os
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
# 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
def get_file_path(
template_dir: str, output_dir: str, filename: str
) -> tuple[str, str]:
"""Constructs the template and output file paths.
Args:
template_dir: The directory containing the template files.
output_dir: The directory where the generated files will be saved.
filename: The name of the file.
Returns:
A tuple containing the template file path and the output file path.
"""
template_file = f"wasm/codegen/{template_dir}/{filename}"
output_file = f"wasm/codegen/{output_dir}/{filename}"
return template_file, output_file
def write_to_file(filepath: str, content: str) -> None:
"""Writes content to a file."""
output_dir = os.path.dirname(filepath)
try:
if output_dir:
os.makedirs(output_dir, exist_ok=True)
with open(filepath, "w") as f:
chars = f.write(content)
print(f"wrote {chars} characters to file '{filepath}'")
except IOError as e:
print(f"Error writing to output file: {filepath} - {e}")
def uppercase_first_letter(input_string: str) -> str:
"""Uppercases the first letter of a string."""
if input_string:
return input_string[0].upper() + input_string[1:]
return input_string
def try_cast_to_scalar_type(value: str) -> int | float | str:
"""Tries to cast a string to an integer, then a float, otherwise returns the original string."""
for type_ in [int, float]:
try:
return type_(value)
except ValueError:
continue
return value
def debug_print(msg: str):
"""Prints a message to the console if STRUCT_DEBUG_MODE is enabled."""
if constants.STRUCT_DEBUG_MODE:
print(msg)
def replace_lines_containing_marker(
lines: list[str],
marker_to_replace: str,
replacement_content: str | list[str],
) -> list[str]:
"""Replaces lines containing a specific marker with new content."""
new_lines = []
replaced = False
for line in lines:
if not replaced and marker_to_replace in line:
indentation = _get_indentation(line)
if isinstance(replacement_content, str):
new_lines.append(indentation + replacement_content)
elif isinstance(replacement_content, list):
for content_line in replacement_content:
if not content_line.strip():
continue
indented_line = (
indentation
+ content_line.replace("\n", "\n" + indentation)
+ "\n"
)
new_lines.append(indented_line)
replaced = True
else:
new_lines.append(line)
return new_lines
def _get_indentation(line: str) -> str:
"""Returns the indentation of the given line as a string of spaces."""
indentation = ""
for char in line:
if char == " ":
indentation += " "
else:
break
return indentation
+37
View File
@@ -0,0 +1,37 @@
# 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()
+474
View File
@@ -0,0 +1,474 @@
# 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.
"""Constants used in the code generation process."""
from introspect import structs as introspect_structs
PRIMITIVE_TYPES = {
# go/keep-sorted start
"char",
"double",
"float",
"int",
"mjtByte",
"mjtMeshBuiltin",
"mjtNum",
"mjtObj", # Adding this to the primitives because it is used as int,
"mjtSize",
"size_t",
"uint64_t",
"uintptr_t",
"unsigned char",
"unsigned int",
"void",
# go/keep-sorted end
}
_PLUGIN_FUNCTIONS = [
# go/keep-sorted start
"mj_getPluginConfig",
"mj_loadAllPluginLibraries",
"mj_loadPluginLibrary",
"mjc_distance",
"mjc_getSDF",
"mjc_gradient",
"mjp_defaultPlugin",
"mjp_defaultResourceProvider",
"mjp_getPlugin",
"mjp_getPluginAtSlot",
"mjp_getResourceProvider",
"mjp_getResourceProviderAtSlot",
"mjp_pluginCount",
"mjp_registerPlugin",
"mjp_registerResourceProvider",
"mjp_resourceProviderCount",
# go/keep-sorted end
]
# Functions that are bound as class methods
_CLASS_METHODS = [
# go/keep-sorted start
"mj_compile",
"mj_copyData",
"mj_copyModel",
"mj_copySpec",
"mj_deleteData",
"mj_deleteModel",
"mj_deleteSpec",
"mj_loadXML",
"mj_makeData",
"mj_makeSpec",
"mj_parseXML", # TODO(manevi): Bind this function.
"mj_parseXMLString",
"mj_recompile", # TODO(manevi): Bind this function.
"mj_saveXML", # TODO(manevi): Bind this function.
"mj_saveXMLString", # TODO(manevi): Bind this function.
# go/keep-sorted end
]
# Omitted because not very useful
_WRITABLE_ERROR = [
"mj_printSchema",
]
# Omitted thread management functions
_THREAD_FUNCTIONS = [
# go/keep-sorted start
"mju_bindThreadPool",
"mju_defaultTask",
"mju_taskJoin",
"mju_threadPoolCreate",
"mju_threadPoolDestroy",
"mju_threadPoolEnqueue",
# go/keep-sorted end
]
# Omitted asset cache functions
_ASSET_CACHE_FUNCTIONS = [
# go/keep-sorted start
"mj_clearCache",
"mj_getCache",
"mj_getCacheCapacity",
"mj_getCacheSize",
"mj_setCacheCapacity",
# go/keep-sorted end
]
# Omitted Virtual Filesystem (VFS) functions
_VFS_FUNCTIONS = [
# go/keep-sorted start
"mj_addBufferVFS",
"mj_addFileVFS",
"mj_defaultVFS",
"mj_deleteFileVFS",
"mj_deleteVFS",
# go/keep-sorted end
]
# Omitted irrelevant visual functions
_VISUAL_FUNCTIONS = [
# go/keep-sorted start
"mjv_averageCamera",
"mjv_copyData",
"mjv_copyModel",
"mjv_defaultScene",
"mjv_freeScene",
"mjv_makeScene",
# go/keep-sorted end
]
_MEMORY_FUNCTIONS = [
# go/keep-sorted start
"mj_freeLastXML",
"mj_freeStack",
"mj_loadModel",
"mj_markStack",
"mj_saveModel",
"mj_stackAllocByte",
"mj_stackAllocInt",
"mj_stackAllocNum",
"mj_warning",
"mjs_bodyToFrame",
"mju_boxQPmalloc",
"mju_clearHandlers",
"mju_error",
"mju_error_i",
"mju_error_s",
"mju_free",
"mju_malloc",
"mju_strncpy",
"mju_warning",
"mju_warning_i",
"mju_warning_s",
# go/keep-sorted end
]
_GETTERS_AND_SETTERS = [
# go/keep-sorted start
"mjs_appendFloatVec",
"mjs_appendIntVec",
"mjs_appendString",
"mjs_getDouble",
"mjs_getPluginAttributes",
"mjs_getString",
"mjs_getUserValue",
"mjs_setBuffer",
"mjs_setDouble",
"mjs_setFloat",
"mjs_setInStringVec",
"mjs_setInt",
"mjs_setPluginAttributes",
"mjs_setString",
"mjs_setStringVec",
"mjs_setUserValue",
# go/keep-sorted end
]
_UTILITY_FUNCTIONS = [
# go/keep-sorted start
"mju_getXMLDependencies",
# go/keep-sorted end
]
# List of functions that should be skipped during the code generation process.
SKIPPED_FUNCTIONS = (
_CLASS_METHODS +
_THREAD_FUNCTIONS +
_MEMORY_FUNCTIONS +
_PLUGIN_FUNCTIONS +
_GETTERS_AND_SETTERS +
_VISUAL_FUNCTIONS +
_ASSET_CACHE_FUNCTIONS +
_VFS_FUNCTIONS +
_WRITABLE_ERROR +
_UTILITY_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 = [
# go/keep-sorted start
"mj_addM",
"mj_angmomMat",
"mj_applyFT",
"mj_constraintUpdate",
"mj_differentiatePos",
"mj_fullM",
"mj_geomDistance",
"mj_getState",
"mj_integratePos",
"mj_jac",
"mj_jacBody",
"mj_jacBodyCom",
"mj_jacDot",
"mj_jacGeom",
"mj_jacPointAxis",
"mj_jacSite",
"mj_jacSubtreeCom",
"mj_mulJacTVec",
"mj_mulJacVec",
"mj_mulM",
"mj_mulM2",
"mj_multiRay",
"mj_normalizeQuat",
"mj_rne",
"mj_saveLastXML",
"mj_setLengthRange",
"mj_setState",
"mj_solveM",
"mj_solveM2",
"mjd_inverseFD",
"mjd_subQuat",
"mjd_transitionFD",
"mju_L1",
"mju_add",
"mju_addScl",
"mju_addTo",
"mju_addToScl",
"mju_band2Dense",
"mju_bandMulMatVec",
"mju_boxQP",
"mju_cholFactor",
"mju_cholFactorBand",
"mju_cholSolve",
"mju_cholSolveBand",
"mju_cholUpdate",
"mju_copy",
"mju_d2n",
"mju_decodePyramid",
"mju_dense2Band",
"mju_dense2sparse",
"mju_dot",
"mju_encodePyramid",
"mju_eye",
"mju_f2n",
"mju_fill",
"mju_insertionSort",
"mju_insertionSortInt",
"mju_isZero",
"mju_mulMatMat",
"mju_mulMatMatT",
"mju_mulMatTMat",
"mju_mulMatTVec",
"mju_mulMatVec",
"mju_mulVecMatVec",
"mju_n2d",
"mju_n2f",
"mju_norm",
"mju_normalize",
"mju_printMatSparse",
"mju_scl",
"mju_sparse2dense",
"mju_sqrMatTD",
"mju_sub",
"mju_subFrom",
"mju_sum",
"mju_symmetrize",
"mju_transpose",
"mju_zero",
# go/keep-sorted end
]
# List of structs that should be skipped during the code generation process.
SKIPPED_STRUCTS = [
# go/keep-sorted start
"mjCache",
"mjSDF",
"mjTask",
"mjThreadPool",
"mjUI",
"mjrContext",
"mjrRect",
"mjuiDef",
"mjuiItem",
"mjuiSection",
"mjuiState",
"mjuiThemeColor",
"mjuiThemeSpacing",
# go/keep-sorted end
]
# Dictionary that maps anonymous structs to their parent struct and field name.
# 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 = {
# go/keep-sorted start
"mjVisualGlobal": {"parent": "mjVisual", "field_name": "global"},
"mjVisualHeadlight": {"parent": "mjVisual", "field_name": "headlight"},
"mjVisualMap": {"parent": "mjVisual", "field_name": "map"},
"mjVisualQuality": {"parent": "mjVisual", "field_name": "quality"},
"mjVisualRgba": {"parent": "mjVisual", "field_name": "rgba"},
"mjVisualScale": {"parent": "mjVisual", "field_name": "scale"},
# go/keep-sorted end
}
# This list is created by subtracting the skipped structs from the list of all
# structs and adding the anonymous structs.
STRUCTS_TO_BIND = 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 = [
# go/keep-sorted start
"mjContact",
"mjSolverStat",
"mjStatistic",
"mjTimerStat",
"mjWarningStat",
"mjsCompiler",
"mjsDefault",
"mjsElement",
"mjsExclude",
"mjsWrap",
"mjvGLCamera",
"mjvLight",
# go/keep-sorted end
]
# 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 = [
# go/keep-sorted start
"contact",
"efc_AR",
"efc_AR_colind",
"efc_AR_rowadr",
"efc_AR_rownnz",
"efc_D",
"efc_J",
"efc_JT",
"efc_JT_colind",
"efc_J_colind",
"efc_J_rowadr",
"efc_J_rownnz",
"efc_J_rowsuper",
"efc_KBIP",
"efc_R",
"efc_aref",
"efc_b",
"efc_diagApprox",
"efc_force",
"efc_frictionloss",
"efc_id",
"efc_island",
"efc_margin",
"efc_pos",
"efc_state",
"efc_type",
"efc_vel",
"iLDiagInv",
"iM_rowadr",
"iM_rownnz",
"iacc",
"iacc_smooth",
"iefc_D",
"iefc_J",
"iefc_JT",
"iefc_JT_colind",
"iefc_JT_rowadr",
"iefc_JT_rownnz",
"iefc_JT_rowsuper",
"iefc_J_colind",
"iefc_J_rowadr",
"iefc_J_rownnz",
"iefc_J_rowsuper",
"iefc_R",
"iefc_aref",
"iefc_force",
"iefc_frictionloss",
"iefc_id",
"iefc_state",
"iefc_type",
"ifrc_constraint",
"ifrc_smooth",
"island_dofadr",
"island_dofnum",
"island_efcadr",
"island_efcind",
"island_efcnum",
"island_idofadr",
"island_iefcadr",
"island_ne",
"island_nefc",
"island_nf",
"island_nv",
"map_efc2iefc",
"map_iefc2efc",
# go/keep-sorted end
]
# Dictionary where keys are the struct names and the values are lists of the
# fields that are manually specified in the structs.h template file.
MANUALLY_ADDED_FIELDS_FROM_TEMPLATE = {
# go/keep-sorted start
"MjData": ["solver", "timer", "warning", "contact"],
"MjSpec": ["option", "visual", "stat", "element", "compiler"],
"MjvScene": [
"model",
"lights",
"camera",
"geoms",
"geomorder",
"flexedgeadr",
"flexedgenum",
"flexvertadr",
"flexvertnum",
"flexfaceadr",
"flexfacenum",
"flexfaceused",
"flexedge",
"flexvert",
"skinfacenum",
"skinvertadr",
"skinvertnum",
"skinvert",
"skinnormal",
"flexface",
"flexnormal",
"flextexcoord",
# go/keep-sorted end
],
}
# Dictionary that maps byte array fields to their corresponding size members.
# 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 = {
"buffer": {"size": "nbuffer"},
"arena": {"size": "narena"},
}
# 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
# 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 = [
"MjData",
"MjModel",
"MjvScene",
"MjSpec",
"MjVisual",
]
+350
View File
@@ -0,0 +1,350 @@
# 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.
"""Helper functions for processing and generating bindings for MuJoCo functions."""
from typing import List, Set, Tuple, cast
from introspect import ast_nodes
from wasm.codegen.helpers import code_builder
from wasm.codegen.helpers import common
from wasm.codegen.helpers import constants
PRIMITIVE_TYPES = constants.PRIMITIVE_TYPES
uppercase_first_letter = common.uppercase_first_letter
def param_is_primitive_value(param: ast_nodes.FunctionParameterDecl) -> bool:
"""Checks if param is a primitive value type."""
if isinstance(param.type, ast_nodes.ValueType):
return param.type.name in PRIMITIVE_TYPES
return False
def param_is_pointer_to_primitive_value(
param: ast_nodes.FunctionParameterDecl,
) -> bool:
"""Checks if param is a pointer to a primitive value."""
return (
isinstance(param.type, ast_nodes.PointerType)
or isinstance(param.type, ast_nodes.ArrayType)
) and (
isinstance(param.type.inner_type, ast_nodes.ValueType)
and param.type.inner_type.name in PRIMITIVE_TYPES
)
def param_is_pointer_to_struct(param: ast_nodes.FunctionParameterDecl) -> bool:
"""Checks if param is a pointer to a struct."""
return (
isinstance(param.type, ast_nodes.PointerType)
or isinstance(param.type, ast_nodes.ArrayType)
) and (
isinstance(param.type.inner_type, ast_nodes.ValueType)
and param.type.inner_type.name not in PRIMITIVE_TYPES
)
def return_is_value_of_type(
func: ast_nodes.FunctionDecl, allowed_types: Set[str]
) -> bool:
"""Checks if func returns an allowed value type."""
return (
isinstance(func.return_type, ast_nodes.ValueType)
and func.return_type.name in allowed_types
)
def return_is_pointer_to_struct(func: ast_nodes.FunctionDecl) -> bool:
"""Checks if func returns a pointer to a struct."""
return (
isinstance(func.return_type, ast_nodes.PointerType)
and isinstance(func.return_type.inner_type, ast_nodes.ValueType)
and func.return_type.inner_type.name not in PRIMITIVE_TYPES
)
def return_is_pointer_to_primitive(func: ast_nodes.FunctionDecl) -> bool:
"""Checks if func returns a pointer to a primitive value."""
return (
isinstance(func.return_type, ast_nodes.PointerType)
and isinstance(func.return_type.inner_type, ast_nodes.ValueType)
and func.return_type.inner_type.name in PRIMITIVE_TYPES
)
def get_const_qualifier(func: ast_nodes.FunctionDecl) -> str:
"""Returns the const qualifier of func's return type."""
if (
isinstance(func.return_type, ast_nodes.PointerType)
and isinstance(func.return_type.inner_type, ast_nodes.ValueType)
and func.return_type.inner_type.is_const
):
return "const "
return ""
def should_be_wrapped(func: ast_nodes.FunctionDecl) -> bool:
"""Checks if a MuJoCo function needs a wrapper function."""
return (
return_is_pointer_to_primitive(func)
or return_is_pointer_to_struct(func)
or any(
param_is_pointer_to_primitive_value(param)
or isinstance(param.type, ast_nodes.ArrayType)
or param_is_pointer_to_struct(param)
for param in func.parameters
)
)
def generate_function_wrapper(func: ast_nodes.FunctionDecl) -> str:
"""Generates C++ code for a wrapper function."""
builder = code_builder.CodeBuilder()
# Build function header
params_unpack_statements = get_params_unpack_statements(func.parameters)
wrapper_params_list = get_params_string(func.parameters)
not_nullable_params = get_params_notnullable(func.parameters)
wrapper_params_str = ", ".join(wrapper_params_list)
ret_type = get_compatible_return_type(func)
builder.line(f"{ret_type} {func.name}_wrapper({wrapper_params_str})")
# Build function body
with builder.block():
invoker_params_list = get_params_string_maybe_with_conversion(
func.parameters
)
invoker_params_str = ", ".join(invoker_params_list)
invoker_call = f"{func.name}({invoker_params_str})"
invoker_statement = get_compatible_return_call(func, invoker_call)
for p in not_nullable_params:
builder.line(f"CHECK_VAL({p});")
for unpack_statement in params_unpack_statements:
builder.line(unpack_statement)
builder.line(f"{invoker_statement};")
return builder.to_string()
def get_params_notnullable(
ast_params: Tuple[ast_nodes.FunctionParameterDecl, ...],
) -> List[str]:
"""Generates list of param names for checking if they aren't null/undefined."""
not_nullable_params = []
for p in ast_params:
if (
isinstance(p.type, (ast_nodes.PointerType, ast_nodes.ArrayType))
and isinstance(p.type.inner_type, ast_nodes.ValueType)
# We only check for char because others are checked in the unpacker
# and we don't want to check twice.
and p.type.inner_type.name == "char"
and not p.nullable
):
not_nullable_params.append(p.name)
return not_nullable_params
def get_params_unpack_statements(
ast_params: Tuple[ast_nodes.FunctionParameterDecl, ...],
) -> List[str]:
"""Generates C++ statements to unpack JS values for pointer/array parameters."""
params_unpack_statements = []
for p in ast_params:
if (
isinstance(p.type, (ast_nodes.PointerType, ast_nodes.ArrayType))
and isinstance(p.type.inner_type, ast_nodes.ValueType)
and p.type.inner_type.name in PRIMITIVE_TYPES
):
if p.type.inner_type.name == "char":
# param is Javascript string
continue
if p.type.inner_type.is_const:
# param is Javascript number[]
params_unpack_statements.append(
f"UNPACK_ARRAY({p.type.inner_type.name}, {p.name});"
)
else:
# param is TypedArray or a WasmBuffer
params_unpack_statements.append(
f"UNPACK_VALUE({p.type.inner_type.name}, {p.name});"
)
return params_unpack_statements
def get_params_string(
parameters: Tuple[ast_nodes.FunctionParameterDecl, ...]
) -> List[str]:
"""Generates a list of C++ parameter declarations as strings."""
result = []
for p in parameters:
if (
isinstance(p.type, ast_nodes.PointerType)
and isinstance(p.type.inner_type, ast_nodes.ValueType)
and p.type.inner_type.name not in PRIMITIVE_TYPES
):
# Pointer to struct parameters
const_qualifier = "const " if p.type.inner_type.is_const else ""
result.append(
f"{const_qualifier}{uppercase_first_letter(p.type.inner_type.name)}&"
f" {p.name}"
)
elif (
isinstance(p.type, ast_nodes.ValueType)
and p.type.name in PRIMITIVE_TYPES
):
# Primitive value parameters
const_qualifier = "const " if p.type.is_const else ""
result.append(f"{const_qualifier}{p.type} {p.name}")
elif (
isinstance(p.type, (ast_nodes.PointerType, ast_nodes.ArrayType))
and isinstance(p.type.inner_type, ast_nodes.ValueType)
and p.type.inner_type.name in PRIMITIVE_TYPES
):
# Pointer to primitive value parameters or arrays
if p.type.inner_type.name == "char":
if p.nullable:
result.append(f"const NullableString& {p.name}")
else:
result.append(f"const String& {p.name}")
elif (
p.type.inner_type.name
in ["int", "float", "double", "mjtNum", "mjtByte"]
and p.type.inner_type.is_const
):
result.append(f"const NumberArray& {p.name}")
else:
result.append(f"const val& {p.name}")
else:
# This case should ideally not be reached if AST is well-formed
# and types are categorized by the helper booleans correctly.
raise TypeError(
"Unable to generate param string. Unhandled parameter type:"
f" {p.type} for param '{p.name}'"
)
return result
def get_params_string_maybe_with_conversion(
ast_params: Tuple[ast_nodes.FunctionParameterDecl, ...],
) -> List[str]:
"""Generates C++ expressions for passing compatible params from JS to MuJoCo C-API functions."""
native_params = []
for p in ast_params:
if param_is_pointer_to_struct(p):
native_params.append(f"{p.name}.get()")
elif param_is_primitive_value(p):
native_params.append(p.name)
elif (
isinstance(p.type, (ast_nodes.PointerType, ast_nodes.ArrayType))
and isinstance(p.type.inner_type, ast_nodes.ValueType)
and p.type.inner_type.name in PRIMITIVE_TYPES
and p.type.inner_type.name != "char"
):
native_params.append(f"{p.name}_.data()")
elif (
isinstance(p.type, (ast_nodes.PointerType, ast_nodes.ArrayType))
and isinstance(p.type.inner_type, ast_nodes.ValueType)
and p.type.inner_type.name == "char"
):
const_qualifier = "const " if p.type.inner_type.is_const else ""
native_params.append(
f"{p.name}.as<{const_qualifier}std::string>().data()"
)
else:
raise TypeError(
f"Unhandled parameter type for conversion: {p.type} for param"
f" '{p.name}'"
)
return native_params
def get_compatible_return_call(
func: ast_nodes.FunctionDecl, invoker: str
) -> str:
"""Generates embind compatible return value conversion."""
if return_is_value_of_type(func, {"void"}):
return invoker
if isinstance(func.return_type, ast_nodes.PointerType) and isinstance(
func.return_type.inner_type, ast_nodes.ValueType
):
if func.return_type.inner_type.name == "char":
return f"return std::string({invoker})"
elif func.return_type.inner_type.name == "mjString":
return f"return *{invoker}"
if return_is_pointer_to_struct(func):
return get_converted_struct_to_class(func, invoker)
if return_is_value_of_type(func, PRIMITIVE_TYPES):
return f"return {invoker}"
raise RuntimeError(
"Failed to calculate return value conversion for function"
f" {func.name} that returns '{func.return_type}'"
)
def get_compatible_return_type(func: ast_nodes.FunctionDecl) -> str:
"""Creates embind compatible return type."""
if (
isinstance(func.return_type, ast_nodes.PointerType)
and isinstance(func.return_type.inner_type, ast_nodes.ValueType)
and func.return_type.inner_type.name in ["char", "mjString"]
):
return "std::string"
if (
isinstance(func.return_type, ast_nodes.PointerType)
and isinstance(func.return_type.inner_type, ast_nodes.ValueType)
and func.return_type.inner_type.name not in PRIMITIVE_TYPES
):
const_qualifier = get_const_qualifier(func)
return f"""{const_qualifier}std::optional<{uppercase_first_letter(func.return_type.inner_type.name)}>"""
if (
isinstance(func.return_type, ast_nodes.ValueType)
and func.return_type.name in PRIMITIVE_TYPES
):
return f"{func.return_type.name}"
return "val"
def get_converted_struct_to_class(
func: ast_nodes.FunctionDecl, invoker: str
) -> str:
"""Generates a C++ function invocation for a struct return-type function."""
const_qualifier = get_const_qualifier(func)
return_type = cast(ast_nodes.PointerType, func.return_type)
struct_name = cast(ast_nodes.ValueType, return_type.inner_type).name
class_constructor = uppercase_first_letter(struct_name)
return_str = f"{class_constructor}(result)"
return f"""{const_qualifier}{struct_name}* result = {invoker};
if (result == nullptr) {{
return std::nullopt;
}}
return {return_str}"""
def is_excluded_function_name(func_name: str) -> bool:
"""Checks if a function name should be excluded from direct binding."""
return (
func_name.startswith("mjr_")
or func_name.startswith("mjui_")
or func_name in constants.SKIPPED_FUNCTIONS
)
+236
View File
@@ -0,0 +1,236 @@
# 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 typing import TypeAlias
from absl.testing import absltest
from introspect import ast_nodes
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.func_ret_ptr_struct = FunctionDecl(
"func_ps", PointerType(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
)
)
self.assertFalse(
function_utils.return_is_value_of_type(
FunctionDecl("func_s", ValueType("MyStruct"), [], "doc"),
PrimitiveTypes,
)
)
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 = FunctionParameterDecl("prim_v", ValueType("int"))
param_arr = FunctionParameterDecl(
"arr_v", ArrayType(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_struct = FunctionParameterDecl(
name="p_struct", type=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 = FunctionParameterDecl(
"a_struct", ArrayType(self.struct_type, extents=(5,))
)
param_ptr_to_struct = FunctionParameterDecl(
"p_struct", PointerType(self.struct_type)
)
param_ptr_to_ptr = FunctionParameterDecl(
"p_ptr", 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 = FunctionDecl(
name="get_data",
return_type=PointerType(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(
name="get_id",
return_type=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 = (
FunctionParameterDecl(
name="mat",
type=PointerType(
inner_type=ValueType(name="mjtNum", is_const=True),
),
),
FunctionParameterDecl(
name="nr",
type=ValueType(name="int"),
),
)
func = FunctionDecl(
name="get_id",
return_type=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 = FunctionParameterDecl(
name="my_struct",
type=PointerType(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(
name="s",
type=PointerType(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(
name="noop",
return_type=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 = FunctionDecl(
name="get_name",
return_type=PointerType(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 = FunctionDecl(
name="get_struct",
return_type=PointerType(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()
@@ -0,0 +1,195 @@
# 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.
"""Code builder for struct constructor code."""
from typing import List, cast
from introspect import ast_nodes
from introspect import structs as introspect_structs
from wasm.codegen.helpers import code_builder
from wasm.codegen.helpers import common
from wasm.codegen.helpers import constants
from wasm.codegen.helpers import structs_wrappers_data
def _has_nested_wrapper_members(
struct_info: ast_nodes.StructDecl
) -> bool:
"""Checks if the struct contains other wrapped structs as direct members."""
for field in struct_info.fields:
struct_field = cast(ast_nodes.StructFieldDecl, field)
if isinstance(struct_field.type, ast_nodes.ValueType):
if struct_field.type.name in constants.STRUCTS_TO_BIND:
return True
if isinstance(struct_field.type, ast_nodes.ArrayType):
if isinstance(struct_field.type.inner_type, ast_nodes.ValueType):
if struct_field.type.inner_type.name in constants.STRUCTS_TO_BIND:
return True
if isinstance(struct_field.type, ast_nodes.PointerType):
if isinstance(struct_field.type.inner_type, ast_nodes.ValueType):
if struct_field.type.inner_type.name in constants.STRUCTS_TO_BIND:
return True
return False
def _build_struct_header_internal(
struct_name: str,
wrapped_fields: List[structs_wrappers_data.WrappedFieldData],
fields_with_init: List[structs_wrappers_data.WrappedFieldData],
use_shallow_copy: bool = False,
is_mjs: bool = False,
):
"""Builds the C++ header file code for a struct."""
wrapper_name = common.uppercase_first_letter(struct_name)
builder = code_builder.CodeBuilder()
with builder.block(f"struct {wrapper_name}"):
if not is_mjs:
builder.line(f"{wrapper_name}();")
builder.line(f"{wrapper_name}(const {wrapper_name} &);")
builder.line(f"{wrapper_name} &operator=(const {wrapper_name} &);")
builder.line(f"explicit {wrapper_name}({struct_name} *ptr);")
builder.line(f"~{wrapper_name}();")
if use_shallow_copy:
builder.line(f"std::unique_ptr<{wrapper_name}> copy();")
for field in wrapped_fields:
if field.definition and field not in fields_with_init:
for line in field.definition.splitlines():
builder.line(line)
builder.line(f"{struct_name}* get() const {{ return ptr_; }}")
builder.line(f"void set({struct_name}* ptr) {{ ptr_ = ptr; }}")
builder.line("")
builder.line("private:")
builder.line(f"{struct_name}* ptr_;")
builder.line("bool owned_ = false;")
if is_mjs and fields_with_init:
builder.line("")
builder.line("public:")
for field in fields_with_init:
if field.definition:
builder.line(f"{field.definition}")
return builder.to_string()+";"
def build_struct_header(
struct_name: str,
use_shallow_copy: bool = False,
fields_with_init: List[structs_wrappers_data.WrappedFieldData] = [],
wrapped_fields: List[structs_wrappers_data.WrappedFieldData] = [],
):
"""Builds the C++ header file code for a struct."""
struct_info = introspect_structs.STRUCTS.get(struct_name)
if struct_name.startswith("mjs"):
return _build_struct_header_internal(
struct_name,
wrapped_fields,
fields_with_init,
use_shallow_copy,
is_mjs=True,
)
if (
(
common.uppercase_first_letter(struct_name)
not in constants.HARDCODED_WRAPPER_STRUCTS
)
and struct_info
and not _has_nested_wrapper_members(struct_info)
):
return _build_struct_header_internal(
struct_name, wrapped_fields, [], use_shallow_copy, is_mjs=False
)
return ""
def build_struct_source(
struct_name: str,
mj_default_func: str | None = None,
fields_with_init: List[structs_wrappers_data.WrappedFieldData] = [],
use_shallow_copy: bool = False,
):
"""Builds the C++ .cc file code for a struct."""
wrapper_name = common.uppercase_first_letter(struct_name)
is_mjs_struct = "Mjs" in wrapper_name
builder = code_builder.CodeBuilder()
fields_init = ""
if fields_with_init:
fields_init = "".join(
field_with_init.initialization
for field_with_init in fields_with_init
)
# constructor passing native ptr
builder.line(
f"{wrapper_name}::{wrapper_name}({struct_name} *ptr) :"
f" ptr_(ptr){fields_init} {{}}"
)
# constructor with default values
if not is_mjs_struct:
with builder.block(
f"{wrapper_name}::{wrapper_name}() : ptr_(new"
f" {struct_name}){fields_init}"
):
builder.line("owned_ = true;")
if mj_default_func:
builder.line(f"{mj_default_func}(ptr_);")
# copy constructor
if use_shallow_copy and not is_mjs_struct:
with builder.block(
f"{wrapper_name}::{wrapper_name}(const"
f" {wrapper_name} &other)"
+ (f" : {wrapper_name}()" if not is_mjs_struct else "")
):
builder.line("*ptr_ = *other.get();")
if fields_with_init:
for field_with_init in fields_with_init:
if field_with_init.ptr_copy_reset is not None:
builder.line(field_with_init.ptr_copy_reset)
# assignment operator
with builder.block(
f"{wrapper_name}&"
f" {wrapper_name}::operator=(const"
f" {wrapper_name} &other)"
):
with builder.block("if (this == &other)"):
builder.line("return *this;")
builder.line("*ptr_ = *other.get();")
if fields_with_init:
for field_with_init in fields_with_init:
if field_with_init.ptr_copy_reset is not None:
builder.line(field_with_init.ptr_copy_reset)
builder.line("return *this;")
# destructor
if is_mjs_struct:
builder.line(f"{wrapper_name}::~{wrapper_name}() {{}}")
else:
with builder.block(f"{wrapper_name}::~{wrapper_name}()"):
builder.line("if (owned_ && ptr_) delete ptr_;")
# copy function
if use_shallow_copy:
with builder.block(
f"std::unique_ptr<{wrapper_name}>"
f" {wrapper_name}::copy()"
):
builder.line(
f"return std::make_unique<{wrapper_name}>(*this);"
)
return builder.to_string()
@@ -0,0 +1,140 @@
# 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
StructFieldDecl = ast_nodes.StructFieldDecl
ValueType = ast_nodes.ValueType
PointerType = ast_nodes.PointerType
ArrayType = ast_nodes.ArrayType
StructDecl = ast_nodes.StructDecl
class StructConstructorCodeBuilderTest(absltest.TestCase):
def test_constructor_code_with_default_function(self):
self.assertEqual(
struct_constructor_code_builder.build_struct_source(
"mjLROpt", "mj_defaultLROpt"
),
"""
MjLROpt::MjLROpt(mjLROpt *ptr) : ptr_(ptr) {}
MjLROpt::MjLROpt() : ptr_(new mjLROpt) {
owned_ = true;
mj_defaultLROpt(ptr_);
}
MjLROpt::~MjLROpt() {
if (owned_ && ptr_) delete ptr_;
}
""".strip(),
)
def test_constructor_code_without_default_function(self):
self.assertEqual(
struct_constructor_code_builder.build_struct_source("mjLROpt"),
"""
MjLROpt::MjLROpt(mjLROpt *ptr) : ptr_(ptr) {}
MjLROpt::MjLROpt() : ptr_(new mjLROpt) {
owned_ = true;
}
MjLROpt::~MjLROpt() {
if (owned_ && ptr_) delete ptr_;
}
""".strip(),
)
def test_constructor_code_with_fields_with_init(self):
field_with_init = StructFieldDecl(
name="element",
type=PointerType(
inner_type=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",
"mjs_defaultTexture",
[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", use_shallow_copy=True
),
"""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> MjvLight::copy() {
return std::make_unique<MjvLight>(*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()
@@ -0,0 +1,100 @@
# 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.
"""Class to build the C++ code for a struct field wrapper."""
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:
"""Builds the C++ code for a primitive type field wrapper."""
if not isinstance(field.type, ValueType):
raise ValueError(f"{field.type} must be ValueType.")
builder = code_builder.CodeBuilder()
# build getter for primitive type field
with builder.block(f"{field.type.name} {field.name}() const"):
builder.line(f"return ptr_->{field.name};")
# build setter for primitive type field
with builder.block(f"void set_{field.name}({field.type.name} value)"):
builder.line(f"ptr_->{field.name} = value;")
return builder.to_string()
def build_memory_view_definition(
field: StructFieldDecl, array_size_str: str, ptr_expr: str
) -> str:
"""Builds the C++ code for a pointer type field wrapper."""
builder = code_builder.CodeBuilder()
with builder.block(f"emscripten::val {field.name}() const"):
builder.line(
"return"
f" emscripten::val(emscripten::typed_memory_view({array_size_str},"
f" {ptr_expr}));"
)
return builder.to_string()
def build_string_field_definition(field: StructFieldDecl) -> str:
"""Builds the C++ code for a string type field wrapper."""
builder = code_builder.CodeBuilder()
with builder.block(f"mjString {field.name}() const"):
builder.line(
f'return (ptr_ && ptr_->{field.name}) ? *(ptr_->{field.name}) : "";'
)
with builder.block(f"void set_{field.name}(const mjString& value)"):
with builder.block(f"if (ptr_ && ptr_->{field.name})"):
builder.line(f"*(ptr_->{field.name}) = value;")
return builder.to_string()
def build_mjvec_pointer_definition(
field: StructFieldDecl, vector_type: str
) -> str:
"""Builds the C++ code for a mjVec type field wrapper."""
ptr_field_expr = f"*(ptr_->{field.name})"
if vector_type == "mjByteVec":
vector_type = "std::vector<uint8_t>"
ptr_field_expr = (
f"*(reinterpret_cast<std::vector<uint8_t>*>(ptr_->{field.name}))"
)
builder = code_builder.CodeBuilder()
with builder.block(f"{vector_type} &{field.name}() const"):
builder.line(f"return {ptr_field_expr};")
return builder.to_string()
def build_simple_property_binding(
field: StructFieldDecl,
struct_wrapper_name: str,
add_setter: bool = False,
add_return_value_policy_as_ref: bool = False,
) -> str:
"""Builds the C++ code for a simple property binding."""
builder = code_builder.CodeBuilder()
setter_txt = ""
if add_setter:
setter_txt = f", &{struct_wrapper_name}::set_{field.name}"
if add_return_value_policy_as_ref:
as_reference_txt = ", reference()"
else:
as_reference_txt = ""
builder.line(
f'.property("{field.name}",'
f" &{struct_wrapper_name}::{field.name}{setter_txt}{as_reference_txt})"
)
return builder.to_string()
@@ -0,0 +1,171 @@
# 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
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(
name="ngeom",
type=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 = StructFieldDecl(
name="geom_rgba",
type=PointerType(
inner_type=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 = StructFieldDecl(
name="string_field",
type=PointerType(
inner_type=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 = StructFieldDecl(
name="vector_field",
type=PointerType(
inner_type=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 = StructFieldDecl(
name="vector_field",
type=PointerType(
inner_type=ValueType(name="mjByteVec"),
),
doc="",
)
self.assertEqual(
struct_field_code_builder.build_mjvec_pointer_definition(
field, "mjByteVec"
),
"""
std::vector<uint8_t> &vector_field() const {
return *(reinterpret_cast<std::vector<uint8_t>*>(ptr_->vector_field));
}""".strip(),
)
def test_simple_property_binding(self):
field = StructFieldDecl(
name="ngeom",
type=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 = StructFieldDecl(
name="ngeom",
type=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 = StructFieldDecl(
name="ngeom",
type=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()
@@ -0,0 +1,358 @@
# 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.
"""Class to handle the different struct field types, and provide the c++ code for the wrappers and bindings."""
import math
from typing import Tuple, Union, cast
from introspect import ast_nodes
from wasm.codegen.helpers import common
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
class StructFieldHandler:
"""Class to handle the different struct field types, and provide the c++ code for the definitions and bindings."""
def __init__(
self,
field: StructFieldDecl,
struct_wrapper_name: str,
):
self.field = field
self.struct_wrapper_name = struct_wrapper_name
self.simple_property_binding = (
struct_field_code_builder.build_simple_property_binding(
self.field, self.struct_wrapper_name
)
)
self.manually_added_fields = (
constants.MANUALLY_ADDED_FIELDS_FROM_TEMPLATE.get(
self.struct_wrapper_name, {}
)
)
def generate(self) -> WrappedFieldData:
"""Generates the C++ definition and binding code for the struct field."""
field_type = self.field.type
if isinstance(field_type, ValueType) and (
field_type.name in constants.PRIMITIVE_TYPES
or field_type.name.startswith("mjt")
):
return self._handle_primitive()
elif isinstance(field_type, PointerType):
return self._handle_pointer()
elif isinstance(field_type, ArrayType):
return self._handle_array()
elif isinstance(field_type, ValueType) and field_type.name.startswith("mj"):
return self._handle_mj_struct()
elif isinstance(field_type, AnonymousStructDecl):
return self._handle_anonymous_struct()
return self._undefined()
def _handle_primitive(self) -> WrappedFieldData:
"""Handles the generation of C++ definition and binding code for primitive fields."""
return WrappedFieldData(
definition=(
struct_field_code_builder.build_primitive_type_definition(
self.field
)
),
binding=struct_field_code_builder.build_simple_property_binding(
self.field,
self.struct_wrapper_name,
add_setter=True,
add_return_value_policy_as_ref=True,
),
is_primitive_or_fixed_size=True,
)
def _handle_pointer(self) -> WrappedFieldData:
"""Handles the generation of C++ definition and binding code for pointer fields."""
if not isinstance(self.field.type, PointerType):
raise ValueError(
f"Expected PointerType, got {type(self.field.type)} for field"
f" {self.field.name}"
)
field_type: PointerType = self.field.type
inner_type_name = (
field_type.inner_type.name
if isinstance(field_type.inner_type, ValueType)
else ""
)
ptr_field_expr = f"ptr_->{self.field.name}"
array_size_str = ""
if self.field.array_extent:
array_size_str = parse_array_extent(
self.field.array_extent, self.struct_wrapper_name, self.field.name
)
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})"
)
# 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"],)
array_size_str = parse_array_extent(
extent, self.struct_wrapper_name, self.field.name
)
elif inner_type_name == "mjString":
return WrappedFieldData(
definition=struct_field_code_builder.build_string_field_definition(
self.field
),
binding=struct_field_code_builder.build_simple_property_binding(
self.field,
self.struct_wrapper_name,
add_setter=True,
add_return_value_policy_as_ref=True,
),
)
elif inner_type_name.startswith("mj") and inner_type_name.endswith("Vec"):
return WrappedFieldData(
definition=struct_field_code_builder.build_mjvec_pointer_definition(
self.field, inner_type_name
),
binding=struct_field_code_builder.build_simple_property_binding(
self.field,
self.struct_wrapper_name,
add_setter=False,
add_return_value_policy_as_ref=True,
),
)
elif inner_type_name in constants.PRIMITIVE_TYPES:
return self._get_manual_definition(
comment_type="primitive pointer field with complex extents"
)
if (
inner_type_name.startswith("mj")
and inner_type_name not in constants.PRIMITIVE_TYPES
):
debug_print(
f"\tcomplex pointer type: needs manual wrapper: {self.field.name}"
)
# it's a pointer to a single struct,
# like the `element` field in mjs structs
# and the struct is not manually added
if (
not self.field.array_extent
and self.struct_wrapper_name
not in constants.MANUALLY_ADDED_FIELDS_FROM_TEMPLATE.keys()
):
ptr_field = cast(PointerType, self.field.type)
wrapper_field_name = common.uppercase_first_letter(
cast(ValueType, ptr_field.inner_type).name
)
return WrappedFieldData(
definition=f"{wrapper_field_name} {self.field.name};",
binding=struct_field_code_builder.build_simple_property_binding(
self.field,
self.struct_wrapper_name,
add_setter=False,
add_return_value_policy_as_ref=True,
),
initialization=f", {self.field.name}(ptr_->{self.field.name})",
)
else:
debug_print(
"\tcomplex pointer type with array extent: needs manual wrapper:"
f" {self.field.name}"
)
return self._get_manual_definition(comment_type="complex pointer field")
return WrappedFieldData(
definition=(
struct_field_code_builder.build_memory_view_definition(
self.field, array_size_str, ptr_field_expr
)
),
binding=self.simple_property_binding,
)
def _handle_array(self) -> WrappedFieldData:
"""Handles the generation of C++ definition and binding code for array fields."""
field_type = self.field.type
if not isinstance(field_type, ArrayType):
raise ValueError(
f"Expected ArrayType, got {type(field_type)} for field"
f" {self.field.name}"
)
inner_type = field_type.inner_type
size = math.prod(field_type.extents)
if isinstance(inner_type, ValueType):
if inner_type.name in constants.PRIMITIVE_TYPES:
ptr_expr = f"ptr_->{self.field.name}"
if len(field_type.extents) > 1:
# for multi-dimensional arrays, we need to cast the field
# to a pointer, so embind can correctly interpret the memory
# view
ptr_expr = f"reinterpret_cast<{inner_type.name}*>({ptr_expr})"
return WrappedFieldData(
definition=(
struct_field_code_builder.build_memory_view_definition(
self.field, str(size), ptr_expr
)
),
binding=self.simple_property_binding,
is_primitive_or_fixed_size=True,
)
elif inner_type.name.startswith("mj") and not inner_type.name.startswith(
"mjt"
):
debug_print(f"\tarray to vector wrapper needed: {self.field.name}")
return self._get_manual_definition(comment_type="array field")
debug_print(f"\tNOT IMPLEMENTED ARRAY field: {self.field.name}")
return 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:
"""Handles the generation of C++ definition and binding code for mj struct fields."""
if (
isinstance(self.field.type, ValueType)
and self.field.name not in self.manually_added_fields
and self.field.type.name in constants.STRUCTS_TO_BIND
):
# TODO(manevi): Find a better way to do this instead of checking the
# struct wrapper name.
definition = ""
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(
definition=definition,
binding=struct_field_code_builder.build_simple_property_binding(
self.field,
self.struct_wrapper_name,
add_setter=False,
add_return_value_policy_as_ref=True,
),
initialization=f", {self.field.name}(&ptr_->{self.field.name})",
ptr_copy_reset=f"{self.field.name}.set(&ptr_->{self.field.name});",
is_primitive_or_fixed_size=True,
)
return self._get_manual_definition(comment_type="struct field")
def _handle_anonymous_struct(self) -> WrappedFieldData:
"""Handles the generation of C++ definition and binding code for anonymous struct fields."""
anonymous_struct_name = ""
for name, value in constants.ANONYMOUS_STRUCTS.items():
if (
common.uppercase_first_letter(value["parent"])
== self.struct_wrapper_name
and value["field_name"] == self.field.name
):
anonymous_struct_name = name
break
if (
isinstance(self.field.type, AnonymousStructDecl)
and self.field.name not in self.manually_added_fields
and anonymous_struct_name in constants.STRUCTS_TO_BIND
):
return WrappedFieldData(
binding=struct_field_code_builder.build_simple_property_binding(
self.field,
self.struct_wrapper_name,
add_setter=False,
add_return_value_policy_as_ref=True,
),
initialization=f", {self.field.name}(&ptr_->{self.field.name})",
ptr_copy_reset=f"{self.field.name}.set(&ptr_->{self.field.name});",
is_primitive_or_fixed_size=True,
)
return self._get_manual_definition(comment_type="anonymous struct field")
def _undefined(self) -> WrappedFieldData:
"""This function adds a TODO comment for fields that are not handled by this class yet."""
return 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:
"""Helper method to generate a comment as a definition for manually added fields."""
if self.field.name in self.manually_added_fields:
return WrappedFieldData(
definition=(
f"// {comment_type} is defined manually. {self.field.name}"
),
binding=self.simple_property_binding,
)
return WrappedFieldData(
definition=(
f"// TODO: Define {comment_type} manually for {self.field.name}"
),
binding=f"// TODO: {self.simple_property_binding}",
)
def parse_array_extent(
extents: Tuple[Union[str, int], ...], wrapper_name: str, field_name: str
) -> str:
"""Parses the array extent of a field, returning a string representing the resolved extents."""
if not extents:
return ""
return " * ".join(
resolve_extent(extent, wrapper_name, field_name) for extent in extents
)
def resolve_extent(
extent: Union[str, int], wrapper_name: str, field_name: str
) -> str:
"""Resolves the extent of an array, handling integers and references to other struct fields.
Args:
extent: The extent to resolve, can be an int or a string referencing a
field.
wrapper_name: The name of the struct wrapper.
field_name: The name of the field being processed.
Returns:
A string representing the resolved extent, either as a number or a field
reference.
"""
if isinstance(extent, int):
return str(extent)
# if starts with mj, it's a mujoco constant,
# so we don't need to get a parent struct ptr
if extent.startswith("mj"):
return str(extent)
if wrapper_name == "MjData" and field_name not in constants.MJDATA_SIZES:
var_name = "model"
else:
var_name = "ptr_"
return f"{var_name}->{extent}"
@@ -0,0 +1,241 @@
# 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
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(
name='ngeom',
type=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 = StructFieldDecl(
name='geom_rgba',
type=PointerType(
inner_type=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 = StructFieldDecl(
name='buffer',
type=PointerType(
inner_type=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<uint8_t*>(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 = StructFieldDecl(
name='element',
type=PointerType(
inner_type=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 = StructFieldDecl(
name='gravity',
type=ArrayType(
inner_type=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 = StructFieldDecl(
name='multi_dim_array',
type=ArrayType(
inner_type=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<float*>(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()
+201
View File
@@ -0,0 +1,201 @@
# 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.
"""Parser for MuJoCo structs."""
import collections
from typing import Dict, List
from introspect import ast_nodes
from introspect import structs
from wasm.codegen.helpers import common
from wasm.codegen.helpers import constants
from wasm.codegen.helpers import struct_constructor_code_builder
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]:
"""Generates WASM bindings for MuJoCo structs."""
for struct_name, wrap_data in wrapped_structs.items():
if struct_name in introspect_structs:
struct_fields = introspect_structs[struct_name].fields
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
):
raise RuntimeError(f"Anonymous struct not found: {struct_name}")
struct_fields = anonymous_struct.type.fields
else:
raise RuntimeError(f"Struct not found: {struct_name}")
debug_print(f"Wrapping struct: {struct_name}")
fields_with_init: List[WrappedFieldData] = []
for field in struct_fields:
field_gen = 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:
wrap_data.use_shallow_copy = False
if field_gen.initialization:
fields_with_init.append(field_gen)
wrap_data.wrapped_fields.append(field_gen)
wrap_data.wrapped_header = (
struct_constructor_code_builder.build_struct_header(
struct_name,
wrap_data.use_shallow_copy,
fields_with_init,
wrap_data.wrapped_fields,
)
)
wrap_data.wrapped_source = (
struct_constructor_code_builder.build_struct_source(
struct_name,
get_default_func_name(struct_name),
fields_with_init,
wrap_data.use_shallow_copy,
)
)
return wrapped_structs
def _get_anonymous_struct_field(
anonymous_structs_key: str,
) -> 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"]]
target_field = next(
(
f
for f in parent_decl.fields
if hasattr(f, "name")
and f.name == info["field_name"]
and hasattr(f, "type")
and isinstance(f.type, AnonymousStructDecl)
),
None,
)
return target_field
def get_default_func_name(struct_name: str) -> str:
"""Returns the default function name for the given struct."""
if (
struct_name in constants.ANONYMOUS_STRUCTS.keys()
or struct_name in constants.NO_DEFAULT_CONSTRUCTORS
or (
common.uppercase_first_letter(struct_name)
in constants.MANUALLY_ADDED_FIELDS_FROM_TEMPLATE.keys()
)
):
return ""
elif struct_name.startswith("mjs"):
return f"mjs_default{struct_name.removeprefix('mjs')}"
elif struct_name.startswith("mjv"):
return f"mjv_default{struct_name.removeprefix('mjv')}"
else:
return f"mj_default{struct_name.removeprefix('mj')}"
def _get_field_struct_type(field_type):
"""Extracts the base struct name if the field type is a struct or pointer to a struct."""
if isinstance(field_type, ast_nodes.ValueType):
return field_type.name
if isinstance(field_type, ast_nodes.PointerType):
if isinstance(field_type.inner_type, ast_nodes.ValueType):
return field_type.inner_type.name
return None
def sort_structs_by_dependency(struct_names: List[str]) -> List[str]:
"""Sorts structs based on their field dependencies using topological sort.
Structs with no dependencies on other structs in the list come first.
If struct A has a field of type struct B, B must come before A in the
sorted list.
Args:
struct_names: A list of struct names to sort.
Returns:
A new list of struct names sorted by dependency.
Raises:
RuntimeError: If a cyclic dependency is detected.
"""
adj = collections.defaultdict(list)
in_degree = collections.defaultdict(int)
struct_set = set(struct_names)
sorted_struct_names = sorted(struct_names)
for struct_name in sorted_struct_names:
if struct_name not in introspect_structs:
# Skip anonymous or other structs not in the main introspect map
continue
struct_decl = introspect_structs[struct_name]
for field in struct_decl.fields:
if isinstance(field, ast_nodes.AnonymousStructDecl):
continue
field_type_name = _get_field_struct_type(field.type)
if (
field_type_name
and field_type_name != struct_name
and field_type_name in struct_set
):
if struct_name not in adj[field_type_name]:
adj[field_type_name].append(struct_name)
in_degree[struct_name] += 1
queue = collections.deque(
[name for name in sorted_struct_names if in_degree[name] == 0]
)
sorted_list = []
while queue:
u = queue.popleft()
sorted_list.append(u)
for v in adj[u]:
in_degree[v] -= 1
if in_degree[v] == 0:
queue.append(v)
if len(sorted_list) == len(struct_names):
return sorted_list
else:
remaining = set(struct_names) - set(sorted_list)
raise RuntimeError(
"Cycle detected in struct dependencies, involving: "
f"{', '.join(sorted(list(remaining)))}"
)
+110
View File
@@ -0,0 +1,110 @@
# 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(
structs_wrappers_data.create_wrapped_structs_set_up_data([
"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(
structs_wrappers_data.create_wrapped_structs_set_up_data(
["mjFakeStruct"]
)
)
with self.assertRaises(RuntimeError):
structs_parser.generate_wasm_bindings(
structs_wrappers_data.create_wrapped_structs_set_up_data(
["mjFakeAnonymousStruct"]
)
)
def test_get_default_func_name_mjv(self):
self.assertEqual(
structs_parser.get_default_func_name("mjvPerturb"), "mjv_defaultPerturb"
)
def test_get_default_func_name(self):
self.assertEqual(
structs_parser.get_default_func_name("mjOption"), "mj_defaultOption"
)
if __name__ == "__main__":
absltest.main()
@@ -0,0 +1,75 @@
# 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.
"""Classes used by the parser to generate the structs wrappers and bindings statements."""
import dataclasses
from typing import Dict, List
from wasm.codegen.helpers import common
@dataclasses.dataclass
class WrappedFieldData:
"""Data class for struct field definition and binding."""
# Line for struct field binding
binding: str
# Line for struct field definition
definition: str | None = None
# Initialization code for fields that require it
initialization: str | None = None
# Statement to reset the inner pointer when copying the field
ptr_copy_reset: str | None = None
# Whether the field is a primitive or fixed size
is_primitive_or_fixed_size: bool = False
@dataclasses.dataclass
class WrappedStructData:
"""Data class for struct wrapper definition and binding."""
# Name of wrapper struct
wrap_name: str
# List of WrappedFieldData for this struct
wrapped_fields: List[WrappedFieldData]
# Struct header code
wrapped_header: str
# Struct source code
wrapped_source: str
# Whether to use shallow copy for this struct
use_shallow_copy: bool = True
def create_wrapped_structs_set_up_data(
struct_names: List[str],
) -> Dict[str, WrappedStructData]:
"""Creates a dictionary of WrappedStructData for the given struct names."""
return {
name: WrappedStructData(
wrap_name=common.uppercase_first_letter(name),
wrapped_fields=[],
wrapped_header="",
wrapped_source="",
)
for name in struct_names
}