Implement cast_ray_predicate to allow filtering the colliders with a …

…function (#297)

As I was writing a xpbd backend for bevy_mod_picking, I found that a ray cast with a predicate was missing to implement the functionality efficiently. 

I was inspired by the predicate from bevy_rapier: https://github.com/dimforge/bevy_rapier/blob/c6bcce4695d596a7a9c8e91748d4dbb3d31f6d13/src/pipeline/query_filter.rs#L12

# Objective

- Implement a ray cast function with the support for a predicate function to filter the colliders

## Solution

- Added support for a predicate starting at SpatialQuery and down the pipeline
- I added a small (and fun to play with) example to illustrate one possible use case
- I chose not to change existing functionality. Therefore there is some code duplication in QueryPipelineAsCompositeShapeWithPredicate. This could be unified by introducing an Option in QueryPipelineAsCompositeShape, but that would create breaking changes at many places.

---------

Co-authored-by: hendrikd <[email protected]>
Co-authored-by: Joona Aalto <[email protected]>

crates/bevy_xpbd_3d/examples/cast_ray_predicate.rs

@@ -0,0 +1,200 @@
+#![allow(clippy::unnecessary_cast)]
+
+use bevy::{pbr::NotShadowReceiver, prelude::*};
+use bevy_xpbd_3d::{math::*, prelude::*};
+use examples_common_3d::XpbdExamplePlugin;
+
+fn main() {
+    App::new()
+        .add_plugins((DefaultPlugins, XpbdExamplePlugin))
+        .insert_resource(ClearColor(Color::rgb(0.05, 0.05, 0.1)))
+        .insert_resource(Msaa::Sample4)
+        .add_systems(Startup, setup)
+        .add_systems(Update, (movement, reset_colors, raycast).chain())
+        .run();
+}
+
+/// The acceleration used for movement.
+#[derive(Component)]
+struct MovementAcceleration(Scalar);
+
+#[derive(Component)]
+struct RayIndicator;
+
+/// If to be ignored by raycast
+#[derive(Component)]
+struct OutOfGlass(bool);
+
+const CUBE_COLOR: Color = Color::rgba(0.2, 0.7, 0.9, 1.0);
+const CUBE_COLOR_GLASS: Color = Color::rgba(0.2, 0.7, 0.9, 0.5);
+
+fn setup(
+    mut commands: Commands,
+    mut materials: ResMut<Assets<StandardMaterial>>,
+    mut meshes: ResMut<Assets<Mesh>>,
+) {
+    let cube_mesh = meshes.add(Mesh::from(shape::Cube { size: 1.0 }));
+
+    // Ground
+    commands.spawn((
+        PbrBundle {
+            mesh: cube_mesh.clone(),
+            material: materials.add(Color::rgb(0.7, 0.7, 0.8).into()),
+            transform: Transform::from_xyz(0.0, -2.0, 0.0).with_scale(Vec3::new(100.0, 1.0, 100.0)),
+            ..default()
+        },
+        RigidBody::Static,
+        Collider::cuboid(1.0, 1.0, 1.0),
+    ));
+
+    let cube_size = 2.0;
+
+    // Spawn cube stacks
+    for x in -1..2 {
+        for y in -1..2 {
+            for z in -1..2 {
+                let position = Vec3::new(x as f32, y as f32 + 5.0, z as f32) * (cube_size + 0.05);
+                let material: StandardMaterial = if x == -1 {
+                    CUBE_COLOR_GLASS.into()
+                } else {
+                    CUBE_COLOR.into()
+                };
+                commands.spawn((
+                    PbrBundle {
+                        mesh: cube_mesh.clone(),
+                        material: materials.add(material.clone()),
+                        transform: Transform::from_translation(position)
+                            .with_scale(Vec3::splat(cube_size as f32)),
+                        ..default()
+                    },
+                    RigidBody::Dynamic,
+                    Collider::cuboid(1.0, 1.0, 1.0),
+                    MovementAcceleration(10.0),
+                    OutOfGlass(x == -1),
+                ));
+            }
+        }
+    }
+
+    // raycast indicator
+    commands.spawn((
+        PbrBundle {
+            mesh: cube_mesh.clone(),
+            material: materials.add(Color::rgb(1.0, 0.0, 0.0).into()),
+            transform: Transform::from_xyz(-500.0, 2.0, 0.0)
+                .with_scale(Vec3::new(1000.0, 0.1, 0.1)),
+            ..default()
+        },
+        RayIndicator,
+        NotShadowReceiver,
+    ));
+
+    // Directional light
+    commands.spawn(DirectionalLightBundle {
+        directional_light: DirectionalLight {
+            illuminance: 20_000.0,
+            shadows_enabled: true,
+            ..default()
+        },
+        transform: Transform::default().looking_at(Vec3::new(-1.0, -2.5, -1.5), Vec3::Y),
+        ..default()
+    });
+
+    // Camera
+    commands.spawn(Camera3dBundle {
+        transform: Transform::from_translation(Vec3::new(0.0, 12.0, 40.0))
+            .looking_at(Vec3::Y * 5.0, Vec3::Y),
+        ..default()
+    });
+}
+
+fn movement(
+    time: Res<Time>,
+    keyboard_input: Res<Input<KeyCode>>,
+    mut query: Query<(&MovementAcceleration, &mut LinearVelocity)>,
+) {
+    // Precision is adjusted so that the example works with
+    // both the `f32` and `f64` features. Otherwise you don't need this.
+    let delta_time = time.delta_seconds_f64().adjust_precision();
+
+    for (movement_acceleration, mut linear_velocity) in &mut query {
+        let up = keyboard_input.any_pressed([KeyCode::W, KeyCode::Up]);
+        let down = keyboard_input.any_pressed([KeyCode::S, KeyCode::Down]);
+        let left = keyboard_input.any_pressed([KeyCode::A, KeyCode::Left]);
+        let right = keyboard_input.any_pressed([KeyCode::D, KeyCode::Right]);
+
+        let horizontal = right as i8 - left as i8;
+        let vertical = down as i8 - up as i8;
+        let direction =
+            Vector::new(horizontal as Scalar, 0.0, vertical as Scalar).normalize_or_zero();
+
+        // Move in input direction
+        if direction != Vector::ZERO {
+            linear_velocity.x += direction.x * movement_acceleration.0 * delta_time;
+            linear_velocity.z += direction.z * movement_acceleration.0 * delta_time;
+        }
+    }
+}
+
+fn reset_colors(
+    mut materials: ResMut<Assets<StandardMaterial>>,
+    cubes: Query<(&Handle<StandardMaterial>, &OutOfGlass)>,
+) {
+    for (material_handle, out_of_glass) in cubes.iter() {
+        if let Some(material) = materials.get_mut(material_handle) {
+            if out_of_glass.0 {
+                material.base_color = CUBE_COLOR_GLASS;
+            } else {
+                material.base_color = CUBE_COLOR;
+            }
+        }
+    }
+}
+
+fn raycast(
+    query: SpatialQuery,
+    mut materials: ResMut<Assets<StandardMaterial>>,
+    cubes: Query<(&Handle<StandardMaterial>, &OutOfGlass)>,
+    mut indicator_transform: Query<&mut Transform, With<RayIndicator>>,
+) {
+    let origin = Vector {
+        x: -200.0,
+        y: 2.0,
+        z: 0.0,
+    };
+    let direction = Vector {
+        x: 1.0,
+        y: 0.0,
+        z: 0.0,
+    };
+
+    let mut ray_indicator_transform = indicator_transform.single_mut();
+
+    if let Some(ray_hit_data) = query.cast_ray_predicate(
+        origin,
+        direction,
+        Scalar::MAX,
+        true,
+        SpatialQueryFilter::new(),
+        &|entity| {
+            if let Ok((_, out_of_glass)) = cubes.get(entity) {
+                return !out_of_glass.0; // only look at cubes not out of glass
+            }
+            true // if the collider has no OutOfGlass component, then check it nevertheless
+        },
+    ) {
+        // set color of hit object to red
+        if let Ok((material_handle, _)) = cubes.get(ray_hit_data.entity) {
+            if let Some(material) = materials.get_mut(material_handle) {
+                material.base_color = Color::RED;
+            }
+        }
+
+        // set length of ray indicator to look more like a laser
+        let contact_point = (origin + direction * ray_hit_data.time_of_impact).x;
+        let target_scale = 1000.0 + contact_point * 2.0;
+        ray_indicator_transform.scale.x = target_scale as f32;
+    } else {
+        ray_indicator_transform.scale.x = 2000.0;
+    }
+}

src/plugins/collision/mod.rs

@@ -229,8 +229,8 @@ impl Collisions {
     /// The order of the entities does not matter.
     pub fn remove_collision_pair(&mut self, entity1: Entity, entity2: Entity) -> Option<Contacts> {
         self.0
-            .remove(&(entity1, entity2))
-            .or_else(|| self.0.remove(&(entity2, entity1)))
+            .swap_remove(&(entity1, entity2))
+            .or_else(|| self.0.swap_remove(&(entity2, entity1)))
     }
 
     /// Removes all collisions that involve the given entity.

src/plugins/spatial_query/pipeline.rs

@@ -58,6 +58,19 @@ impl SpatialQueryPipeline {
         }
     }
 
+    pub(crate) fn as_composite_shape_with_predicate<'a>(
+        &'a self,
+        query_filter: SpatialQueryFilter,
+        predicate: &'a dyn Fn(Entity) -> bool,
+    ) -> QueryPipelineAsCompositeShapeWithPredicate {
+        QueryPipelineAsCompositeShapeWithPredicate {
+            pipeline: self,
+            colliders: &self.colliders,
+            query_filter,
+            predicate,
+        }
+    }
+
     /// Updates the associated acceleration structures with a new set of entities.
     pub fn update<'a>(
         &mut self,
@@ -172,6 +185,48 @@ impl SpatialQueryPipeline {
             })
     }
 
+    /// Casts a [ray](spatial_query#raycasting) and computes the closest [hit](RayHitData) with a collider.
+    /// If there are no hits, `None` is returned.
+    ///
+    /// ## Arguments
+    ///
+    /// - `origin`: Where the ray is cast from.
+    /// - `direction`: What direction the ray is cast in.
+    /// - `max_time_of_impact`: The maximum distance that the ray can travel.
+    /// - `solid`: If true and the ray origin is inside of a collider, the hit point will be the ray origin itself.
+    /// Otherwise, the collider will be treated as hollow, and the hit point will be at the collider's boundary.
+    /// - `query_filter`: A [`SpatialQueryFilter`] that determines which colliders are taken into account in the query.
+    /// - `predicate`: A function with which the colliders are filtered. Given the Entity it should return false, if the
+    /// entity should be ignored.
+    ///
+    /// See also: [`SpatialQuery::cast_ray`]
+    pub fn cast_ray_predicate(
+        &self,
+        origin: Vector,
+        direction: Vector,
+        max_time_of_impact: Scalar,
+        solid: bool,
+        query_filter: SpatialQueryFilter,
+        predicate: &dyn Fn(Entity) -> bool,
+    ) -> Option<RayHitData> {
+        let pipeline_shape = self.as_composite_shape_with_predicate(query_filter, predicate);
+        let ray = parry::query::Ray::new(origin.into(), direction.into());
+        let mut visitor = RayCompositeShapeToiAndNormalBestFirstVisitor::new(
+            &pipeline_shape,
+            &ray,
+            max_time_of_impact,
+            solid,
+        );
+
+        self.qbvh
+            .traverse_best_first(&mut visitor)
+            .map(|(_, (entity_index, hit))| RayHitData {
+                entity: self.entity_from_index(entity_index),
+                time_of_impact: hit.toi,
+                normal: hit.normal.into(),
+            })
+    }
+
     /// Casts a [ray](spatial_query#raycasting) and computes all [hits](RayHitData) until `max_hits` is reached.
     ///
     /// Note that the order of the results is not guaranteed, and if there are more hits than `max_hits`,
@@ -715,6 +770,49 @@ impl<'a> TypedSimdCompositeShape for QueryPipelineAsCompositeShape<'a> {
     }
 }
 
+pub(crate) struct QueryPipelineAsCompositeShapeWithPredicate<'a, 'b> {
+    colliders: &'a HashMap<Entity, (Isometry<Scalar>, Collider, CollisionLayers)>,
+    pipeline: &'a SpatialQueryPipeline,
+    query_filter: SpatialQueryFilter,
+    predicate: &'b dyn Fn(Entity) -> bool,
+}
+
+impl<'a, 'b> TypedSimdCompositeShape for QueryPipelineAsCompositeShapeWithPredicate<'a, 'b> {
+    type PartShape = dyn Shape;
+    type PartId = u32;
+    type QbvhStorage = DefaultStorage;
+
+    fn map_typed_part_at(
+        &self,
+        shape_id: Self::PartId,
+        mut f: impl FnMut(Option<&Isometry<Scalar>>, &Self::PartShape),
+    ) {
+        if let Some((entity, (iso, shape, layers))) =
+            self.colliders
+                .get_key_value(&utils::entity_from_index_and_gen(
+                    shape_id,
+                    *self.pipeline.entity_generations.get(&shape_id).unwrap(),
+                ))
+        {
+            if self.query_filter.test(*entity, *layers) && (self.predicate)(*entity) {
+                f(Some(iso), &**shape.shape_scaled());
+            }
+        }
+    }
+
+    fn map_untyped_part_at(
+        &self,
+        shape_id: Self::PartId,
+        f: impl FnMut(Option<&Isometry<Scalar>>, &dyn Shape),
+    ) {
+        self.map_typed_part_at(shape_id, f);
+    }
+
+    fn typed_qbvh(&self) -> &parry::partitioning::GenericQbvh<Self::PartId, Self::QbvhStorage> {
+        &self.pipeline.qbvh
+    }
+}
+
 /// The result of a [point projection](spatial_query#point-projection) on a [collider](Collider).
 #[derive(Debug, Clone, PartialEq)]
 #[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]

src/plugins/spatial_query/system_param.rs

@@ -133,6 +133,68 @@ impl<'w, 's> SpatialQuery<'w, 's> {
             .cast_ray(origin, direction, max_time_of_impact, solid, query_filter)
     }
 
+    /// Casts a [ray](spatial_query#raycasting) and computes the closest [hit](RayHitData) with a collider.
+    /// If there are no hits, `None` is returned.
+    ///
+    /// ## Arguments
+    ///
+    /// - `origin`: Where the ray is cast from.
+    /// - `direction`: What direction the ray is cast in.
+    /// - `max_time_of_impact`: The maximum distance that the ray can travel.
+    /// - `solid`: If true and the ray origin is inside of a collider, the hit point will be the ray origin itself.
+    /// Otherwise, the collider will be treated as hollow, and the hit point will be at the collider's boundary.
+    /// - `query_filter`: A [`SpatialQueryFilter`] that determines which colliders are taken into account in the query.
+    /// - `predicate`: A function with which the colliders are filtered. Given the Entity it should return false, if the
+    /// entity should be ignored.
+    ///
+    /// ## Example
+    ///
+    /// ```
+    /// use bevy::prelude::*;
+    /// # #[cfg(feature = "2d")]
+    /// # use bevy_xpbd_2d::prelude::*;
+    /// # #[cfg(feature = "3d")]
+    /// use bevy_xpbd_3d::prelude::*;
+    ///
+    /// # #[cfg(all(feature = "3d", feature = "f32"))]
+    /// fn print_hits(spatial_query: SpatialQuery) {
+    ///     // Cast ray and print first hit
+    ///     if let Some(first_hit) = spatial_query.cast_ray(
+    ///         Vec3::ZERO,                    // Origin
+    ///         Vec3::X,                       // Direction
+    ///         100.0,                         // Maximum time of impact (travel distance)
+    ///         true,                          // Does the ray treat colliders as "solid"
+    ///         SpatialQueryFilter::default(), // Query filter
+    ///         &|entity| {                    // Predicate
+    ///             if let Some(value) = query.get(entity) {
+    ///                 return value == x;     // ignore if value from query is x
+    ///             }
+    ///             true                       // else check for collision
+    ///         }
+    ///     ) {
+    ///         println!("First hit: {:?}", first_hit);
+    ///     }
+    /// }
+    /// ```
+    pub fn cast_ray_predicate(
+        &self,
+        origin: Vector,
+        direction: Vector,
+        max_time_of_impact: Scalar,
+        solid: bool,
+        query_filter: SpatialQueryFilter,
+        predicate: &dyn Fn(Entity) -> bool,
+    ) -> Option<RayHitData> {
+        self.query_pipeline.cast_ray_predicate(
+            origin,
+            direction,
+            max_time_of_impact,
+            solid,
+            query_filter,
+            predicate,
+        )
+    }
+
     /// Casts a [ray](spatial_query#raycasting) and computes all [hits](RayHitData) until `max_hits` is reached.
     ///
     /// Note that the order of the results is not guaranteed, and if there are more hits than `max_hits`,