Refactor WASM bindings in preparation for codegen directory structure update

PiperOrigin-RevId: 829459535
Change-Id: I4045e163a6e7fcee9d9585131d1e5e30f0b50b89

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)