More simplification to functions.py in WASM bindings

PiperOrigin-RevId: 834643175
Change-Id: Iff873f6f0916e0039ff6a44468afd9b4d8f98f37

wasm/codegen/generated/bindings.cc

@@ -12374,18 +12374,6 @@ std::optional<MjsPlugin> mjs_asPlugin_wrapper(MjsElement& element) {
 EMSCRIPTEN_BINDINGS(mujoco_functions) {
   function("parseXMLString", &parseXMLString_wrapper, take_ownership());
   function("error", &error_wrapper);
-  function("mj_resetCallbacks", &mj_resetCallbacks);
-  function("mj_version", &mj_version);
-  function("mju_bandDiag", &mju_bandDiag);
-  function("mju_springDamper", &mju_springDamper);
-  function("mju_min", &mju_min);
-  function("mju_max", &mju_max);
-  function("mju_clip", &mju_clip);
-  function("mju_sign", &mju_sign);
-  function("mju_round", &mju_round);
-  function("mju_isBad", &mju_isBad);
-  function("mju_Halton", &mju_Halton);
-  function("mju_sigmoid", &mju_sigmoid);
   function("mj_copyBack", &mj_copyBack_wrapper);
   function("mj_saveLastXML", &mj_saveLastXML_wrapper);
   function("mj_step", &mj_step_wrapper);
@@ -12403,6 +12391,7 @@ EMSCRIPTEN_BINDINGS(mujoco_functions) {
   function("mj_resetData", &mj_resetData_wrapper);
   function("mj_resetDataDebug", &mj_resetDataDebug_wrapper);
   function("mj_resetDataKeyframe", &mj_resetDataKeyframe_wrapper);
+  function("mj_resetCallbacks", &mj_resetCallbacks);
   function("mj_setConst", &mj_setConst_wrapper);
   function("mj_setLengthRange", &mj_setLengthRange_wrapper);
   function("mjs_activatePlugin", &mjs_activatePlugin_wrapper);
@@ -12494,6 +12483,7 @@ EMSCRIPTEN_BINDINGS(mujoco_functions) {
   function("mj_local2Global", &mj_local2Global_wrapper);
   function("mj_getTotalmass", &mj_getTotalmass_wrapper);
   function("mj_setTotalmass", &mj_setTotalmass_wrapper);
+  function("mj_version", &mj_version);
   function("mj_versionString", &mj_versionString_wrapper);
   function("mj_multiRay", &mj_multiRay_wrapper);
   function("mj_ray", &mj_ray_wrapper);
@@ -12605,6 +12595,7 @@ EMSCRIPTEN_BINDINGS(mujoco_functions) {
   function("mju_band2Dense", &mju_band2Dense_wrapper);
   function("mju_dense2Band", &mju_dense2Band_wrapper);
   function("mju_bandMulMatVec", &mju_bandMulMatVec_wrapper);
+  function("mju_bandDiag", &mju_bandDiag);
   function("mju_eig3", &mju_eig3_wrapper);
   function("mju_boxQP", &mju_boxQP_wrapper);
   function("mju_muscleGain", &mju_muscleGain_wrapper);
@@ -12612,10 +12603,17 @@ EMSCRIPTEN_BINDINGS(mujoco_functions) {
   function("mju_muscleDynamics", &mju_muscleDynamics_wrapper);
   function("mju_encodePyramid", &mju_encodePyramid_wrapper);
   function("mju_decodePyramid", &mju_decodePyramid_wrapper);
+  function("mju_springDamper", &mju_springDamper);
+  function("mju_min", &mju_min);
+  function("mju_max", &mju_max);
+  function("mju_clip", &mju_clip);
+  function("mju_sign", &mju_sign);
+  function("mju_round", &mju_round);
   function("mju_type2Str", &mju_type2Str_wrapper);
   function("mju_str2Type", &mju_str2Type_wrapper);
   function("mju_writeNumBytes", &mju_writeNumBytes_wrapper);
   function("mju_warningText", &mju_warningText_wrapper);
+  function("mju_isBad", &mju_isBad);
   function("mju_isZero", &mju_isZero_wrapper);
   function("mju_standardNormal", &mju_standardNormal_wrapper);
   function("mju_f2n", &mju_f2n_wrapper);
@@ -12624,6 +12622,8 @@ EMSCRIPTEN_BINDINGS(mujoco_functions) {
   function("mju_n2d", &mju_n2d_wrapper);
   function("mju_insertionSort", &mju_insertionSort_wrapper);
   function("mju_insertionSortInt", &mju_insertionSortInt_wrapper);
+  function("mju_Halton", &mju_Halton);
+  function("mju_sigmoid", &mju_sigmoid);
   function("mjd_transitionFD", &mjd_transitionFD_wrapper);
   function("mjd_inverseFD", &mjd_inverseFD_wrapper);
   function("mjd_subQuat", &mjd_subQuat_wrapper);

wasm/codegen/generators/functions.py

@@ -60,9 +60,12 @@ def get_const_qualifier(func: ast_nodes.FunctionDecl) -> str:
 
 def should_be_wrapped(func: ast_nodes.FunctionDecl) -> bool:
   """Checks if a MuJoCo function needs a wrapper function."""
-  return get_pointer_return_inner_value_type(func) is not None or any(
-      get_inner_value_type(param) for param in func.parameters
-  )
+  if get_pointer_return_inner_value_type(func):
+    return True
+  for param in func.parameters:
+    if get_inner_value_type(param):
+      return True
+  return False
 
 
 def generate_function_wrapper(func: ast_nodes.FunctionDecl) -> str:
@@ -88,11 +91,8 @@ def generate_function_wrapper(func: ast_nodes.FunctionDecl) -> str:
     if bound_check_code:
       builder.line(bound_check_code)
 
-    c_params_list = get_params_string_maybe_with_conversion(func.parameters)
-    c_params_str = ", ".join(c_params_list)
-    c_call = f"{func.name}({c_params_str})"
-    c_statement = get_compatible_return_call(func, c_call)
-    builder.line(f"{c_statement};")
+    for line in get_compatible_return_code(func).splitlines():
+      builder.line(line)
 
   return builder.to_string()
 
@@ -194,7 +194,7 @@ def get_param_string(p: ast_nodes.FunctionParameterDecl) -> str:
 
 def get_params_string_maybe_with_conversion(
     ast_params: Tuple[ast_nodes.FunctionParameterDecl, ...],
-) -> List[str]:
+) -> str:
   """Generates C++ expressions for passing compatible params from JS to MuJoCo C-API functions."""
 
   native_params = []
@@ -217,27 +217,27 @@ def get_params_string_maybe_with_conversion(
           f"Unhandled parameter type for conversion: {p.type} for param"
           f" '{p.name}'"
       )
-  return native_params
+  return ", ".join(native_params)
 
 
-def get_compatible_return_call(
-    func: ast_nodes.FunctionDecl, invoker: str
-) -> str:
+def get_compatible_return_code(func: ast_nodes.FunctionDecl) -> str:
   """Generates embind compatible return value conversion."""
+  c_params = get_params_string_maybe_with_conversion(func.parameters)
+  c_call = f"{func.name}({c_params})"
 
   if isinstance(func.return_type, ast_nodes.ValueType):
     if func.return_type.name == "void":
-      return invoker
+      return f"{c_call};"
     if func.return_type.name in constants.PRIMITIVE_TYPES:
-      return f"return {invoker}"
+      return f"return {c_call};"
 
   if inner_type := get_pointer_return_inner_value_type(func):
     if inner_type.name == "char":
-      return f"return std::string({invoker})"
+      return f"return std::string({c_call});"
     elif inner_type.name == "mjString":
-      return f"return *{invoker}"
+      return f"return *{c_call};"
     elif inner_type.name not in constants.PRIMITIVE_TYPES:
-      return get_converted_struct_to_class(func, invoker)
+      return get_optional_return_code(func, c_call)
 
   raise RuntimeError(
       "Failed to calculate return value conversion for function"
@@ -262,21 +262,22 @@ def get_compatible_return_type(func: ast_nodes.FunctionDecl) -> str:
   return "val"
 
 
-def get_converted_struct_to_class(
-    func: ast_nodes.FunctionDecl, invoker: str
+def get_optional_return_code(
+    func: ast_nodes.FunctionDecl, c_call: str
 ) -> str:
-  """Generates a C++ function invocation for a struct return-type function."""
+  """Generates code to return std::optional of the wrapped struct."""
 
   const_qualifier = get_const_qualifier(func)
   return_type = cast(ast_nodes.PointerType, func.return_type)
   struct_name = cast(ast_nodes.ValueType, return_type.inner_type).name
-  class_constructor = common.capitalize(struct_name)
-  return_str = f"{class_constructor}(result)"
-  return f"""{const_qualifier}{struct_name}* result = {invoker};
-  if (result == nullptr) {{
-    return std::nullopt;
-  }}
-  return {return_str}"""
+
+  builder = code_builder.CodeBuilder()
+  builder.line(f"{const_qualifier}{struct_name}* result = {c_call};")
+  with builder.block("if (result == nullptr)"):
+    builder.line("return std::nullopt;")
+  builder.line(f"return {common.capitalize(struct_name)}(result);")
+
+  return builder.to_string()
 
 
 def is_excluded_function_name(func_name: str) -> bool:
@@ -291,63 +292,26 @@ class Generator:
   """Generates Embind bindings for MuJoCo functions."""
 
   def __init__(self, functions: Mapping[str, ast_nodes.FunctionDecl]):
-    self.direct_bind_functions: List[ast_nodes.FunctionDecl] = []
-    self.wrapper_bind_functions: List[ast_nodes.FunctionDecl] = []
-
-    for func in functions.values():
-      if should_be_wrapped(func):
-        self.wrapper_bind_functions.append(func)
-      else:
-        self.direct_bind_functions.append(func)
-
-  def _generate_wrappers(self) -> str:
-    """Generates Embind bindings for all functions that need wrappers."""
-
-    code = []
-    for func in self.wrapper_bind_functions:
-      if func.name not in constants.MANUAL_WRAPPER_FUNCTIONS:
-        wrapper_code = generate_function_wrapper(func)
-        code.append(wrapper_code)
-
-    return "\n\n".join(code)
-
-  def _generate_direct_bindable_functions(self) -> list[str]:
-    """Generates Embind bindings for all directly bindable functions."""
-
-    result = []
-    for func in self.direct_bind_functions:
-      result.append(self._generate_function_binding(func))
-
-    return result
-
-  def _generate_function_binding(
-      self, func: ast_nodes.FunctionDecl, is_wrapper=False
-  ) -> str:
-    """Generates the Embind code for a single function."""
-
-    js_name, cpp_func = func.name, func.name
-    if is_wrapper:
-      cpp_func += "_wrapper"
-
-    return f'function("{js_name}", &{cpp_func});'
-
-  def _generate_wrapper_bindable_functions(self) -> list[str]:
-    """Generates Embind bindings for all functions that need wrappers."""
-
-    result = []
-    for func in self.wrapper_bind_functions:
-      result.append(self._generate_function_binding(func, True))
-
-    return result
+    self.functions = functions
 
   def generate(self) -> list[tuple[str, list[str]]]:
     """Generates the bindings file for all functions."""
 
-    wrapper_functions = self._generate_wrappers()
-    function_bindings = self._generate_direct_bindable_functions()
-    function_bindings += self._generate_wrapper_bindable_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))
+
+    function_bindings = []
+    for func in self.functions.values():
+      js_name = func.name
+      cpp_func = func.name
+      if should_be_wrapped(func):
+        cpp_func += "_wrapper"
+      function_bindings.append(f'function("{js_name}", &{cpp_func});')
 
     return [
-        ("// {{ WRAPPER_FUNCTIONS }}", [wrapper_functions]),
+        ("// {{ WRAPPER_FUNCTIONS }}", ["\n\n".join(function_wrappers)]),
         ("// {{ FUNCTION_BINDINGS }}", ["\n".join(function_bindings)]),
     ]

wasm/codegen/tests/generators_test.py

@@ -157,17 +157,17 @@ class FunctionUtilsTest(absltest.TestCase):
         type=ast_nodes.PointerType(ast_nodes.ValueType("customstruct")),
     )
     result = functions.get_params_string_maybe_with_conversion((param,))
-    self.assertEqual(result, ["s.get()"])
+    self.assertEqual(result, "s.get()")
 
-  def test_get_compatible_return_call(self):
+  def test_get_compatible_return_code(self):
     func = ast_nodes.FunctionDecl(
         name="noop",
         return_type=ast_nodes.ValueType("void"),
         parameters=tuple(),
         doc="does nothing",
     )
-    result = functions.get_compatible_return_call(func, "noop()")
-    self.assertEqual(result, "noop()")
+    result = functions.get_compatible_return_code(func)
+    self.assertEqual(result, "noop();")
 
   def test_get_compatible_return_type(self):
     func = ast_nodes.FunctionDecl(
@@ -179,14 +179,14 @@ class FunctionUtilsTest(absltest.TestCase):
     result = functions.get_compatible_return_type(func)
     self.assertEqual(result.strip(), "std::string")
 
-  def test_get_converted_struct_to_class(self):
+  def test_get_optional_return_code(self):
     func = ast_nodes.FunctionDecl(
         name="get_struct",
         return_type=ast_nodes.PointerType(ast_nodes.ValueType("mystruct")),
         parameters=tuple(),
         doc="returns struct",
     )
-    result = functions.get_converted_struct_to_class(func, "get_struct()")
+    result = functions.get_optional_return_code(func, "get_struct()")
     self.assertIn("mystruct* result = get_struct();", result)
     self.assertIn("return Mystruct(result)", result)
 
@@ -836,49 +836,5 @@ class EnumsGeneratorTest(absltest.TestCase):
     self.assertEqual(actual_code, expected_code)
 
 
-class FunctionsGeneratorTest(absltest.TestCase):
-
-  def setUp(self):
-    super().setUp()
-    self.generator = functions.Generator({})
-    self.int_type = ast_nodes.ValueType(name="int")
-
-  def test_generate_function_binding_simple_case(self):
-    func_simple_void = ast_nodes.FunctionDecl(
-        name="do_nothing",
-        return_type=ast_nodes.ValueType(name="void"),
-        parameters=tuple(),
-        doc="doc",
-    )
-    self.assertEqual(
-        self.generator._generate_function_binding(func_simple_void),
-        'function("do_nothing", &do_nothing);',
-    )
-
-  def test_generate_direct_bindable_functions_simple_filter(self):
-    direct_bind = ast_nodes.FunctionDecl(
-        name="direct_bind",
-        return_type=self.int_type,
-        parameters=(
-            ast_nodes.FunctionParameterDecl(name="val", type=self.int_type),
-        ),
-        doc="doc",
-    )
-    needs_wrap = ast_nodes.FunctionDecl(
-        name="needs_wrap",
-        return_type=ast_nodes.PointerType(inner_type=self.int_type),
-        parameters=tuple(),
-        doc="doc",
-    )
-    self.generator = functions.Generator({
-        "direct1": direct_bind,
-        "wrapped1": needs_wrap,
-    })
-
-    generated_code = self.generator._generate_direct_bindable_functions()
-    self.assertIn('function("direct_bind", &direct_bind);', generated_code)
-    self.assertNotIn('function("needs_wrap", &needs_wrap);', generated_code)
-
-
 if __name__ == "__main__":
   absltest.main()