Skip to main content

wesl/
package.rs

1//! Shader packaging.
2
3use std::{
4    collections::HashSet,
5    path::{Path, PathBuf},
6};
7
8use crate::{
9    error::{Diagnostic, Error, TomlError},
10    pass,
11    toml_cfg::{self, WeslToml},
12};
13use wgsl_parse::{
14    lexer::{Lexer, Token},
15    syntax::{ModulePath, TranslationUnit},
16};
17use wgsl_types::idents::RESERVED_WORDS;
18
19pub use wesl_core::{StaticPackage, StaticPackageModule};
20
21#[cfg(feature = "package")]
22use quote::{format_ident, quote};
23
24/// WGSL identifier predicate, including reserved words, but excluding keywords.
25pub(crate) fn is_mod_ident(name: &str) -> bool {
26    let mut lex = Lexer::new(name);
27    RESERVED_WORDS.contains(&name)
28        || matches!(
29            (lex.next(), lex.next()),
30            (Some(Ok((_, Token::Ident(_), _))), None)
31        )
32}
33
34// https://github.com/webgpu-tools/wesl-spec/blob/main/Imports.md#reference-level-explanation
35pub(crate) const RESERVED_MOD_NAMES: &[&str] = &[
36    // WGSL keywords
37    "const_assert",
38    "continue",
39    "continuing",
40    "default",
41    "diagnostic",
42    "discard",
43    "else",
44    "enable",
45    "false",
46    "fn",
47    "for",
48    "if",
49    "let",
50    "loop",
51    "override",
52    "requires",
53    "return",
54    "struct",
55    "switch",
56    "true",
57    "var",
58    "while",
59    // WESL keywords
60    "self",
61    "super",
62    "package",
63    "as",
64    "import",
65];
66
67/// A builder that generates code for WESL packages.
68///
69/// WESL packages are bundles of shader files that can be reused in other projects, like
70/// Rust crates. See the `consumer` example to see how to use them.
71///
72/// This type is mainly designed to be used in a build script (`build.rs` file). Add `wesl`
73/// to the build-dependencies of your project and enable the `package` feature flag.
74///
75/// # Usage
76///
77/// ```no_run
78/// // in build.rs
79/// # #[allow(clippy::needless_doctest_main)]
80/// fn main() {
81/// #  #[cfg(feature = "package")]
82///    wesl::package::PackageBuilder::new("my_package")
83///        // read all wesl files in the directory "src/shaders"
84///        .scan_root("src/shaders")
85///        .expect("failed to scan WESL files")
86///        // validation is optional
87///        .validate()
88///        .inspect_err(|e| eprintln!("{e}"))
89///        .expect("validation error")
90///        // write "my_package.rs" in OUT_DIR
91///        .build_artifact()
92///        .expect("failed to build artifact");
93/// }
94/// ```
95///
96/// Then, in your `lib.rs` file, expose the generated module with the [`wesl_core::wesl_pkg`] macro.
97///
98/// ```rust,ignore
99/// wesl::wesl_pkg!(pub my_package);
100/// ```
101///
102/// The package name must be a valid rust identifier, E.g. it must not contain dashes `-`.
103/// Dashes are replaced with underscores `_`.
104pub struct PackageBuilder {
105    name: String,
106    dependencies: Vec<&'static StaticPackage>,
107}
108
109/// The type holding the source code of packages.
110///
111/// This struct is produced by [`PackageBuilder::scan_root`], but one can also create or edit
112/// packages manually by modifying this struct. The final package is produced by calling
113/// [`Self::build_artifact`] or [`Self::codegen`].
114pub struct Package {
115    pub crate_name: String,
116    pub root: PackageModule,
117    pub dependencies: Vec<&'static StaticPackage>,
118}
119
120/// The type holding the source code of individual modules in packages.
121///
122/// See [`Package`].
123#[derive(Debug)]
124pub struct PackageModule {
125    pub name: String,
126    pub source: String,
127    pub submodules: Vec<PackageModule>,
128}
129
130/// Error type for [`PackageBuilder::scan_root`].
131#[derive(Debug, thiserror::Error)]
132pub enum ScanDirectoryError {
133    #[error("Package root directory not found: `{0}`")]
134    RootNotFound(PathBuf),
135    #[error("Module name `{0}` is reserved")]
136    ReservedModName(String),
137    #[error("I/O error while scanning package directory: {0}")]
138    Io(#[from] std::io::Error),
139}
140
141fn cargo_crate_name() -> String {
142    std::env::var("CARGO_PKG_NAME").expect("CARGO_PKG_NAME environment variable is not defined")
143}
144
145impl PackageBuilder {
146    pub fn new(name: &str) -> Self {
147        Self {
148            name: name.replace('-', "_"),
149            dependencies: Vec::new(),
150        }
151    }
152
153    /// Add a package dependency.
154    ///
155    /// Learn more about packages in [`PackageBuilder`].
156    pub fn add_package(mut self, pkg: &'static StaticPackage) -> Self {
157        self.dependencies.push(pkg);
158        self
159    }
160
161    /// Add several package dependencies.
162    ///
163    /// Learn more about packages in [`PackageBuilder`].
164    pub fn add_packages(mut self, pkgs: impl IntoIterator<Item = &'static StaticPackage>) -> Self {
165        for pkg in pkgs {
166            self = self.add_package(pkg);
167        }
168        self
169    }
170
171    // TODO: semantics have changed: the package root module is now in a special `package.wesl` file relative to the root dir.
172    /// Read all files to include in the package, starting from the package root directory.
173    ///
174    /// The package will include all `.wesl` and `.wgsl` files reachable from the package root, recursively.
175    pub fn scan_root(self, pkg_root_dir: impl AsRef<Path>) -> Result<Package, ScanDirectoryError> {
176        fn process_path(path: &Path) -> Result<Option<PackageModule>, ScanDirectoryError> {
177            let path_with_ext_wesl = path.with_extension("wesl");
178            let path_with_ext_wgsl = path.with_extension("wgsl");
179            let path_without_ext = path.with_extension("");
180            let path_filename = path_without_ext
181                .file_stem()
182                .unwrap()
183                .to_string_lossy()
184                .to_string();
185
186            if !path_with_ext_wesl.is_file()
187                && !path_with_ext_wgsl.is_file()
188                && !path_without_ext.is_dir()
189            {
190                return Ok(None);
191            }
192
193            if RESERVED_MOD_NAMES.contains(&path_filename.as_str()) {
194                return Err(ScanDirectoryError::ReservedModName(path_filename));
195            }
196
197            if !is_mod_ident(&path_filename) {
198                // skip file or directory names that are not valid identifiers,
199                // but WGSL reserved words not in RESERVED_MOD_NAMES are allowed.
200                println!(
201                    "cargo::warning=skipped file/dir: not a WGSL ident `{path_filename}` {path:?}"
202                );
203                return Ok(None);
204            }
205
206            // check for source file
207            let source = if path_with_ext_wesl.is_file() {
208                std::fs::read_to_string(&path_with_ext_wesl)?
209            } else if path_with_ext_wgsl.is_file() {
210                std::fs::read_to_string(&path_with_ext_wgsl)?
211            } else {
212                String::new()
213            };
214
215            // check for submodules
216            let mut submodules = Vec::new();
217            if path_without_ext.is_dir() {
218                // use hashset to avoid duplicate entries
219                let mut unique_submodules = HashSet::new();
220                for entry in std::fs::read_dir(&path_without_ext)? {
221                    let Ok(entry) = entry else { continue };
222                    let submodule_path = entry.path().with_extension("");
223                    unique_submodules.insert(submodule_path);
224                }
225                for entry in unique_submodules {
226                    // errors in the top module should be returned
227                    // other errors should only be logged
228                    match process_path(&entry) {
229                        Ok(Some(module)) => submodules.push(module),
230                        Ok(None) => {}
231                        Err(err) => {
232                            println!("cargo::error=error processing submodule {entry:?}: {err}")
233                        }
234                    }
235                }
236            };
237
238            // check for empty module
239            if source.is_empty() && submodules.is_empty() {
240                return Ok(None);
241            }
242
243            let module = PackageModule {
244                name: path_filename,
245                source,
246                submodules,
247            };
248
249            Ok(Some(module))
250        }
251
252        let pkg_root_dir = pkg_root_dir.as_ref().to_path_buf();
253        let potential_module = process_path(&pkg_root_dir)?;
254        let Some(mut module) = potential_module else {
255            return Err(ScanDirectoryError::RootNotFound(pkg_root_dir));
256        };
257        // top level module should be named by package builder and not file path
258        module.name = self.name;
259
260        Ok(Package {
261            crate_name: cargo_crate_name(),
262            root: module,
263            dependencies: self.dependencies,
264        })
265    }
266
267    // TODO: semantics have changed: the package root module is now in a special `package.wesl` file relative to the root dir.
268    /// Reads a wesl.toml file and builds a package from its configuration.
269    ///
270    /// The wesl.toml file should be located in the given directory and contain:
271    /// - `edition` (required), `root`, `include`, `exclude` as top-level keys
272    /// - `[dependencies]` section (optional)
273    ///
274    /// # Example
275    ///
276    /// ```no_run
277    /// # #[cfg(feature = "package")]
278    /// wesl::package::PackageBuilder::new("my_package")
279    ///     .scan_toml(".")  // looks for ./wesl.toml
280    ///     .expect("failed to scan WESL files")
281    ///     .build_artifact()
282    ///     .expect("failed to build artifact");
283    /// ```
284    pub fn scan_toml(self, toml_path: impl AsRef<Path>) -> Result<Package, TomlError> {
285        let mut toml_path = toml_path.as_ref().to_path_buf();
286
287        if toml_path.is_dir() {
288            toml_path.push("wesl.toml");
289        }
290
291        if !toml_path.is_file() {
292            return Err(TomlError::TomlNotFound(toml_path));
293        }
294
295        let dir = toml_path.parent().unwrap(/* SAFETY: toml_path is a file, must have a parent. */);
296
297        let config = WeslToml::from_file(&toml_path)?;
298        let result = toml_cfg::scan_from_config(&self.name, dir, &config)?;
299
300        for warning in &result.warnings {
301            println!("cargo::warning={warning}");
302        }
303
304        Ok(Package {
305            crate_name: cargo_crate_name(),
306            root: result.module,
307            dependencies: self.dependencies,
308        })
309    }
310}
311
312impl PackageModule {
313    fn validate(&self, parent_path: ModulePath) -> Result<(), Error> {
314        let mut path = parent_path.clone();
315        path.push(&self.name);
316
317        eprintln!("INFO: validate {path}");
318
319        let to_diagnostic = |e: Error| {
320            Diagnostic::from(e)
321                .with_module_path(path.clone(), None)
322                .with_source(self.source.clone())
323        };
324        let mut module: TranslationUnit = self
325            .source
326            .parse()
327            .map_err(|e: wgsl_parse::Error| to_diagnostic(e.into()))?;
328        pass::retarget_idents(&mut module);
329        pass::validate_wesl(&module).map_err(|e| to_diagnostic(e.into()))?;
330        for module in &self.submodules {
331            module.validate(path.clone())?;
332        }
333        Ok(())
334    }
335}
336
337impl Package {
338    /// Run [validation][pass::validate_wesl] on each of the scanned files.
339    pub fn validate(self) -> Result<Self, Error> {
340        self.root.validate(ModulePath::new_root())?;
341        Ok(self)
342    }
343}
344
345#[cfg(feature = "package")]
346impl PackageModule {
347    fn codegen(&self) -> proc_macro2::TokenStream {
348        let mod_ident = format_ident!("r#{}", self.name);
349        let name = &self.name;
350        let source = &self.source;
351
352        let submodules = self.submodules.iter().map(|submod| {
353            let name = &submod.name;
354            let ident = format_ident!("r#{}", name);
355            quote! { &#ident::MODULE }
356        });
357
358        let submods = self.submodules.iter().map(|submod| submod.codegen());
359
360        quote! {
361            pub mod #mod_ident {
362                use super::StaticPackageModule;
363                pub static MODULE: StaticPackageModule = StaticPackageModule {
364                    name: #name,
365                    source: #source,
366                    submodules: &[#(#submodules),*]
367                };
368
369                #(#submods)*
370            }
371        }
372    }
373}
374
375#[cfg(feature = "package")]
376impl Package {
377    /// Generate the rust code that holds the packaged wesl files.
378    /// You probably want to use [`Self::build_artifact`] instead.
379    pub fn codegen(&self) -> String {
380        let deps = self.dependencies.iter().map(|dep| {
381            let crate_name = format_ident!("r#{}", dep.crate_name);
382            let mod_name = format_ident!("r#{}", dep.root.name);
383            quote! { &#crate_name::#mod_name::PACKAGE }
384        });
385
386        let crate_name = &self.crate_name;
387        let root_name = &self.root.name;
388        let root_source = &self.root.source;
389
390        let submodules = self.root.submodules.iter().map(|submod| {
391            let name = &submod.name;
392            let ident = format_ident!("r#{}", name);
393            quote! { &#ident::MODULE }
394        });
395
396        let submods = self.root.submodules.iter().map(|submod| submod.codegen());
397
398        let tokens = quote! {
399            pub static PACKAGE: StaticPackage = StaticPackage {
400                crate_name: #crate_name,
401                root: &MODULE,
402                dependencies: &[#(#deps),*],
403            };
404
405            pub static MODULE: StaticPackageModule = StaticPackageModule {
406                name: #root_name,
407                source: #root_source,
408                submodules: &[#(#submodules),*]
409            };
410
411            #(#submods)*
412        };
413
414        tokens.to_string()
415    }
416
417    /// Generate the build artifact that can then be exposed by the [`wesl_core::wesl_pkg`] macro.
418    ///
419    /// This function must be called from a `build.rs` file. Refer to the crate documentation
420    /// for more details.
421    ///
422    /// # Panics
423    /// Panics if the OUT_DIR environment variable is not set. This should not happen if
424    /// ran from a `build.rs` file.
425    pub fn build_artifact(&self) -> std::io::Result<()> {
426        let code = self.codegen();
427        let out_dir = std::path::Path::new(
428            &std::env::var_os("OUT_DIR").expect("OUT_DIR environment variable is not defined"),
429        )
430        .join(format!("{}.rs", self.root.name));
431        std::fs::write(&out_dir, code)?;
432        Ok(())
433    }
434}