Skip to main content

ursus_core/assets/storage/
shader.rs

1use std::collections::HashMap;
2use std::path::PathBuf;
3
4pub enum ShaderSource {
5    File(PathBuf),
6    Bytes(Vec<u8>),
7}
8
9pub struct ShaderDef {
10    pub name: String,
11    pub vert: ShaderSource,
12    pub frag: Option<ShaderSource>,
13}
14
15impl ShaderDef {
16    pub fn from_files(name: impl Into<String>, vert: impl Into<PathBuf>, frag: impl Into<PathBuf>) -> Self {
17        Self { name: name.into(), vert: ShaderSource::File(vert.into()), frag: Some(ShaderSource::File(frag.into())) }
18    }
19
20    pub fn from_files_vert_only(name: impl Into<String>, vert: impl Into<PathBuf>) -> Self {
21        Self { name: name.into(), vert: ShaderSource::File(vert.into()), frag: None }
22    }
23
24    pub fn from_bytes(name: impl Into<String>, vert: Vec<u8>, frag: Vec<u8>) -> Self {
25        Self { name: name.into(), vert: ShaderSource::Bytes(vert), frag: Some(ShaderSource::Bytes(frag)) }
26    }
27
28    pub fn from_bytes_vert_only(name: impl Into<String>, vert: Vec<u8>) -> Self {
29        Self { name: name.into(), vert: ShaderSource::Bytes(vert), frag: None }
30    }
31}
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
34pub struct ShaderHandle(pub u32);
35
36struct CompiledShader {
37    vert_spv: Vec<u8>,
38    frag_spv: Option<Vec<u8>>,
39}
40
41pub struct ShaderRegistry {
42    shaders: Vec<ShaderDef>,
43    by_name: HashMap<String, ShaderHandle>,
44    compiled: HashMap<ShaderHandle, CompiledShader>,
45    version: HashMap<ShaderHandle, u32>,
46}
47
48impl ShaderRegistry {
49    pub fn empty() -> Self {
50        Self { shaders: Vec::new(), by_name: HashMap::new(), compiled: HashMap::new(), version: HashMap::new() }
51    }
52
53    pub fn register(&mut self, def: ShaderDef) -> ShaderHandle {
54        let handle = ShaderHandle(self.shaders.len() as u32);
55        self.by_name.insert(def.name.clone(), handle);
56        self.shaders.push(def);
57        handle
58    }
59
60    pub fn register_if_absent(&mut self, def: ShaderDef) -> ShaderHandle {
61        if let Some(h) = self.by_name(&def.name) {
62            return h;
63        }
64        self.register(def)
65    }
66
67    pub fn load_spv(&mut self, handle: ShaderHandle) -> anyhow::Result<(&[u8], Option<&[u8]>)> {
68        if !self.compiled.contains_key(&handle) {
69            let def = self
70                .shaders
71                .get(handle.0 as usize)
72                .ok_or_else(|| anyhow::anyhow!("ShaderHandle {:?} не найден", handle))?;
73
74            let vert_spv = load_source(&def.vert)
75                .map_err(|e| anyhow::anyhow!("Ошибка загрузки vert шейдера '{}': {}", def.name, e))?;
76            let frag_spv = def
77                .frag
78                .as_ref()
79                .map(load_source)
80                .transpose()
81                .map_err(|e| anyhow::anyhow!("Ошибка загрузки frag шейдера '{}': {}", def.name, e))?;
82
83            self.compiled.insert(handle, CompiledShader { vert_spv, frag_spv });
84            log::info!("Шейдер '{}' загружен", def.name);
85        }
86
87        let compiled = &self.compiled[&handle];
88        Ok((&compiled.vert_spv, compiled.frag_spv.as_deref()))
89    }
90
91    pub fn unload(&mut self, handle: ShaderHandle) {
92        self.compiled.remove(&handle);
93    }
94
95    pub fn reload(&mut self, handle: ShaderHandle) -> anyhow::Result<()> {
96        self.unload(handle);
97        self.load_spv(handle)?;
98        *self.version.entry(handle).or_insert(0) += 1;
99        Ok(())
100    }
101
102    pub fn version(&self, handle: ShaderHandle) -> u32 {
103        self.version.get(&handle).copied().unwrap_or(0)
104    }
105
106    pub fn get(&self, handle: ShaderHandle) -> Option<&ShaderDef> {
107        self.shaders.get(handle.0 as usize)
108    }
109
110    pub fn by_name(&self, name: &str) -> Option<ShaderHandle> {
111        self.by_name.get(name).copied()
112    }
113}
114
115impl Default for ShaderRegistry {
116    fn default() -> Self {
117        Self::empty()
118    }
119}
120
121fn load_source(source: &ShaderSource) -> anyhow::Result<Vec<u8>> {
122    match source {
123        ShaderSource::Bytes(bytes) => Ok(bytes.clone()),
124        ShaderSource::File(path) => {
125            std::fs::read(path).map_err(|e| anyhow::anyhow!("Не удалось прочитать {:?}: {}", path, e))
126        }
127    }
128}