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

Refactor WASM bindings in preparation for codegen directory structure update

Browse Source 
PiperOrigin-RevId: 829459535
Change-Id: I4045e163a6e7fcee9d9585131d1e5e30f0b50b89
This commit is contained in:
Matija Kecman
2025-11-07 09:01:25 -08:00
committed by Copybara-Service
parent 59debb50b1
commit 44220fcc51
8 changed files with 96 additions and 222 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wasm/codegen/binding_builder.py
-76
View File
@@ -1,76 +0,0 @@
# Copyright 2025 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Builds WASM bindings for MuJoCo."""
from introspect import ast_nodes
from introspect import enums as introspect_enums
from introspect import functions as introspect_functions
from wasm.codegen.generators import enums
from wasm.codegen.generators import functions
from wasm.codegen.generators import structs
from wasm.codegen.helpers import common
from wasm.codegen.helpers import constants as _constants
from wasm.codegen.helpers import functions as function_utils
class BindingBuilder:
"""Builds WASM bindings for MuJoCo."""
def __init__(
self,
template_path_cc: str,
):
with open(template_path_cc, "r") as f:
self.content_cc = f.readlines()
self.markers_and_content = []
def set_enums(self):
"""Generates and sets the enum bindings."""
generator = enums.Generator(introspect_enums.ENUMS)
self.markers_and_content += generator.generate()
return self
def set_structs(self):
"""Generates and sets the struct bindings."""
generator = structs.Generator()
self.markers_and_content += generator.generate()
return self
def set_functions(self):
"""Generates and sets the function wrappers and bindings."""
functions_to_bind: dict[str, ast_nodes.FunctionDecl] = {}
for name, func in introspect_functions.FUNCTIONS.items():
if not function_utils.is_excluded_function_name(name):
if name not in _constants.BOUNDCHECK_FUNCS:
functions_to_bind[name] = func
generator = functions.Generator(functions_to_bind)
self.markers_and_content += generator.generate()
return self
def to_string(self) -> str:
for marker, content in self.markers_and_content:
self.content_cc = common.replace_lines_containing_marker(
self.content_cc, marker, content
)
return "".join(self.content_cc)
def build(self, generated_path_cc: str):
"""Writes the generated content to the output files."""
common.write_to_file(generated_path_cc, self.to_string())
wasm/codegen/bindings_diff_test.py
+3 -6
View File
@@ -16,7 +16,7 @@ from pathlib import Path
from absl.testing import absltest
from wasm.codegen import binding_builder
from wasm.codegen import update
ERROR_MESSAGE = """
The file '{}' needs to be updated, please run:
@@ -25,17 +25,14 @@ update.py as described in wasm/README.md""".lstrip()
class BindingsDiffTest(absltest.TestCase):
def setUp(self):
super().setUp()
def test_bindings_source(self):
SCRIPT_DIR = Path(__file__).parent
with open(SCRIPT_DIR / 'generated/bindings.cc', 'r') as f:
self.generated_src = f.read()
self.template_path_cc = SCRIPT_DIR / 'templates/bindings.cc'
self.builder = binding_builder.BindingBuilder(self.template_path_cc)
self.builder = update.BindingBuilder(self.template_path_cc)
def test_bindings_source(self):
generator_output = (
self.builder.set_enums().set_structs().set_functions().to_string()
)
wasm/codegen/generators/constants.py
-34
View File
@@ -1,34 +0,0 @@
# Copyright 2025 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Generator for the constants."""
from wasm.codegen.helpers import common
# TODO(manevi): Delete this file and use the genrule to handle the file copying
class Generator:
"""Generator for the constants."""
def run(self):
"""Runs the generator."""
template_cc_file, output_cc_file = common.get_file_path(
"templates", "generated", "constants.cc"
)
with open(template_cc_file, "r") as f_template:
template_content = f_template.read()
with open(output_cc_file, "w") as f_output:
f_output.write(template_content)
wasm/codegen/generators/functions.py
+1 -1
View File
@@ -14,7 +14,7 @@
"""Generates Embind bindings for MuJoCo functions."""
from typing import List, Mapping, Optional
from typing import List, Mapping
from introspect import ast_nodes
wasm/codegen/helpers/common.py
+18 -75
View File
@@ -15,36 +15,6 @@
"""Utility functions for code generation."""
import os
import pathlib
from wasm.codegen.helpers import constants
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 = 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
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:
@@ -71,53 +41,26 @@ def uppercase_first_letter(input_string: str) -> str:
return input_string[:1].upper() + input_string[1:]
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 replace_lines_containing_marker(
lines: list[str],
marker_to_replace: str,
replacement_content: str | list[str],
marker: str,
content: 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"
for i, line in enumerate(lines):
if marker in line:
indent = line[: len(line) - len(line.lstrip(" "))]
replacement_lines = []
for text in content:
if text.strip():
# Prepend indent to ensure the first replacement line matches the
# indentation of the marker and also ensure that text containing
# newlines is also indented correctly.
# TODO(matijak): This is working around an upstream problem, we should
# make it a precondition that content elements do not contain newlines
# and fix callers to ensure that.
replacement_lines.append(
indent + text.replace("\n", f"\n{indent}") + "\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
return lines[:i] + replacement_lines + lines[i + 1 :]
return lines
wasm/codegen/helpers/helpers_test.py
-5
View File
@@ -64,11 +64,6 @@ class CommonUtilsTest(absltest.TestCase):
common.uppercase_first_letter(" leading space"), " leading space"
)
def test_try_cast_to_scalar_type(self):
self.assertEqual(common.try_cast_to_scalar_type("123"), 123)
self.assertEqual(common.try_cast_to_scalar_type("123.456"), 123.456)
self.assertEqual(common.try_cast_to_scalar_type("abc"), "abc")
class FunctionUtilsTest(absltest.TestCase):
wasm/codegen/helpers/structs.py
+8 -11
View File
@@ -187,12 +187,6 @@ def _generate_field_data(
binding=_simple_property_binding(f, w),
is_primitive_or_fixed_size=True,
)
else:
return WrappedFieldData(
definition=f"// TODO: NOT IMPLEMENTED ARRAY wrapper for {f.name}",
typename=_get_field_struct_type(f.type),
binding=f"// TODO: NOT IMPLEMENTED ARRAY binding for {f.name}",
)
elif isinstance(f.type, ast_nodes.PointerType):
@@ -280,10 +274,12 @@ def _generate_field_data(
binding=_simple_property_binding(f, w),
)
# SHOULD NOT OCCUR
print("Error: field {f.name} not properly handled")
return WrappedFieldData(
definition=f"// TODO: UNDEFINED definition for {f.name}",
definition=f"// Error: field {f.name} not properly handled.",
typename=_get_field_struct_type(f.type),
binding=f"// TODO: UNDEFINED binding for {f.name}",
binding=f"// Error: field {f.name} not properly handled.",
)
@@ -512,6 +508,8 @@ def _build_struct_bindings(
):
"""Builds the C++ bindings for a struct."""
w = common.uppercase_first_letter(struct_name)
is_mjs = w.startswith("Mjs")
builder = code_builder.CodeBuilder()
with builder.block(
header_line=f'emscripten::class_<{w}>("{w}")', braces=False
@@ -528,9 +526,8 @@ def _build_struct_bindings(
builder.line(".constructor<const MjSpec &>()")
elif w == "MjvScene":
builder.line(".constructor<MjModel *, int>()")
is_mjs = w.startswith("Mjs")
if not is_mjs and w not in ["MjData", "MjModel", "MjSpec"]:
builder.line(".constructor<>()")
elif not is_mjs:
builder.line(".constructor<>()")
shallow_copy = use_shallow_copy(wrapped_fields)
wasm/codegen/update.py
+66 -14
View File
@@ -12,25 +12,77 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""Generates WASM bindings for MuJoCo's API.
"""Generates Javascript/TypeScript bindings for MuJoCo."""
This script leverages MuJoCo's introspect dicts to gather information
about its internal structures and then uses a code generation framework to
produce corresponding WASM bindings.
"""
import os
from wasm.codegen import binding_builder
from introspect import ast_nodes
from introspect import enums as introspect_enums
from introspect import functions as introspect_functions
from wasm.codegen.generators import enums
from wasm.codegen.generators import functions
from wasm.codegen.generators import structs
from wasm.codegen.helpers import common
from wasm.codegen.helpers import constants as _constants
from wasm.codegen.helpers import functions as function_utils
def generate_all_bindings():
"""Generates WASM bindings for MuJoCo."""
template_path_cc, generated_path_cc = common.get_file_path(
"templates", "generated", "bindings.cc"
)
builder = binding_builder.BindingBuilder(template_path_cc)
builder.set_enums().set_structs().set_functions().build(generated_path_cc)
class BindingBuilder:
"""Builds WASM bindings for MuJoCo."""
def __init__(
self,
template_path_cc: str,
):
with open(template_path_cc, "r") as f:
self.content_cc = f.readlines()
self.markers_and_content = []
def set_enums(self):
"""Generates and sets the enum bindings."""
generator = enums.Generator(introspect_enums.ENUMS)
self.markers_and_content += generator.generate()
return self
def set_structs(self):
"""Generates and sets the struct bindings."""
generator = structs.Generator()
self.markers_and_content += generator.generate()
return self
def set_functions(self):
"""Generates and sets the function wrappers and bindings."""
functions_to_bind: dict[str, ast_nodes.FunctionDecl] = {}
for name, func in introspect_functions.FUNCTIONS.items():
if not function_utils.is_excluded_function_name(name):
if name not in _constants.BOUNDCHECK_FUNCS:
functions_to_bind[name] = func
generator = functions.Generator(functions_to_bind)
self.markers_and_content += generator.generate()
return self
def to_string(self) -> str:
for marker, content in self.markers_and_content:
self.content_cc = common.replace_lines_containing_marker(
self.content_cc, marker, content
)
return "".join(self.content_cc)
def build(self, generated_path_cc: str):
common.write_to_file(generated_path_cc, self.to_string())
if __name__ == "__main__":
generate_all_bindings()
template_file = "wasm/codegen/templates/bindings.cc"
generated_file = "wasm/codegen/generated/bindings.cc"
builder = BindingBuilder(template_file)
builder.set_enums()
builder.set_structs()
builder.set_functions()
builder.build(generated_file)