ranim_render/
resource.rs

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
14/// A render resource.
15pub(crate) trait GpuResource {
16    fn new(ctx: &WgpuContext) -> Self
17    where
18        Self: Sized;
19}
20
21/// A storage for pipelines
22#[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// MARK: RenderTextures
53/// Texture resources used for rendering
54#[allow(unused)]
55pub struct RenderTextures {
56    width: u32,
57    height: u32,
58    pub render_texture: ReadbackWgpuTexture,
59    // multisample_texture: wgpu::Texture,
60    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    /// Bind group for depth texture (used in OIT resolve)
65    pub(crate) depth_bind_group: wgpu::BindGroup,
66    // pub(crate) multisample_view: wgpu::TextureView,
67    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 multisample_texture = ctx.device.create_texture(&wgpu::TextureDescriptor {
113        //     label: Some("Multisample Texture"),
114        //     size: wgpu::Extent3d {
115        //         width: width as u32,
116        //         height: height as u32,
117        //         depth_or_array_layers: 1,
118        //     },
119        //     mip_level_count: 1,
120        //     sample_count: 4,
121        //     dimension: wgpu::TextureDimension::D2,
122        //     format,
123        //     usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
124        //     view_formats: &[
125        //         wgpu::TextureFormat::Rgba8UnormSrgb,
126        //         wgpu::TextureFormat::Rgba8Unorm,
127        //     ],
128        // });
129        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 multisample_view = multisample_texture.create_view(&wgpu::TextureViewDescriptor {
157        //     format: Some(format),
158        //     ..Default::default()
159        // });
160        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        // Create depth bind group for OIT resolve
170        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            // multisample_texture,
185            depth_stencil_texture,
186            render_view,
187            linear_render_view,
188            depth_texture_view,
189            depth_bind_group,
190            // multisample_view,
191            depth_stencil_view,
192            output_dirty: true,
193            depth_dirty: true,
194        }
195    }
196
197    /// Mark textures as dirty after rendering.
198    pub fn mark_dirty(&mut self) {
199        self.output_dirty = true;
200        self.depth_dirty = true;
201    }
202
203    /// Start async readback of the output texture (non-blocking).
204    pub fn start_readback(&mut self, ctx: &WgpuContext) {
205        self.render_texture.start_readback(ctx);
206        self.output_dirty = false;
207    }
208
209    /// Finish a pending async readback, copying data into the CPU-side buffer.
210    pub fn finish_readback(&mut self, ctx: &WgpuContext) {
211        self.render_texture.finish_readback(ctx);
212    }
213
214    /// Try to finish a pending readback without blocking.
215    /// Returns `true` if the readback completed (or there was nothing pending),
216    /// `false` if the GPU hasn't finished yet.
217    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
259/// A handle to a render packet.
260///
261/// In its inner is an [`Arc`] reference count of the [`RenderInstanceKey`].
262pub 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//
276// MARK: RenderPool
277#[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    /// When called, all instances not referenced are recorded into the `last_frame_dropped` map.
354    /// An will be cleaned in the next call.
355    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        // println!("dropped {}", self.last_frame_dropped.len());
363        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}