1use std::{
2 any::{Any, TypeId},
3 collections::HashMap,
4 sync::{
5 Arc, RwLock,
6 atomic::{AtomicBool, Ordering},
7 },
8};
9
10use image::{ImageBuffer, Luma, Rgba};
11
12use crate::utils::{ReadbackWgpuTexture, WgpuContext};
13
14pub(crate) trait GpuResource {
16 fn new(ctx: &WgpuContext) -> Self
17 where
18 Self: Sized;
19}
20
21#[derive(bevy_ecs::prelude::Resource, Default)]
23pub struct PipelinesPool {
24 inner: RwLock<HashMap<TypeId, Arc<dyn Any + Send + Sync>>>,
25}
26
27impl PipelinesPool {
28 pub(crate) fn get_or_init<P: GpuResource + Send + Sync + 'static>(
29 &self,
30 ctx: &WgpuContext,
31 ) -> Arc<P> {
32 let id = std::any::TypeId::of::<P>();
33 {
34 let inner = self.inner.read().unwrap();
35 if let Some(pipeline) = inner.get(&id) {
36 return pipeline.clone().downcast::<P>().unwrap();
37 }
38 }
39 let mut inner = self.inner.write().unwrap();
40 inner
41 .entry(id)
42 .or_insert_with(|| {
43 let pipeline = P::new(ctx);
44 Arc::new(pipeline)
45 })
46 .clone()
47 .downcast::<P>()
48 .unwrap()
49 }
50}
51
52#[derive(Clone)]
54pub(crate) struct RenderTextureState(Arc<RenderTextureStateInner>);
55
56struct RenderTextureStateInner {
57 output_dirty: AtomicBool,
58 depth_dirty: AtomicBool,
59}
60
61impl Default for RenderTextureState {
62 fn default() -> Self {
63 Self(Arc::new(RenderTextureStateInner {
64 output_dirty: AtomicBool::new(true),
65 depth_dirty: AtomicBool::new(true),
66 }))
67 }
68}
69
70impl RenderTextureState {
71 pub(crate) fn mark_dirty(&self) {
72 self.0.output_dirty.store(true, Ordering::Release);
73 self.0.depth_dirty.store(true, Ordering::Release);
74 }
75}
76
77#[allow(unused)]
79pub struct RenderTextures {
80 width: u32,
81 height: u32,
82 pub render_texture: ReadbackWgpuTexture,
83 pub depth_stencil_texture: ReadbackWgpuTexture,
85 pub render_view: wgpu::TextureView,
86 pub linear_render_view: wgpu::TextureView,
87 pub depth_texture_view: wgpu::TextureView,
88 pub(crate) depth_bind_group: wgpu::BindGroup,
90 pub(crate) depth_stencil_view: wgpu::TextureView,
92
93 state: RenderTextureState,
94}
95
96pub(crate) const OUTPUT_TEXTURE_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Rgba8UnormSrgb;
97impl RenderTextures {
98 pub fn width(&self) -> u32 {
99 self.width
100 }
101
102 pub fn height(&self) -> u32 {
103 self.height
104 }
105
106 pub fn ratio(&self) -> f32 {
107 self.width as f32 / self.height as f32
108 }
109
110 pub(crate) fn new(ctx: &WgpuContext, width: u32, height: u32) -> Self {
111 let format = OUTPUT_TEXTURE_FORMAT;
112 let render_texture = ReadbackWgpuTexture::new(
113 ctx,
114 &wgpu::TextureDescriptor {
115 label: Some("Target Texture"),
116 size: wgpu::Extent3d {
117 width,
118 height,
119 depth_or_array_layers: 1,
120 },
121 mip_level_count: 1,
122 sample_count: 1,
123 dimension: wgpu::TextureDimension::D2,
124 format,
125 usage: wgpu::TextureUsages::RENDER_ATTACHMENT
126 | wgpu::TextureUsages::COPY_SRC
127 | wgpu::TextureUsages::COPY_DST
128 | wgpu::TextureUsages::TEXTURE_BINDING,
129 view_formats: &[
130 wgpu::TextureFormat::Rgba8UnormSrgb,
131 wgpu::TextureFormat::Rgba8Unorm,
132 ],
133 },
134 );
135 let depth_stencil_texture = ReadbackWgpuTexture::new(
153 ctx,
154 &wgpu::TextureDescriptor {
155 label: Some("Depth Stencil Texture"),
156 size: wgpu::Extent3d {
157 width,
158 height,
159 depth_or_array_layers: 1,
160 },
161 mip_level_count: 1,
162 sample_count: 1,
163 dimension: wgpu::TextureDimension::D2,
164 format: wgpu::TextureFormat::Depth32Float,
165 usage: wgpu::TextureUsages::RENDER_ATTACHMENT
166 | wgpu::TextureUsages::COPY_SRC
167 | wgpu::TextureUsages::TEXTURE_BINDING,
168 view_formats: &[],
169 },
170 );
171 let render_view = render_texture.create_view(&wgpu::TextureViewDescriptor {
172 format: Some(format),
173 ..Default::default()
174 });
175 let linear_render_view = render_texture.create_view(&wgpu::TextureViewDescriptor {
176 format: Some(wgpu::TextureFormat::Rgba8Unorm),
177 ..Default::default()
178 });
179 let depth_stencil_view =
184 depth_stencil_texture.create_view(&wgpu::TextureViewDescriptor::default());
185
186 let depth_texture_view = depth_stencil_texture.create_view(&wgpu::TextureViewDescriptor {
187 label: Some("Depth Texture View"),
188 aspect: wgpu::TextureAspect::DepthOnly,
189 ..Default::default()
190 });
191
192 use crate::pipelines::OITResolvePipeline;
194 let depth_bind_group = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
195 label: Some("Depth Texture Bind Group"),
196 layout: &OITResolvePipeline::depth_bind_group_layout(ctx),
197 entries: &[wgpu::BindGroupEntry {
198 binding: 0,
199 resource: wgpu::BindingResource::TextureView(&depth_texture_view),
200 }],
201 });
202
203 Self {
204 width,
205 height,
206 render_texture,
207 depth_stencil_texture,
209 render_view,
210 linear_render_view,
211 depth_texture_view,
212 depth_bind_group,
213 depth_stencil_view,
215 state: RenderTextureState::default(),
216 }
217 }
218
219 pub fn mark_dirty(&self) {
221 self.state.mark_dirty();
222 }
223
224 pub(crate) fn state(&self) -> RenderTextureState {
225 self.state.clone()
226 }
227
228 pub fn start_readback(&mut self, ctx: &WgpuContext) {
230 self.render_texture.start_readback(ctx);
231 self.state.0.output_dirty.store(false, Ordering::Release);
232 }
233
234 pub fn finish_readback(&mut self, ctx: &WgpuContext) {
236 self.render_texture.finish_readback(ctx);
237 }
238
239 pub fn try_finish_readback(&mut self, ctx: &WgpuContext) -> bool {
243 self.render_texture.try_finish_readback(ctx)
244 }
245
246 pub fn get_rendered_texture_data(&mut self, ctx: &WgpuContext) -> &[u8] {
247 if !self.state.0.output_dirty.load(Ordering::Acquire) {
248 return self.render_texture.texture_data();
249 }
250 self.state.0.output_dirty.store(false, Ordering::Release);
251 self.render_texture.update_texture_data(ctx)
252 }
253
254 pub fn get_rendered_texture_img_buffer(
255 &mut self,
256 ctx: &WgpuContext,
257 ) -> ImageBuffer<Rgba<u8>, &[u8]> {
258 ImageBuffer::from_raw(self.width, self.height, self.get_rendered_texture_data(ctx)).unwrap()
259 }
260
261 pub fn get_depth_texture_data(&mut self, ctx: &WgpuContext) -> &[f32] {
262 if !self.state.0.depth_dirty.load(Ordering::Acquire) {
263 return bytemuck::cast_slice(self.depth_stencil_texture.texture_data());
264 }
265 self.state.0.depth_dirty.store(false, Ordering::Release);
266 bytemuck::cast_slice(self.depth_stencil_texture.update_texture_data(ctx))
267 }
268
269 pub fn get_depth_texture_img_buffer(
270 &mut self,
271 ctx: &WgpuContext,
272 ) -> ImageBuffer<Luma<u8>, Vec<u8>> {
273 let data = self
274 .get_depth_texture_data(ctx)
275 .iter()
276 .map(|&d| (d.clamp(0.0, 1.0) * 255.0) as u8)
277 .collect::<Vec<_>>();
278 ImageBuffer::from_raw(self.width, self.height, data).unwrap()
279 }
280}