1use std::{
2 any::{Any, TypeId},
3 collections::HashMap,
4 sync::{Arc, RwLock},
5};
6
7use image::{ImageBuffer, Luma, Rgba};
8
9use crate::{
10 primitives::{Primitive, RenderResource},
11 utils::{ReadbackWgpuTexture, WgpuContext},
12};
13
14pub(crate) trait GpuResource {
16 fn new(ctx: &WgpuContext) -> Self
17 where
18 Self: Sized;
19}
20
21#[derive(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#[allow(unused)]
55pub struct RenderTextures {
56 width: u32,
57 height: u32,
58 pub render_texture: ReadbackWgpuTexture,
59 pub depth_stencil_texture: ReadbackWgpuTexture,
61 pub render_view: wgpu::TextureView,
62 pub linear_render_view: wgpu::TextureView,
63 pub depth_texture_view: wgpu::TextureView,
64 pub(crate) depth_bind_group: wgpu::BindGroup,
66 pub(crate) depth_stencil_view: wgpu::TextureView,
68
69 output_dirty: bool,
70 depth_dirty: bool,
71}
72
73pub(crate) const OUTPUT_TEXTURE_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Rgba8UnormSrgb;
74impl RenderTextures {
75 pub fn width(&self) -> u32 {
76 self.width
77 }
78
79 pub fn height(&self) -> u32 {
80 self.height
81 }
82
83 pub fn ratio(&self) -> f32 {
84 self.width as f32 / self.height as f32
85 }
86
87 pub(crate) fn new(ctx: &WgpuContext, width: u32, height: u32) -> Self {
88 let format = OUTPUT_TEXTURE_FORMAT;
89 let render_texture = ReadbackWgpuTexture::new(
90 ctx,
91 &wgpu::TextureDescriptor {
92 label: Some("Target Texture"),
93 size: wgpu::Extent3d {
94 width,
95 height,
96 depth_or_array_layers: 1,
97 },
98 mip_level_count: 1,
99 sample_count: 1,
100 dimension: wgpu::TextureDimension::D2,
101 format,
102 usage: wgpu::TextureUsages::RENDER_ATTACHMENT
103 | wgpu::TextureUsages::COPY_SRC
104 | wgpu::TextureUsages::COPY_DST
105 | wgpu::TextureUsages::TEXTURE_BINDING,
106 view_formats: &[
107 wgpu::TextureFormat::Rgba8UnormSrgb,
108 wgpu::TextureFormat::Rgba8Unorm,
109 ],
110 },
111 );
112 let depth_stencil_texture = ReadbackWgpuTexture::new(
130 ctx,
131 &wgpu::TextureDescriptor {
132 label: Some("Depth Stencil Texture"),
133 size: wgpu::Extent3d {
134 width,
135 height,
136 depth_or_array_layers: 1,
137 },
138 mip_level_count: 1,
139 sample_count: 1,
140 dimension: wgpu::TextureDimension::D2,
141 format: wgpu::TextureFormat::Depth32Float,
142 usage: wgpu::TextureUsages::RENDER_ATTACHMENT
143 | wgpu::TextureUsages::COPY_SRC
144 | wgpu::TextureUsages::TEXTURE_BINDING,
145 view_formats: &[],
146 },
147 );
148 let render_view = render_texture.create_view(&wgpu::TextureViewDescriptor {
149 format: Some(format),
150 ..Default::default()
151 });
152 let linear_render_view = render_texture.create_view(&wgpu::TextureViewDescriptor {
153 format: Some(wgpu::TextureFormat::Rgba8Unorm),
154 ..Default::default()
155 });
156 let depth_stencil_view =
161 depth_stencil_texture.create_view(&wgpu::TextureViewDescriptor::default());
162
163 let depth_texture_view = depth_stencil_texture.create_view(&wgpu::TextureViewDescriptor {
164 label: Some("Depth Texture View"),
165 aspect: wgpu::TextureAspect::DepthOnly,
166 ..Default::default()
167 });
168
169 use crate::pipelines::OITResolvePipeline;
171 let depth_bind_group = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
172 label: Some("Depth Texture Bind Group"),
173 layout: &OITResolvePipeline::depth_bind_group_layout(ctx),
174 entries: &[wgpu::BindGroupEntry {
175 binding: 0,
176 resource: wgpu::BindingResource::TextureView(&depth_texture_view),
177 }],
178 });
179
180 Self {
181 width,
182 height,
183 render_texture,
184 depth_stencil_texture,
186 render_view,
187 linear_render_view,
188 depth_texture_view,
189 depth_bind_group,
190 depth_stencil_view,
192 output_dirty: true,
193 depth_dirty: true,
194 }
195 }
196
197 pub fn mark_dirty(&mut self) {
199 self.output_dirty = true;
200 self.depth_dirty = true;
201 }
202
203 pub fn start_readback(&mut self, ctx: &WgpuContext) {
205 self.render_texture.start_readback(ctx);
206 self.output_dirty = false;
207 }
208
209 pub fn finish_readback(&mut self, ctx: &WgpuContext) {
211 self.render_texture.finish_readback(ctx);
212 }
213
214 pub fn try_finish_readback(&mut self, ctx: &WgpuContext) -> bool {
218 self.render_texture.try_finish_readback(ctx)
219 }
220
221 pub fn get_rendered_texture_data(&mut self, ctx: &WgpuContext) -> &[u8] {
222 if !self.output_dirty {
223 return self.render_texture.texture_data();
224 }
225 self.output_dirty = false;
226 self.render_texture.update_texture_data(ctx)
227 }
228
229 pub fn get_rendered_texture_img_buffer(
230 &mut self,
231 ctx: &WgpuContext,
232 ) -> ImageBuffer<Rgba<u8>, &[u8]> {
233 ImageBuffer::from_raw(self.width, self.height, self.get_rendered_texture_data(ctx)).unwrap()
234 }
235
236 pub fn get_depth_texture_data(&mut self, ctx: &WgpuContext) -> &[f32] {
237 if !self.depth_dirty {
238 return bytemuck::cast_slice(self.depth_stencil_texture.texture_data());
239 }
240 self.depth_dirty = false;
241 bytemuck::cast_slice(self.depth_stencil_texture.update_texture_data(ctx))
242 }
243
244 pub fn get_depth_texture_img_buffer(
245 &mut self,
246 ctx: &WgpuContext,
247 ) -> ImageBuffer<Luma<u8>, Vec<u8>> {
248 let data = self
249 .get_depth_texture_data(ctx)
250 .iter()
251 .map(|&d| (d.clamp(0.0, 1.0) * 255.0) as u8)
252 .collect::<Vec<_>>();
253 ImageBuffer::from_raw(self.width, self.height, data).unwrap()
254 }
255}
256
257slotmap::new_key_type! { pub struct RenderInstanceKey; }
258
259pub struct Handle<T> {
263 key: Arc<RenderInstanceKey>,
264 _phantom: std::marker::PhantomData<T>,
265}
266
267impl<T> Clone for Handle<T> {
268 fn clone(&self) -> Self {
269 Self {
270 key: self.key.clone(),
271 _phantom: std::marker::PhantomData,
272 }
273 }
274}
275#[derive(Default)]
278pub struct RenderPool {
279 #[allow(clippy::type_complexity)]
280 inner: slotmap::SlotMap<
281 RenderInstanceKey,
282 (
283 Arc<RenderInstanceKey>,
284 TypeId,
285 Box<dyn Any + Send + Sync + 'static>,
286 ),
287 >,
288 last_frame_dropped: HashMap<TypeId, Vec<RenderInstanceKey>>,
289}
290
291impl RenderPool {
292 pub fn new() -> Self {
293 Self::default()
294 }
295
296 pub fn get_packet<T: 'static>(&self, handle: &Handle<T>) -> &T {
297 self.get(*handle.key)
298 .map(|x| x.downcast_ref::<T>().unwrap())
299 .unwrap()
300 }
301
302 pub fn alloc_packet<P: Primitive>(
303 &mut self,
304 ctx: &WgpuContext,
305 data: &P,
306 ) -> Handle<P::RenderPacket> {
307 let key = self.alloc(ctx, data);
308 Handle {
309 key,
310 _phantom: std::marker::PhantomData,
311 }
312 }
313
314 pub fn show(&self) {
315 self.inner
316 .iter()
317 .enumerate()
318 .for_each(|(idx, (_, (k, _, _)))| {
319 print!("{idx}: {}, ", Arc::strong_count(k));
320 });
321 println!();
322 }
323
324 fn get(&self, key: RenderInstanceKey) -> Option<&(dyn Any + Send + Sync + 'static)> {
325 self.inner.get(key).map(|x| x.2.as_ref())
326 }
327
328 fn alloc<P: Primitive>(&mut self, ctx: &WgpuContext, data: &P) -> Arc<RenderInstanceKey> {
329 let last_frame_dropped = self
330 .last_frame_dropped
331 .entry(TypeId::of::<P::RenderPacket>())
332 .or_default();
333 if let Some(key) = last_frame_dropped.pop() {
334 let entry = self.inner.get_mut(key).unwrap();
335 let key = entry.0.clone();
336 (entry.2.as_mut() as &mut dyn Any)
337 .downcast_mut::<P::RenderPacket>()
338 .unwrap()
339 .update(ctx, data);
340 key
341 } else {
342 let handle = self.inner.insert_with_key(|key| {
343 (
344 Arc::new(key),
345 TypeId::of::<P::RenderPacket>(),
346 Box::new(P::RenderPacket::init(ctx, data)),
347 )
348 });
349 self.inner.get(handle).unwrap().0.clone()
350 }
351 }
352
353 pub fn clean(&mut self) {
356 self.inner.retain(|key, (_, t_id, _)| {
357 self.last_frame_dropped
358 .get(t_id)
359 .map(|x| !x.contains(&key))
360 .unwrap_or(true)
361 });
362 self.last_frame_dropped.clear();
364 self.inner
365 .iter()
366 .filter(|(_, (key, _, _))| Arc::strong_count(key) == 1)
367 .for_each(|(key, (_, t_id, _))| {
368 self.last_frame_dropped.entry(*t_id).or_default().push(key);
369 });
370 }
371}