Change WASM generators from classes to functions.

PiperOrigin-RevId: 836257289
Change-Id: I7a3c1d7bb545353a3d278330582e33c0551792cb

wasm/codegen/generators/binding_builder.py

@@ -17,6 +17,7 @@ from introspect import enums as introspect_enums
 from introspect import functions as introspect_functions
 
 from wasm.codegen.generators import common
+from wasm.codegen.generators import constants
 from wasm.codegen.generators import enums
 from wasm.codegen.generators import functions
 from wasm.codegen.generators import structs
@@ -36,14 +37,12 @@ class BindingBuilder:
 
   def set_enums(self):
     """Generates and sets the enum bindings."""
-    generator = enums.Generator(introspect_enums.ENUMS)
+    self.markers_and_content += enums.generate(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 += structs.generate(constants.STRUCTS_TO_BIND)
-    self.markers_and_content += generator.generate()
     return self
 
   def set_functions(self):
@@ -54,8 +53,7 @@ class BindingBuilder:
       if not functions.is_excluded_function_name(name):
         functions_to_bind[name] = func
 
-    generator = functions.Generator(functions_to_bind)
+    self.markers_and_content += functions.generate(functions_to_bind)
-    self.markers_and_content += generator.generate()
     return self
 
   def to_string(self) -> str:

wasm/codegen/generators/enums.py

@@ -21,18 +21,14 @@ from introspect import ast_nodes
 from wasm.codegen.generators import code_builder
 
 
-class Generator:
+def generate(
-  """Generates Embind code for MuJoCo enums."""
+    enums: Mapping[str, ast_nodes.EnumDecl],
-
+) -> list[tuple[str, list[str]]]:
-  def __init__(self, enums: Mapping[str, ast_nodes.EnumDecl]):
-    self.enums = enums
-
-  def generate(self) -> list[tuple[str, list[str]]]:
   """Generates all Embind code for the provided enums."""
 
   builder = code_builder.CodeBuilder()
   with builder.block('EMSCRIPTEN_BINDINGS(mujoco_enums)'):
-      for e in self.enums.values():
+    for e in enums.values():
       if e.values:  # Skip empty enums.
         with builder.block(f'enum_<{e.name}>("{e.name}")', braces=False):
           names = list(e.values.keys())

wasm/codegen/generators/functions.py

@@ -14,7 +14,7 @@
 
 """Helper functions for processing and generating bindings for MuJoCo functions."""
 
-from typing import List, Mapping, Set, Tuple, cast
+from typing import Mapping, Tuple, cast
 
 from introspect import ast_nodes
 
@@ -288,30 +288,24 @@ def is_excluded_function_name(func_name: str) -> bool:
   )
 
 
-class Generator:
+def generate(
+    functions: Mapping[str, ast_nodes.FunctionDecl],
+) -> list[tuple[str, list[str]]]:
   """Generates Embind bindings for MuJoCo functions."""
-
+  wrapper_functions = []
-  def __init__(self, functions: Mapping[str, ast_nodes.FunctionDecl]):
+  for func in functions.values():
-    self.functions = functions
-
-  def generate(self) -> list[tuple[str, list[str]]]:
-    """Generates the bindings file for all functions."""
-
-    function_wrappers = []
-    for func in self.functions.values():
     if should_be_wrapped(func):
       if func.name not in constants.MANUAL_WRAPPER_FUNCTIONS:
-          function_wrappers.append(generate_function_wrapper(func))
+        wrapper_functions.append(generate_function_wrapper(func))
+  wrapper_content = "\n\n".join(wrapper_functions)
 
   function_bindings = []
-    for func in self.functions.values():
+  for func in functions.values():
-      js_name = func.name
+    suffix = "_wrapper" if should_be_wrapped(func) else ""
-      cpp_func = func.name
+    function_bindings.append(f'function("{func.name}", &{func.name}{suffix});')
-      if should_be_wrapped(func):
+  bindings_content = "\n".join(function_bindings)
-        cpp_func += "_wrapper"
-      function_bindings.append(f'function("{js_name}", &{cpp_func});')
 
   return [
-        ("// {{ WRAPPER_FUNCTIONS }}", ["\n\n".join(function_wrappers)]),
+      ("// {{ WRAPPER_FUNCTIONS }}", [wrapper_content]),
-        ("// {{ FUNCTION_BINDINGS }}", ["\n".join(function_bindings)]),
+      ("// {{ FUNCTION_BINDINGS }}", [bindings_content]),
   ]

wasm/codegen/generators/structs.py

@@ -719,17 +719,12 @@ def sort_structs_by_dependency(
     )
 
 
-class Generator:
+def generate(struct_to_bind: List[str]) -> list[tuple[str, list[str]]]:
-  """Generates C++ code for binding and wrapping MuJoCo structs."""
-
-  def generate(self) -> list[tuple[str, list[str]]]:
   """Generates C++ header file for binding and wrapping MuJoCo structs."""
 
   # Traverse the introspect dictionary to get the field
   # wrapper/bindings statements set up for each struct
-    self.structs_to_bind_data = generate_wasm_bindings(
+  structs_to_bind_data = generate_wasm_bindings(struct_to_bind)
-        constants.STRUCTS_TO_BIND
-    )
 
   autogenned_struct_definitions = []
   markers_and_content = []
@@ -746,10 +741,10 @@ class Generator:
   ))
 
   # Sort by struct name by dependency to ensure deterministic output order
-    sorted_struct_names = sort_structs_by_dependency(self.structs_to_bind_data)
+  sorted_struct_names = sort_structs_by_dependency(structs_to_bind_data)
 
   for struct_name in sorted_struct_names:
-      struct_data = self.structs_to_bind_data[struct_name]
+    struct_data = structs_to_bind_data[struct_name]
     if struct_data.wrapped_header:
       autogenned_struct_definitions.append(struct_data.wrapped_header + "\n")
     else:
@@ -768,7 +763,7 @@ class Generator:
   autogenned_struct_source = []
   autogenned_struct_bindings = []
   for struct_name in sorted_struct_names:
-      struct_data = self.structs_to_bind_data[struct_name]
+    struct_data = structs_to_bind_data[struct_name]
     if struct_data.wrapped_source:
       autogenned_struct_source.append(struct_data.wrapped_source + "\n")
     autogenned_struct_bindings.append(struct_data.bindings)

wasm/codegen/tests/generators_test.py

@@ -19,7 +19,6 @@ from introspect import ast_nodes
 
 from wasm.codegen.generators import code_builder
 from wasm.codegen.generators import common
-from wasm.codegen.generators import constants
 from wasm.codegen.generators import enums
 from wasm.codegen.generators import functions
 from wasm.codegen.generators import structs
@@ -801,7 +800,20 @@ class EnumsGeneratorTest(absltest.TestCase):
 
   def test_generate_enum_bindings(self):
 
-    generator = enums.Generator({
+    expected_code = """
+EMSCRIPTEN_BINDINGS(mujoco_enums) {
+  enum_<TestEnum>("TestEnum")
+    .value("FIRST_VAL", FIRST_VAL)
+    .value("SECOND_VAL", SECOND_VAL)
+    .value("THIRD_VAL", THIRD_VAL);
+
+  enum_<AnotherEnum>("AnotherEnum")
+    .value("ALPHA", ALPHA)
+    .value("BETA", BETA);
+
+}""".strip()
+
+    markers_and_content = enums.generate({
         "TestEnum": ast_nodes.EnumDecl(
             name="TestEnum",
             declname="enum TestEnum_",
@@ -818,21 +830,6 @@ class EnumsGeneratorTest(absltest.TestCase):
             values={},
         ),
     })
-
-    expected_code = """
-EMSCRIPTEN_BINDINGS(mujoco_enums) {
-  enum_<TestEnum>("TestEnum")
-    .value("FIRST_VAL", FIRST_VAL)
-    .value("SECOND_VAL", SECOND_VAL)
-    .value("THIRD_VAL", THIRD_VAL);
-
-  enum_<AnotherEnum>("AnotherEnum")
-    .value("ALPHA", ALPHA)
-    .value("BETA", BETA);
-
-}""".strip()
-
-    markers_and_content = generator.generate()
     actual_code = "\n\n".join(markers_and_content[0][1])
 
     self.assertEqual(actual_code, expected_code)