Skip to main content

vk_graph_hot/
shader.rs

1//! Hot-reload shader descriptions and builder APIs.
2
3pub use shaderc::{OptimizationLevel, SourceLanguage, SpirvVersion};
4
5use {
6    super::{compile_shader, guess_shader_source_language},
7    derive_builder::{Builder, UninitializedFieldError},
8    log::{Level, debug, error, log_enabled},
9    notify::{RecommendedWatcher, RecursiveMode, Watcher},
10    shaderc::{CompileOptions, EnvVersion, ShaderKind, TargetEnv},
11    std::path::{Path, PathBuf},
12    vk_graph::driver::{DriverError, ash::vk, shader::SpecializationMap},
13};
14
15/// Describes a shader program which runs on some pipeline stage.
16///
17/// _NOTE:_ When compiled on Apple platforms the macro `MOLTEN_VK` will be defined automatically.
18/// This may be used to handle any differences introduced by SPIRV-Cross translation to Metal
19/// Shading Language (MSL) at runtime.
20#[allow(missing_docs)]
21#[derive(Builder, Clone, Debug)]
22#[builder(
23    build_fn(private, name = "fallible_build", error = "UninitializedFieldError"),
24    derive(Clone, Debug),
25    pattern = "owned"
26)]
27pub struct HotShader {
28    /// The name of the entry point which will be executed by this shader.
29    ///
30    /// The default value is `main`.
31    #[builder(default = "\"main\".to_owned()", setter(into))]
32    pub entry_name: String,
33
34    /// Macro definitions.
35    #[builder(default, setter(strip_option))]
36    pub macro_definitions: Option<Vec<(String, Option<String>)>>,
37
38    /// Sets the optimization level.
39    #[builder(default, setter(strip_option))]
40    pub optimization_level: Option<OptimizationLevel>,
41
42    /// Shader source code path.
43    #[builder(setter(custom))]
44    pub path: PathBuf,
45
46    /// Sets the source language.
47    #[builder(default, setter(strip_option))]
48    pub source_language: Option<SourceLanguage>,
49
50    /// Data about Vulkan specialization constants.
51    ///
52    /// # Examples
53    ///
54    /// Basic usage (GLSL):
55    ///
56    /// ```glsl
57    /// // fire.comp
58    /// #version 460 core
59    ///
60    /// // Defaults to 6 if not set using HotShader specialization_info!
61    /// layout(constant_id = 0) const uint MY_COUNT = 6;
62    ///
63    /// layout(set = 0, binding = 0) uniform sampler2D my_samplers[MY_COUNT];
64    ///
65    /// void main()
66    /// {
67    ///     // Code uses MY_COUNT number of my_samplers here
68    /// }
69    /// ```
70    ///
71    /// ```no_run
72    /// use vk_graph::driver::shader::SpecializationMap;
73    /// use vk_graph_hot::HotShader;
74    ///
75    /// // We instead specify 42 for MY_COUNT:
76    /// let shader = HotShader::new_compute("shaders/fire.comp")
77    ///     .specialization(
78    ///         SpecializationMap::new(42u32.to_ne_bytes())
79    ///             .constant(0, 0, 4)
80    ///     );
81    /// ```
82    #[builder(default, setter(strip_option))]
83    pub specialization: Option<SpecializationMap>,
84
85    /// The shader stage this structure applies to.
86    pub stage: vk::ShaderStageFlags,
87
88    /// Sets the target SPIR-V version.
89    #[builder(default, setter(strip_option))]
90    pub target_spirv: Option<SpirvVersion>,
91
92    /// Sets the compiler mode to treat all warnings as errors.
93    #[builder(default)]
94    pub warnings_as_errors: bool,
95}
96
97impl HotShader {
98    /// Specifies a shader with the given `stage` and shader code values.
99    #[allow(clippy::new_ret_no_self)]
100    pub fn new(stage: vk::ShaderStageFlags, path: impl AsRef<Path>) -> HotShaderBuilder {
101        HotShaderBuilder::new(stage, path)
102    }
103
104    /// Creates a new ray tracing shader.
105    pub fn new_any_hit(path: impl AsRef<Path>) -> HotShaderBuilder {
106        Self::new(vk::ShaderStageFlags::ANY_HIT_KHR, path)
107    }
108
109    /// Creates a new ray tracing shader.
110    pub fn new_callable(path: impl AsRef<Path>) -> HotShaderBuilder {
111        Self::new(vk::ShaderStageFlags::CALLABLE_KHR, path)
112    }
113
114    /// Creates a new ray tracing shader.
115    pub fn new_closest_hit(path: impl AsRef<Path>) -> HotShaderBuilder {
116        Self::new(vk::ShaderStageFlags::CLOSEST_HIT_KHR, path)
117    }
118
119    /// Creates a new compute shader.
120    pub fn new_compute(path: impl AsRef<Path>) -> HotShaderBuilder {
121        Self::new(vk::ShaderStageFlags::COMPUTE, path)
122    }
123
124    /// Creates a new fragment shader.
125    pub fn new_fragment(path: impl AsRef<Path>) -> HotShaderBuilder {
126        Self::new(vk::ShaderStageFlags::FRAGMENT, path)
127    }
128
129    /// Creates a new geometry shader.
130    pub fn new_geometry(path: impl AsRef<Path>) -> HotShaderBuilder {
131        Self::new(vk::ShaderStageFlags::GEOMETRY, path)
132    }
133
134    /// Creates a new ray tracing shader.
135    pub fn new_intersection(path: impl AsRef<Path>) -> HotShaderBuilder {
136        Self::new(vk::ShaderStageFlags::INTERSECTION_KHR, path)
137    }
138
139    /// Creates a new mesh shader.
140    pub fn new_mesh(path: impl AsRef<Path>) -> HotShaderBuilder {
141        Self::new(vk::ShaderStageFlags::MESH_EXT, path)
142    }
143
144    /// Creates a new ray tracing shader.
145    pub fn new_miss(path: impl AsRef<Path>) -> HotShaderBuilder {
146        Self::new(vk::ShaderStageFlags::MISS_KHR, path)
147    }
148
149    /// Creates a new ray tracing shader.
150    pub fn new_ray_gen(path: impl AsRef<Path>) -> HotShaderBuilder {
151        Self::new(vk::ShaderStageFlags::RAYGEN_KHR, path)
152    }
153
154    /// Creates a new mesh task shader.
155    pub fn new_task(path: impl AsRef<Path>) -> HotShaderBuilder {
156        Self::new(vk::ShaderStageFlags::TASK_EXT, path)
157    }
158
159    /// Creates a new tessellation control shader.
160    pub fn new_tessellation_ctrl(path: impl AsRef<Path>) -> HotShaderBuilder {
161        Self::new(vk::ShaderStageFlags::TESSELLATION_CONTROL, path)
162    }
163
164    /// Creates a new tessellation evaluation shader.
165    pub fn new_tessellation_eval(path: impl AsRef<Path>) -> HotShaderBuilder {
166        Self::new(vk::ShaderStageFlags::TESSELLATION_EVALUATION, path)
167    }
168
169    /// Creates a new vertex shader.
170    pub fn new_vertex(path: impl AsRef<Path>) -> HotShaderBuilder {
171        Self::new(vk::ShaderStageFlags::VERTEX, path)
172    }
173
174    pub(super) fn compile_and_watch(
175        &self,
176        watcher: &mut RecommendedWatcher,
177    ) -> Result<Vec<u8>, DriverError> {
178        let shader_kind = if self.stage == vk::ShaderStageFlags::empty() {
179            None
180        } else {
181            Some(match self.stage {
182                vk::ShaderStageFlags::ANY_HIT_KHR => ShaderKind::AnyHit,
183                vk::ShaderStageFlags::CALLABLE_KHR => ShaderKind::Callable,
184                vk::ShaderStageFlags::CLOSEST_HIT_KHR => ShaderKind::ClosestHit,
185                vk::ShaderStageFlags::COMPUTE => ShaderKind::Compute,
186                vk::ShaderStageFlags::FRAGMENT => ShaderKind::Fragment,
187                vk::ShaderStageFlags::GEOMETRY => ShaderKind::Geometry,
188                vk::ShaderStageFlags::INTERSECTION_KHR => ShaderKind::Intersection,
189                vk::ShaderStageFlags::MISS_KHR => ShaderKind::Miss,
190                vk::ShaderStageFlags::RAYGEN_KHR => ShaderKind::RayGeneration,
191                vk::ShaderStageFlags::TASK_EXT => ShaderKind::Task,
192                vk::ShaderStageFlags::TESSELLATION_CONTROL => ShaderKind::TessControl,
193                vk::ShaderStageFlags::TESSELLATION_EVALUATION => ShaderKind::TessEvaluation,
194                vk::ShaderStageFlags::VERTEX => ShaderKind::Vertex,
195                _ => {
196                    error!(
197                        "unsupported shader stage for shaderc kind inference: {:?}",
198                        self.stage
199                    );
200                    return Err(DriverError::Unsupported);
201                }
202            })
203        };
204
205        let mut additional_opts = CompileOptions::new().map_err(|err| {
206            error!("unable to initialize compiler options: {err:?}");
207
208            DriverError::Unsupported
209        })?;
210
211        if let Some(macro_definitions) = &self.macro_definitions {
212            for (name, value) in macro_definitions {
213                additional_opts.add_macro_definition(name, value.as_deref());
214            }
215        }
216
217        additional_opts.set_target_env(TargetEnv::Vulkan, EnvVersion::Vulkan1_2 as _);
218
219        if let Some(language) = self.source_language.or_else(|| {
220            let language = guess_shader_source_language(&self.path);
221
222            if let Some(language) = language {
223                debug!("Guessed source language: {:?}", language);
224            }
225
226            language
227        }) {
228            additional_opts.set_source_language(language);
229        }
230
231        additional_opts.set_target_spirv(self.target_spirv.unwrap_or(SpirvVersion::V1_5));
232
233        if let Some(level) = self.optimization_level {
234            additional_opts.set_optimization_level(level);
235        }
236
237        if self.warnings_as_errors {
238            additional_opts.set_warnings_as_errors();
239        }
240
241        let res = compile_shader(
242            &self.path,
243            &self.entry_name,
244            shader_kind,
245            Some(&additional_opts),
246        )
247        .map_err(|err| {
248            if !log_enabled!(Level::Error) {
249                panic!("unable to compile shader {}: {err}", self.path.display());
250            }
251
252            error!("unable to compile shader {}: {err}", self.path.display());
253
254            DriverError::InvalidData
255        })?;
256
257        for path in res.files_included {
258            watcher
259                .watch(&path, RecursiveMode::NonRecursive)
260                .map_err(|err| {
261                    error!("unable to watch file: {err}");
262
263                    DriverError::Unsupported
264                })?;
265        }
266
267        Ok(res.spirv_code)
268    }
269
270    /// Creates a shader using a shader kind inferred from the source code.
271    pub fn from_path(path: impl AsRef<Path>) -> HotShaderBuilder {
272        HotShaderBuilder::default().path(path)
273    }
274}
275
276impl From<HotShaderBuilder> for HotShader {
277    fn from(builder: HotShaderBuilder) -> HotShader {
278        builder.build()
279    }
280}
281
282// HACK: https://github.com/colin-kiegel/rust-derive-builder/issues/56
283impl HotShaderBuilder {
284    /// Specifies a shader with the given `stage` and shader path values.
285    pub fn new(stage: vk::ShaderStageFlags, path: impl AsRef<Path>) -> Self {
286        Self::default().stage(stage).path(path)
287    }
288
289    /// Builds a new `HotShader`.
290    pub fn build(self) -> HotShader {
291        self.try_build().expect("invalid hot shader")
292    }
293
294    /// Builds a new `HotShader`, returning an error if required fields are missing.
295    pub fn try_build(self) -> Result<HotShader, UninitializedFieldError> {
296        let this = self;
297
298        #[cfg(target_os = "macos")]
299        let this = this.macro_definition("MOLTEN_VK", Some("1".to_string()));
300
301        let mut this = this;
302
303        if this.stage.is_none() {
304            this.stage = Some(vk::ShaderStageFlags::empty());
305        }
306
307        this.fallible_build()
308    }
309
310    /// Defines a single macro.
311    pub fn macro_definition(
312        mut self,
313        key: impl Into<String>,
314        value: impl Into<Option<String>>,
315    ) -> Self {
316        let macro_definitions = self
317            .macro_definitions
318            .get_or_insert_with(|| Some(Vec::new()))
319            .get_or_insert_with(Vec::new);
320        macro_definitions.push((key.into(), value.into()));
321
322        self
323    }
324
325    /// Shader source code path.
326    pub fn path(mut self, path: impl AsRef<Path>) -> Self {
327        self.path = Some(path.as_ref().to_owned());
328        self
329    }
330}
331
332#[cfg(test)]
333mod test {
334    use super::*;
335
336    #[test]
337    fn new_compute_sets_stage_and_path() {
338        let shader = HotShader::new_compute("shader.comp").build();
339
340        assert_eq!(shader.stage, vk::ShaderStageFlags::COMPUTE);
341        assert_eq!(shader.path, PathBuf::from("shader.comp"));
342        assert_eq!(shader.entry_name, "main");
343    }
344
345    #[test]
346    fn from_path_defaults_to_empty_stage_for_inference() {
347        let shader = HotShader::from_path("shader.glsl").build();
348
349        assert!(shader.stage.is_empty());
350        assert_eq!(shader.path, PathBuf::from("shader.glsl"));
351    }
352
353    #[test]
354    fn macro_definition_accumulates_values() {
355        let shader = HotShader::new_fragment("shader.frag")
356            .macro_definition("FOO", Some("1".to_owned()))
357            .macro_definition("BAR", None::<String>)
358            .build();
359
360        assert_eq!(
361            shader.macro_definitions,
362            Some(vec![
363                ("FOO".to_owned(), Some("1".to_owned())),
364                ("BAR".to_owned(), None),
365            ])
366        );
367    }
368
369    #[test]
370    fn try_build_requires_path() {
371        let err = HotShaderBuilder::default().try_build().unwrap_err();
372
373        assert_eq!(err.field_name(), "path");
374    }
375}