`bevy_mod_picking` raycast backend with support for multiple raycasting sets

picking_raycast_multiple.rs

// This is a rough & dirty fork of `bevy_picking_raycast` that supports multiple raycast sets.
// This is needed when you have multiple cameras you want to raycast from at the same time.
// The key change is making the `Backend` generic for raycast sets.  It's harder to work
// with as it requires an understanding of `bevy_mod_raycast`.
//
// Other changes: removed `Picking` component management; added enable/disable via config

//! A raycasting backend for `bevy_mod_picking` that uses `bevy_mod_raycast` for raycasting.

use bevy::{prelude::*, utils::HashMap, window::PrimaryWindow};
use bevy_mod_picking::picking_core::backend::{prelude::*, PickingBackend};
use bevy_mod_raycast::{Ray3d, RaycastSource};

use std::marker::PhantomData;

pub trait PickingSet: 'static + Send + Sync + Reflect + Clone {}
impl<T: 'static + Send + Sync + Reflect + Clone> PickingSet for T {}

#[derive(Clone)]
pub struct Backend<T: PickingSet> {
    enabled: bool,
    phantom: PhantomData<T>,
}

impl<T: PickingSet> Backend<T> {
    pub fn disabled() -> Self {
        Self {
            enabled: false,
            phantom: PhantomData,
        }
    }
}

impl<T: PickingSet> Default for Backend<T> {
    fn default() -> Self {
        Self {
            enabled: true,
            phantom: PhantomData,
        }
    }
}

impl<T: PickingSet> PickingBackend for Backend<T> {}

impl<T: PickingSet> bevy::app::Plugin for Backend<T> {
    fn build(&self, app: &mut App) {
        app.insert_resource(Config::<T> {
            enabled: self.enabled,
            phantom: PhantomData,
        })
        .add_systems(
            (
                build_rays_from_pointers::<T>.run_if(enabled::<T>),
                spawn_raycast_sources::<T>.run_if(enabled::<T>),
            )
                .chain()
                .in_set(PickSet::PostInput),
        )
        .add_systems(
            (
                bevy_mod_raycast::update_raycast::<T>.run_if(enabled::<T>),
                update_hits::<T>.run_if(enabled::<T>),
            )
                .chain()
                .in_set(PickSet::Backend),
        );
    }
}

#[derive(Resource, Debug)]
pub struct Config<T: PickingSet> {
    pub enabled: bool,
    phantom: PhantomData<T>,
}

/// check to see if picking is enabled
fn enabled<T: PickingSet>(config: Res<Config<T>>) -> bool {
    config.enabled
}

/// Marks a camera that should be used for picking with [`bevy_mod_raycast`].
#[derive(Debug, Default, Clone, Component, Reflect)]
pub struct RaycastPickCamera<T: PickingSet> {
    #[reflect(ignore)]
    /// Maps the pointers visible to this [`RaycastPickCamera`] to their corresponding ray. We need
    /// to create a map because many pointers may be visible to this camera.
    ray_map: HashMap<PointerId, Ray3d>,
    phantom: PhantomData<T>,
}

// --
//
// TODO:
//
// The following design, where we need to add children to the cameras, only exists because
// `bevy_mod_raycast` only supports raycasting via components. Ideally, we would be able to run
// raycasts on demand without needing to supply them as components on entities.
//
// --

/// Builds rays and updates raycasting [`RaycastPickCamera`]s from [`PointerLocation`]s.
pub fn build_rays_from_pointers<T: PickingSet>(
    pointers: Query<(&PointerId, &PointerLocation)>,
    windows: Query<&Window>,
    primary_window: Query<Entity, With<PrimaryWindow>>,
    images: Res<Assets<Image>>,
    mut picking_cameras: Query<(&Camera, &GlobalTransform, &mut RaycastPickCamera<T>)>,
) {
    picking_cameras.iter_mut().for_each(|(_, _, mut pick_cam)| {
        pick_cam.ray_map.clear();
    });

    for (pointer_id, pointer_location) in &pointers {
        let pointer_location = match pointer_location.location() {
            Some(l) => l,
            None => continue,
        };
        picking_cameras
            .iter_mut()
            .filter(|(camera, _, _)| {
                pointer_location.is_in_viewport(camera, &windows, &primary_window, &images)
            })
            .for_each(|(camera, transform, mut source)| {
                let pointer_pos = pointer_location.position;
                if let Some(ray) = Ray3d::from_screenspace(pointer_pos, camera, transform) {
                    source.ray_map.insert(*pointer_id, ray);
                }
            });
    }
}

/// A newtype, used solely to mark the [`RaycastSource`] children on the [`RaycastPickCamera`] so we
/// know what pointer they are associated with.
#[derive(Component)]
struct PointerMarker(PointerId);

/// Using the rays in each [`RaycastPickCamera`], updates their child [`RaycastSource`]s.
pub fn spawn_raycast_sources<T: PickingSet>(
    mut commands: Commands,
    picking_cameras: Query<(Entity, &RaycastPickCamera<T>)>,
    child_sources: Query<Entity, With<RaycastSource<T>>>,
) {
    child_sources
        .iter()
        .for_each(|pick_source| commands.entity(pick_source).despawn_recursive());

    picking_cameras.iter().for_each(|(entity, pick_cam)| {
        pick_cam.ray_map.iter().for_each(|(pointer, ray)| {
            let mut new_source = RaycastSource::<T>::default();
            new_source.ray = Some(*ray);
            let pointer_marker = PointerMarker(*pointer);
            let new_child = commands.spawn((new_source, pointer_marker)).id();
            commands.entity(entity).add_child(new_child);
        })
    })
}

/// Produces [`PointerHits`]s from [`RaycastSource`] intersections.
fn update_hits<T: PickingSet>(
    pick_cameras: Query<(Entity, &Camera), With<RaycastPickCamera<T>>>,
    mut pick_sources: Query<(&PointerMarker, &RaycastSource<T>, &Parent)>,
    mut output_events: EventWriter<PointerHits>,
) {
    pick_sources
        .iter_mut()
        .filter_map(|(pointer, pick_source, parent)| {
            pick_cameras
                .get(parent.get())
                .map(|(entity, camera)| (pointer, pick_source, entity, camera))
                .ok()
        })
        .for_each(|(pointer_marker, pick_source, cam_entity, camera)| {
            let under_cursor: Vec<(Entity, HitData)> = pick_source
                .intersections()
                .iter()
                .map(|(entity, intersection)| {
                    (
                        *entity,
                        HitData {
                            camera: cam_entity,
                            depth: intersection.distance(),
                            position: Some(intersection.position()),
                            normal: Some(intersection.normal()),
                        },
                    )
                })
                .collect();

            if !under_cursor.is_empty() {
                output_events.send(PointerHits {
                    pointer: pointer_marker.0,
                    picks: under_cursor,
                    order: camera.order,
                });
            }
        });
}