chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Simplified struct code generation for WASM bindings

Browse Source 
PiperOrigin-RevId: 827878780
Change-Id: I97ab732eb11c63a6aaf7fa95add0b51bc0a56525
This commit is contained in:
Matija Kecman
2025-11-04 02:50:09 -08:00
committed by Copybara-Service
parent c22d94e470
commit 3d59b20b65
6 changed files with 122 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wasm/codegen/generators/structs.py
+1 -8
View File
@@ -18,23 +18,16 @@ from typing import Optional
from wasm.codegen.helpers import constants
from wasm.codegen.helpers import structs_parser
from wasm.codegen.helpers import structs_wrappers_data
class Generator:
"""Generates C++ code for binding and wrapping MuJoCo structs."""
def __init__(self):
# Set up the correct input dict based on the structs we want to bind
# and already have a wrapper manually created in the template/bindings.cc
wrapped_structs = structs_wrappers_data.create_wrapped_structs_set_up_data(
constants.STRUCTS_TO_BIND
)
# Traverse the introspect dictionary to get the field
# wrapper/bindings statements set up for each struct
self.structs_to_bind_data = structs_parser.generate_wasm_bindings(
wrapped_structs
constants.STRUCTS_TO_BIND
)
def generate_header(
wasm/codegen/helpers/struct_constructor_code_builder.py
+60 -25
View File
@@ -25,9 +25,7 @@ from wasm.codegen.helpers import constants
from wasm.codegen.helpers import structs_wrappers_data
def _has_nested_wrapper_members(
struct_info: ast_nodes.StructDecl
) -> bool:
def _has_nested_wrapper_members(struct_info: ast_nodes.StructDecl) -> bool:
"""Checks if the struct contains other wrapped structs as direct members."""
for field in struct_info.fields:
struct_field = cast(ast_nodes.StructFieldDecl, field)
@@ -49,10 +47,12 @@ def _build_struct_header_internal(
struct_name: str,
wrapped_fields: List[structs_wrappers_data.WrappedFieldData],
fields_with_init: List[structs_wrappers_data.WrappedFieldData],
use_shallow_copy: bool = False,
is_mjs: bool = False,
):
"""Builds the C++ header file code for a struct."""
shallow_copy = use_shallow_copy(wrapped_fields)
wrapper_name = common.uppercase_first_letter(struct_name)
builder = code_builder.CodeBuilder()
with builder.block(f"struct {wrapper_name}"):
@@ -64,7 +64,7 @@ def _build_struct_header_internal(
builder.line(f"explicit {wrapper_name}({struct_name} *ptr);")
builder.line(f"~{wrapper_name}();")
if use_shallow_copy:
if shallow_copy:
builder.line(f"std::unique_ptr<{wrapper_name}> copy();")
for field in wrapped_fields:
@@ -87,24 +87,62 @@ def _build_struct_header_internal(
for field in fields_with_init:
if field.definition:
builder.line(f"{field.definition}")
return builder.to_string()+";"
return builder.to_string() + ";"
def _get_default_func_name(struct_name: str) -> str:
"""Returns the default function name for the given struct."""
if (
struct_name in constants.ANONYMOUS_STRUCTS.keys()
or struct_name in constants.NO_DEFAULT_CONSTRUCTORS
or (
common.uppercase_first_letter(struct_name)
in constants.MANUALLY_ADDED_FIELDS_FROM_TEMPLATE.keys()
)
):
return ""
elif struct_name.startswith("mjs"):
return f"mjs_default{struct_name.removeprefix('mjs')}"
elif struct_name.startswith("mjv"):
return f"mjv_default{struct_name.removeprefix('mjv')}"
else:
return f"mj_default{struct_name.removeprefix('mj')}"
def _find_fields_with_init(
wrapped_fields: List[structs_wrappers_data.WrappedFieldData],
) -> List[structs_wrappers_data.WrappedFieldData]:
"""Finds the fields with initialization in the wrapped fields list."""
fields_with_init = []
for field in wrapped_fields:
if field.initialization:
fields_with_init.append(field)
return fields_with_init
def use_shallow_copy(
wrapped_fields: List[structs_wrappers_data.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,
use_shallow_copy: bool = False,
fields_with_init: List[structs_wrappers_data.WrappedFieldData] = [],
wrapped_fields: List[structs_wrappers_data.WrappedFieldData] = [],
wrapped_fields: List[structs_wrappers_data.WrappedFieldData],
):
"""Builds the C++ header file code for a struct."""
struct_info = introspect_structs.STRUCTS.get(struct_name)
if struct_name.startswith("mjs"):
fields_with_init = _find_fields_with_init(wrapped_fields)
return _build_struct_header_internal(
struct_name,
wrapped_fields,
fields_with_init,
use_shallow_copy,
is_mjs=True,
)
@@ -117,27 +155,28 @@ def build_struct_header(
and not _has_nested_wrapper_members(struct_info)
):
return _build_struct_header_internal(
struct_name, wrapped_fields, [], use_shallow_copy, is_mjs=False
struct_name, wrapped_fields, [], is_mjs=False
)
return ""
def build_struct_source(
struct_name: str,
mj_default_func: str | None = None,
fields_with_init: List[structs_wrappers_data.WrappedFieldData] = [],
use_shallow_copy: bool = False,
wrapped_fields: List[structs_wrappers_data.WrappedFieldData],
):
"""Builds the C++ .cc file code for a struct."""
wrapper_name = common.uppercase_first_letter(struct_name)
is_mjs_struct = "Mjs" in wrapper_name
builder = code_builder.CodeBuilder()
fields_with_init = _find_fields_with_init(wrapped_fields)
shallow_copy = use_shallow_copy(wrapped_fields)
mj_default_func = _get_default_func_name(struct_name)
fields_init = ""
if fields_with_init:
fields_init = "".join(
field_with_init.initialization
for field_with_init in fields_with_init
field_with_init.initialization for field_with_init in fields_with_init
)
# constructor passing native ptr
builder.line(
@@ -154,10 +193,9 @@ def build_struct_source(
if mj_default_func:
builder.line(f"{mj_default_func}(ptr_);")
# copy constructor
if use_shallow_copy and not is_mjs_struct:
if shallow_copy and not is_mjs_struct:
with builder.block(
f"{wrapper_name}::{wrapper_name}(const"
f" {wrapper_name} &other)"
f"{wrapper_name}::{wrapper_name}(const {wrapper_name} &other)"
+ (f" : {wrapper_name}()" if not is_mjs_struct else "")
):
builder.line("*ptr_ = *other.get();")
@@ -186,12 +224,9 @@ def build_struct_source(
with builder.block(f"{wrapper_name}::~{wrapper_name}()"):
builder.line("if (owned_ && ptr_) delete ptr_;")
# copy function
if use_shallow_copy:
if shallow_copy:
with builder.block(
f"std::unique_ptr<{wrapper_name}>"
f" {wrapper_name}::copy()"
f"std::unique_ptr<{wrapper_name}> {wrapper_name}::copy()"
):
builder.line(
f"return std::make_unique<{wrapper_name}>(*this);"
)
builder.line(f"return std::make_unique<{wrapper_name}>(*this);")
return builder.to_string()
wasm/codegen/helpers/struct_constructor_code_builder_test.py
+25 -16
View File
@@ -16,37 +16,49 @@ from absl.testing import absltest
from introspect import ast_nodes
from wasm.codegen.helpers import struct_constructor_code_builder
from wasm.codegen.helpers import struct_field_handler
from wasm.codegen.helpers import structs_parser
class StructConstructorCodeBuilderTest(absltest.TestCase):
def test_constructor_code_with_default_function(self):
wrapped_structs = structs_parser.generate_wasm_bindings(["mjLROpt"])
self.assertEqual(
struct_constructor_code_builder.build_struct_source(
"mjLROpt", "mj_defaultLROpt"
),
wrapped_structs["mjLROpt"].wrapped_source,
"""
MjLROpt::MjLROpt(mjLROpt *ptr) : ptr_(ptr) {}
MjLROpt::MjLROpt() : ptr_(new mjLROpt) {
owned_ = true;
mj_defaultLROpt(ptr_);
}
MjLROpt::MjLROpt(const MjLROpt &other) : MjLROpt() {
*ptr_ = *other.get();
}
MjLROpt& MjLROpt::operator=(const MjLROpt &other) {
if (this == &other) {
return *this;
}
*ptr_ = *other.get();
return *this;
}
MjLROpt::~MjLROpt() {
if (owned_ && ptr_) delete ptr_;
}
std::unique_ptr<MjLROpt> MjLROpt::copy() {
return std::make_unique<MjLROpt>(*this);
}
""".strip(),
)
def test_constructor_code_without_default_function(self):
wrapped_structs = structs_parser.generate_wasm_bindings(["mjsElement"])
self.assertEqual(
struct_constructor_code_builder.build_struct_source("mjLROpt"),
wrapped_structs["mjsElement"].wrapped_source,
"""
MjLROpt::MjLROpt(mjLROpt *ptr) : ptr_(ptr) {}
MjLROpt::MjLROpt() : ptr_(new mjLROpt) {
owned_ = true;
}
MjLROpt::~MjLROpt() {
if (owned_ && ptr_) delete ptr_;
MjsElement::MjsElement(mjsElement *ptr) : ptr_(ptr) {}
MjsElement::~MjsElement() {}
std::unique_ptr<MjsElement> MjsElement::copy() {
return std::make_unique<MjsElement>(*this);
}
""".strip(),
)
@@ -65,7 +77,6 @@ MjLROpt::~MjLROpt() {
self.assertEqual(
struct_constructor_code_builder.build_struct_source(
"mjsTexture",
"mjs_defaultTexture",
[wrapped_field_data],
),
"""
@@ -76,9 +87,7 @@ MjsTexture::~MjsTexture() {}
def test_constructor_code_with_shallow_copy(self):
self.assertEqual(
struct_constructor_code_builder.build_struct_source(
"mjvLight", use_shallow_copy=True
),
struct_constructor_code_builder.build_struct_source("mjvLight", []),
"""MjvLight::MjvLight(mjvLight *ptr) : ptr_(ptr) {}
MjvLight::MjvLight() : ptr_(new mjvLight) {
owned_ = true;
@@ -104,14 +113,14 @@ std::unique_ptr<MjvLight> MjvLight::copy() {
def test_build_struct_header_with_nested_wrappers(self):
self.assertEqual(
struct_constructor_code_builder.build_struct_header("mjData"),
struct_constructor_code_builder.build_struct_header("mjData", []),
"",
)
def test_build_struct_header_basic_struct(self):
self.assertEqual(
struct_constructor_code_builder.build_struct_header("mjLROpt"),
struct_constructor_code_builder.build_struct_header("mjLROpt", []),
"""
struct MjLROpt {
MjLROpt();
wasm/codegen/helpers/structs_parser.py
+27 -45
View File
@@ -33,11 +33,14 @@ introspect_structs = structs.STRUCTS
def generate_wasm_bindings(
wrapped_structs: Dict[str, structs_wrappers_data.WrappedStructData],
structs_to_bind: List[str],
) -> Dict[str, structs_wrappers_data.WrappedStructData]:
"""Generates WASM bindings for MuJoCo structs."""
for struct_name, wrap_data in wrapped_structs.items():
wrapped_structs: Dict[str, structs_wrappers_data.WrappedStructData] = {}
for struct_name in structs_to_bind:
wrapped_name = common.uppercase_first_letter(struct_name)
if struct_name in introspect_structs:
struct_fields = introspect_structs[struct_name].fields
elif struct_name in constants.ANONYMOUS_STRUCTS:
@@ -52,35 +55,33 @@ def generate_wasm_bindings(
debug_print(f"Wrapping struct: {struct_name}")
fields_with_init: List[structs_wrappers_data.WrappedFieldData] = []
wrapped_fields: List[structs_wrappers_data.WrappedFieldData] = []
for field in struct_fields:
field_gen = struct_field_handler.StructFieldHandler(
field, wrap_data.wrap_name
wrapped_field = struct_field_handler.StructFieldHandler(
field, wrapped_name
).generate()
# If the struct has at least one non-primitive or fixed size field
# we avoid shallow copy to avoid uninitialized memory.
if not field_gen.is_primitive_or_fixed_size:
wrap_data.use_shallow_copy = False
if field_gen.initialization:
fields_with_init.append(field_gen)
wrap_data.wrapped_fields.append(field_gen)
wrapped_fields.append(wrapped_field)
wrap_data.wrapped_header = (
struct_constructor_code_builder.build_struct_header(
struct_name,
wrap_data.use_shallow_copy,
fields_with_init,
wrap_data.wrapped_fields,
)
wrapped_header = struct_constructor_code_builder.build_struct_header(
struct_name,
wrapped_fields,
)
wrap_data.wrapped_source = (
struct_constructor_code_builder.build_struct_source(
struct_name,
get_default_func_name(struct_name),
fields_with_init,
wrap_data.use_shallow_copy,
)
wrapped_source = struct_constructor_code_builder.build_struct_source(
struct_name,
wrapped_fields,
)
wrap_data = structs_wrappers_data.WrappedStructData(
wrap_name=wrapped_name,
wrapped_fields=wrapped_fields,
wrapped_header=wrapped_header,
wrapped_source=wrapped_source,
use_shallow_copy=struct_constructor_code_builder.use_shallow_copy(
wrapped_fields
),
)
wrapped_structs[struct_name] = wrap_data
return wrapped_structs
@@ -104,25 +105,6 @@ def _get_anonymous_struct_field(
return target_field
def get_default_func_name(struct_name: str) -> str:
"""Returns the default function name for the given struct."""
if (
struct_name in constants.ANONYMOUS_STRUCTS.keys()
or struct_name in constants.NO_DEFAULT_CONSTRUCTORS
or (
common.uppercase_first_letter(struct_name)
in constants.MANUALLY_ADDED_FIELDS_FROM_TEMPLATE.keys()
)
):
return ""
elif struct_name.startswith("mjs"):
return f"mjs_default{struct_name.removeprefix('mjs')}"
elif struct_name.startswith("mjv"):
return f"mjv_default{struct_name.removeprefix('mjv')}"
else:
return f"mj_default{struct_name.removeprefix('mj')}"
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):
wasm/codegen/helpers/structs_parser_test.py
+9 -29
View File
@@ -25,15 +25,13 @@ class StructsParserTest(absltest.TestCase):
def setUp(self):
super().setUp()
self.wrapped_structs = structs_parser.generate_wasm_bindings(
structs_wrappers_data.create_wrapped_structs_set_up_data([
"mjModel",
"mjData",
"mjVisualGlobal",
"mjVisualQuality",
"mjVisual",
])
)
self.wrapped_structs = structs_parser.generate_wasm_bindings([
"mjModel",
"mjData",
"mjVisualGlobal",
"mjVisualQuality",
"mjVisual",
])
def test_sort_structs_by_dependency(self):
mock_introspect_structs = {
@@ -83,27 +81,9 @@ class StructsParserTest(absltest.TestCase):
def test_generate_wasm_bindings_with_error(self):
with self.assertRaises(RuntimeError):
structs_parser.generate_wasm_bindings(
structs_wrappers_data.create_wrapped_structs_set_up_data(
["mjFakeStruct"]
)
)
structs_parser.generate_wasm_bindings(["mjFakeStruct"])
with self.assertRaises(RuntimeError):
structs_parser.generate_wasm_bindings(
structs_wrappers_data.create_wrapped_structs_set_up_data(
["mjFakeAnonymousStruct"]
)
)
def test_get_default_func_name_mjv(self):
self.assertEqual(
structs_parser.get_default_func_name("mjvPerturb"), "mjv_defaultPerturb"
)
def test_get_default_func_name(self):
self.assertEqual(
structs_parser.get_default_func_name("mjOption"), "mj_defaultOption"
)
structs_parser.generate_wasm_bindings(["mjFakeAnonymousStruct"])
if __name__ == "__main__":
wasm/codegen/helpers/structs_wrappers_data.py
-14
View File
@@ -59,17 +59,3 @@ class WrappedStructData:
# Whether to use shallow copy for this struct
use_shallow_copy: bool = True
def create_wrapped_structs_set_up_data(
struct_names: List[str],
) -> Dict[str, WrappedStructData]:
"""Creates a dictionary of WrappedStructData for the given struct names."""
return {
name: WrappedStructData(
wrap_name=common.uppercase_first_letter(name),
wrapped_fields=[],
wrapped_header="",
wrapped_source="",
)
for name in struct_names
}