10fe5ff07f
PiperOrigin-RevId: 834264723 Change-Id: Id33d7568a1f9f0cfb1271cf9dd2d8713c00eee34
787 lines
25 KiB
Python
787 lines
25 KiB
Python
# 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.
|
|
|
|
"""Generates Embind bindings for MuJoCo structs."""
|
|
|
|
import collections
|
|
import dataclasses
|
|
import math
|
|
from typing import Dict, List, Tuple, Union, cast
|
|
|
|
from introspect import ast_nodes
|
|
from introspect import structs as introspect_structs
|
|
|
|
from wasm.codegen.generators import code_builder
|
|
from wasm.codegen.generators import common
|
|
from wasm.codegen.generators import constants
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class WrappedFieldData:
|
|
"""Data class for struct field definition and binding."""
|
|
|
|
# Line for struct field binding
|
|
binding: str = ""
|
|
|
|
# Line for struct field definition
|
|
definition: str = ""
|
|
|
|
# Used for constructor fields that are class types and need to be initialized
|
|
# with a pointer to the corresponding member in the native struct
|
|
# (e.g., "alt(&ptr_->alt)").
|
|
ptr_initialization: str = ""
|
|
|
|
# Statement to reset the inner pointer when copying the field
|
|
ptr_copy_reset: str = ""
|
|
|
|
# Whether the field is a primitive or fixed size
|
|
is_primitive_or_fixed_size: bool = False
|
|
|
|
# Underlying type of the field. If non-empty, used to determine the order in
|
|
# which structs are written in the bindings.h file.
|
|
typename: str = ""
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class WrappedStructData:
|
|
"""Data class for struct wrapper definition and binding."""
|
|
|
|
# Name of wrapper struct
|
|
wrap_name: str
|
|
|
|
# List of WrappedFieldData for this struct
|
|
wrapped_fields: List[WrappedFieldData]
|
|
|
|
# Struct header code
|
|
wrapped_header: str
|
|
|
|
# Struct source code
|
|
wrapped_source: str
|
|
|
|
# Struct bindings code
|
|
bindings: str = ""
|
|
|
|
|
|
def _simple_property_binding(
|
|
field: ast_nodes.StructFieldDecl,
|
|
struct_wrapper_name: str,
|
|
setter: bool = False,
|
|
reference: bool = False,
|
|
) -> str:
|
|
"""Builds the C++ code for a simple property binding."""
|
|
f = field
|
|
w = struct_wrapper_name
|
|
setter_txt = f", &{w}::set_{f.name}" if setter else ""
|
|
reference_txt = ", reference()" if reference else ""
|
|
return f'.property("{f.name}", &{w}::{f.name}{setter_txt}{reference_txt})'
|
|
|
|
|
|
def _generate_field_data(
|
|
field: ast_nodes.StructFieldDecl, struct_wrapper_name: str
|
|
) -> WrappedFieldData:
|
|
"""Generates the C++ definition and binding code for the struct field."""
|
|
f = field
|
|
w = struct_wrapper_name
|
|
s = common.lowercase_first_letter(w)
|
|
|
|
if f.name in constants.MANUAL_FIELDS.get(w, []):
|
|
# Note: Manually handled MjModel fields are special cased so that a
|
|
# by-reference embind return value policy is used.
|
|
return WrappedFieldData(
|
|
typename=_get_field_struct_type(f.type),
|
|
definition=f"// {f.name} field is handled manually in template file struct declaration", # pylint: disable=line-too-long
|
|
binding=_simple_property_binding(f, w, reference=(w == "MjModel")),
|
|
)
|
|
|
|
if f.name in constants.SKIPPED_FIELDS.get(w, []):
|
|
return WrappedFieldData(
|
|
typename="",
|
|
definition=f"// {f.name} field is skipped.",
|
|
binding=f"// {f.name} field is skipped.",
|
|
)
|
|
|
|
if isinstance(f.type, ast_nodes.ValueType) and (
|
|
f.type.name in constants.PRIMITIVE_TYPES or f.type.name.startswith("mjt")
|
|
):
|
|
|
|
builder = code_builder.CodeBuilder()
|
|
with builder.function(f"{f.type.name} {f.name}() const"):
|
|
builder.line(f"return ptr_->{f.name};")
|
|
with builder.function(f"void set_{f.name}({f.type.name} value)"):
|
|
builder.line(f"ptr_->{f.name} = value;")
|
|
|
|
return WrappedFieldData(
|
|
definition=builder.to_string(),
|
|
typename=_get_field_struct_type(f.type),
|
|
binding=_simple_property_binding(f, w, setter=True, reference=True),
|
|
is_primitive_or_fixed_size=True,
|
|
)
|
|
|
|
elif isinstance(f.type, ast_nodes.ValueType) and f.type.name.startswith("mj"):
|
|
return WrappedFieldData(
|
|
definition=f"{common.uppercase_first_letter(f.type.name)} {f.name};",
|
|
typename=_get_field_struct_type(f.type),
|
|
binding=_simple_property_binding(f, w, setter=False, reference=True),
|
|
ptr_initialization=f"{f.name}(&ptr_->{f.name})",
|
|
ptr_copy_reset=f"{f.name}.set(&ptr_->{f.name});",
|
|
is_primitive_or_fixed_size=True,
|
|
)
|
|
|
|
elif isinstance(f.type, ast_nodes.AnonymousStructDecl):
|
|
|
|
anonymous_struct_name = ""
|
|
for name, value in constants.ANONYMOUS_STRUCTS.items():
|
|
if value["parent"] == s and value["field_name"] == f.name:
|
|
anonymous_struct_name = name
|
|
break
|
|
|
|
if not anonymous_struct_name:
|
|
raise RuntimeError(
|
|
f"Anonymous struct for field {f.name} in {w} not found in"
|
|
" ANONYMOUS_STRUCTS."
|
|
)
|
|
|
|
if anonymous_struct_name in constants.STRUCTS_TO_BIND:
|
|
return WrappedFieldData(
|
|
binding=_simple_property_binding(f, w, setter=False, reference=True),
|
|
typename=anonymous_struct_name,
|
|
definition=(
|
|
f"{common.uppercase_first_letter(anonymous_struct_name)} {f.name};"
|
|
),
|
|
ptr_initialization=f"{f.name}(&ptr_->{f.name})",
|
|
ptr_copy_reset=f"{f.name}.set(&ptr_->{f.name});",
|
|
is_primitive_or_fixed_size=True,
|
|
)
|
|
|
|
elif isinstance(f.type, ast_nodes.ArrayType):
|
|
|
|
inner_type = f.type.inner_type
|
|
size = math.prod(f.type.extents)
|
|
inner_type_name = (
|
|
f.type.inner_type.name
|
|
if isinstance(f.type.inner_type, ast_nodes.ValueType)
|
|
else ""
|
|
)
|
|
if (
|
|
isinstance(inner_type, ast_nodes.ValueType)
|
|
and inner_type.name in constants.PRIMITIVE_TYPES
|
|
):
|
|
ptr_expr = f"ptr_->{f.name}"
|
|
if len(f.type.extents) > 1:
|
|
# for multi-dimensional arrays, we need to cast the field
|
|
# to a pointer, so embind can correctly interpret the memory
|
|
# view
|
|
ptr_expr = f"reinterpret_cast<{inner_type.name}*>({ptr_expr})"
|
|
|
|
builder = code_builder.CodeBuilder()
|
|
with builder.function(f"emscripten::val {f.name}() const"):
|
|
builder.line(
|
|
"return"
|
|
f" emscripten::val(emscripten::typed_memory_view({str(size)},"
|
|
f" {ptr_expr}));"
|
|
)
|
|
|
|
return WrappedFieldData(
|
|
definition=builder.to_string(),
|
|
typename=_get_field_struct_type(f.type),
|
|
binding=_simple_property_binding(f, w),
|
|
is_primitive_or_fixed_size=True,
|
|
)
|
|
elif inner_type_name.startswith("mj"):
|
|
return WrappedFieldData(
|
|
definition=(
|
|
f"std::vector<{common.uppercase_first_letter(inner_type_name)}>"
|
|
f" {f.name};"
|
|
),
|
|
ptr_initialization=f"{f.name}(&ptr_->{f.name})",
|
|
typename=_get_field_struct_type(f.type),
|
|
binding=_simple_property_binding(f, w, reference=True),
|
|
)
|
|
|
|
elif isinstance(f.type, ast_nodes.PointerType):
|
|
inner_type_name = (
|
|
f.type.inner_type.name
|
|
if isinstance(f.type.inner_type, ast_nodes.ValueType)
|
|
else ""
|
|
)
|
|
ptr_field_expr = f"ptr_->{f.name}"
|
|
array_size_str = ""
|
|
|
|
if f.array_extent:
|
|
array_size_str = parse_array_extent(f.array_extent, w, f.name)
|
|
elif f.name in constants.BYTE_FIELDS.keys():
|
|
# for byte fields, we need to cast the pointer to uint8_t*
|
|
# so embind can correctly interpret the memory view
|
|
ptr_field_expr = f"static_cast<uint8_t*>({ptr_field_expr})"
|
|
# for these byte fields, there is no array_extent, so we add the size of
|
|
# in the config file based in the documentation
|
|
extent = (constants.BYTE_FIELDS[f.name]["size"],)
|
|
array_size_str = parse_array_extent(extent, w, f.name)
|
|
elif inner_type_name == "mjString":
|
|
|
|
builder = code_builder.CodeBuilder()
|
|
with builder.function(f"mjString {f.name}() const"):
|
|
builder.line(
|
|
f'return (ptr_ && ptr_->{f.name}) ? *(ptr_->{f.name}) : "";'
|
|
)
|
|
with builder.function(f"void set_{f.name}(const mjString& value)"):
|
|
with builder.block(f"if (ptr_ && ptr_->{f.name})"):
|
|
builder.line(f"*(ptr_->{f.name}) = value;")
|
|
|
|
return WrappedFieldData(
|
|
definition=builder.to_string(),
|
|
typename=_get_field_struct_type(f.type),
|
|
binding=_simple_property_binding(f, w, setter=True, reference=True),
|
|
)
|
|
elif inner_type_name.startswith("mj") and inner_type_name.endswith("Vec"):
|
|
ptr_field_expr_vec = f"*(ptr_->{f.name})"
|
|
vector_type = inner_type_name
|
|
if vector_type == "mjByteVec":
|
|
vector_type = "std::vector<uint8_t>"
|
|
ptr_field_expr_vec = (
|
|
f"*(reinterpret_cast<std::vector<uint8_t>*>(ptr_->{f.name}))"
|
|
)
|
|
|
|
builder = code_builder.CodeBuilder()
|
|
with builder.function(f"{vector_type} &{f.name}() const"):
|
|
builder.line(f"return {ptr_field_expr_vec};")
|
|
|
|
return WrappedFieldData(
|
|
definition=builder.to_string(),
|
|
typename=_get_field_struct_type(f.type),
|
|
binding=_simple_property_binding(f, w, setter=False, reference=True),
|
|
)
|
|
|
|
if (
|
|
inner_type_name.startswith("mj")
|
|
and inner_type_name not in constants.PRIMITIVE_TYPES
|
|
and not f.array_extent
|
|
and w not in constants.MANUAL_FIELDS.keys()
|
|
):
|
|
ptr_field = cast(ast_nodes.PointerType, f.type)
|
|
wrapper_field_name = common.uppercase_first_letter(
|
|
cast(ast_nodes.ValueType, ptr_field.inner_type).name
|
|
)
|
|
return WrappedFieldData(
|
|
definition=f"{wrapper_field_name} {f.name};",
|
|
typename=_get_field_struct_type(f.type),
|
|
binding=_simple_property_binding(f, w, setter=False, reference=True),
|
|
ptr_initialization=f"{f.name}(ptr_->{f.name})",
|
|
)
|
|
|
|
builder = code_builder.CodeBuilder()
|
|
with builder.function(f"emscripten::val {f.name}() const"):
|
|
builder.line(
|
|
"return"
|
|
f" emscripten::val(emscripten::typed_memory_view({array_size_str},"
|
|
f" {ptr_field_expr}));"
|
|
)
|
|
|
|
return WrappedFieldData(
|
|
definition=builder.to_string(),
|
|
typename=_get_field_struct_type(f.type),
|
|
binding=_simple_property_binding(f, w),
|
|
)
|
|
|
|
raise RuntimeError(f"Field {f.name} from struct {w} not properly handled")
|
|
|
|
|
|
def _default_function_statement(struct_name: str) -> str:
|
|
"""Returns the default function name for the given struct."""
|
|
if struct_name == "mjvGeom":
|
|
f = "mjv_initGeom(ptr_, mjGEOM_NONE, nullptr, nullptr, nullptr, nullptr);"
|
|
return f
|
|
elif struct_name in constants.ANONYMOUS_STRUCTS.keys():
|
|
return ""
|
|
elif struct_name in constants.NO_DEFAULT_CONSTRUCTORS:
|
|
return ""
|
|
elif struct_name.startswith("mjs"):
|
|
return f"mjs_default{struct_name.removeprefix('mjs')}(ptr_);"
|
|
elif struct_name.startswith("mjv"):
|
|
return f"mjv_default{struct_name.removeprefix('mjv')}(ptr_);"
|
|
elif struct_name.startswith("mj"):
|
|
return f"mj_default{struct_name.removeprefix('mj')}(ptr_);"
|
|
|
|
return ""
|
|
|
|
|
|
def _delete_ptr_statement(struct_name: str) -> str:
|
|
"""Returns the delete function name for the given struct."""
|
|
if struct_name == "mjVFS":
|
|
return "mj_deleteVFS(ptr_);"
|
|
|
|
return "delete ptr_;"
|
|
|
|
|
|
def _find_member_inits(
|
|
wrapped_fields: List[WrappedFieldData],
|
|
) -> List[WrappedFieldData]:
|
|
"""Finds the fields with ptr_initialization in the wrapped fields list."""
|
|
member_inits = []
|
|
for field in wrapped_fields:
|
|
if field.ptr_initialization:
|
|
member_inits.append(field)
|
|
return member_inits
|
|
|
|
|
|
def use_shallow_copy(
|
|
wrapped_fields: List[WrappedFieldData],
|
|
) -> bool:
|
|
"""Returns true if the struct fields can be shallow copied."""
|
|
for field in wrapped_fields:
|
|
if not field.is_primitive_or_fixed_size:
|
|
return False
|
|
return True
|
|
|
|
|
|
def build_struct_header(
|
|
struct_name: str,
|
|
wrapped_fields: List[WrappedFieldData],
|
|
):
|
|
"""Builds the C++ header file code for a struct."""
|
|
s = struct_name
|
|
w = common.uppercase_first_letter(s)
|
|
|
|
if w in constants.MANUAL_STRUCTS:
|
|
return ""
|
|
|
|
if (
|
|
s not in constants.ANONYMOUS_STRUCTS and
|
|
s not in introspect_structs.STRUCTS
|
|
):
|
|
raise RuntimeError(f"Struct {s} not found in introspect structs")
|
|
|
|
not_mjs = not w.startswith("Mjs")
|
|
member_inits = _find_member_inits(wrapped_fields)
|
|
shallow_copy = not_mjs and use_shallow_copy(wrapped_fields)
|
|
builder = code_builder.CodeBuilder()
|
|
|
|
with builder.struct(f"{w}"):
|
|
# destructor
|
|
if not_mjs:
|
|
builder.line(f"~{w}();")
|
|
|
|
# default constructor
|
|
if not_mjs and w not in ["MjData", "MjModel"]:
|
|
builder.line(f"{w}();")
|
|
|
|
# constructor passing native ptr
|
|
if w != "MjData":
|
|
builder.line(f"explicit {w}({s} *ptr);")
|
|
else:
|
|
# special handling for MjData
|
|
builder.line("MjData(MjModel *m);")
|
|
builder.line("explicit MjData(const MjModel &, const MjData &);")
|
|
builder.line("std::vector<MjContact> contact() const;")
|
|
|
|
# copy constructor
|
|
if shallow_copy or w in ["MjSpec", "MjModel"]:
|
|
builder.line(f"{w}(const {w} &);")
|
|
|
|
# assignment operator
|
|
if shallow_copy or w == "MjSpec":
|
|
builder.line(f"{w} &operator=(const {w} &);")
|
|
|
|
# explicit copy function
|
|
if shallow_copy or w in ["MjSpec", "MjData", "MjModel"]:
|
|
builder.line(f"std::unique_ptr<{w}> copy();")
|
|
|
|
# C struct getter/setter
|
|
builder.line(f"{s}* get() const;")
|
|
builder.line(f"void set({s}* ptr);")
|
|
|
|
# field declarations
|
|
for field in wrapped_fields:
|
|
if field.definition and field not in member_inits:
|
|
for line in field.definition.splitlines():
|
|
builder.line(line)
|
|
|
|
# define private struct members
|
|
builder.private()
|
|
builder.line(f"{s}* ptr_;")
|
|
if not_mjs and w not in ["MjData", "MjModel"]:
|
|
builder.line("bool owned_ = false;")
|
|
|
|
# define public struct members
|
|
if member_inits:
|
|
builder.public()
|
|
for f in member_inits:
|
|
if f.definition:
|
|
builder.line(f"{f.definition}")
|
|
|
|
if w == "MjData":
|
|
builder.line("mjModel *model;")
|
|
|
|
return builder.to_string() + ";"
|
|
|
|
|
|
def build_struct_source(
|
|
struct_name: str,
|
|
wrapped_fields: List[WrappedFieldData],
|
|
):
|
|
"""Builds the C++ .cc file code for a struct."""
|
|
# These structs require specific function calls for creation and/or deletion
|
|
# which, for now, are hardcoded in the template file.
|
|
if struct_name in [
|
|
"mjData",
|
|
"mjModel",
|
|
"mjvScene",
|
|
"mjSpec",
|
|
]:
|
|
return ""
|
|
|
|
s = struct_name
|
|
w = common.uppercase_first_letter(s)
|
|
is_mjs = w.startswith("Mjs")
|
|
|
|
member_inits = _find_member_inits(wrapped_fields)
|
|
shallow_copy = use_shallow_copy(wrapped_fields)
|
|
|
|
fields_init = ""
|
|
if member_inits:
|
|
fields_init = ", " + ", ".join(
|
|
mjs_field.ptr_initialization for mjs_field in member_inits
|
|
)
|
|
|
|
builder = code_builder.CodeBuilder()
|
|
|
|
# constructor passing native ptr
|
|
with builder.function(f"{w}::{w}({s} *ptr) : ptr_(ptr){fields_init}"):
|
|
pass
|
|
|
|
# destructor
|
|
if not is_mjs:
|
|
with builder.function(f"{w}::~{w}()"):
|
|
with builder.block("if (owned_ && ptr_)"):
|
|
delete_ptr = _delete_ptr_statement(s)
|
|
builder.line(delete_ptr)
|
|
|
|
if not is_mjs:
|
|
# default constructor
|
|
with builder.function(f"{w}::{w}() : ptr_(new {s}){fields_init}"):
|
|
builder.line("owned_ = true;")
|
|
default_func = _default_function_statement(s)
|
|
if default_func:
|
|
builder.line(default_func)
|
|
|
|
if shallow_copy and not is_mjs:
|
|
# copy constructor
|
|
with builder.function(f"{w}::{w}(const {w} &other) : {w}()"):
|
|
builder.line("*ptr_ = *other.get();")
|
|
for f in member_inits:
|
|
if f.ptr_copy_reset is not None:
|
|
builder.line(f.ptr_copy_reset)
|
|
|
|
# assignment operator
|
|
with builder.function(f"{w}& {w}::operator=(const {w} &other)"):
|
|
with builder.block("if (this == &other)"):
|
|
builder.line("return *this;")
|
|
builder.line("*ptr_ = *other.get();")
|
|
for f in member_inits:
|
|
if f.ptr_copy_reset is not None:
|
|
builder.line(f.ptr_copy_reset)
|
|
builder.line("return *this;")
|
|
|
|
# explicit copy function
|
|
with builder.function(f"std::unique_ptr<{w}> {w}::copy()"):
|
|
builder.line(f"return std::make_unique<{w}>(*this);")
|
|
|
|
# C struct getter/setter
|
|
with builder.function(f"{s}* {w}::get() const"):
|
|
builder.line("return ptr_;")
|
|
with builder.function(f"void {w}::set({s}* ptr)"):
|
|
builder.line("ptr_ = ptr;")
|
|
|
|
return builder.to_string()
|
|
|
|
|
|
def _build_struct_bindings(
|
|
struct_name: str,
|
|
wrapped_fields: List[WrappedFieldData],
|
|
):
|
|
"""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
|
|
):
|
|
if w == "MjData":
|
|
builder.line(".constructor<MjModel *>()")
|
|
builder.line(".constructor<const MjModel &, const MjData &>()")
|
|
elif w == "MjModel":
|
|
builder.line(
|
|
'.class_function("loadFromXML", &loadFromXML, take_ownership())'
|
|
)
|
|
builder.line(".constructor<const MjModel &>()")
|
|
elif w == "MjSpec":
|
|
builder.line(".constructor<const MjSpec &>()")
|
|
elif w == "MjvScene":
|
|
builder.line(".constructor<MjModel *, int>()")
|
|
builder.line(".constructor<>()")
|
|
elif not is_mjs:
|
|
builder.line(".constructor<>()")
|
|
|
|
shallow_copy = use_shallow_copy(wrapped_fields)
|
|
if shallow_copy and not is_mjs:
|
|
builder.line(f'.function("copy", &{w}::copy, take_ownership())')
|
|
|
|
for field in wrapped_fields[:-1]:
|
|
if field.binding:
|
|
builder.line(field.binding)
|
|
if wrapped_fields:
|
|
builder.line(f"{wrapped_fields[-1].binding};")
|
|
|
|
return builder.to_string()
|
|
|
|
|
|
def parse_array_extent(
|
|
extents: Tuple[Union[str, int], ...], wrapper_name: str, field_name: str
|
|
) -> str:
|
|
"""Parses the array extent of a field, returning a string representing the resolved extents."""
|
|
if not extents:
|
|
return ""
|
|
return " * ".join(
|
|
resolve_extent(extent, wrapper_name, field_name) for extent in extents
|
|
)
|
|
|
|
|
|
def resolve_extent(
|
|
extent: Union[str, int], wrapper_name: str, field_name: str
|
|
) -> str:
|
|
"""Resolves the extent of an array, handling integers and references to other struct fields.
|
|
|
|
Args:
|
|
extent: The extent to resolve, can be an int or a string referencing a
|
|
field.
|
|
wrapper_name: The name of the struct wrapper.
|
|
field_name: The name of the field being processed.
|
|
|
|
Returns:
|
|
A string representing the resolved extent, either as a number or a field
|
|
reference.
|
|
"""
|
|
if isinstance(extent, int):
|
|
return str(extent)
|
|
# if starts with mj, it's a mujoco constant,
|
|
# so we don't need to get a parent struct ptr
|
|
if extent.startswith("mj"):
|
|
return str(extent)
|
|
if wrapper_name == "MjData" and field_name not in constants.MJDATA_SIZES:
|
|
var_name = "model"
|
|
else:
|
|
var_name = "ptr_"
|
|
return f"{var_name}->{extent}"
|
|
|
|
|
|
def generate_wasm_bindings(
|
|
structs_to_bind: List[str],
|
|
) -> Dict[str, WrappedStructData]:
|
|
"""Generates WASM bindings for MuJoCo structs."""
|
|
|
|
wrapped_structs: Dict[str, WrappedStructData] = {}
|
|
for struct_name in structs_to_bind:
|
|
s = struct_name
|
|
w = common.uppercase_first_letter(s)
|
|
|
|
if s in introspect_structs.STRUCTS:
|
|
struct_fields = introspect_structs.STRUCTS[s].fields
|
|
elif s in constants.ANONYMOUS_STRUCTS:
|
|
anonymous_struct = _get_anonymous_struct_field(s)
|
|
if not anonymous_struct or not isinstance(
|
|
anonymous_struct.type, ast_nodes.AnonymousStructDecl
|
|
):
|
|
raise RuntimeError(f"Anonymous struct not found: {s}")
|
|
struct_fields = anonymous_struct.type.fields
|
|
else:
|
|
raise RuntimeError(f"Struct not found: {s}")
|
|
|
|
wrapped_fields: List[WrappedFieldData] = []
|
|
for field in struct_fields:
|
|
wrapped_fields.append(_generate_field_data(field, w))
|
|
|
|
wrap_data = WrappedStructData(
|
|
wrap_name=w,
|
|
wrapped_fields=wrapped_fields,
|
|
wrapped_header=build_struct_header(s, wrapped_fields),
|
|
wrapped_source=build_struct_source(s, wrapped_fields),
|
|
bindings=_build_struct_bindings(s, wrapped_fields),
|
|
)
|
|
|
|
wrapped_structs[s] = wrap_data
|
|
|
|
return wrapped_structs
|
|
|
|
|
|
def _get_anonymous_struct_field(
|
|
anonymous_structs_key: str,
|
|
) -> ast_nodes.StructFieldDecl | None:
|
|
"""Looks up the given key in the anonymous_structs dict and generates bindings for its fields."""
|
|
info = constants.ANONYMOUS_STRUCTS[anonymous_structs_key]
|
|
parent_decl = introspect_structs.STRUCTS[info["parent"]]
|
|
target_field = next(
|
|
(
|
|
f
|
|
for f in parent_decl.fields
|
|
if hasattr(f, "name")
|
|
and f.name == info["field_name"]
|
|
and hasattr(f, "type")
|
|
and isinstance(f.type, ast_nodes.AnonymousStructDecl)
|
|
),
|
|
None,
|
|
)
|
|
return target_field
|
|
|
|
|
|
def _get_field_struct_type(field_type):
|
|
"""Extracts the base struct name if the field type is a struct or pointer to a struct."""
|
|
if isinstance(field_type, ast_nodes.ValueType):
|
|
return field_type.name
|
|
if isinstance(field_type, ast_nodes.PointerType):
|
|
if isinstance(field_type.inner_type, ast_nodes.ValueType):
|
|
return field_type.inner_type.name
|
|
return None
|
|
|
|
|
|
def sort_structs_by_dependency(
|
|
struct_wrappers: dict[str, WrappedStructData],
|
|
) -> List[str]:
|
|
"""Sorts structs based on their field dependencies using topological sort.
|
|
|
|
Structs with no dependencies on other structs in the list come first.
|
|
Struct A has a dependency on struct B if struct A has a field where the
|
|
underlying_type is B. Note that this definition is stricter than the C++
|
|
struct dependency criterion where forward declarations can be used to
|
|
eliminate dependencies A and B if A only has a pointer to B.
|
|
|
|
Args:
|
|
struct_wrappers: A dictionary mapping struct names to their
|
|
WrappedStructData.
|
|
|
|
Returns:
|
|
A new list of struct names sorted by dependency.
|
|
|
|
Raises:
|
|
RuntimeError: If a cyclic dependency is detected.
|
|
"""
|
|
adj = collections.defaultdict(list)
|
|
in_degree = collections.defaultdict(int)
|
|
struct_names = struct_wrappers.keys()
|
|
struct_set = set(struct_names)
|
|
sorted_struct_names = sorted(struct_names)
|
|
|
|
for struct_name in sorted_struct_names:
|
|
if struct_name == "mjData":
|
|
adj["mjModel"].append("mjData")
|
|
in_degree["mjData"] += 1
|
|
for field in struct_wrappers[struct_name].wrapped_fields:
|
|
|
|
field_type_name = field.typename
|
|
if (
|
|
field_type_name
|
|
and field_type_name != struct_name
|
|
and field_type_name in struct_set
|
|
):
|
|
if struct_name not in adj[field_type_name]:
|
|
adj[field_type_name].append(struct_name)
|
|
in_degree[struct_name] += 1
|
|
|
|
queue = collections.deque(
|
|
[name for name in sorted_struct_names if in_degree[name] == 0]
|
|
)
|
|
sorted_list = []
|
|
|
|
while queue:
|
|
u = queue.popleft()
|
|
sorted_list.append(u)
|
|
for v in adj[u]:
|
|
in_degree[v] -= 1
|
|
if in_degree[v] == 0:
|
|
queue.append(v)
|
|
|
|
if len(sorted_list) == len(struct_names):
|
|
return sorted_list
|
|
else:
|
|
remaining = set(struct_names) - set(sorted_list)
|
|
raise RuntimeError(
|
|
"Cycle detected in struct dependencies, involving: "
|
|
f"{', '.join(sorted(list(remaining)))}"
|
|
)
|
|
|
|
|
|
class Generator:
|
|
"""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(
|
|
constants.STRUCTS_TO_BIND
|
|
)
|
|
|
|
autogenned_struct_definitions = []
|
|
markers_and_content = []
|
|
|
|
typedefs = []
|
|
for type_name in sorted(constants.ANONYMOUS_STRUCTS):
|
|
s = constants.ANONYMOUS_STRUCTS[type_name]
|
|
typedefs.append(
|
|
f"using {type_name} = decltype(::{s['parent']}::{s['field_name']});"
|
|
)
|
|
markers_and_content.append((
|
|
"// {{ ANONYMOUS_STRUCT_TYPEDEFS }}",
|
|
typedefs,
|
|
))
|
|
|
|
# Sort by struct name by dependency to ensure deterministic output order
|
|
sorted_struct_names = sort_structs_by_dependency(self.structs_to_bind_data)
|
|
|
|
for struct_name in sorted_struct_names:
|
|
struct_data = self.structs_to_bind_data[struct_name]
|
|
if struct_data.wrapped_header:
|
|
autogenned_struct_definitions.append(struct_data.wrapped_header + "\n")
|
|
else:
|
|
markers_and_content.append((
|
|
f"// INSERT-GENERATED-{struct_data.wrap_name}-DEFINITIONS",
|
|
[
|
|
l.definition if l.definition else ""
|
|
for l in struct_data.wrapped_fields
|
|
],
|
|
))
|
|
markers_and_content.append((
|
|
"// {{ AUTOGENNED_STRUCTS_HEADER }}",
|
|
autogenned_struct_definitions,
|
|
))
|
|
|
|
autogenned_struct_source = []
|
|
autogenned_struct_bindings = []
|
|
for struct_name in sorted_struct_names:
|
|
struct_data = self.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)
|
|
|
|
markers_and_content.append((
|
|
"// {{ AUTOGENNED_STRUCTS_SOURCE }}",
|
|
autogenned_struct_source,
|
|
))
|
|
markers_and_content.append((
|
|
"// {{ AUTOGENNED_STRUCTS_BINDINGS }}",
|
|
autogenned_struct_bindings,
|
|
))
|
|
return markers_and_content
|