Skip to main content

cvkg_render_gpu/kvasir/
registry.rs

1use super::resource::{ResourceDescriptor, ResourceId, ResourceKind, ResourceLifetime};
2use std::collections::HashMap;
3
4pub struct TextureResource {
5    pub texture: Option<wgpu::Texture>,
6    pub view: wgpu::TextureView,
7    pub lifetime: ResourceLifetime,
8}
9
10#[derive(Default)]
11pub struct ResourceRegistry {
12    textures: HashMap<ResourceId, TextureResource>,
13    pool: HashMap<(wgpu::TextureFormat, u32, u32), Vec<TextureResource>>,
14    next_id: u32,
15}
16
17impl ResourceRegistry {
18    pub fn new() -> Self {
19        Self {
20            textures: HashMap::new(),
21            pool: HashMap::new(),
22            next_id: 10000,
23        }
24    }
25
26    pub fn register_external_texture(
27        &mut self,
28        id: ResourceId,
29        texture: wgpu::Texture,
30        view: wgpu::TextureView,
31    ) {
32        self.textures.insert(
33            id,
34            TextureResource {
35                texture: Some(texture),
36                view,
37                lifetime: ResourceLifetime::Frame,
38            },
39        );
40    }
41
42    pub fn alias(&mut self, alias_id: ResourceId, actual_id: ResourceId) {
43        if let Some(res) = self.textures.get(&actual_id) {
44            let cloned = TextureResource {
45                texture: res.texture.clone(),
46                view: res.view.clone(),
47                lifetime: ResourceLifetime::Frame,
48            };
49            self.textures.insert(alias_id, cloned);
50        }
51    }
52
53    pub fn alias_view(&mut self, alias_id: ResourceId, view: wgpu::TextureView) {
54        self.textures.insert(
55            alias_id,
56            TextureResource {
57                texture: None,
58                view,
59                lifetime: ResourceLifetime::Frame,
60            },
61        );
62    }
63
64    pub fn allocate_offscreen(
65        &mut self,
66        device: &wgpu::Device,
67        target_id: u64,
68        size: [u32; 2],
69    ) -> ResourceId {
70        let desc = ResourceDescriptor {
71            label: Some(format!("Offscreen {}", target_id)),
72            kind: ResourceKind::Image {
73                format: wgpu::TextureFormat::Bgra8UnormSrgb,
74                width: size[0].max(1),
75                height: size[1].max(1),
76                mip_level_count: 1,
77                usage: wgpu::TextureUsages::RENDER_ATTACHMENT
78                    | wgpu::TextureUsages::TEXTURE_BINDING,
79            },
80            lifetime: ResourceLifetime::Frame,
81        };
82        let id = ResourceId(1000 + target_id as u32);
83
84        if let ResourceKind::Image {
85            format,
86            width,
87            height,
88            mip_level_count,
89            usage,
90        } = &desc.kind
91        {
92            let pool_key = (*format, *width, *height);
93            let pooled_tex = if let Some(pool_list) = self.pool.get_mut(&pool_key) {
94                pool_list.pop()
95            } else {
96                None
97            };
98
99            let tex_res = if let Some(mut t) = pooled_tex {
100                t.lifetime = desc.lifetime;
101                t
102            } else {
103                let texture = device.create_texture(&wgpu::TextureDescriptor {
104                    label: desc.label.as_deref(),
105                    size: wgpu::Extent3d {
106                        width: *width,
107                        height: *height,
108                        depth_or_array_layers: 1,
109                    },
110                    mip_level_count: *mip_level_count,
111                    sample_count: 1,
112                    dimension: wgpu::TextureDimension::D2,
113                    format: *format,
114                    usage: *usage,
115                    view_formats: &[],
116                });
117                let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
118                TextureResource {
119                    texture: Some(texture),
120                    view,
121                    lifetime: desc.lifetime,
122                }
123            };
124
125            self.textures.insert(id, tex_res);
126        }
127        id
128    }
129
130    pub fn allocate_image(
131        &mut self,
132        device: &wgpu::Device,
133        desc: &ResourceDescriptor,
134    ) -> ResourceId {
135        let id = ResourceId(self.next_id);
136        self.next_id += 1;
137
138        if let ResourceKind::Image {
139            format,
140            width,
141            height,
142            mip_level_count,
143            usage,
144        } = &desc.kind
145        {
146            let pool_key = (*format, *width, *height);
147            let pooled_tex = if let Some(pool_list) = self.pool.get_mut(&pool_key) {
148                pool_list.pop()
149            } else {
150                None
151            };
152
153            let tex_res = if let Some(mut t) = pooled_tex {
154                t.lifetime = desc.lifetime;
155                t
156            } else {
157                let texture = device.create_texture(&wgpu::TextureDescriptor {
158                    label: desc.label.as_deref(),
159                    size: wgpu::Extent3d {
160                        width: *width,
161                        height: *height,
162                        depth_or_array_layers: 1,
163                    },
164                    mip_level_count: *mip_level_count,
165                    sample_count: 1,
166                    dimension: wgpu::TextureDimension::D2,
167                    format: *format,
168                    usage: *usage,
169                    view_formats: &[],
170                });
171                let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
172                TextureResource {
173                    texture: Some(texture),
174                    view,
175                    lifetime: desc.lifetime,
176                }
177            };
178
179            self.textures.insert(id, tex_res);
180        } else {
181            panic!("allocate_image called with non-Image descriptor");
182        }
183        id
184    }
185
186    pub fn get_texture_view(&self, id: ResourceId) -> Option<wgpu::TextureView> {
187        self.textures.get(&id).map(|res| res.view.clone())
188    }
189
190    pub fn get_texture(&self, id: ResourceId) -> Option<wgpu::Texture> {
191        self.textures.get(&id).and_then(|res| res.texture.clone())
192    }
193
194    pub fn remove_image(&mut self, id: ResourceId) {
195        self.textures.remove(&id);
196    }
197
198    pub fn evict_frame_resources(&mut self) {
199        // Move transient frame resources back into the pool instead of destroying them
200        let mut to_remove = Vec::new();
201        for (id, res) in &self.textures {
202            if res.lifetime == ResourceLifetime::Frame {
203                to_remove.push(*id);
204            }
205        }
206        for id in to_remove {
207            if let Some(res) = self.textures.remove(&id)
208                && let Some(tex) = res.texture
209            {
210                let size = tex.size();
211                let format = tex.format();
212                let pool_key = (format, size.width, size.height);
213                let pool_list = self.pool.entry(pool_key).or_default();
214                if pool_list.len() < 4 {
215                    pool_list.push(TextureResource {
216                        texture: Some(tex),
217                        view: res.view,
218                        lifetime: res.lifetime,
219                    });
220                }
221            }
222        }
223    }
224}