use core::cmp::Ordering;
use arrayvec::ArrayVec;
use hashbrown::HashMap;
use unlit_ecs::{Has, Scope};
use unlit_wgpu::instance_stream::InstanceBuffer;
use unlit_wgpu::resources::{ResHandle, Virtual};
use unlit_wgpu::specialize::{PipelineVariant, SurfaceKey, Variants};
use unlit_wgpu::pipeline::{GLOBAL_GROUP, MATERIAL_GROUP, MESH_GROUP};
use unlit_wgpu::scene::{DrawEntry, DrawRange, MAX_VERTEX_BUFFERS, Scene, VertexBufferBinding};
use crate::bounds::FrustumPlanes;
use crate::components::{Camera, GpuMaterial, GpuMesh, GpuRenderPipeline, ZSortedDrawing};
use crate::culling::{VisibleMesh, collect_visible};
use crate::pipeline::{
DrawContext, FamilyContext, GlobalResources, InstanceContext, InstanceData,
RegisteredRenderPipeline, RenderPipelineFactory, RenderPipelineId, RenderPipelineKey,
};
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub(crate) struct FamilyId(u32);
impl FamilyId {
pub(crate) fn new(index: u32) -> Self {
Self(index)
}
pub(crate) fn as_usize(self) -> usize {
self.0 as usize
}
}
pub(crate) struct Family<K, F, I>
where
K: RenderPipelineKey,
F: RenderPipelineFactory<<K::Variant as PipelineVariant<wgpu::RenderPipeline>>::Descriptor>,
I: InstanceData,
{
variants: Variants<wgpu::RenderPipeline, K::Variant>,
factory: F,
registered: Vec<RenderPipelineId>,
instances: I,
stream: InstanceBuffer,
records: Vec<u8>,
}
impl<K, F, I> Family<K, F, I>
where
K: RenderPipelineKey,
F: RenderPipelineFactory<<K::Variant as PipelineVariant<wgpu::RenderPipeline>>::Descriptor>,
I: InstanceData,
{
pub(crate) fn new(device: &wgpu::Device, factory: F, instances: I) -> Self {
Self {
variants: Variants::new(device),
factory,
registered: Vec::new(),
stream: InstanceBuffer::new(instances.stream()),
instances,
records: Vec::new(),
}
}
}
pub(crate) struct VisibleEntry {
pub(crate) mesh: VisibleMesh,
pub(crate) family: FamilyId,
pub(crate) pipeline_id: RenderPipelineId,
pub(crate) instance_index: u32,
pub(crate) sort_key: u64,
pub(crate) depth: f32,
pub(crate) z_sorted: bool,
pub(crate) handles_key: DrawHandlesKey,
}
pub(crate) struct FamilyFrame<'a> {
pub(crate) world: Scope<'a>,
pub(crate) meshes: &'a [VisibleMesh],
pub(crate) camera: &'a Camera,
pub(crate) surface: SurfaceKey,
pub(crate) device: &'a wgpu::Device,
pub(crate) resources: GlobalResources,
}
struct Resolver<'a, 'f, K, F>
where
K: RenderPipelineKey,
F: RenderPipelineFactory<<K::Variant as PipelineVariant<wgpu::RenderPipeline>>::Descriptor>,
{
family: FamilyId,
frame: &'a FamilyFrame<'f>,
variants: &'a mut Variants<wgpu::RenderPipeline, K::Variant>,
factory: &'a F,
registered: &'a mut Vec<RenderPipelineId>,
visible: &'a mut Vec<VisibleEntry>,
register: &'a mut dyn FnMut(RegisteredRenderPipeline) -> RenderPipelineId,
}
impl<K, F> Resolver<'_, '_, K, F>
where
K: RenderPipelineKey,
F: RenderPipelineFactory<<K::Variant as PipelineVariant<wgpu::RenderPipeline>>::Descriptor>,
{
fn collect(&mut self) {
profiling::scope!("scene.resolve.family");
let world = &self.frame.world;
for &candidate in self.frame.meshes {
let Ok((pipeline, gpu_mesh, material, z_sorted)) = world.query_one::<(
&GpuRenderPipeline<K>,
&GpuMesh,
Option<&GpuMaterial>,
Has<ZSortedDrawing>,
)>(candidate.entity) else {
continue;
};
let (sort_key, material_bg) = match &material {
Some(material) => (
material.bind_group_id.index() as u64,
Some(material.bind_group_id.clone()),
),
None => (0, None),
};
self.push(
candidate,
&pipeline,
&gpu_mesh,
material.as_deref(),
sort_key,
material_bg,
z_sorted,
);
}
}
#[expect(
clippy::too_many_arguments,
reason = "one field per resolved draw component"
)]
fn push(
&mut self,
mesh: VisibleMesh,
pipeline: &GpuRenderPipeline<K>,
gpu_mesh: &GpuMesh,
material: Option<&GpuMaterial>,
sort_key: u64,
material_bg: Option<ResHandle<wgpu::BindGroup>>,
z_sorted: bool,
) {
let draw = DrawContext {
surface: self.frame.surface,
device: self.frame.device,
world: self.frame.world,
entity: mesh.entity,
mesh: gpu_mesh,
material,
};
let variant = pipeline.key().variant(&draw);
let ordinal = self.variants.specialize(variant);
let pipeline_id = match self.registered.get(ordinal as usize) {
Some(&id) => id,
None => {
let context = FamilyContext {
device: self.frame.device,
resources: self.frame.resources.clone(),
};
let desc = self
.factory
.descriptor(&context, self.variants.get(ordinal));
let id = (self.register)(desc);
self.registered.push(id);
id
}
};
let centre = mesh.world_from_local.transform_point3(gpu_mesh.aabb.center);
let depth = self.frame.camera.view_depth(centre);
let handles_key = DrawHandlesKey {
mesh: gpu_mesh.parts.root.clone(),
shape: DrawShape {
indexed: gpu_mesh.indexed,
count: gpu_mesh.count,
first: gpu_mesh.first,
base_vertex: gpu_mesh.base_vertex,
},
material: material_bg,
};
self.visible.push(VisibleEntry {
mesh,
family: self.family,
pipeline_id,
instance_index: 0,
sort_key,
depth,
z_sorted,
handles_key,
});
}
}
pub(crate) trait AnyFamily {
fn collect_and_resolve(
&mut self,
family: FamilyId,
frame: &FamilyFrame<'_>,
visible: &mut Vec<VisibleEntry>,
register: &mut dyn FnMut(RegisteredRenderPipeline) -> RenderPipelineId,
);
fn begin_instances(&mut self);
fn push_instance(&mut self, context: &mut InstanceContext<'_>) -> u32;
fn upload_instances(&mut self, device: &wgpu::Device, encoder: &mut wgpu::CommandEncoder);
fn instance_binding(&self) -> Option<(u32, wgpu::Buffer)>;
}
impl<K, F, I> AnyFamily for Family<K, F, I>
where
K: RenderPipelineKey,
F: RenderPipelineFactory<<K::Variant as PipelineVariant<wgpu::RenderPipeline>>::Descriptor>,
I: InstanceData,
{
fn collect_and_resolve(
&mut self,
family: FamilyId,
frame: &FamilyFrame<'_>,
visible: &mut Vec<VisibleEntry>,
register: &mut dyn FnMut(RegisteredRenderPipeline) -> RenderPipelineId,
) {
Resolver::<K, F> {
family,
frame,
variants: &mut self.variants,
factory: &self.factory,
registered: &mut self.registered,
visible,
register,
}
.collect();
}
fn begin_instances(&mut self) {
self.records.clear();
}
fn push_instance(&mut self, context: &mut InstanceContext<'_>) -> u32 {
let stride = self.instances.stream().array_stride as usize;
if stride == 0 {
return 0;
}
let index = (self.records.len() / stride) as u32;
let start = self.records.len();
self.records.resize(start + stride, 0);
self.instances.write(context, &mut self.records[start..]);
index
}
fn upload_instances(&mut self, device: &wgpu::Device, encoder: &mut wgpu::CommandEncoder) {
self.stream.upload(device, encoder, &self.records);
}
fn instance_binding(&self) -> Option<(u32, wgpu::Buffer)> {
let desc = self.instances.stream();
if desc.array_stride == 0 {
return None;
}
self.stream
.buffer()
.map(|buffer| (desc.slot, buffer.clone()))
}
}
pub(crate) struct SceneFrame<'a> {
pub(crate) families: &'a mut [Box<dyn AnyFamily>],
pub(crate) meshes: &'a mut Vec<VisibleMesh>,
pub(crate) visible: &'a mut Vec<VisibleEntry>,
pub(crate) register: &'a mut dyn FnMut(RegisteredRenderPipeline) -> RenderPipelineId,
}
pub(crate) fn collect_and_sort_visible(
frame: SceneFrame<'_>,
scope: Scope<'_>,
camera: &Camera,
surface: SurfaceKey,
device: &wgpu::Device,
resources: GlobalResources,
) {
let SceneFrame {
families,
meshes,
visible,
register,
} = frame;
let frustum = FrustumPlanes::from_clip_from_world(camera.clip_from_world());
{
profiling::scope!("scene.cull");
collect_visible(&scope, &frustum, meshes);
}
visible.clear();
let frame = FamilyFrame {
world: scope,
meshes,
camera,
surface,
device,
resources,
};
{
profiling::scope!("scene.resolve");
for (index, family) in families.iter_mut().enumerate() {
family.collect_and_resolve(FamilyId::new(index as u32), &frame, visible, register);
}
}
profiling::scope!("scene.sort");
visible.sort_unstable_by(|a, b| match (a.z_sorted, b.z_sorted) {
(false, false) => a
.pipeline_id
.cmp(&b.pipeline_id)
.then_with(|| a.sort_key.cmp(&b.sort_key)),
(true, true) => {
b.depth.partial_cmp(&a.depth).unwrap_or(Ordering::Equal)
}
(false, true) => Ordering::Less,
(true, false) => Ordering::Greater,
});
}
#[derive(Clone, PartialEq, Eq, Hash)]
pub(crate) struct DrawHandlesKey {
pub(crate) mesh: ResHandle<Virtual>,
pub(crate) shape: DrawShape,
pub(crate) material: Option<ResHandle<wgpu::BindGroup>>,
}
pub(crate) struct RenderPipelineHandles {
pub(crate) pipeline: wgpu::RenderPipeline,
pub(crate) global: Option<wgpu::BindGroup>,
}
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub(crate) struct DrawShape {
pub(crate) indexed: bool,
pub(crate) count: u32,
pub(crate) first: u32,
pub(crate) base_vertex: u32,
}
pub(crate) struct EntryHandles {
pub(crate) mesh_bg: Option<wgpu::BindGroup>,
pub(crate) material_bg: Option<wgpu::BindGroup>,
pub(crate) vertex_buffers: ArrayVec<VertexBufferBinding, MAX_VERTEX_BUFFERS>,
pub(crate) index_buffer: Option<(wgpu::Buffer, wgpu::IndexFormat)>,
}
fn same_draw(a: &VisibleEntry, b: &VisibleEntry) -> bool {
!a.z_sorted && !b.z_sorted && a.pipeline_id == b.pipeline_id && a.handles_key == b.handles_key
}
pub(crate) fn assemble_scene(
scene: &mut Scene,
visible: &[VisibleEntry],
pipelines: &[RenderPipelineHandles],
handles: &HashMap<DrawHandlesKey, EntryHandles>,
instances: &[Option<(u32, wgpu::Buffer)>],
) {
let mut start = 0;
while start < visible.len() {
let mut end = start + 1;
while end < visible.len() && same_draw(&visible[end - 1], &visible[end]) {
end += 1;
}
let entry = &visible[start];
let shape = entry.handles_key.shape;
let handle = handles
.get(&entry.handles_key)
.expect("every visible entry's handles were interned");
let pipeline = &pipelines[entry.pipeline_id.as_usize()];
let instance = &instances[entry.family.as_usize()];
let base = entry.instance_index;
let count = (end - start) as u32;
let first = shape.first;
let range = if shape.indexed {
DrawRange::indexed(first..first + shape.count)
.with_base_vertex(shape.base_vertex as i32)
.with_instances(base..base + count)
} else {
DrawRange::vertices(first..first + shape.count).with_instances(base..base + count)
};
let mut draw = DrawEntry::new(&pipeline.pipeline, range);
if let Some(global) = &pipeline.global {
draw = draw.with_bind_group(GLOBAL_GROUP, global);
}
if let Some(material_bg) = &handle.material_bg {
draw = draw.with_bind_group(MATERIAL_GROUP, material_bg);
}
if let Some(mesh_bg) = &handle.mesh_bg {
draw = draw.with_bind_group(MESH_GROUP, mesh_bg);
}
for (slot, buffer, range) in &handle.vertex_buffers {
draw = draw.with_vertex_buffer_range(*slot, buffer, range.clone());
}
if let Some((buffer, format)) = &handle.index_buffer {
draw = draw.with_index_buffer(buffer, *format);
}
if let Some((slot, buffer)) = instance {
draw = draw.with_vertex_buffer(*slot, buffer);
}
scene.draws.push(draw);
start = end;
}
}