use bevy::prelude::*;
use bevy_rapier3d::prelude::*;
use crate::conversions;
use robocomp::cleanup_manager::Cleanup;
use robocomp::rc::{
BelongsToRcGroup, BelongsToRcLink, BelongsToRcRobotRoot, BelongsToRcSceneRoot, RcCcd,
RcCollider, RcColliderConstruction, RcColliderIsProcessed, RcDisableSleep, RcGroup,
RcGroupIsProcessed, RcJoint, RcJointCategory, RcJointIsProcessed, RcJointKind,
RcJointLocalAnchor1, RcJointLocalAnchor2, RcJointMotorControlConfig, RcLink, RcLinkIsProcessed,
RcPrismaticJointLimits, RcPrismaticJointMotor, RcRevoluteJointLimits, RcRevoluteJointMotor,
RcRobotGraph, RcRobotGraphJoint, RcRobotGraphLink, RcRobotRoot, RcRobotRootIsProcessed,
RcSceneRoot, RcSceneRootIsProcessed, RcSceneRootReady, RcSceneRootWRobotRootsDetached,
RcSphericalJointLimits, RcSphericalJointMotor,
};
pub struct PreProcessorPlugin;
impl Plugin for PreProcessorPlugin {
fn build(&self, app: &mut App) {
app .configure_sets(Update, PreProcessorPluginSystemSet)
.add_systems(
Update,
(
process_group_markers,
process_scene_root_markers,
process_root_markers,
process_collider_markers,
process_link_markers,
process_joint_markers,
check_processing_complete,
detach_root_markers_and_flatten_hierarchy,
)
.chain()
.in_set(PreProcessorPluginSystemSet),
)
.add_systems(Startup, || info!("PreProcessorPlugin started..."));
}
}
#[derive(Debug, Hash, PartialEq, Eq, Clone, SystemSet)]
pub struct PreProcessorPluginSystemSet;
pub fn process_group_markers(
group_markers: Query<Entity, (With<RcGroup>, Without<RcGroupIsProcessed>, Without<Cleanup>)>,
children: Query<&Children, Without<Cleanup>>,
mut commands: Commands,
) {
for group_ent in group_markers.iter() {
debug!("Processing RcGroup: {:?}", group_ent);
let num_of_recursive_children = children.iter_descendants(group_ent).count();
if num_of_recursive_children == 0 {
continue;
}
for child_ent in children.iter_descendants(group_ent) {
commands
.entity(child_ent)
.insert(BelongsToRcGroup(group_ent));
}
commands.entity(group_ent).insert(RcGroupIsProcessed);
}
}
pub fn process_scene_root_markers(
scene_root_markers: Query<
Entity,
(
With<RcSceneRoot>,
Without<RcSceneRootIsProcessed>,
Without<RcSceneRootReady>,
Without<Cleanup>,
),
>,
children: Query<&Children, Without<Cleanup>>,
mut commands: Commands,
) {
for scene_root_ent in scene_root_markers.iter() {
debug!("Processing RcSceneRoot: {:?}", scene_root_ent);
let num_of_recursive_children = children.iter_descendants(scene_root_ent).count();
if num_of_recursive_children == 0 {
continue;
}
for child_ent in children.iter_descendants(scene_root_ent) {
commands
.entity(child_ent)
.insert(BelongsToRcSceneRoot(scene_root_ent));
}
commands
.entity(scene_root_ent)
.insert(RcSceneRootIsProcessed);
}
}
pub fn process_root_markers(
mut commands: Commands,
robot_roots: Query<
Entity,
(
With<RcRobotRoot>,
With<BelongsToRcSceneRoot>,
Without<RcRobotRootIsProcessed>,
),
>,
children: Query<&Children, Without<Cleanup>>,
) {
for root_ent in robot_roots.iter() {
for child_ent in children.iter_descendants(root_ent) {
commands
.entity(child_ent)
.insert(BelongsToRcRobotRoot(root_ent));
}
commands.entity(root_ent).insert((
RcRobotGraph::default(),
BelongsToRcRobotRoot(root_ent),
RcRobotRootIsProcessed,
));
}
}
pub fn process_collider_markers(
mut commands: Commands,
child_mesh_nodes: Query<(Entity, &Mesh3d, &ChildOf)>,
collider_markers: Query<
(
Entity,
&BelongsToRcRobotRoot,
&RcCollider,
Option<&RcColliderConstruction>,
),
(
With<RcCollider>,
With<BelongsToRcSceneRoot>,
Without<RcColliderIsProcessed>,
),
>,
robot_roots: Query<
Entity,
(
With<RcRobotRoot>,
With<BelongsToRcSceneRoot>,
With<RcRobotRootIsProcessed>,
),
>,
meshes: Res<Assets<Mesh>>,
) {
if collider_markers.is_empty() {
return;
}
for (child_ent, mesh_3d, child_of) in child_mesh_nodes.iter() {
let Ok((
collider_ent,
&BelongsToRcRobotRoot(root_ent),
collider_marker,
collider_construction,
)) = collider_markers.get(child_of.parent())
else {
continue;
};
if !robot_roots.contains(root_ent) {
warn!("Root not found for collider entity: {:?}", child_ent);
continue;
}
let Some(mesh) = meshes.get(mesh_3d) else {
continue;
};
let collider_construction = match collider_construction {
Some(collider_construction) => match collider_construction {
RcColliderConstruction::ConvexHull => ComputedColliderShape::ConvexHull,
RcColliderConstruction::Trimesh => {
ComputedColliderShape::TriMesh(TriMeshFlags::all())
}
RcColliderConstruction::ConvexDecomposition => {
ComputedColliderShape::ConvexDecomposition(VHACDParameters::default())
}
},
None => ComputedColliderShape::ConvexHull,
};
let Some(collider) = Collider::from_bevy_mesh(mesh, &collider_construction) else {
warn!(
"Failed to create collider from mesh for entity: {:?}",
child_ent
);
continue;
};
let RcCollider {
density: collider_density,
keep_mesh,
..
} = collider_marker;
commands
.entity(collider_ent)
.insert((collider, ColliderMassProperties::Density(*collider_density)));
if !*keep_mesh {
commands.entity(child_ent).insert(Cleanup::Recursive);
}
commands.entity(collider_ent).insert(RcColliderIsProcessed);
}
}
pub fn process_link_markers(
mut commands: Commands,
mut robot_roots: Query<
&mut RcRobotGraph,
(
With<RcRobotRoot>,
With<BelongsToRcSceneRoot>,
With<RcRobotRootIsProcessed>,
),
>,
link_markers: Query<
(
Entity,
&BelongsToRcRobotRoot,
&RcLink,
Option<&RcCcd>,
Option<&RcDisableSleep>,
),
(
With<RcLink>,
With<BelongsToRcSceneRoot>,
Without<RcLinkIsProcessed>,
),
>,
children: Query<&Children, Without<Cleanup>>,
) {
if link_markers.is_empty() {
return;
}
for (link_ent, &BelongsToRcRobotRoot(root_ent), link_marker, rc_ccd, disable_sleep) in
link_markers.iter()
{
let Ok(mut robot_graph) = robot_roots.get_mut(root_ent) else {
warn!("Root not found for link entity: {:?}", link_ent);
continue;
};
let RcLink { rigid_body, .. } = link_marker;
commands
.entity(link_ent)
.insert((conversions::rigid_body_from_rc(rigid_body.clone()), Velocity::zero()));
if rc_ccd.is_some() {
commands.entity(link_ent).insert(Ccd::enabled());
}
if disable_sleep.is_some() {
commands.entity(link_ent).insert(Sleeping::disabled());
}
for child in children.iter_descendants(link_ent) {
commands.entity(child).insert(BelongsToRcLink(link_ent));
}
robot_graph.add_link(RcRobotGraphLink {
link_name: link_marker.name.clone(),
link_entity: link_ent,
});
commands.entity(link_ent).insert(RcLinkIsProcessed);
}
}
pub fn process_joint_markers(
mut commands: Commands,
mut robot_roots: Query<
&mut RcRobotGraph,
(
With<RcRobotRoot>,
With<BelongsToRcSceneRoot>,
With<RcRobotRootIsProcessed>,
),
>,
joint_anchors: Query<
(
Entity,
&ChildOf,
&GlobalTransform,
Has<RcJointLocalAnchor1>,
Has<RcJointLocalAnchor2>,
),
(
Or<(With<RcJointLocalAnchor1>, With<RcJointLocalAnchor2>)>,
With<BelongsToRcSceneRoot>,
),
>,
joint_markers: Query<
(
Entity,
&BelongsToRcRobotRoot,
&RcJoint,
&RcJointKind,
Option<&RcJointMotorControlConfig>,
Option<&RcRevoluteJointMotor>,
Option<&RcRevoluteJointLimits>,
Option<&RcPrismaticJointMotor>,
Option<&RcPrismaticJointLimits>,
Option<&RcSphericalJointMotor>,
Option<&RcSphericalJointLimits>,
),
(
With<RcJoint>,
With<BelongsToRcSceneRoot>,
Without<RcJointIsProcessed>,
Without<Cleanup>,
),
>,
links: Query<
(Entity, &BelongsToRcRobotRoot, &RcLink, &GlobalTransform),
(With<RcLinkIsProcessed>, With<BelongsToRcSceneRoot>),
>,
) {
if joint_markers.is_empty() {
return;
}
for (
joint_ent,
&BelongsToRcRobotRoot(root_ent),
joint_marker,
joint_kind,
joint_motor_control_config,
revolute_joint_motor,
revolute_joint_limits,
prismatic_joint_motor,
prismatic_joint_limits,
spherical_joint_motor,
spherical_joint_limits,
) in joint_markers.iter()
{
let Ok(mut robot_graph) = robot_roots.get_mut(root_ent) else {
warn!("Root not found for joint entity: {:?}", joint_ent);
continue;
};
let (joint_parent_name, joint_child_name) = match joint_kind {
RcJointKind::Fixed { parent, child } => (parent, child),
RcJointKind::Revolute { parent, child, .. } => (parent, child),
RcJointKind::Prismatic { parent, child, .. } => (parent, child),
RcJointKind::Spherical { parent, child } => (parent, child),
};
let Some((parent_link, _, _, parent_link_glob_transform)) =
links
.iter()
.find(|&(_, &BelongsToRcRobotRoot(cur_root_ent), ref link, _)| {
cur_root_ent == root_ent && link.name == *joint_parent_name
})
else {
warn!("Parent link not found for joint: {:?}", joint_ent);
continue;
};
let Some((child_link, _, _, child_link_glob_transform)) =
links
.iter()
.find(|&(_, &BelongsToRcRobotRoot(cur_root_ent), ref link, _)| {
cur_root_ent == root_ent && link.name == *joint_child_name
})
else {
warn!("Child link not found for joint: {:?}", joint_ent);
continue;
};
let Some(local_anchor_1_glob_translation) = joint_anchors
.iter()
.find(|(_, child_of, _, is_anchor1, _)| *is_anchor1 && child_of.parent() == joint_ent)
.map(|(_, _, global_transform, _, _)| global_transform.translation())
else {
warn!("Local anchor 1 not found for joint: {:?}", joint_ent);
continue;
};
let Some(local_anchor_2_glob_translation) = joint_anchors
.iter()
.find(|(_, child_of, _, _, is_anchor2)| *is_anchor2 && child_of.parent() == joint_ent)
.map(|(_, _, global_transform, _, _)| global_transform.translation())
else {
warn!("Local anchor 2 not found for joint: {:?}", joint_ent);
continue;
};
debug!(
"Global Anchor 1: {:?}, Global Anchor 2: {:?}",
local_anchor_1_glob_translation, local_anchor_2_glob_translation
);
let local_anchor_1 = parent_link_glob_transform
.to_matrix()
.inverse()
.transform_point3(local_anchor_1_glob_translation);
let local_anchor_2 = child_link_glob_transform
.to_matrix()
.inverse()
.transform_point3(local_anchor_2_glob_translation);
debug!(
"Local Anchor 1: {:?}, Local Anchor 2: {:?}",
local_anchor_1, local_anchor_2
);
let joint: TypedJoint = match joint_kind {
RcJointKind::Fixed { .. } => {
let joint = FixedJointBuilder::new()
.local_anchor1(local_anchor_1)
.local_anchor2(local_anchor_2);
debug!("Fixed Joint created: {:?}", joint_ent);
joint.into()
}
RcJointKind::Revolute { axis, .. } => {
let mut joint = RevoluteJointBuilder::new(*axis)
.local_anchor1(local_anchor_1)
.local_anchor2(local_anchor_2);
if let Some(RcRevoluteJointMotor {
model,
velocity,
position,
max_force,
}) = revolute_joint_motor
{
joint = joint.motor_model(conversions::rapier_motor_model_from_rd(&model));
if let Some(velocity) = velocity {
let (_, damping) = conversions::stiffness_damping_from_rd(&model);
joint = joint.motor_velocity(velocity.target, damping);
}
if let Some(position) = position {
let (stiffness, damping) =
conversions::stiffness_damping_from_rd(&model);
joint = joint.motor_position(position.target, stiffness, damping);
}
if let Some(max_force) = max_force {
joint = joint.motor_max_force(*max_force);
}
}
if let Some(RcRevoluteJointLimits(limits)) = revolute_joint_limits {
joint = joint.limits([limits.min, limits.max]);
}
debug!("Revolute Joint created: {:?}", joint_ent);
joint.into()
}
RcJointKind::Prismatic { axis, .. } => {
let mut joint = PrismaticJointBuilder::new(*axis)
.local_anchor1(local_anchor_1)
.local_anchor2(local_anchor_2);
if let Some(RcPrismaticJointMotor {
model,
velocity,
position,
max_force,
}) = prismatic_joint_motor
{
joint = joint.motor_model(conversions::rapier_motor_model_from_rd(&model));
if let Some(velocity) = velocity {
let (_, damping) = conversions::stiffness_damping_from_rd(&model);
joint = joint.motor_velocity(velocity.target, damping);
}
if let Some(position) = position {
let (stiffness, damping) =
conversions::stiffness_damping_from_rd(&model);
joint = joint.motor_position(position.target, stiffness, damping);
}
if let Some(max_force) = max_force {
joint = joint.motor_max_force(*max_force);
}
}
if let Some(RcPrismaticJointLimits(limits)) = prismatic_joint_limits {
joint = joint.limits([limits.min, limits.max]);
}
debug!("Prismatic Joint created: {:?}", joint_ent);
joint.into()
}
RcJointKind::Spherical { .. } => {
let mut joint = SphericalJointBuilder::new()
.local_anchor1(local_anchor_1)
.local_anchor2(local_anchor_2);
if let Some(RcSphericalJointMotor {
model_ang_x,
model_ang_y,
model_ang_z,
velocity_ang_x,
velocity_ang_y,
velocity_ang_z,
position_ang_x,
position_ang_y,
position_ang_z,
max_force_ang_x,
max_force_ang_y,
max_force_ang_z,
}) = spherical_joint_motor
{
joint = joint
.motor_model(
JointAxis::AngX,
conversions::rapier_motor_model_from_rd(&model_ang_x),
)
.motor_model(
JointAxis::AngY,
conversions::rapier_motor_model_from_rd(&model_ang_y),
)
.motor_model(
JointAxis::AngZ,
conversions::rapier_motor_model_from_rd(&model_ang_z),
);
if let Some(velocity_ang_x) = velocity_ang_x {
let (_, damping) =
conversions::stiffness_damping_from_rd(&model_ang_x);
joint = joint.motor_velocity(JointAxis::AngX, velocity_ang_x.target, damping);
}
if let Some(velocity_ang_y) = velocity_ang_y {
let (_, damping) =
conversions::stiffness_damping_from_rd(&model_ang_y);
joint = joint.motor_velocity(JointAxis::AngY, velocity_ang_y.target, damping);
}
if let Some(velocity_ang_z) = velocity_ang_z {
let (_, damping) =
conversions::stiffness_damping_from_rd(&model_ang_z);
joint = joint.motor_velocity(JointAxis::AngZ, velocity_ang_z.target, damping);
}
if let Some(position_ang_x) = position_ang_x {
let (stiffness, damping) =
conversions::stiffness_damping_from_rd(&model_ang_x);
joint = joint.motor_position(
JointAxis::AngX,
position_ang_x.target,
stiffness,
damping,
);
}
if let Some(position_ang_y) = position_ang_y {
let (stiffness, damping) =
conversions::stiffness_damping_from_rd(&model_ang_y);
joint = joint.motor_position(
JointAxis::AngY,
position_ang_y.target,
stiffness,
damping,
);
}
if let Some(position_ang_z) = position_ang_z {
let (stiffness, damping) =
conversions::stiffness_damping_from_rd(&model_ang_z);
joint = joint.motor_position(
JointAxis::AngZ,
position_ang_z.target,
stiffness,
damping,
);
}
if let Some(max_force_ang_x) = max_force_ang_x {
joint = joint.motor_max_force(JointAxis::AngX, *max_force_ang_x);
}
if let Some(max_force_ang_y) = max_force_ang_y {
joint = joint.motor_max_force(JointAxis::AngY, *max_force_ang_y);
}
if let Some(max_force_ang_z) = max_force_ang_z {
joint = joint.motor_max_force(JointAxis::AngZ, *max_force_ang_z);
}
}
if let Some(RcSphericalJointLimits {
ang_x,
ang_y,
ang_z,
}) = spherical_joint_limits
{
if let Some(ang_x) = ang_x {
joint = joint.limits(JointAxis::AngX, [ang_x.min, ang_x.max]);
}
if let Some(ang_y) = ang_y {
joint = joint.limits(JointAxis::AngY, [ang_y.min, ang_y.max]);
}
if let Some(ang_z) = ang_z {
joint = joint.limits(JointAxis::AngZ, [ang_z.min, ang_z.max]);
}
}
debug!("Spherical Joint created: {:?}", joint_ent);
joint.into()
}
};
let mut child_ent_commands = commands.entity(child_link);
match joint_marker.category {
RcJointCategory::Impulse => {
child_ent_commands.insert(ImpulseJoint::new(parent_link, joint));
}
RcJointCategory::Multibody => {
child_ent_commands.insert(MultibodyJoint::new(parent_link, joint));
}
}
if let Some(joint_motor_control_config) = joint_motor_control_config {
child_ent_commands.insert(joint_motor_control_config.clone());
}
robot_graph.add_joint(RcRobotGraphJoint {
joint_name: joint_marker.name.clone(),
joint_entity: child_link,
from_link: joint_parent_name.clone(),
to_link: joint_child_name.clone(),
});
commands.entity(joint_ent).insert(Cleanup::Recursive);
}
}
pub fn check_processing_complete(
mut commands: Commands,
scene_root_markers: Query<
(Entity, Option<&RcSceneRootIsProcessed>),
(With<RcSceneRoot>, Without<RcSceneRootReady>),
>,
root_markers: Query<
(
Entity,
&BelongsToRcSceneRoot,
Option<&RcRobotRootIsProcessed>,
),
With<RcRobotRoot>,
>,
collider_markers: Query<
(
Entity,
&BelongsToRcSceneRoot,
Option<&RcColliderIsProcessed>,
),
With<RcCollider>,
>,
link_markers: Query<(Entity, &BelongsToRcSceneRoot, Option<&RcLinkIsProcessed>), With<RcLink>>,
joint_markers: Query<
(Entity, &BelongsToRcSceneRoot, Option<&RcJointIsProcessed>),
With<RcJoint>,
>,
) {
for (scene_root_ent, scene_root_is_processed) in scene_root_markers.iter() {
if scene_root_is_processed.is_none() {
continue;
}
let all_root_markers_processed = root_markers
.iter()
.filter(|&(_, &BelongsToRcSceneRoot(cur_scene_root), _)| {
cur_scene_root == scene_root_ent
})
.all(|(_, _, is_processed)| is_processed.is_some());
if !all_root_markers_processed {
continue;
}
let all_collider_markers_processed = collider_markers
.iter()
.filter(|&(_, &BelongsToRcSceneRoot(cur_scene_root), _)| {
cur_scene_root == scene_root_ent
})
.all(|(_, _, is_processed)| is_processed.is_some());
if !all_collider_markers_processed {
continue;
}
let all_link_markers_processed = link_markers
.iter()
.filter(|&(_, &BelongsToRcSceneRoot(cur_scene_root), _)| {
cur_scene_root == scene_root_ent
})
.all(|(_, _, is_processed)| is_processed.is_some());
if !all_link_markers_processed {
continue;
}
let all_joint_markers_processed = joint_markers
.iter()
.filter(|&(_, &BelongsToRcSceneRoot(cur_scene_root), _)| {
cur_scene_root == scene_root_ent
})
.all(|(_, _, is_processed)| is_processed.is_some());
if !all_joint_markers_processed {
continue;
}
debug!("scene_root_ent: {} is ready!!", scene_root_ent);
commands.entity(scene_root_ent).insert(RcSceneRootReady);
}
}
pub fn detach_root_markers_and_flatten_hierarchy(
mut commands: Commands,
scene_root_markers: Query<
Entity,
(
With<RcSceneRoot>,
With<RcSceneRootIsProcessed>,
With<RcSceneRootReady>,
Without<Cleanup>,
Without<RcSceneRootWRobotRootsDetached>,
),
>,
root_markers: Query<
(Entity, &BelongsToRcSceneRoot),
(With<RcRobotRoot>, With<RcRobotRootIsProcessed>),
>,
children: Query<(Entity, &ChildOf), Without<Cleanup>>,
) {
for scene_root in scene_root_markers.iter() {
let root_markers_belonging_to_scene_root = root_markers
.iter()
.filter(|&(_, &BelongsToRcSceneRoot(cur_scene_root))| cur_scene_root == scene_root);
if root_markers_belonging_to_scene_root.count() == 0 {
continue;
}
for (root_ent, &BelongsToRcSceneRoot(cur_scene_root)) in root_markers.iter() {
if cur_scene_root != scene_root {
continue;
}
for (child_ent, child_of) in children.iter() {
if child_of.parent() != root_ent {
continue;
}
commands.entity(child_ent).remove_parent_in_place();
}
}
commands
.entity(scene_root)
.insert(RcSceneRootWRobotRootsDetached);
}
}