Refactor WASM bindings in preparation for codegen directory structure update

PiperOrigin-RevId: 829459535
Change-Id: I4045e163a6e7fcee9d9585131d1e5e30f0b50b89

wasm/codegen/binding_builder.py

@@ -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

@@ -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):
+  def test_bindings_source(self):
-    super().setUp()
-
     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

@@ -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

@@ -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

@@ -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,
+    marker: str,
-    replacement_content: str | list[str],
+    content: list[str],
 ) -> list[str]:
   """Replaces lines containing a specific marker with new content."""
-
+  for i, line in enumerate(lines):
-  new_lines = []
+    if marker in line:
-  replaced = False
+      indent = line[: len(line) - len(line.lstrip(" "))]
-  for line in lines:
+      replacement_lines = []
-    if not replaced and marker_to_replace in line:
+      for text in content:
-      indentation = _get_indentation(line)
+        if text.strip():
-      if isinstance(replacement_content, str):
+          # Prepend indent to ensure the first replacement line matches the
-        new_lines.append(indentation + replacement_content)
+          # indentation of the marker and also ensure that text containing
-      elif isinstance(replacement_content, list):
+          # newlines is also indented correctly.
-        for content_line in replacement_content:
+          # TODO(matijak): This is working around an upstream problem, we should
-          if not content_line.strip():
+          # make it a precondition that content elements do not contain newlines
-            continue
+          # and fix callers to ensure that.
-          indented_line = (
+          replacement_lines.append(
-              indentation
+              indent + text.replace("\n", f"\n{indent}") + "\n"
-              + content_line.replace("\n", "\n" + indentation)
-              + "\n"
           )
-          new_lines.append(indented_line)
+      return lines[:i] + replacement_lines + lines[i + 1 :]
-      replaced = True
+  return lines
-    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

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):
 

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<const MjSpec &>()")
     elif w == "MjvScene":
       builder.line(".constructor<MjModel *, int>()")
-
+      builder.line(".constructor<>()")
-    is_mjs = w.startswith("Mjs")
+    elif not is_mjs:
-    if not is_mjs and w not in ["MjData", "MjModel", "MjSpec"]:
       builder.line(".constructor<>()")
 
     shallow_copy = use_shallow_copy(wrapped_fields)

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
+import os
-about its internal structures and then uses a code generation framework to
-produce corresponding WASM bindings.
-"""
 
-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():
+class BindingBuilder:
-  """Generates WASM bindings for MuJoCo."""
+  """Builds WASM bindings for MuJoCo."""
-  template_path_cc, generated_path_cc = common.get_file_path(
+
-      "templates", "generated", "bindings.cc"
+  def __init__(
-  )
+      self,
-  builder = binding_builder.BindingBuilder(template_path_cc)
+      template_path_cc: str,
-  builder.set_enums().set_structs().set_functions().build(generated_path_cc)
+  ):
+    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)