Skip to main content

ranim_render/
resource.rs

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
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(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// MARK: RenderTextures
53#[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/// Texture resources used for rendering
78#[allow(unused)]
79pub struct RenderTextures {
80    width: u32,
81    height: u32,
82    pub render_texture: ReadbackWgpuTexture,
83    // multisample_texture: wgpu::Texture,
84    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    /// Bind group for depth texture (used in OIT resolve)
89    pub(crate) depth_bind_group: wgpu::BindGroup,
90    // pub(crate) multisample_view: wgpu::TextureView,
91    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 multisample_texture = ctx.device.create_texture(&wgpu::TextureDescriptor {
136        //     label: Some("Multisample Texture"),
137        //     size: wgpu::Extent3d {
138        //         width: width as u32,
139        //         height: height as u32,
140        //         depth_or_array_layers: 1,
141        //     },
142        //     mip_level_count: 1,
143        //     sample_count: 4,
144        //     dimension: wgpu::TextureDimension::D2,
145        //     format,
146        //     usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
147        //     view_formats: &[
148        //         wgpu::TextureFormat::Rgba8UnormSrgb,
149        //         wgpu::TextureFormat::Rgba8Unorm,
150        //     ],
151        // });
152        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 multisample_view = multisample_texture.create_view(&wgpu::TextureViewDescriptor {
180        //     format: Some(format),
181        //     ..Default::default()
182        // });
183        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        // Create depth bind group for OIT resolve
193        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            // multisample_texture,
208            depth_stencil_texture,
209            render_view,
210            linear_render_view,
211            depth_texture_view,
212            depth_bind_group,
213            // multisample_view,
214            depth_stencil_view,
215            state: RenderTextureState::default(),
216        }
217    }
218
219    /// Mark textures as dirty after rendering.
220    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    /// Start async readback of the output texture (non-blocking).
229    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    /// Finish a pending async readback, copying data into the CPU-side buffer.
235    pub fn finish_readback(&mut self, ctx: &WgpuContext) {
236        self.render_texture.finish_readback(ctx);
237    }
238
239    /// Try to finish a pending readback without blocking.
240    /// Returns `true` if the readback completed (or there was nothing pending),
241    /// `false` if the GPU hasn't finished yet.
242    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}