From 44220fcc510d5a28421526ee6caede0ccc2e2c0e Mon Sep 17 00:00:00 2001 From: Matija Kecman Date: Fri, 7 Nov 2025 09:01:25 -0800 Subject: [PATCH] Refactor WASM bindings in preparation for codegen directory structure update PiperOrigin-RevId: 829459535 Change-Id: I4045e163a6e7fcee9d9585131d1e5e30f0b50b89 --- wasm/codegen/binding_builder.py | 76 ----------------------- wasm/codegen/bindings_diff_test.py | 9 +-- wasm/codegen/generators/constants.py | 34 ---------- wasm/codegen/generators/functions.py | 2 +- wasm/codegen/helpers/common.py | 93 ++++++---------------------- wasm/codegen/helpers/helpers_test.py | 5 -- wasm/codegen/helpers/structs.py | 19 +++--- wasm/codegen/update.py | 80 +++++++++++++++++++----- 8 files changed, 96 insertions(+), 222 deletions(-) delete mode 100644 wasm/codegen/binding_builder.py delete mode 100644 wasm/codegen/generators/constants.py diff --git a/wasm/codegen/binding_builder.py b/wasm/codegen/binding_builder.py deleted file mode 100644 index 1f61766a..00000000 --- a/wasm/codegen/binding_builder.py +++ /dev/null @@ -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()) diff --git a/wasm/codegen/bindings_diff_test.py b/wasm/codegen/bindings_diff_test.py index edbfd732..cf77f017 100644 --- a/wasm/codegen/bindings_diff_test.py +++ b/wasm/codegen/bindings_diff_test.py @@ -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() ) diff --git a/wasm/codegen/generators/constants.py b/wasm/codegen/generators/constants.py deleted file mode 100644 index a7550ca2..00000000 --- a/wasm/codegen/generators/constants.py +++ /dev/null @@ -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) diff --git a/wasm/codegen/generators/functions.py b/wasm/codegen/generators/functions.py index 5e7d8f8f..bdfbcaa7 100644 --- a/wasm/codegen/generators/functions.py +++ b/wasm/codegen/generators/functions.py @@ -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 diff --git a/wasm/codegen/helpers/common.py b/wasm/codegen/helpers/common.py index c809b6ac..a056c325 100644 --- a/wasm/codegen/helpers/common.py +++ b/wasm/codegen/helpers/common.py @@ -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 diff --git a/wasm/codegen/helpers/helpers_test.py b/wasm/codegen/helpers/helpers_test.py index 57f63f2f..a539f854 100644 --- a/wasm/codegen/helpers/helpers_test.py +++ b/wasm/codegen/helpers/helpers_test.py @@ -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): diff --git a/wasm/codegen/helpers/structs.py b/wasm/codegen/helpers/structs.py index ed4b87d7..810e6b89 100644 --- a/wasm/codegen/helpers/structs.py +++ b/wasm/codegen/helpers/structs.py @@ -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()") elif w == "MjvScene": builder.line(".constructor()") - - 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) diff --git a/wasm/codegen/update.py b/wasm/codegen/update.py index b32684a3..776e2d2e 100644 --- a/wasm/codegen/update.py +++ b/wasm/codegen/update.py @@ -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)