Mujoco_WASM/wasm/codegen/generators/functions.py

# 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.

"""Helper functions for processing and generating bindings for MuJoCo functions."""

from typing import List, Mapping, Set, Tuple, cast

from introspect import ast_nodes

from wasm.codegen.generators import code_builder
from wasm.codegen.generators import common
from wasm.codegen.generators import constants


PRIMITIVE_TYPES = constants.PRIMITIVE_TYPES
uppercase_first_letter = common.uppercase_first_letter


def param_is_primitive_value(param: ast_nodes.FunctionParameterDecl) -> bool:
  """Checks if param is a primitive value type."""
  if isinstance(param.type, ast_nodes.ValueType):
    return param.type.name in PRIMITIVE_TYPES
  return False


def param_is_pointer_to_primitive_value(
    param: ast_nodes.FunctionParameterDecl,
) -> bool:
  """Checks if param is a pointer to a primitive value."""
  return (
      isinstance(param.type, ast_nodes.PointerType)
      or isinstance(param.type, ast_nodes.ArrayType)
  ) and (
      isinstance(param.type.inner_type, ast_nodes.ValueType)
      and param.type.inner_type.name in PRIMITIVE_TYPES
  )


def param_is_pointer_to_struct(param: ast_nodes.FunctionParameterDecl) -> bool:
  """Checks if param is a pointer to a struct."""
  return (
      isinstance(param.type, ast_nodes.PointerType)
      or isinstance(param.type, ast_nodes.ArrayType)
  ) and (
      isinstance(param.type.inner_type, ast_nodes.ValueType)
      and param.type.inner_type.name not in PRIMITIVE_TYPES
  )


def return_is_value_of_type(
    func: ast_nodes.FunctionDecl, allowed_types: Set[str]
) -> bool:
  """Checks if func returns an allowed value type."""
  return (
      isinstance(func.return_type, ast_nodes.ValueType)
      and func.return_type.name in allowed_types
  )


def return_is_pointer_to_struct(func: ast_nodes.FunctionDecl) -> bool:
  """Checks if func returns a pointer to a struct."""
  return (
      isinstance(func.return_type, ast_nodes.PointerType)
      and isinstance(func.return_type.inner_type, ast_nodes.ValueType)
      and func.return_type.inner_type.name not in PRIMITIVE_TYPES
  )


def return_is_pointer_to_primitive(func: ast_nodes.FunctionDecl) -> bool:
  """Checks if func returns a pointer to a primitive value."""
  return (
      isinstance(func.return_type, ast_nodes.PointerType)
      and isinstance(func.return_type.inner_type, ast_nodes.ValueType)
      and func.return_type.inner_type.name in PRIMITIVE_TYPES
  )


def get_const_qualifier(func: ast_nodes.FunctionDecl) -> str:
  """Returns the const qualifier of func's return type."""
  if (
      isinstance(func.return_type, ast_nodes.PointerType)
      and isinstance(func.return_type.inner_type, ast_nodes.ValueType)
      and func.return_type.inner_type.is_const
  ):
    return "const "
  return ""


def should_be_wrapped(func: ast_nodes.FunctionDecl) -> bool:
  """Checks if a MuJoCo function needs a wrapper function."""
  return (
      return_is_pointer_to_primitive(func)
      or return_is_pointer_to_struct(func)
      or any(
          param_is_pointer_to_primitive_value(param)
          or isinstance(param.type, ast_nodes.ArrayType)
          or param_is_pointer_to_struct(param)
          for param in func.parameters
      )
  )


def generate_function_wrapper(func: ast_nodes.FunctionDecl) -> str:
  """Generates C++ code for a wrapper function."""

  bound_check_code = constants.FUNCTION_BOUNDS_CHECKS.get(func.name, "")
  wrapper_parameters = [
      p for p in func.parameters if f"{p.name} = " not in bound_check_code
  ]
  wrapper_params = ", ".join([get_param_string(p) for p in wrapper_parameters])
  ret_type = get_compatible_return_type(func)
  builder = code_builder.CodeBuilder()
  with builder.function(f"{ret_type} {func.name}_wrapper({wrapper_params})"):

    for p in wrapper_parameters:
      if c_notnullable := get_param_notnullable(p):
        builder.line(f"CHECK_VAL({c_notnullable});")

    for p in wrapper_parameters:
      if c_unpack := get_param_unpack_statement(p):
        builder.line(c_unpack)

    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};")

  return builder.to_string()


def get_param_notnullable(
    p: ast_nodes.FunctionParameterDecl,
) -> str:
  """Generates list of param names for checking if they aren't null/undefined."""

  if (
      isinstance(p.type, (ast_nodes.PointerType, ast_nodes.ArrayType))
      and isinstance(p.type.inner_type, ast_nodes.ValueType)
      # We only check for char because others are checked in the unpacker
      # and we don't want to check twice.
      and p.type.inner_type.name == "char"
      and not p.nullable
  ):
    return p.name
  return ""


def get_param_unpack_statement(
    p: ast_nodes.FunctionParameterDecl,
) -> str | None:
  """Generates C++ statements to unpack JS values for pointer/array parameters."""

  if (
      isinstance(p.type, (ast_nodes.PointerType, ast_nodes.ArrayType))
      and isinstance(p.type.inner_type, ast_nodes.ValueType)
      and p.type.inner_type.name in PRIMITIVE_TYPES
  ):
    if p.type.inner_type.name == "char":
      # param is Javascript string
      return ""
    if p.type.inner_type.is_const:
      # param is Javascript number[]
      if p.nullable:
        return f"UNPACK_NULLABLE_ARRAY({p.type.inner_type.name}, {p.name});"
      else:
        return f"UNPACK_ARRAY({p.type.inner_type.name}, {p.name});"
    else:
      # param is TypedArray or a WasmBuffer
      if p.nullable:
        return f"UNPACK_NULLABLE_VALUE({p.type.inner_type.name}, {p.name});"
      else:
        return f"UNPACK_VALUE({p.type.inner_type.name}, {p.name});"
  return None


def get_param_string(
    p: ast_nodes.FunctionParameterDecl
) -> str:
  """Generates a list of C++ parameter declarations as strings."""

  if (
      isinstance(p.type, ast_nodes.PointerType)
      and isinstance(p.type.inner_type, ast_nodes.ValueType)
      and p.type.inner_type.name not in PRIMITIVE_TYPES
  ):
    # Pointer to struct parameters
    const_qualifier = "const " if p.type.inner_type.is_const else ""
    return (
        f"{const_qualifier}{uppercase_first_letter(p.type.inner_type.name)}&"
        f" {p.name}"
    )
  elif (
      isinstance(p.type, ast_nodes.ValueType)
      and p.type.name in PRIMITIVE_TYPES
  ):
    # Primitive value parameters
    const_qualifier = "const " if p.type.is_const else ""
    return f"{const_qualifier}{p.type} {p.name}"
  elif (
      isinstance(p.type, (ast_nodes.PointerType, ast_nodes.ArrayType))
      and isinstance(p.type.inner_type, ast_nodes.ValueType)
      and p.type.inner_type.name in PRIMITIVE_TYPES
  ):
    # Pointer to primitive value parameters or arrays
    if p.type.inner_type.name == "char":
      if p.nullable:
        return f"const NullableString& {p.name}"
      else:
        return f"const String& {p.name}"
    elif (
        p.type.inner_type.name
        in ["int", "float", "double", "mjtNum", "mjtByte"]
        and p.type.inner_type.is_const
    ):
      return f"const NumberArray& {p.name}"
    else:
      return f"const val& {p.name}"
  else:
    # This case should ideally not be reached if AST is well-formed
    # and types are categorized by the helper booleans correctly.
    raise TypeError(
        "Unable to generate param string. Unhandled parameter type:"
        f" {p.type} for param '{p.name}'"
    )


def get_params_string_maybe_with_conversion(
    ast_params: Tuple[ast_nodes.FunctionParameterDecl, ...],
) -> List[str]:
  """Generates C++ expressions for passing compatible params from JS to MuJoCo C-API functions."""

  native_params = []
  for p in ast_params:
    if param_is_pointer_to_struct(p):
      native_params.append(f"{p.name}.get()")
    elif param_is_primitive_value(p):
      native_params.append(p.name)
    elif (
        isinstance(p.type, (ast_nodes.PointerType, ast_nodes.ArrayType))
        and isinstance(p.type.inner_type, ast_nodes.ValueType)
        and p.type.inner_type.name in PRIMITIVE_TYPES
        and p.type.inner_type.name != "char"
    ):
      native_params.append(f"{p.name}_.data()")
    elif (
        isinstance(p.type, (ast_nodes.PointerType, ast_nodes.ArrayType))
        and isinstance(p.type.inner_type, ast_nodes.ValueType)
        and p.type.inner_type.name == "char"
    ):
      const_qualifier = "const " if p.type.inner_type.is_const else ""
      native_params.append(
          f"{p.name}.as<{const_qualifier}std::string>().data()"
      )
    else:
      raise TypeError(
          f"Unhandled parameter type for conversion: {p.type} for param"
          f" '{p.name}'"
      )
  return native_params


def get_compatible_return_call(
    func: ast_nodes.FunctionDecl, invoker: str
) -> str:
  """Generates embind compatible return value conversion."""

  if return_is_value_of_type(func, {"void"}):
    return invoker
  if isinstance(func.return_type, ast_nodes.PointerType) and isinstance(
      func.return_type.inner_type, ast_nodes.ValueType
  ):
    if func.return_type.inner_type.name == "char":
      return f"return std::string({invoker})"
    elif func.return_type.inner_type.name == "mjString":
      return f"return *{invoker}"
  if return_is_pointer_to_struct(func):
    return get_converted_struct_to_class(func, invoker)
  if return_is_value_of_type(func, PRIMITIVE_TYPES):
    return f"return {invoker}"
  raise RuntimeError(
      "Failed to calculate return value conversion for function"
      f" {func.name} that returns '{func.return_type}'"
  )


def get_compatible_return_type(func: ast_nodes.FunctionDecl) -> str:
  """Creates embind compatible return type."""

  if (
      isinstance(func.return_type, ast_nodes.PointerType)
      and isinstance(func.return_type.inner_type, ast_nodes.ValueType)
      and func.return_type.inner_type.name in ["char", "mjString"]
  ):
    return "std::string"
  if (
      isinstance(func.return_type, ast_nodes.PointerType)
      and isinstance(func.return_type.inner_type, ast_nodes.ValueType)
      and func.return_type.inner_type.name not in PRIMITIVE_TYPES
  ):
    const_qualifier = get_const_qualifier(func)
    return f"""{const_qualifier}std::optional<{uppercase_first_letter(func.return_type.inner_type.name)}>"""
  if (
      isinstance(func.return_type, ast_nodes.ValueType)
      and func.return_type.name in PRIMITIVE_TYPES
  ):
    return f"{func.return_type.name}"
  return "val"


def get_converted_struct_to_class(
    func: ast_nodes.FunctionDecl, invoker: str
) -> str:
  """Generates a C++ function invocation for a struct return-type function."""

  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 = uppercase_first_letter(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}"""


def is_excluded_function_name(func_name: str) -> bool:
  """Checks if a function name should be excluded from direct binding."""
  return (
      func_name.startswith("mjr_")
      or func_name.startswith("mjui_")
      or func_name in constants.SKIPPED_FUNCTIONS
  )


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

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

    return [
        ("// {{ WRAPPER_FUNCTIONS }}", [wrapper_functions]),
        ("// {{ FUNCTION_BINDINGS }}", ["\n".join(function_bindings)]),
    ]