use std::marker::PhantomData;
use bevy_app::{App, Plugin};
use bevy_asset::{Asset, AssetId, Handle};
use bevy_derive::{Deref, DerefMut};
use bevy_ecs::{
entity::EntityHashMap,
prelude::Entity,
query::{QueryFilter, QueryItem, ReadOnlyQueryData},
system::{lifetimeless::Read, Query, ResMut, Resource},
};
use crate::{prelude::ViewVisibility, Extract, ExtractSchedule, RenderApp};
pub trait ExtractInstance: Send + Sync + Sized + 'static {
type QueryData: ReadOnlyQueryData;
type QueryFilter: QueryFilter;
fn extract(item: QueryItem<'_, Self::QueryData>) -> Option<Self>;
}
#[derive(Default)]
pub struct ExtractInstancesPlugin<EI>
where
EI: ExtractInstance,
{
only_extract_visible: bool,
marker: PhantomData<fn() -> EI>,
}
#[derive(Resource, Deref, DerefMut)]
pub struct ExtractedInstances<EI>(EntityHashMap<EI>)
where
EI: ExtractInstance;
impl<EI> Default for ExtractedInstances<EI>
where
EI: ExtractInstance,
{
fn default() -> Self {
Self(Default::default())
}
}
impl<EI> ExtractInstancesPlugin<EI>
where
EI: ExtractInstance,
{
pub fn new() -> Self {
Self {
only_extract_visible: false,
marker: PhantomData,
}
}
pub fn extract_visible() -> Self {
Self {
only_extract_visible: true,
marker: PhantomData,
}
}
}
impl<EI> Plugin for ExtractInstancesPlugin<EI>
where
EI: ExtractInstance,
{
fn build(&self, app: &mut App) {
if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
render_app.init_resource::<ExtractedInstances<EI>>();
if self.only_extract_visible {
render_app.add_systems(ExtractSchedule, extract_visible::<EI>);
} else {
render_app.add_systems(ExtractSchedule, extract_all::<EI>);
}
}
}
}
fn extract_all<EI>(
mut extracted_instances: ResMut<ExtractedInstances<EI>>,
query: Extract<Query<(Entity, EI::QueryData), EI::QueryFilter>>,
) where
EI: ExtractInstance,
{
extracted_instances.clear();
for (entity, other) in &query {
if let Some(extract_instance) = EI::extract(other) {
extracted_instances.insert(entity, extract_instance);
}
}
}
fn extract_visible<EI>(
mut extracted_instances: ResMut<ExtractedInstances<EI>>,
query: Extract<Query<(Entity, &ViewVisibility, EI::QueryData), EI::QueryFilter>>,
) where
EI: ExtractInstance,
{
extracted_instances.clear();
for (entity, view_visibility, other) in &query {
if view_visibility.get() {
if let Some(extract_instance) = EI::extract(other) {
extracted_instances.insert(entity, extract_instance);
}
}
}
}
impl<A> ExtractInstance for AssetId<A>
where
A: Asset,
{
type QueryData = Read<Handle<A>>;
type QueryFilter = ();
fn extract(item: QueryItem<'_, Self::QueryData>) -> Option<Self> {
Some(item.id())
}
}