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