1use std::{marker::PhantomData, ops::Range, sync::PoisonError};
17
18use bevy_asset::prelude::*;
19use bevy_core_pipeline::tonemapping::{
20 DebandDither, Tonemapping, TonemappingLuts, get_lut_bind_group_layout_entries, get_lut_bindings,
21};
22use bevy_ecs::{
23 entity::EntityHashMap,
24 prelude::*,
25 query::ROQueryItem,
26 system::{
27 StaticSystemParam, SystemParamItem, SystemState,
28 lifetimeless::{Read, SRes},
29 },
30};
31use bevy_image::BevyDefault;
32use bevy_render::{
33 Extract,
34 prelude::*,
35 render_asset::RenderAssets,
36 render_phase::{
37 DrawFunctions, PhaseItem, PhaseItemExtraIndex, RenderCommand, RenderCommandResult, SetItemPipeline,
38 TrackedRenderPass, ViewSortedRenderPhases,
39 },
40 render_resource::{
41 BindGroup, BindGroupEntry, BindGroupLayout, BindingResource, BlendState, Buffer, BufferAddress, BufferBinding,
42 BufferDescriptor, BufferId, BufferInitDescriptor, BufferSize, BufferUsages, ColorTargetState, ColorWrites,
43 CompareFunction, DepthBiasState, DepthStencilState, FragmentState, FrontFace, IndexFormat, MultisampleState,
44 PipelineCache, PolygonMode, PrimitiveState, PrimitiveTopology, RenderPipelineDescriptor, SamplerId, ShaderDefVal,
45 ShaderStages, ShaderType, SpecializedRenderPipeline, SpecializedRenderPipelines, StencilFaceState, StencilState,
46 TextureFormat, TextureViewId, VertexBufferLayout, VertexState, VertexStepMode, binding_types::uniform_buffer,
47 },
48 renderer::{RenderDevice, RenderQueue},
49 sync_world::MainEntity,
50 texture::{FallbackImage, GpuImage},
51 view::{ExtractedView, ViewTarget, ViewUniform, ViewUniformOffset, ViewUniforms},
52};
53use bevy_utils::default;
54use bytemuck::cast_slice;
55use fixedbitset::FixedBitSet;
56use vec_belt::VecBelt;
57
58use crate::vertex::{DrawItems, Vertex};
59
60#[derive(Resource)]
63pub struct VertexPipeline<T: Vertex> {
64 view_layout: BindGroupLayout,
65 vertex_prop: T::PipelineProp,
66}
67
68impl<T: Vertex> VertexPipeline<T> {
69 #[inline]
72 pub const fn vertex_prop(&self) -> &T::PipelineProp {
73 &self.vertex_prop
74 }
75}
76
77impl<T: Vertex> FromWorld for VertexPipeline<T> {
78 fn from_world(world: &mut World) -> Self {
79 let device = world.resource::<RenderDevice>();
80
81 let [lut_texture, lut_sampler] = get_lut_bind_group_layout_entries();
82 let view_layout = device.create_bind_group_layout("hephae_view_layout", &[
83 uniform_buffer::<ViewUniform>(true).build(0, ShaderStages::VERTEX_FRAGMENT),
84 lut_texture.build(1, ShaderStages::FRAGMENT),
85 lut_sampler.build(2, ShaderStages::FRAGMENT),
86 ]);
87
88 let mut state = SystemState::<T::PipelineParam>::new(world);
89 let vertex_prop = T::init_pipeline(state.get_mut(world));
90 state.apply(world);
91
92 Self {
93 view_layout,
94 vertex_prop,
95 }
96 }
97}
98
99#[derive(Resource)]
101pub struct PipelineShader<T: Vertex>(pub(crate) Handle<Shader>, PhantomData<fn() -> T>);
102impl<T: Vertex> PipelineShader<T> {
103 #[inline]
105 pub fn shader(&self) -> AssetId<Shader> {
106 self.0.id()
107 }
108}
109
110pub fn load_shader<T: Vertex>(mut commands: Commands, server: Res<AssetServer>) {
112 commands.insert_resource(PipelineShader::<T>(server.load(T::SHADER), PhantomData));
113}
114
115pub fn extract_shader<T: Vertex>(mut commands: Commands, shader: Extract<Option<Res<PipelineShader<T>>>>) {
118 if let Some(ref shader) = *shader {
119 if shader.is_changed() {
120 commands.insert_resource(PipelineShader::<T>(shader.0.clone_weak(), PhantomData));
121 }
122 }
123}
124
125#[derive(Eq, PartialEq, Hash, Copy, Clone)]
130pub struct ViewKey {
131 pub hdr: bool,
133 pub msaa: u8,
135 pub tonemapping: Option<Tonemapping>,
137 pub dither: bool,
139 pub shader: AssetId<Shader>,
142}
143
144impl<T: Vertex> SpecializedRenderPipeline for VertexPipeline<T> {
145 type Key = (ViewKey, T::PipelineKey);
146
147 fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
148 let (view_key, key) = key;
149 let mut defs = Vec::new();
150 if let Some(tonemapping) = view_key.tonemapping {
151 defs.extend([
152 "TONEMAP_IN_SHADER".into(),
153 ShaderDefVal::UInt("TONEMAPPING_LUT_TEXTURE_BINDING_INDEX".into(), 1),
154 ShaderDefVal::UInt("TONEMAPPING_LUT_SAMPLER_BINDING_INDEX".into(), 2),
155 match tonemapping {
156 Tonemapping::None => "TONEMAP_METHOD_NONE",
157 Tonemapping::Reinhard => "TONEMAP_METHOD_REINHARD",
158 Tonemapping::ReinhardLuminance => "TONEMAP_METHOD_REINHARD_LUMINANCE",
159 Tonemapping::AcesFitted => "TONEMAP_METHOD_ACES_FITTED",
160 Tonemapping::AgX => "TONEMAP_METHOD_AGX",
161 Tonemapping::SomewhatBoringDisplayTransform => "TONEMAP_METHOD_SOMEWHAT_BORING_DISPLAY_TRANSFORM",
162 Tonemapping::TonyMcMapface => "TONEMAP_METHOD_TONY_MC_MAPFACE",
163 Tonemapping::BlenderFilmic => "TONEMAP_METHOD_BLENDER_FILMIC",
164 }
165 .into(),
166 ]);
167
168 if view_key.dither {
169 defs.push("DEBAND_DITHER".into());
170 }
171 }
172
173 let format = match view_key.hdr {
174 true => ViewTarget::TEXTURE_FORMAT_HDR,
175 false => TextureFormat::bevy_default(),
176 };
177
178 let mut desc = RenderPipelineDescriptor {
179 label: Some("hephae_pipeline_descriptor".into()),
180 layout: [self.view_layout.clone()].into(),
181 push_constant_ranges: Vec::new(),
182 vertex: VertexState {
183 shader: Handle::Weak(view_key.shader),
184 shader_defs: defs.clone(),
185 entry_point: "vertex".into(),
186 buffers: [VertexBufferLayout {
187 array_stride: size_of::<T>() as BufferAddress,
188 step_mode: VertexStepMode::Vertex,
189 attributes: T::ATTRIBUTES.into(),
190 }]
191 .into(),
192 },
193 primitive: PrimitiveState {
194 topology: PrimitiveTopology::TriangleList,
195 strip_index_format: None,
196 front_face: FrontFace::Ccw,
197 cull_mode: None,
198 unclipped_depth: false,
199 polygon_mode: PolygonMode::Fill,
200 conservative: false,
201 },
202 depth_stencil: T::DEPTH_FORMAT.map(|format| DepthStencilState {
203 format,
204 depth_write_enabled: false,
205 depth_compare: CompareFunction::GreaterEqual,
206 stencil: StencilState {
207 front: StencilFaceState::IGNORE,
208 back: StencilFaceState::IGNORE,
209 read_mask: 0,
210 write_mask: 0,
211 },
212 bias: DepthBiasState {
213 constant: 0,
214 slope_scale: 0.,
215 clamp: 0.,
216 },
217 }),
218 multisample: MultisampleState {
219 count: 1 << view_key.msaa,
220 mask: !0,
221 alpha_to_coverage_enabled: false,
222 },
223 fragment: Some(FragmentState {
224 shader: Handle::Weak(view_key.shader),
225 shader_defs: defs,
226 entry_point: "fragment".into(),
227 targets: [Some(ColorTargetState {
228 format,
229 blend: Some(BlendState::ALPHA_BLENDING),
230 write_mask: ColorWrites::ALL,
231 })]
232 .into(),
233 }),
234 zero_initialize_workgroup_memory: false,
235 };
236
237 T::specialize_pipeline(key, &self.vertex_prop, &mut desc);
238 desc
239 }
240}
241
242#[derive(Resource)]
244pub struct DrawBuffers<T: Vertex> {
245 pub(crate) vertices: VecBelt<T>,
246 pub(crate) indices: VecBelt<u32>,
247 vertex_buffer: Buffer,
248}
249
250impl<T: Vertex> FromWorld for DrawBuffers<T> {
251 #[inline]
252 fn from_world(world: &mut World) -> Self {
253 let device = world.resource::<RenderDevice>();
254 Self {
255 vertices: VecBelt::new(4096),
256 indices: VecBelt::new(6144),
257 vertex_buffer: device.create_buffer(&BufferDescriptor {
258 label: Some("hephae_vertex_buffer"),
259 size: (4096 * size_of::<T>()) as BufferAddress,
260 usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
261 mapped_at_creation: false,
262 }),
263 }
264 }
265}
266
267#[derive(Component)]
269pub struct ViewBatches<T: Vertex>(pub EntityHashMap<(T::BatchProp, Range<u32>)>);
270impl<T: Vertex> Default for ViewBatches<T> {
271 #[inline]
272 fn default() -> Self {
273 Self(default())
274 }
275}
276
277#[derive(Component)]
279pub struct ViewBindGroup<T: Vertex> {
280 bind_group: BindGroup,
281 last_buffer: BufferId,
282 last_lut_texture: TextureViewId,
283 last_lut_sampler: SamplerId,
284 _marker: PhantomData<fn() -> T>,
285}
286
287#[derive(Component)]
289pub struct ViewIndexBuffer<T: Vertex> {
290 indices: Vec<u32>,
291 index_buffer: Option<Buffer>,
292 _marker: PhantomData<fn() -> T>,
293}
294
295impl<T: Vertex> Default for ViewIndexBuffer<T> {
296 #[inline]
297 fn default() -> Self {
298 Self {
299 indices: Vec::with_capacity(6144),
300 index_buffer: None,
301 _marker: PhantomData,
302 }
303 }
304}
305
306#[derive(Component)]
307pub(crate) struct VisibleDrawers<T: Vertex>(pub VecBelt<Entity>, PhantomData<fn() -> T>);
308impl<T: Vertex> Default for VisibleDrawers<T> {
309 #[inline]
310 fn default() -> Self {
311 Self(VecBelt::new(1024), PhantomData)
312 }
313}
314
315pub(crate) fn queue_vertices<T: Vertex>(
316 draw_functions: Res<DrawFunctions<T::Item>>,
317 pipeline: Res<VertexPipeline<T>>,
318 shader: Res<PipelineShader<T>>,
319 mut pipelines: ResMut<SpecializedRenderPipelines<VertexPipeline<T>>>,
320 pipeline_cache: Res<PipelineCache>,
321 mut transparent_phases: ResMut<ViewSortedRenderPhases<T::Item>>,
322 mut views: Query<(
323 Entity,
324 &mut VisibleDrawers<T>,
325 &ExtractedView,
326 &Msaa,
327 Option<&Tonemapping>,
328 Option<&DebandDither>,
329 )>,
330 mut items: Query<(Entity, &MainEntity, &mut DrawItems<T>)>,
331 mut iterated: Local<FixedBitSet>,
332) {
333 let draw_function = draw_functions.read().id::<DrawRequests<T>>();
334 for item in &mut views {
335 let (view_entity, mut visible_drawers, view, &msaa, tonemapping, dither): (
336 Entity,
337 Mut<VisibleDrawers<T>>,
338 &ExtractedView,
339 &Msaa,
340 Option<&Tonemapping>,
341 Option<&DebandDither>,
342 ) = item;
343
344 let Some(transparent_phase) = transparent_phases.get_mut(&view_entity) else {
345 continue;
346 };
347
348 let view_key = ViewKey {
349 hdr: view.hdr,
350 msaa: msaa.samples().trailing_zeros() as u8,
351 tonemapping: (!view.hdr).then_some(tonemapping.copied()).flatten(),
352 dither: !view.hdr && dither.copied().unwrap_or_default() == DebandDither::Enabled,
353 shader: shader.0.id(),
354 };
355
356 iterated.clear();
357 visible_drawers.0.clear(|entities| {
358 let mut iter = items.iter_many_mut(entities);
359 while let Some((e, &main_e, mut items)) = iter.fetch_next() {
360 let index = e.index() as usize;
361 if iterated[index] {
362 continue;
363 }
364
365 iterated.grow_and_insert(index);
366 transparent_phase.items.extend(
367 items
368 .0
369 .get_mut()
370 .unwrap_or_else(PoisonError::into_inner)
371 .iter_mut()
372 .enumerate()
373 .map(|(i, &mut (.., layer, ref key))| {
374 T::create_item(
375 layer,
376 (e, main_e),
377 pipelines.specialize(&pipeline_cache, &pipeline, (view_key, key.clone())),
378 draw_function,
379 i,
380 )
381 }),
382 );
383 }
384 });
385 }
386}
387
388pub(crate) fn prepare_indices<T: Vertex>(
389 mut param_set: ParamSet<(
390 (
391 Res<RenderDevice>,
392 Res<RenderQueue>,
393 ResMut<DrawBuffers<T>>,
394 ResMut<ViewSortedRenderPhases<T::Item>>,
395 Query<(Entity, &mut ViewIndexBuffer<T>)>,
396 Query<&mut DrawItems<T>>,
397 ),
398 StaticSystemParam<T::BatchParam>,
399 Query<&mut ViewBatches<T>>,
400 )>,
401 mut batched_entities: Local<Vec<(Entity, Entity, T::PipelineKey, Range<u32>)>>,
402 mut batched_results: Local<Vec<(Entity, Entity, T::BatchProp, Range<u32>)>>,
403) {
404 let (device, queue, buffers, mut transparent_phases, mut views, mut items): (
405 Res<RenderDevice>,
406 Res<RenderQueue>,
407 ResMut<DrawBuffers<T>>,
408 ResMut<ViewSortedRenderPhases<T::Item>>,
409 Query<(Entity, &mut ViewIndexBuffer<T>)>,
410 Query<&mut DrawItems<T>>,
411 ) = param_set.p0();
412
413 batched_entities.clear();
414
415 let buffers = buffers.into_inner();
416 buffers.vertices.clear(|vertices| {
417 let contents = cast_slice::<T, u8>(&vertices);
418 if (buffers.vertex_buffer.size() as usize) < contents.len() {
419 buffers.vertex_buffer = device.create_buffer_with_data(&BufferInitDescriptor {
420 label: Some("hephae_vertex_buffer"),
421 contents,
422 usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
423 });
424 } else if let Some(len) = BufferSize::new(contents.len() as u64) {
425 queue
426 .write_buffer_with(&buffers.vertex_buffer, 0, len)
427 .unwrap()
428 .copy_from_slice(contents);
429 }
430 });
431
432 buffers.indices.clear(|indices| {
433 for (view_entity, view_indices) in &mut views {
434 let Some(transparent_phase) = transparent_phases.get_mut(&view_entity) else {
435 continue;
436 };
437
438 let view_indices = view_indices.into_inner();
439 view_indices.indices.clear();
440
441 let mut batch_item_index = 0;
442 let mut batch_index_range = 0;
443 let mut batch_key = None::<T::PipelineKey>;
444
445 for item_index in 0..transparent_phase.items.len() {
446 let item = &mut transparent_phase.items[item_index];
447 let Ok(mut items) = items.get_mut(item.entity()) else {
448 batch_key = None;
449 continue;
450 };
451
452 let Some((range, .., key)) =
453 items.0.get_mut().unwrap_or_else(PoisonError::into_inner).get(
454 std::mem::replace(item.batch_range_and_extra_index_mut().1, PhaseItemExtraIndex::NONE).0 as usize,
455 )
456 else {
457 continue;
458 };
459
460 view_indices.indices.extend(&indices[range.clone()]);
461 if match batch_key {
462 None => true,
463 Some(ref batch_key) => batch_key != key,
464 } {
465 batch_item_index = item_index;
466 batched_entities.push((view_entity, item.entity(), key.clone(), batch_index_range..batch_index_range));
467 }
468
469 batch_index_range = view_indices.indices.len() as u32;
470 transparent_phase.items[batch_item_index].batch_range_mut().end += 1;
471 batched_entities.last_mut().unwrap().3.end = batch_index_range;
472
473 batch_key = Some(key.clone());
474 }
475
476 let contents = cast_slice::<u32, u8>(&view_indices.indices);
477 if view_indices
478 .index_buffer
479 .as_ref()
480 .is_none_or(|index_buffer| (index_buffer.size() as usize) < contents.len())
481 {
482 view_indices.index_buffer = Some(device.create_buffer_with_data(&BufferInitDescriptor {
483 label: Some("hephae_index_buffer"),
484 contents,
485 usage: BufferUsages::INDEX | BufferUsages::COPY_DST,
486 }));
487 } else if let Some(len) = BufferSize::new(contents.len() as u64) {
488 queue
489 .write_buffer_with(view_indices.index_buffer.as_ref().unwrap(), 0, len)
490 .unwrap()
491 .copy_from_slice(contents);
492 }
493 }
494 });
495
496 for mut item in &mut items {
497 item.0.get_mut().unwrap_or_else(PoisonError::into_inner).clear();
498 }
499
500 let mut param = param_set.p1();
501 batched_results.extend(batched_entities.drain(..).map(|(view_entity, batch_entity, key, range)| {
502 (view_entity, batch_entity, T::create_batch(&mut param, key), range)
503 }));
504
505 drop(param);
506
507 let mut batches = param_set.p2();
508 for mut view_batches in &mut batches {
509 view_batches.0.clear();
510 }
511
512 for (view_entity, batch_entity, prop, range) in batched_results.drain(..) {
513 let Ok(mut view_batches) = batches.get_mut(view_entity) else {
514 continue;
515 };
516
517 view_batches.0.insert(batch_entity, (prop, range));
518 }
519}
520
521pub(crate) fn prepare_view_bind_groups<T: Vertex>(
523 mut commands: Commands,
524 pipeline: Res<VertexPipeline<T>>,
525 render_device: Res<RenderDevice>,
526 view_uniforms: Res<ViewUniforms>,
527 mut views: Query<(Entity, &Tonemapping, Option<&mut ViewBindGroup<T>>)>,
528 tonemapping_luts: Res<TonemappingLuts>,
529 images: Res<RenderAssets<GpuImage>>,
530 fallback_image: Res<FallbackImage>,
531) {
532 let Some(buffer) = view_uniforms.uniforms.buffer() else {
533 return;
534 };
535
536 let view_binding = BindingResource::Buffer(BufferBinding {
537 buffer,
538 offset: 0,
539 size: Some(ViewUniform::min_size()),
540 });
541
542 for (entity, &tonemapping, bind_group) in &mut views {
543 let (lut_texture, lut_sampler) = get_lut_bindings(&images, &tonemapping_luts, &tonemapping, &fallback_image);
544 let create_bind_group = || ViewBindGroup::<T> {
545 bind_group: render_device.create_bind_group("hephae_view_bind_group", &pipeline.view_layout, &[
546 BindGroupEntry {
547 binding: 0,
548 resource: view_binding.clone(),
549 },
550 BindGroupEntry {
551 binding: 1,
552 resource: BindingResource::TextureView(lut_texture),
553 },
554 BindGroupEntry {
555 binding: 2,
556 resource: BindingResource::Sampler(lut_sampler),
557 },
558 ]),
559 last_buffer: buffer.id(),
560 last_lut_texture: lut_texture.id(),
561 last_lut_sampler: lut_sampler.id(),
562 _marker: PhantomData,
563 };
564
565 if let Some(mut bind_group) = bind_group {
566 if bind_group.last_buffer != buffer.id() ||
567 bind_group.last_lut_texture != lut_texture.id() ||
568 bind_group.last_lut_sampler != lut_sampler.id()
569 {
570 *bind_group = create_bind_group();
571 }
572 } else {
573 commands.entity(entity).insert(create_bind_group());
574 }
575 }
576}
577
578pub type DrawRequests<T> = (
580 SetItemPipeline,
581 SetViewBindGroup<T, 0>,
582 <T as Vertex>::RenderCommand,
583 DrawBatch<T>,
584);
585
586pub struct SetViewBindGroup<T: Vertex, const I: usize>(PhantomData<fn() -> T>);
588impl<P: PhaseItem, T: Vertex, const I: usize> RenderCommand<P> for SetViewBindGroup<T, I> {
589 type Param = ();
590 type ViewQuery = (Read<ViewUniformOffset>, Read<ViewBindGroup<T>>);
591 type ItemQuery = ();
592
593 #[inline]
594 fn render<'w>(
595 _: &P,
596 (view_uniform, view_bind_group): ROQueryItem<'w, Self::ViewQuery>,
597 _: Option<ROQueryItem<'w, Self::ItemQuery>>,
598 _: SystemParamItem<'w, '_, Self::Param>,
599 pass: &mut TrackedRenderPass<'w>,
600 ) -> RenderCommandResult {
601 pass.set_bind_group(I, &view_bind_group.bind_group, &[view_uniform.offset]);
602 RenderCommandResult::Success
603 }
604}
605
606pub struct DrawBatch<T: Vertex>(PhantomData<fn() -> T>);
608impl<P: PhaseItem, T: Vertex> RenderCommand<P> for DrawBatch<T> {
609 type Param = SRes<DrawBuffers<T>>;
610 type ViewQuery = (Read<ViewBatches<T>>, Read<ViewIndexBuffer<T>>);
611 type ItemQuery = ();
612
613 #[inline]
614 fn render<'w>(
615 item: &P,
616 (batches, index_buffer): ROQueryItem<'w, Self::ViewQuery>,
617 _: Option<ROQueryItem<'w, Self::ItemQuery>>,
618 buffers: SystemParamItem<'w, '_, Self::Param>,
619 pass: &mut TrackedRenderPass<'w>,
620 ) -> RenderCommandResult {
621 let Some((.., range)) = batches.0.get(&item.entity()) else {
622 return RenderCommandResult::Skip;
623 };
624
625 let Some(ref index_buffer) = index_buffer.index_buffer else {
626 return RenderCommandResult::Skip;
627 };
628
629 let buffers = buffers.into_inner();
630 pass.set_vertex_buffer(0, buffers.vertex_buffer.slice(..));
631 pass.set_index_buffer(index_buffer.slice(..), 0, IndexFormat::Uint32);
632 pass.draw_indexed(range.clone(), 0, 0..1);
633
634 RenderCommandResult::Success
635 }
636}