Generate MjVisual, MjModel, MjData and MjSpec structs declaration

PiperOrigin-RevId: 834264723
Change-Id: Id33d7568a1f9f0cfb1271cf9dd2d8713c00eee34
This commit is contained in:
Google DeepMind
2025-11-19 06:21:24 -08:00
committed by Copybara-Service
parent 07d8bd5b80
commit 10fe5ff07f
4 changed files with 1442 additions and 1509 deletions
File diff suppressed because it is too large Load Diff
+1 -7
View File
@@ -235,11 +235,7 @@ SKIPPED_STRUCTS: List[str] = [
# or some of their fields need to be handled manually for now;
# making their wrapper constructors/destructors non-trivial.
MANUAL_STRUCTS: List[str] = [
"MjData",
"MjModel",
"MjvScene",
"MjSpec",
"MjVisual",
]
# Dictionary that maps anonymous structs to their parent struct and field name.
@@ -364,9 +360,7 @@ SKIPPED_FIELDS: Dict[str, List[str]] = {}
# Fields handled manually in template file struct declaration.
MANUAL_FIELDS: Dict[str, List[str]] = {
# go/keep-sorted start
"MjData": ["solver", "timer", "warning", "contact"],
"MjModel": ["opt", "vis", "stat"],
"MjSpec": ["option", "visual", "stat", "element", "compiler"],
"MjData": ["contact"],
"MjvScene": [
# go/keep-sorted start
"camera",
+55 -18
View File
@@ -129,7 +129,6 @@ def _generate_field_data(
)
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),
@@ -156,7 +155,10 @@ def _generate_field_data(
if anonymous_struct_name in constants.STRUCTS_TO_BIND:
return WrappedFieldData(
binding=_simple_property_binding(f, w, setter=False, reference=True),
typename=_get_field_struct_type(f.type),
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,
@@ -166,7 +168,11 @@ def _generate_field_data(
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
@@ -192,9 +198,18 @@ def _generate_field_data(
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)
@@ -279,13 +294,7 @@ def _generate_field_data(
binding=_simple_property_binding(f, w),
)
# SHOULD NOT OCCUR
print("Error: field {f.name} not properly handled")
return WrappedFieldData(
definition=f"// Error: field {f.name} not properly handled.",
typename=_get_field_struct_type(f.type),
binding=f"// Error: field {f.name} not properly handled.",
)
raise RuntimeError(f"Field {f.name} from struct {w} not properly handled")
def _default_function_statement(struct_name: str) -> str:
@@ -353,42 +362,67 @@ def build_struct_header(
):
raise RuntimeError(f"Struct {s} not found in introspect structs")
is_mjs = w.startswith("Mjs")
not_mjs = not w.startswith("Mjs")
member_inits = _find_member_inits(wrapped_fields)
shallow_copy = use_shallow_copy(wrapped_fields)
shallow_copy = not_mjs and use_shallow_copy(wrapped_fields)
builder = code_builder.CodeBuilder()
with builder.struct(f"{w}"):
builder.line(f"explicit {w}({s} *ptr);")
if not is_mjs:
# destructor
if not_mjs:
builder.line(f"~{w}();")
# default constructor
if not_mjs and w not in ["MjData", "MjModel"]:
builder.line(f"{w}();")
if shallow_copy and not is_mjs:
# 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 is_mjs:
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() + ";"
@@ -649,6 +683,9 @@ def sort_structs_by_dependency(
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
+13 -99
View File
@@ -109,72 +109,9 @@ std::vector<WrapperType> InitWrapperArray(ArrayType* array, SizeType size) {
// {{ AUTOGENNED_STRUCTS_HEADER }}
struct MjVisual {
MjVisual();
explicit MjVisual(mjVisual *ptr_);
MjVisual(const MjVisual &);
MjVisual &operator=(const MjVisual &);
~MjVisual();
std::unique_ptr<MjVisual> copy();
mjVisual* get() const;
void set(mjVisual* ptr);
// INSERT-GENERATED-MjVisual-DEFINITIONS
private:
mjVisual* ptr_;
bool owned_ = false;
public:
MjVisualGlobal global;
MjVisualQuality quality;
MjVisualHeadlight headlight;
MjVisualMap map;
MjVisualScale scale;
MjVisualRgba rgba;
};
struct MjModel {
explicit MjModel(mjModel *m);
explicit MjModel(const MjModel &other);
~MjModel();
std::unique_ptr<MjModel> copy();
mjModel* get() const;
void set(mjModel* ptr);
// INSERT-GENERATED-MjModel-DEFINITIONS
private:
mjModel* ptr_;
public:
MjOption opt;
MjStatistic stat;
MjVisual vis;
};
struct MjData {
MjData(MjModel *m);
explicit MjData(const MjModel &, const MjData &);
~MjData();
std::vector<MjContact> contact() const;
std::unique_ptr<MjData> copy();
mjData* get() const;
void set(mjData* ptr);
// INSERT-GENERATED-MjData-DEFINITIONS
private:
mjData* ptr_;
public:
mjModel *model;
std::vector<MjSolverStat> solver;
std::vector<MjTimerStat> timer;
std::vector<MjWarningStat> warning;
};
struct MjvScene {
MjvScene();
MjvScene(MjModel *m, int maxgeom);
// MjvScene(const MjvScene &);
~MjvScene();
std::unique_ptr<MjvScene> copy();
int GetSumFlexFaces() const;
@@ -268,29 +205,6 @@ struct MjvScene {
std::vector<MjvGLCamera> camera;
};
struct MjSpec {
MjSpec();
explicit MjSpec(mjSpec *ptr);
MjSpec(const MjSpec &);
MjSpec &operator=(const MjSpec &);
~MjSpec();
std::unique_ptr<MjSpec> copy();
mjSpec* get() const;
void set(mjSpec* ptr);
// INSERT-GENERATED-MjSpec-DEFINITIONS
private:
mjSpec* ptr_;
bool owned_ = false;
public:
MjOption option;
MjVisual visual;
MjStatistic stat;
MjsCompiler compiler;
MjsElement element;
};
val get_mjDISABLESTRING() { return MakeValArray(mjDISABLESTRING); }
val get_mjENABLESTRING() { return MakeValArray(mjENABLESTRING); }
val get_mjTIMERSTRING() { return MakeValArray(mjTIMERSTRING); }
@@ -348,13 +262,13 @@ EMSCRIPTEN_BINDINGS(mujoco_enums) {
// {{ AUTOGENNED_STRUCTS_SOURCE }}
// =============== MjModel =============== //
MjModel::MjModel(mjModel *m)
: ptr_(m), opt(&m->opt), stat(&m->stat), vis(&m->vis) {}
MjModel::MjModel(mjModel *ptr)
: ptr_(ptr), opt(&ptr->opt), vis(&ptr->vis), stat(&ptr->stat) {}
MjModel::MjModel(const MjModel &other)
: ptr_(mj_copyModel(nullptr, other.get())),
opt(&ptr_->opt),
stat(&ptr_->stat),
vis(&ptr_->vis) {}
vis(&ptr_->vis),
stat(&ptr_->stat) {}
MjModel::~MjModel() {
if (ptr_) {
mj_deleteModel(ptr_);
@@ -462,30 +376,30 @@ std::vector<MjvGeom> MjvScene::geoms() const {
MjSpec::MjSpec()
: ptr_(mj_makeSpec()),
element(ptr_->element),
compiler(&ptr_->compiler),
option(&ptr_->option),
visual(&ptr_->visual),
stat(&ptr_->stat),
compiler(&ptr_->compiler),
element(ptr_->element) {
stat(&ptr_->stat) {
owned_ = true;
mjs_defaultSpec(ptr_);
};
MjSpec::MjSpec(mjSpec *ptr)
: ptr_(ptr),
element(ptr_->element),
compiler(&ptr_->compiler),
option(&ptr_->option),
visual(&ptr_->visual),
stat(&ptr_->stat),
compiler(&ptr_->compiler),
element(ptr_->element) {}
stat(&ptr_->stat) {}
MjSpec::MjSpec(const MjSpec &other)
: ptr_(mj_copySpec(other.get())),
element(ptr_->element),
compiler(&ptr_->compiler),
option(&ptr_->option),
visual(&ptr_->visual),
stat(&ptr_->stat),
compiler(&ptr_->compiler),
element(ptr_->element) {
stat(&ptr_->stat) {
owned_ = true;
}