1use std::{
2 collections::HashMap,
3 hash::{DefaultHasher, Hash, Hasher},
4};
5
6use crate::utils::ArcRef;
7
8use super::{
9 manager::{GraphicsPipelineDesc, VertexAttributeLayout},
10 super::{
11 GPUInner,
12 texture::{Texture, TextureSampler, BlendState},
13 shader::{
14 bind_group_manager::BindGroupCreateInfo,
15 GraphicsShader,
16 IndexBufferSize,
17 ShaderBindingType,
18 ShaderCullMode,
19 ShaderFrontFace,
20 ShaderPollygonMode,
21 ShaderTopology,
22 graphics::GraphicsShaderType,
23 types::ShaderReflect
24 },
25 buffer::Buffer,
26 command::{
27 BindGroupAttachment,
28 utils::BindGroupType,
29 renderpass::IntermediateRenderPipeline,
30 },
31 }
32};
33
34#[derive(Debug, Clone, Hash)]
35pub struct RenderPipeline {
36 pub(crate) bind_group: Vec<(u32, wgpu::BindGroup)>,
37 pub(crate) pipeline_desc: GraphicsPipelineDesc,
38 pub(crate) index_format: Option<IndexBufferSize>,
39}
40
41#[derive(Debug, Clone)]
42pub struct RenderPipelineBuilder {
43 pub(crate) gpu: ArcRef<GPUInner>,
44 pub(crate) attachments: Vec<BindGroupAttachment>,
45 pub(crate) shader: Option<IntermediateRenderPipeline>,
46 pub(crate) blend: Option<wgpu::BlendState>,
47 pub(crate) color_write_mask: Option<wgpu::ColorWrites>,
48 pub(crate) shader_reflection: Option<Vec<ShaderReflect>>,
49}
50
51impl RenderPipelineBuilder {
52 pub(crate) fn new(gpu: ArcRef<GPUInner>) -> Self {
53 Self {
54 gpu,
55 attachments: Vec::new(),
56 shader: None,
57 blend: None,
58 color_write_mask: None,
59 shader_reflection: None,
60 }
61 }
62
63 #[inline]
64 pub fn set_blend(mut self, blend: Option<&BlendState>) -> Self {
65 match blend {
66 Some(blend) => {
67 self.blend = Some(blend.clone().into());
68 self.color_write_mask = Some(blend.clone().into());
69 }
70 None => {
71 self.blend = None;
72 self.color_write_mask = None;
73 }
74 }
75
76 self
77 }
78
79 #[inline]
80 pub fn set_shader(self, shader: Option<&GraphicsShader>) -> Self {
81 self.set_shader_with_options(shader, None, None, None, None, None)
82 }
83
84 #[inline]
85 pub fn set_shader_with_options(
86 mut self,
87 shader: Option<&GraphicsShader>,
88 topology: Option<ShaderTopology>,
89 cull_mode: Option<ShaderCullMode>,
90 front_face: Option<ShaderFrontFace>,
91 polygon_mode: Option<ShaderPollygonMode>,
92 index_format: Option<IndexBufferSize>,
93 ) -> Self {
94 match shader {
95 Some(shader) => {
96 let shader_inner = shader.inner.borrow();
97 let (vertex_shader, fragment_shader) = match &shader_inner.ty {
98 GraphicsShaderType::GraphicsSplit {
99 vertex_module,
100 fragment_module,
101 } => (vertex_module.clone(), fragment_module.clone()),
102 GraphicsShaderType::GraphicsSingle { module } => (module.clone(), module.clone()),
103 };
104
105 let layout = shader_inner.bind_group_layouts.clone();
106
107 let vertex_reflect = shader_inner.reflection.iter().find(|r| {
108 matches!(r, ShaderReflect::Vertex { .. })
109 || matches!(r, ShaderReflect::VertexFragment { .. })
110 });
111
112 let fragment_reflect = shader_inner.reflection.iter().find(|r| {
113 matches!(r, ShaderReflect::Fragment { .. })
114 || matches!(r, ShaderReflect::VertexFragment { .. })
115 });
116
117 let vertex_entry_point = match vertex_reflect {
118 Some(ShaderReflect::Vertex { entry_point, .. }) => Some(entry_point),
119 Some(ShaderReflect::VertexFragment {
120 vertex_entry_point, ..
121 }) => Some(vertex_entry_point),
122 _ => None,
123 };
124
125 let fragment_entry_point = match fragment_reflect {
126 Some(ShaderReflect::Fragment { entry_point, .. }) => Some(entry_point),
127 Some(ShaderReflect::VertexFragment {
128 fragment_entry_point,
129 ..
130 }) => Some(fragment_entry_point),
131 _ => None,
132 };
133
134 #[cfg(any(debug_assertions, feature = "enable-release-validation"))]
135 {
136 if vertex_entry_point.is_none() {
137 panic!("Vertex shader entry point is not found in shader reflection");
138 }
139
140 if fragment_entry_point.is_none() {
141 panic!("Fragment shader entry point is not found in shader reflection");
142 }
143 }
144
145 let vertex_entry_point = vertex_entry_point.unwrap();
146 let fragment_entry_point = fragment_entry_point.unwrap();
147
148 let attrib_inner = shader.attrib.borrow();
149 let shader_binding = IntermediateRenderPipeline {
150 shader: (vertex_shader, fragment_shader),
151 vertex_attribute: (attrib_inner.stride, attrib_inner.attributes.clone()),
152 shader_entry: (vertex_entry_point.clone(), fragment_entry_point.clone()),
153 layout: layout,
154 topology: topology.unwrap_or(attrib_inner.topology),
155 cull_mode: cull_mode.into(),
156 front_face: front_face.unwrap_or(attrib_inner.front_face),
157 polygon_mode: polygon_mode.unwrap_or(attrib_inner.polygon_mode),
158 index_format: index_format.or_else(|| attrib_inner.index.clone()),
159 };
160
161 self.shader = Some(shader_binding);
162 self.shader_reflection = Some(shader_inner.reflection.clone());
163 }
164 None => {
165 self.shader = None;
166 self.shader_reflection = None;
167 }
168 }
169
170 self
171 }
172
173 #[inline]
174 pub fn set_attachment_sampler(
175 mut self,
176 group: u32,
177 binding: u32,
178 sampler: Option<&TextureSampler>,
179 ) -> Self {
180 match sampler {
181 Some(sampler) => {
182 let attachment = {
183 let gpu_inner = self.gpu.borrow();
184
185 BindGroupAttachment {
186 group,
187 binding,
188 attachment: BindGroupType::Sampler(
189 sampler.make_wgpu(gpu_inner.device()),
190 ),
191 }
192 };
193
194 self.insert_or_replace_attachment(group, binding, attachment);
195 }
196 None => {
197 self.remove_attachment(group, binding);
198 }
199 }
200
201 self
202 }
203
204 #[inline]
205 pub fn set_attachment_texture(
206 mut self,
207 group: u32,
208 binding: u32,
209 texture: Option<&Texture>,
210 ) -> Self {
211 match texture {
212 Some(texture) => {
213 let attachment = {
214 BindGroupAttachment {
215 group,
216 binding,
217 attachment: BindGroupType::Texture(
218 texture.inner.borrow().wgpu_view.clone(),
219 ),
220 }
221 };
222
223 self.insert_or_replace_attachment(group, binding, attachment);
224 }
225 None => {
226 self.remove_attachment(group, binding);
227 }
228 }
229
230 self
231 }
232
233 #[inline]
234 pub fn set_attachment_texture_storage(
235 mut self,
236 group: u32,
237 binding: u32,
238 texture: Option<&Texture>,
239 ) -> Self {
240 match texture {
241 Some(texture) => {
242 let inner = texture.inner.borrow();
243 let attachment = BindGroupAttachment {
244 group,
245 binding,
246 attachment: BindGroupType::TextureStorage(inner.wgpu_view.clone()),
247 };
248
249 self.insert_or_replace_attachment(group, binding, attachment);
250 }
251 None => {
252 self.remove_attachment(group, binding);
253 }
254 }
255
256 self
257 }
258
259 #[inline]
260 pub fn set_attachment_uniform(
261 mut self,
262 group: u32,
263 binding: u32,
264 buffer: Option<&Buffer>,
265 ) -> Self {
266 match buffer {
267 Some(buffer) => {
268 let inner = buffer.inner.borrow();
269 let attachment = BindGroupAttachment {
270 group,
271 binding,
272 attachment: BindGroupType::Uniform(inner.buffer.clone()),
273 };
274
275 self.insert_or_replace_attachment(group, binding, attachment);
276 }
277 None => {
278 self.remove_attachment(group, binding);
279 }
280 }
281
282 self
283 }
284
285 #[inline]
286 pub fn set_attachment_uniform_vec<T>(
287 mut self,
288 group: u32,
289 binding: u32,
290 buffer: Option<Vec<T>>,
291 ) -> Self
292 where
293 T: bytemuck::Pod + bytemuck::Zeroable,
294 {
295 match buffer {
296 Some(buffer) => {
297 let attachment = {
298 let mut inner = self.gpu.borrow_mut();
299
300 let buffer = inner.create_buffer_with(&buffer, wgpu::BufferUsages::COPY_DST);
301 BindGroupAttachment {
302 group,
303 binding,
304 attachment: BindGroupType::Uniform(buffer),
305 }
306 };
307
308 self.insert_or_replace_attachment(group, binding, attachment);
309 }
310 None => {
311 self.remove_attachment(group, binding);
312 }
313 }
314
315 self
316 }
317
318 #[inline]
319 pub fn set_attachment_uniform_raw<T>(
320 mut self,
321 group: u32,
322 binding: u32,
323 buffer: Option<&[T]>,
324 ) -> Self
325 where
326 T: bytemuck::Pod + bytemuck::Zeroable,
327 {
328 match buffer {
329 Some(buffer) => {
330 let mut inner = self.gpu.borrow_mut();
331
332 let buffer = inner.create_buffer_with(&buffer, wgpu::BufferUsages::COPY_DST);
333 let attachment = BindGroupAttachment {
334 group,
335 binding,
336 attachment: BindGroupType::Uniform(buffer),
337 };
338
339 drop(inner);
340
341 self.insert_or_replace_attachment(group, binding, attachment);
342 }
343 None => {
344 self.remove_attachment(group, binding);
345 }
346 }
347
348 self
349 }
350
351 #[inline]
352 pub fn set_attachment_storage(
353 mut self,
354 group: u32,
355 binding: u32,
356 buffer: Option<&Buffer>,
357 ) -> Self {
358 match buffer {
359 Some(buffer) => {
360 let inner = buffer.inner.borrow();
361 let attachment = BindGroupAttachment {
362 group,
363 binding,
364 attachment: BindGroupType::Storage(inner.buffer.clone()),
365 };
366
367 self.insert_or_replace_attachment(group, binding, attachment);
368 }
369 None => {
370 self.remove_attachment(group, binding);
371 }
372 }
373
374 self
375 }
376
377 #[inline]
378 pub fn set_attachment_storage_raw<T>(
379 mut self,
380 group: u32,
381 binding: u32,
382 buffer: Option<&[T]>,
383 ) -> Self
384 where
385 T: bytemuck::Pod + bytemuck::Zeroable,
386 {
387 match buffer {
388 Some(buffer) => {
389 let mut inner = self.gpu.borrow_mut();
390
391 let buffer = inner.create_buffer_with(&buffer, wgpu::BufferUsages::COPY_DST);
392 let attachment = BindGroupAttachment {
393 group,
394 binding,
395 attachment: BindGroupType::Storage(buffer),
396 };
397
398 drop(inner);
399
400 self.insert_or_replace_attachment(group, binding, attachment);
401 }
402 None => {
403 self.remove_attachment(group, binding);
404 }
405 }
406
407 self
408 }
409
410 #[inline]
411 pub fn set_attachment_storage_vec<T>(
412 mut self,
413 group: u32,
414 binding: u32,
415 buffer: Option<Vec<T>>,
416 ) -> Self
417 where
418 T: bytemuck::Pod + bytemuck::Zeroable,
419 {
420 match buffer {
421 Some(buffer) => {
422 let mut inner = self.gpu.borrow_mut();
423
424 let buffer = inner.create_buffer_with(&buffer, wgpu::BufferUsages::COPY_DST);
425 let attachment = BindGroupAttachment {
426 group,
427 binding,
428 attachment: BindGroupType::Storage(buffer),
429 };
430
431 drop(inner);
432
433 self.insert_or_replace_attachment(group, binding, attachment);
434 }
435 None => {
436 self.remove_attachment(group, binding);
437 }
438 }
439
440 self
441 }
442
443 #[inline]
444 pub(crate) fn remove_attachment(&mut self, group: u32, binding: u32) {
445 self.attachments
446 .retain(|a| a.group != group || a.binding != binding);
447 }
448
449 pub(crate) fn insert_or_replace_attachment(
450 &mut self,
451 group: u32,
452 binding: u32,
453 attachment: BindGroupAttachment,
454 ) {
455 let index = self
456 .attachments
457 .iter()
458 .position(|a| a.group == group && a.binding == binding);
459
460 if let Some(index) = index {
461 self.attachments[index] = attachment;
462 } else {
463 self.attachments.push(attachment);
464 }
465 }
466
467 pub fn build(self) -> Result<RenderPipeline, RenderPipelineError> {
468 if self.shader.is_none() {
469 return Err(RenderPipelineError::ShaderNotSet);
470 }
471
472 let shader_binding = self.shader.unwrap();
473 for attachment in &self.attachments {
474 let r#type = self
475 .shader_reflection
476 .as_ref()
477 .unwrap()
478 .iter()
479 .find_map(|b| {
480 let bindings = match b {
481 ShaderReflect::Vertex { bindings, .. }
482 | ShaderReflect::Fragment { bindings, .. }
483 | ShaderReflect::VertexFragment { bindings, .. } => bindings,
484 _ => return None,
485 };
486
487 bindings.iter().find_map(|shaderbinding| {
488 if shaderbinding.group == attachment.group
489 && shaderbinding.binding == attachment.binding
490 {
491 Some(shaderbinding)
492 } else {
493 None
494 }
495 })
496 });
497
498 if r#type.is_none() {
499 return Err(RenderPipelineError::AttachmentNotSet(
500 attachment.group,
501 attachment.binding,
502 ));
503 }
504
505 let r#type = r#type.unwrap();
506
507 if !match r#type.ty {
508 ShaderBindingType::UniformBuffer(_) => {
509 matches!(attachment.attachment, BindGroupType::Uniform(_))
510 }
511 ShaderBindingType::StorageBuffer(_, _) => {
512 matches!(attachment.attachment, BindGroupType::Storage(_))
513 }
514 ShaderBindingType::StorageTexture(_) => {
515 matches!(attachment.attachment, BindGroupType::TextureStorage(_))
516 }
517 ShaderBindingType::Sampler(_) => {
518 matches!(attachment.attachment, BindGroupType::Sampler(_))
519 }
520 ShaderBindingType::Texture(_) => {
521 matches!(attachment.attachment, BindGroupType::Texture(_))
522 }
523 ShaderBindingType::PushConstant(_) => {
524 matches!(attachment.attachment, BindGroupType::Uniform(_))
525 }
526 } {
527 return Err(RenderPipelineError::InvalidAttachmentType(
528 attachment.group,
529 attachment.binding,
530 r#type.ty,
531 ));
532 }
533 }
534
535 let bind_group_hash_key = {
536 let mut hasher = DefaultHasher::new();
537 hasher.write_u64(0u64); for attachment in &self.attachments {
540 attachment.group.hash(&mut hasher);
541 attachment.binding.hash(&mut hasher);
542 match &attachment.attachment {
543 BindGroupType::Uniform(uniform) => {
544 uniform.hash(&mut hasher);
545 }
546 BindGroupType::Texture(texture) => {
547 texture.hash(&mut hasher);
548 }
549 BindGroupType::TextureStorage(texture) => texture.hash(&mut hasher),
550 BindGroupType::Sampler(sampler) => sampler.hash(&mut hasher),
551 BindGroupType::Storage(storage) => storage.hash(&mut hasher),
552 }
553 }
554
555 hasher.finish()
556 };
557
558 let bind_group_attachments = {
559 let mut gpu_inner = self.gpu.borrow_mut();
560
561 match gpu_inner.get_bind_group(bind_group_hash_key) {
562 Some(bind_group) => bind_group,
563 None => {
564 let mut bind_group_attachments: HashMap<u32, Vec<wgpu::BindGroupEntry>> =
565 self.attachments.iter().fold(HashMap::new(), |mut map, e| {
566 let (group, binding, attachment) = (e.group, e.binding, &e.attachment);
567 let entry = match attachment {
568 BindGroupType::Uniform(buffer) => wgpu::BindGroupEntry {
569 binding,
570 resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
571 buffer,
572 offset: 0,
573 size: None,
574 }),
575 },
576 BindGroupType::Texture(texture) => wgpu::BindGroupEntry {
577 binding,
578 resource: wgpu::BindingResource::TextureView(texture),
579 },
580 BindGroupType::Sampler(sampler) => wgpu::BindGroupEntry {
581 binding,
582 resource: wgpu::BindingResource::Sampler(sampler),
583 },
584 BindGroupType::Storage(buffer) => wgpu::BindGroupEntry {
585 binding,
586 resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
587 buffer,
588 offset: 0,
589 size: None,
590 }),
591 },
592 BindGroupType::TextureStorage(texture) => wgpu::BindGroupEntry {
593 binding,
594 resource: wgpu::BindingResource::TextureView(texture),
595 },
596 };
597
598 map.entry(group).or_insert_with(Vec::new).push(entry);
599 map
600 });
601
602 for entries in bind_group_attachments.values_mut() {
606 entries.sort_by_key(|e| e.binding);
607 }
608
609 let bind_group = bind_group_attachments
610 .iter()
611 .map(|(group, entries)| {
612 let layout = shader_binding
613 .layout
614 .iter()
615 .find(|l| l.group == *group)
616 .unwrap();
617
618 (layout, entries.as_slice())
619 })
620 .collect::<Vec<_>>();
621
622 let create_info = BindGroupCreateInfo {
623 entries: bind_group,
624 };
625
626 gpu_inner.create_bind_group(bind_group_hash_key, create_info)
627 }
628 }
629 };
630
631 let attribute = &shader_binding.vertex_attribute;
632 let vertex_desc = VertexAttributeLayout {
633 stride: attribute.0 as wgpu::BufferAddress,
634 step_mode: wgpu::VertexStepMode::Vertex,
635 attributes: attribute.1.clone(),
636 };
637
638 let primitive_state = wgpu::PrimitiveState {
639 topology: shader_binding.topology.into(),
640 strip_index_format: None,
641 front_face: shader_binding.front_face.into(),
642 cull_mode: shader_binding.cull_mode.map(|c| c.into()),
643 polygon_mode: shader_binding.polygon_mode.into(),
644 unclipped_depth: false,
645 conservative: false,
646 };
647
648 let layout = shader_binding
649 .layout
650 .iter()
651 .map(|l| l.layout.clone())
652 .collect::<Vec<_>>();
653
654 let pipeline_desc = GraphicsPipelineDesc {
655 shaders: shader_binding.shader.clone(),
656 entry_point: shader_binding.shader_entry.clone(),
657 render_target: vec![(
658 wgpu::TextureFormat::Rgba8UnormSrgb,
659 self.blend.clone(),
660 self.color_write_mask.clone(),
661 )],
662 depth_stencil: None,
663 vertex_desc,
664 primitive_state,
665 bind_group_layout: layout,
666 msaa_count: 1,
667 };
668
669 Ok(RenderPipeline {
670 bind_group: bind_group_attachments,
671 pipeline_desc,
672 index_format: shader_binding.index_format,
673 })
674 }
675}
676
677#[derive(Debug, Clone, Copy)]
678pub enum RenderPipelineError {
679 ShaderNotSet,
680 InvalidShaderType,
681 AttachmentNotSet(u32, u32),
682 InvalidAttachmentType(u32, u32, ShaderBindingType),
683}