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