bevy_shader/
shader_cache.rs

1use crate::shader::*;
2use alloc::borrow::Cow;
3use alloc::sync::Arc;
4use bevy_asset::AssetId;
5use bevy_platform::collections::{hash_map::EntryRef, HashMap, HashSet};
6use core::hash::Hash;
7use thiserror::Error;
8use tracing::debug;
9use tracing::warn;
10
11pub(crate) fn wesl_module_path(import_path: &ShaderImport) -> Option<wesl::syntax::ModulePath> {
12    match import_path {
13        ShaderImport::Custom(name) => {
14            let mut segments = name.split("::");
15            let package = segments.next().filter(|s| !s.is_empty())?;
16            let components = segments
17                .map(|s| (!s.is_empty()).then(|| s.to_string()))
18                .collect::<Option<Vec<_>>>()?;
19            Some(wesl::syntax::ModulePath {
20                origin: wesl::syntax::PathOrigin::Package(package.to_string()),
21                components,
22            })
23        }
24        // `ModulePath::from_path` would strip anything after a `.` as an extension.
25        ShaderImport::AssetPath(path) => {
26            let components: Vec<String> = path
27                .split('/')
28                .filter(|component| !component.is_empty())
29                .map(str::to_string)
30                .collect();
31            (!components.is_empty()).then_some(wesl::syntax::ModulePath {
32                origin: wesl::syntax::PathOrigin::Absolute,
33                components,
34            })
35        }
36    }
37}
38
39fn is_module_not_found(error: &wesl::Error) -> bool {
40    match error {
41        wesl::Error::ResolveError(wesl::error::ResolveError::ModuleNotFound(..))
42        | wesl::Error::ImportError(wesl::error::ImportError::ResolveError(
43            wesl::error::ResolveError::ModuleNotFound(..),
44        )) => true,
45        wesl::Error::Error(diagnostic) => is_module_not_found(&diagnostic.error),
46        _ => false,
47    }
48}
49
50/// Fully composed source code of a shader module, with all shader defs applied.
51///
52/// This is roughly equivalent to [`wgpu::ShaderSource`](https://docs.rs/wgpu/latest/wgpu/enum.ShaderSource.html),
53/// but with less variants and more concrete types instead of [`Cow`].
54///
55/// This source will be parsed and validated by the renderer.
56///
57/// Any necessary shader translation (e.g. from WGSL to SPIR-V or vice versa)
58/// must be done internally by the renderer.
59#[derive(Clone, Debug)]
60pub enum ShaderCacheSource<'a> {
61    /// SPIR-V module represented as a slice of words.
62    SpirV(&'a [u8]),
63    /// WGSL module as a string slice.
64    Wgsl(String),
65}
66
67/// An id of a pipeline, typically in the [`PipelineCache`](https://docs.rs/bevy/latest/bevy/render/render_resource/struct.PipelineCache.html)
68/// Typically corresponds to a unique combination of [`Shader`] and [`ShaderDefVal`]s.
69pub type CachedPipelineId = usize;
70
71struct ShaderData<ShaderModule> {
72    pipelines: HashSet<CachedPipelineId>,
73    processed_shaders: HashMap<Box<[ShaderDefVal]>, Arc<ShaderModule>>,
74    resolved_imports: HashMap<ShaderImport, AssetId<Shader>>,
75    dependents: HashSet<AssetId<Shader>>,
76}
77
78impl<T> Default for ShaderData<T> {
79    fn default() -> Self {
80        Self {
81            pipelines: Default::default(),
82            processed_shaders: Default::default(),
83            resolved_imports: Default::default(),
84            dependents: Default::default(),
85        }
86    }
87}
88
89/// A cache for shaders and shader imports, with asset state-tracking for
90/// waiting to load shaders until all imports are resolved.
91///
92/// Note that the `RenderDevice` generic parameter is a means by which
93/// to avoid a cyclic dependency with `bevy_render`, while also permitting
94/// alternative rendering implementations. The actual processing of the
95/// shader source into a usable compiled module is left to the renderer.
96pub struct ShaderCache<ShaderModule, RenderDevice> {
97    device: RenderDevice,
98    data: HashMap<AssetId<Shader>, ShaderData<ShaderModule>>,
99    load_module: fn(
100        &RenderDevice,
101        ShaderCacheSource,
102        &ValidateShader,
103    ) -> Result<ShaderModule, ShaderCacheError>,
104    module_path_to_asset_id: HashMap<wesl::syntax::ModulePath, AssetId<Shader>>,
105    shaders: HashMap<AssetId<Shader>, Shader>,
106    import_path_shaders: HashMap<ShaderImport, AssetId<Shader>>,
107    waiting_on_import: HashMap<ShaderImport, Vec<AssetId<Shader>>>,
108    missing_import_logged: HashSet<AssetId<Shader>>,
109}
110
111/// A compile time shader value definition to be inlined into the shader source.
112/// Variant tuples contain the name of the definition, and the value.
113#[expect(missing_docs, reason = "Enum variants are self-explanatory")]
114#[derive(serde::Serialize, serde::Deserialize, Clone, PartialEq, Eq, Debug, Hash)]
115pub enum ShaderDefVal {
116    Bool(Cow<'static, str>, bool),
117    Int(Cow<'static, str>, i32),
118    UInt(Cow<'static, str>, u32),
119}
120
121impl From<&'static str> for ShaderDefVal {
122    fn from(key: &'static str) -> Self {
123        ShaderDefVal::Bool(key.into(), true)
124    }
125}
126
127impl From<String> for ShaderDefVal {
128    fn from(key: String) -> Self {
129        ShaderDefVal::Bool(key.into(), true)
130    }
131}
132
133impl ShaderDefVal {
134    /// Returns the value of the define as a string.
135    pub fn value_as_string(&self) -> String {
136        match self {
137            ShaderDefVal::Bool(_, def) => def.to_string(),
138            ShaderDefVal::Int(_, def) => def.to_string(),
139            ShaderDefVal::UInt(_, def) => def.to_string(),
140        }
141    }
142}
143
144impl<ShaderModule, RenderDevice> ShaderCache<ShaderModule, RenderDevice> {
145    /// Creates a new [`ShaderCache`] with the given shader module loading
146    /// function. `load_module` is responsible for actually compiling shader
147    /// source into a module usable by the render device.
148    pub fn new(
149        device: RenderDevice,
150        load_module: fn(
151            &RenderDevice,
152            ShaderCacheSource,
153            &ValidateShader,
154        ) -> Result<ShaderModule, ShaderCacheError>,
155    ) -> Self {
156        Self {
157            device,
158            load_module,
159            data: Default::default(),
160            module_path_to_asset_id: Default::default(),
161            shaders: Default::default(),
162            import_path_shaders: Default::default(),
163            waiting_on_import: Default::default(),
164            missing_import_logged: Default::default(),
165        }
166    }
167
168    /// Attempts to retrieve or create a compiled shader module for the given
169    /// shader id and shader definitions.
170    ///
171    /// The provided `pipeline` is tracked so it may later be reported "dirty"
172    /// when a shader is removed or replaced.
173    ///
174    /// Note that the cache is keyed by `id` and `shader_defs`, meaning providing
175    /// the same `shader_defs` in a different order, or with redundancies, will
176    /// not result in cache hits, and thus require re-composing the module and
177    /// calling `load_module` again.
178    pub fn get(
179        &mut self,
180        pipeline: CachedPipelineId,
181        id: AssetId<Shader>,
182        shader_defs: &[ShaderDefVal],
183    ) -> Result<Arc<ShaderModule>, ShaderCacheError> {
184        let shader = self
185            .shaders
186            .get(&id)
187            .ok_or(ShaderCacheError::ShaderNotLoaded(id))?;
188
189        let data = self.data.entry(id).or_default();
190
191        // Wesl imports are scanned as both module and item candidates.
192        if !matches!(shader.source, Source::Wesl(_)) {
193            let n_asset_imports = shader
194                .imports
195                .iter()
196                .filter(|import| matches!(import, ShaderImport::AssetPath(_)))
197                .count();
198            let n_resolved_asset_imports = data
199                .resolved_imports
200                .keys()
201                .filter(|import| matches!(import, ShaderImport::AssetPath(_)))
202                .count();
203            if n_asset_imports != n_resolved_asset_imports {
204                return Err(ShaderCacheError::ShaderImportNotYetAvailable);
205            }
206        }
207
208        data.pipelines.insert(pipeline);
209
210        let mut wesl_dependencies: Vec<AssetId<Shader>> = Vec::new();
211
212        let module = match data.processed_shaders.entry_ref(shader_defs) {
213            EntryRef::Occupied(entry) => entry.into_mut(),
214            EntryRef::Vacant(entry) => {
215                debug!(
216                    "processing shader {}, with shader defs {:?}",
217                    id, shader_defs
218                );
219                let shader_source = match &shader.source {
220                    Source::SpirV(data) => ShaderCacheSource::SpirV(data.as_ref()),
221                    Source::Wesl(_) => {
222                        if let Some(module_path) = wesl_module_path(&shader.import_path) {
223                            let mut compiler_options = wesl::CompileOptions {
224                                imports: true,
225                                condcomp: true,
226                                visibility: false,
227                                ..Default::default()
228                            };
229
230                            let mut closure_defs = Vec::new();
231                            let mut visited = HashSet::new();
232                            let mut to_visit = vec![id];
233                            while let Some(dep_id) = to_visit.pop() {
234                                if !visited.insert(dep_id) {
235                                    continue;
236                                }
237                                let Some(dep) = self.shaders.get(&dep_id) else {
238                                    continue;
239                                };
240                                if dep_id != id {
241                                    closure_defs.extend(dep.shader_defs.iter());
242                                }
243                                for import in &dep.imports {
244                                    if let Some(import_id) = self.import_path_shaders.get(import) {
245                                        to_visit.push(*import_id);
246                                    }
247                                }
248                            }
249                            let mut constants = alloc::collections::BTreeMap::new();
250                            for shader_def in closure_defs
251                                .into_iter()
252                                .chain(shader_defs.iter())
253                                .chain(shader.shader_defs.iter())
254                            {
255                                match shader_def {
256                                    ShaderDefVal::Bool(key, value) => {
257                                        compiler_options
258                                            .features
259                                            .flags
260                                            .insert(key.to_string(), (*value).into());
261                                    }
262                                    ShaderDefVal::Int(key, value) => {
263                                        compiler_options
264                                            .features
265                                            .flags
266                                            .insert(key.to_string(), true.into());
267                                        constants.insert(key.as_ref(), value.to_string());
268                                    }
269                                    ShaderDefVal::UInt(key, value) => {
270                                        compiler_options
271                                            .features
272                                            .flags
273                                            .insert(key.to_string(), true.into());
274                                        constants.insert(key.as_ref(), value.to_string());
275                                    }
276                                }
277                            }
278                            let constants_source: String = constants
279                                .iter()
280                                .map(|(name, value)| format!("const {name} = {value};\n"))
281                                .collect();
282
283                            let shader_resolver = ShaderResolver::new(
284                                &self.module_path_to_asset_id,
285                                &self.shaders,
286                                &constants_source,
287                            );
288
289                            let compiled =
290                                wesl::compile(&module_path, &compiler_options, &shader_resolver)
291                                    .map_err(|error| {
292                                        // We use render_plain to avoid rendering ANSI codes
293                                        // Workaround for: https://github.com/tokio-rs/tracing/issues/3378
294                                        if is_module_not_found(&error) {
295                                            if self.missing_import_logged.insert(id) {
296                                                warn!(
297                                                    "Shader `{}` has an unresolved import:\n{}",
298                                                    shader.path,
299                                                    error.diagnostic().render_plain()
300                                                );
301                                            }
302                                            ShaderCacheError::ShaderImportNotYetAvailable
303                                        } else {
304                                            ShaderCacheError::ProcessShaderError(
305                                                error.diagnostic().render_plain(),
306                                            )
307                                        }
308                                    })?;
309
310                            for used in &compiled.modules {
311                                let used_path = match &used.path.origin {
312                                    wesl::syntax::PathOrigin::Package(pkg) if pkg.contains('/') => {
313                                        Cow::Owned(wesl::syntax::ModulePath {
314                                            origin: wesl::syntax::PathOrigin::Package(
315                                                pkg.rsplit('/').next().unwrap().to_string(),
316                                            ),
317                                            components: used.path.components.clone(),
318                                        })
319                                    }
320                                    _ => Cow::Borrowed(&used.path),
321                                };
322                                if let Some(dep_id) =
323                                    self.module_path_to_asset_id.get(used_path.as_ref())
324                                    && *dep_id != id
325                                {
326                                    wesl_dependencies.push(*dep_id);
327                                }
328                            }
329
330                            ShaderCacheSource::Wgsl(compiled.to_string())
331                        } else {
332                            return Err(ShaderCacheError::ProcessShaderError(format!(
333                                "Wesl shader `{}` has a malformed import path `{:?}`",
334                                shader.path, shader.import_path
335                            )));
336                        }
337                    }
338                    Source::Wgsl(wgsl_source) => ShaderCacheSource::Wgsl(wgsl_source.to_string()),
339                };
340
341                let shader_module =
342                    (self.load_module)(&self.device, shader_source, &shader.validate_shader)?;
343
344                entry.insert(Arc::new(shader_module))
345            }
346        };
347        let module = module.clone();
348
349        for dep_id in wesl_dependencies {
350            self.data.entry(dep_id).or_default().dependents.insert(id);
351        }
352
353        Ok(module)
354    }
355
356    fn clear(&mut self, id: AssetId<Shader>) -> Vec<CachedPipelineId> {
357        let mut shaders_to_clear = vec![id];
358        let mut visited = HashSet::new();
359        let mut pipelines_to_queue = Vec::new();
360        while let Some(handle) = shaders_to_clear.pop() {
361            if !visited.insert(handle) {
362                continue;
363            }
364            if let Some(data) = self.data.get_mut(&handle) {
365                data.processed_shaders.clear();
366                pipelines_to_queue.extend(data.pipelines.iter().copied());
367                shaders_to_clear.extend(data.dependents.iter().copied());
368            }
369        }
370
371        pipelines_to_queue
372    }
373
374    /// Inserts and possibly replaces a shader at the given asset id.
375    ///
376    /// Returns a vec of which cached pipelines depended on it
377    /// (directly or indirectly via a shader import) and thus must be recompiled.
378    pub fn set_shader(&mut self, id: AssetId<Shader>, shader: Shader) -> Vec<CachedPipelineId> {
379        self.missing_import_logged.clear();
380        let mut pipelines_to_queue = self.clear(id);
381        let path = &shader.import_path;
382        self.import_path_shaders.insert(path.clone(), id);
383        if let Some(waiting_shaders) = self.waiting_on_import.get_mut(path) {
384            for waiting_shader in waiting_shaders.drain(..) {
385                // resolve waiting shader import
386                let data = self.data.entry(waiting_shader).or_default();
387                data.resolved_imports.insert(path.clone(), id);
388                // add waiting shader as dependent of this shader
389                let data = self.data.entry(id).or_default();
390                data.dependents.insert(waiting_shader);
391            }
392        }
393
394        for import in shader.imports.iter() {
395            if let Some(import_id) = self.import_path_shaders.get(import).copied() {
396                // resolve import because it is currently available
397                let data = self.data.entry(id).or_default();
398                data.resolved_imports.insert(import.clone(), import_id);
399                // add this shader as a dependent of the import
400                let data = self.data.entry(import_id).or_default();
401                data.dependents.insert(id);
402            } else {
403                let waiting = self.waiting_on_import.entry(import.clone()).or_default();
404                waiting.push(id);
405            }
406        }
407
408        if let Source::Wesl(_) = shader.source {
409            match wesl_module_path(&shader.import_path) {
410                Some(module_path) => {
411                    if let Some(existing) = self.module_path_to_asset_id.get(&module_path).copied()
412                        && existing != id
413                    {
414                        warn!(
415                            "Shader module path `{module_path}` was registered to a different \
416                            shader, replacing with `{}`",
417                            shader.path
418                        );
419                        pipelines_to_queue.extend(self.clear(existing));
420                    }
421                    self.module_path_to_asset_id.insert(module_path, id);
422                }
423                None => warn!(
424                    "Shader `{}` has a malformed import path `{:?}`",
425                    shader.path, shader.import_path
426                ),
427            }
428        }
429        self.shaders.insert(id, shader);
430        pipelines_to_queue
431    }
432
433    /// Removes the shader with the given asset id.
434    ///
435    /// Returns a vec of which cached pipelines depended on it
436    /// (directly or indirectly via a shader import) and thus must be recompiled.
437    pub fn remove(&mut self, id: AssetId<Shader>) -> Vec<CachedPipelineId> {
438        let pipelines_to_queue = self.clear(id);
439        if let Some(shader) = self.shaders.remove(&id) {
440            if self.import_path_shaders.get(&shader.import_path) == Some(&id) {
441                self.import_path_shaders.remove(&shader.import_path);
442            }
443            if let Source::Wesl(_) = shader.source
444                && let Some(module_path) = wesl_module_path(&shader.import_path)
445                && self.module_path_to_asset_id.get(&module_path) == Some(&id)
446            {
447                self.module_path_to_asset_id.remove(&module_path);
448            }
449        }
450
451        pipelines_to_queue
452    }
453}
454
455/// A Wesl import resolver. Maps module paths to actual Wesl shader source.
456pub struct ShaderResolver<'a> {
457    module_path_to_asset_id: &'a HashMap<wesl::syntax::ModulePath, AssetId<Shader>>,
458    shaders: &'a HashMap<AssetId<Shader>, Shader>,
459    constants_source: &'a str,
460}
461
462impl<'a> ShaderResolver<'a> {
463    /// Creates a shader resolver with the given map of module paths to shader asset ids,
464    /// map of shader asset ids to shader source, and the source of the virtual
465    /// `constants` module. This resolver is not meant to be long living.
466    pub fn new(
467        module_path_to_asset_id: &'a HashMap<wesl::syntax::ModulePath, AssetId<Shader>>,
468        shaders: &'a HashMap<AssetId<Shader>, Shader>,
469        constants_source: &'a str,
470    ) -> Self {
471        Self {
472            module_path_to_asset_id,
473            shaders,
474            constants_source,
475        }
476    }
477}
478
479impl<'a> wesl::Resolver for ShaderResolver<'a> {
480    fn resolve_source(
481        &self,
482        module_path: &wesl::syntax::ModulePath,
483    ) -> Result<Cow<'_, str>, wesl::error::ResolveError> {
484        let module_path = self.canonical_path(module_path);
485        if module_path.origin == wesl::syntax::PathOrigin::Package("constants".to_string())
486            && module_path.components.is_empty()
487        {
488            return Ok(Cow::Borrowed(self.constants_source));
489        }
490        let asset_id = self
491            .module_path_to_asset_id
492            .get(&module_path)
493            .ok_or_else(|| {
494                wesl::error::ResolveError::ModuleNotFound(
495                    module_path.clone(),
496                    "no shader registered for this module path".to_string(),
497                )
498            })?;
499
500        let shader = self.shaders.get(asset_id).ok_or_else(|| {
501            wesl::error::ResolveError::ModuleNotFound(
502                module_path.clone(),
503                "shader asset not loaded".to_string(),
504            )
505        })?;
506        Ok(Cow::Borrowed(shader.source.as_str()))
507    }
508
509    fn canonical_path(&self, module_path: &wesl::syntax::ModulePath) -> wesl::syntax::ModulePath {
510        match &module_path.origin {
511            wesl::syntax::PathOrigin::Package(pkg) if pkg.contains('/') => {
512                wesl::syntax::ModulePath {
513                    origin: wesl::syntax::PathOrigin::Package(
514                        pkg.rsplit('/').next().unwrap().to_string(),
515                    ),
516                    components: module_path.components.clone(),
517                }
518            }
519            _ => module_path.clone(),
520        }
521    }
522
523    fn display_name(&self, module_path: &wesl::syntax::ModulePath) -> Option<String> {
524        let module_path = self.canonical_path(module_path);
525        let asset_id = self.module_path_to_asset_id.get(&module_path)?;
526        let shader = self.shaders.get(asset_id)?;
527        Some(shader.path.clone())
528    }
529}
530
531/// Type of error returned by a `PipelineCache` when the creation of a GPU pipeline object failed.
532#[expect(missing_docs, reason = "Enum variants are self-explanatory")]
533#[derive(Error, Debug)]
534pub enum ShaderCacheError {
535    #[error(
536        "Pipeline could not be compiled because the following shader could not be loaded: {0:?}"
537    )]
538    ShaderNotLoaded(AssetId<Shader>),
539    #[error("Failed to process shader:\n{0}")]
540    ProcessShaderError(String),
541    #[error("Shader import not yet available.")]
542    ShaderImportNotYetAvailable,
543    #[error("Could not create shader module: {0}")]
544    CreateShaderModule(String),
545}
546
547#[cfg(test)]
548mod tests {
549    use super::*;
550
551    fn test_cache() -> ShaderCache<String, ()> {
552        ShaderCache::new((), |_, source, _| match source {
553            ShaderCacheSource::Wgsl(wgsl) => Ok(wgsl),
554            _ => panic!("expected wgsl output"),
555        })
556    }
557
558    #[test]
559    fn import_resolution() {
560        let mut cache = test_cache();
561
562        let (maths_id, lighting_id, root_id) = set_test_shaders(&mut cache);
563
564        let compiled = cache
565            .get(0, root_id, &[ShaderDefVal::Bool("BRIGHT".into(), true)])
566            .unwrap();
567        assert!(compiled.contains("fn fragment"));
568        assert!(compiled.contains("* 2.0"));
569        assert!(compiled.contains("+ 0.1"));
570
571        let (maths, lighting, _) = test_shaders();
572        assert!(cache.set_shader(lighting_id, lighting).contains(&0));
573        let _ = cache.get(0, root_id, &[ShaderDefVal::Bool("BRIGHT".into(), true)]);
574        assert!(cache.set_shader(maths_id, maths).contains(&0));
575    }
576
577    #[test]
578    fn constants_module() {
579        let mut cache = test_cache();
580        let mut shader = Shader::from_wesl(
581            r#"
582@group(constants::MATERIAL_BIND_GROUP) @binding(0) var<uniform> scale: f32;
583var<uniform> batch: array<vec4<f32>, constants::BATCH_SIZE>;
584
585@fragment
586fn fragment() -> @location(0) vec4<f32> {
587    return vec4<f32>(scale) + batch[0];
588}
589"#,
590            "shaders/constants_user.wesl",
591        );
592        shader.shader_defs = vec![ShaderDefVal::UInt("BATCH_SIZE".into(), 4)];
593        let id = AssetId::Uuid {
594            uuid: bevy_asset::uuid::Uuid::from_u128(7),
595        };
596        cache.set_shader(id, shader);
597
598        let compiled = cache
599            .get(
600                0,
601                id,
602                &[ShaderDefVal::UInt("MATERIAL_BIND_GROUP".into(), 2)],
603            )
604            .unwrap();
605        assert!(compiled.contains("= 2;"));
606        assert!(compiled.contains("= 4;"));
607    }
608
609    #[test]
610    fn compile_error_display_names() {
611        let mut cache = test_cache();
612        let (maths, _, root) = test_shaders();
613        let (maths_id, lighting_id, root_id) = test_ids();
614        let broken_lighting = Shader::from_wesl(
615            "fn brighten(x: f32) -> f32 { return x + ; }",
616            "embedded://bevy_pbr/render/lighting.wesl",
617        );
618        cache.set_shader(maths_id, maths);
619        cache.set_shader(lighting_id, broken_lighting);
620        cache.set_shader(root_id, root);
621
622        let error = cache
623            .get(0, root_id, &[ShaderDefVal::Bool("BRIGHT".into(), true)])
624            .expect_err("syntax error");
625        let ShaderCacheError::ProcessShaderError(message) = error else {
626            panic!("expected ProcessShaderError, got: {error:?}");
627        };
628        assert!(message.contains("embedded://bevy_pbr/render/lighting.wesl"));
629    }
630
631    #[test]
632    fn import_retry() {
633        let mut cache = test_cache();
634        let (maths, lighting, root) = test_shaders();
635        let (maths_id, lighting_id, root_id) = test_ids();
636
637        cache.set_shader(root_id, root);
638        assert!(matches!(
639            cache.get(0, root_id, &[]),
640            Err(ShaderCacheError::ShaderImportNotYetAvailable)
641        ));
642
643        cache.set_shader(lighting_id, lighting);
644        cache.set_shader(maths_id, maths);
645        cache.get(0, root_id, &[]).unwrap();
646    }
647
648    fn test_shaders() -> (Shader, Shader, Shader) {
649        let maths = Shader::from_wesl(
650            "fn double(x: f32) -> f32 { return x * 2.0; }",
651            "embedded://bevy_render/maths.wesl",
652        );
653        let lighting = Shader::from_wesl(
654            "fn brighten(x: f32) -> f32 { return x + 0.1; }",
655            "embedded://bevy_pbr/render/lighting.wesl",
656        );
657        let root = Shader::from_wesl(
658            r#"
659import bevy_pbr::render::lighting::brighten;
660
661@fragment
662fn fragment() -> @location(0) vec4<f32> {
663    var value = bevy_render::maths::double(1.0);
664    @if(BRIGHT) { value = brighten(value); }
665    return vec4<f32>(value);
666}
667"#,
668            "shaders/root.wesl",
669        );
670        (maths, lighting, root)
671    }
672
673    fn test_ids() -> (AssetId<Shader>, AssetId<Shader>, AssetId<Shader>) {
674        let id = |n| AssetId::Uuid {
675            uuid: bevy_asset::uuid::Uuid::from_u128(n),
676        };
677        (id(1), id(2), id(3))
678    }
679
680    #[test]
681    fn library_def_scoping() {
682        let mut cache = test_cache();
683        let id = |n| AssetId::Uuid {
684            uuid: bevy_asset::uuid::Uuid::from_u128(n),
685        };
686
687        let mut lib_a = Shader::from_wesl(
688            "var<uniform> batch_a: array<vec4<f32>, constants::BATCH_SIZE>;",
689            "embedded://bevy_a/bindings.wesl",
690        );
691        lib_a.shader_defs = vec![ShaderDefVal::UInt("BATCH_SIZE".into(), 3)];
692        let mut lib_b = Shader::from_wesl(
693            "var<uniform> batch_b: array<vec4<f32>, constants::BATCH_SIZE>;",
694            "embedded://bevy_b/bindings.wesl",
695        );
696        lib_b.shader_defs = vec![ShaderDefVal::UInt("BATCH_SIZE".into(), 7)];
697
698        let root_a = Shader::from_wesl(
699            r#"
700import bevy_a::bindings::batch_a;
701@fragment
702fn fragment() -> @location(0) vec4<f32> { return batch_a[0]; }
703"#,
704            "shaders/root_a.wesl",
705        );
706        let root_b = Shader::from_wesl(
707            r#"
708import bevy_b::bindings::batch_b;
709@fragment
710fn fragment() -> @location(0) vec4<f32> { return batch_b[0]; }
711"#,
712            "shaders/root_b.wesl",
713        );
714
715        cache.set_shader(id(1), lib_a);
716        cache.set_shader(id(2), lib_b);
717        cache.set_shader(id(3), root_a);
718        cache.set_shader(id(4), root_b);
719
720        let compiled_a = cache.get(0, id(3), &[]).unwrap();
721        let compiled_b = cache.get(1, id(4), &[]).unwrap();
722        assert!(compiled_a.contains("= 3;") && !compiled_a.contains("= 7;"));
723        assert!(compiled_b.contains("= 7;") && !compiled_b.contains("= 3;"));
724    }
725
726    #[test]
727    fn cyclic_import_invalidation() {
728        let mut cache = test_cache();
729        let id = |n| AssetId::Uuid {
730            uuid: bevy_asset::uuid::Uuid::from_u128(n),
731        };
732        let module_a = Shader::from_wesl(
733            "import bevy_cycle::b::from_b;\nfn from_a() -> f32 { return 1.0; }",
734            "embedded://bevy_cycle/a.wesl",
735        );
736        let module_b = Shader::from_wesl(
737            "import bevy_cycle::a::from_a;\nfn from_b() -> f32 { return 2.0; }",
738            "embedded://bevy_cycle/b.wesl",
739        );
740        cache.set_shader(id(1), module_a);
741        cache.set_shader(id(2), module_b);
742        let module_a = Shader::from_wesl(
743            "import bevy_cycle::b::from_b;\nfn from_a() -> f32 { return 3.0; }",
744            "embedded://bevy_cycle/a.wesl",
745        );
746        cache.set_shader(id(1), module_a);
747    }
748
749    fn set_test_shaders(
750        cache: &mut ShaderCache<String, ()>,
751    ) -> (AssetId<Shader>, AssetId<Shader>, AssetId<Shader>) {
752        let (maths, lighting, root) = test_shaders();
753        let (maths_id, lighting_id, root_id) = test_ids();
754        cache.set_shader(maths_id, maths);
755        cache.set_shader(lighting_id, lighting);
756        cache.set_shader(root_id, root);
757        (maths_id, lighting_id, root_id)
758    }
759}
bevy_shader/shader_cache.rs

bevy_shader/
shader_cache.rs
