wesl/
frontend.rs

1//! WESL compiler frontend.
2//!
3//! Types and functions in this module are small wrappers that make it easy to use WESL in typical scenarios, "batteries included".
4
5use std::{
6    collections::HashSet,
7    path::{Path, PathBuf},
8};
9
10use wesl_core::StaticPackage;
11use wgsl_parse::{
12    SyntaxNode,
13    syntax::{Ident, ModulePath, PathOrigin, TranslationUnit, Visibility},
14};
15
16use crate::{
17    SyntaxUtil,
18    error::{Diagnostic, Error, ResolveError, UsageError},
19    mangler::{self, Mangler},
20    pass::{self, CompilerDriver, Features, Module, UsedItems},
21    resolver::{Constants, Resolver, StandardResolver},
22    sourcemap::{BasicSourceMap, SourceMapper},
23};
24
25/// Compilation options used by [`Compiler`].
26#[derive(Clone, Debug, PartialEq)]
27pub struct CompileOptions {
28    /// Toggle [WESL Imports](https://github.com/webgpu-tools/wesl-spec/blob/main/Imports.md).
29    ///
30    /// If disabled:
31    ///
32    /// * The compiler will silently remove all import statements and inline paths.
33    /// * Validation will not trigger an error if referencing an imported item.
34    pub imports: bool,
35    /// Toggle [WESL Conditional Translation](https://github.com/webgpu-tools/wesl-spec/blob/main/ConditionalTranslation.md).
36    ///
37    /// See [`Self::features`] to enable/disable each feature flag.
38    pub condcomp: bool,
39    /// Toggle [WESL Visibility](https://github.com/webgpu-tools/wesl-spec/blob/main/Visibility.md)
40    ///
41    /// When disabled, WESL simply ignores `public` and `private` keywords on declarations and treats
42    /// all declarations as public.
43    pub visibility: bool,
44    /// Toggle generics. Generics are super experimental, don't expect anything from it.
45    ///
46    /// Requires the `generics` crate feature flag.
47    pub generics: bool,
48    /// Enable stripping (aka. Dead Code Elimination).
49    ///
50    /// By default, all declarations reachable by entrypoint functions, const_asserts and
51    /// pipeline-overridable constants in the main module are kept.
52    /// See [`Self::keep`] and [`Self::keep_main`] to control what gets stripped.
53    ///
54    /// Stripping can have side-effects: modules are loaded only if statically accessed,
55    /// and `const_assert` statements are not always preserved.
56    /// Refer to the WESL docs to learn more.
57    pub strip: bool,
58    /// Enable lowering/polyfills. This transforms the output code in various ways.
59    ///
60    /// See [`pass::lower`] for the list of transforms.
61    pub lower: bool,
62    /// Enable validation of individual WESL modules and of the final output.
63    ///
64    /// This will catch *some* errors, not all.
65    /// See [`pass::validate_wesl`] and [`pass::validate_wgsl`] for the list of validations.
66    ///
67    /// Requires the `eval` crate feature flag.
68    pub validate: bool,
69    /// Enable sourcemapping, which provides better error diagnostics.
70    pub sourcemap: bool,
71    /// Declaration name mangling scheme.
72    pub mangler: ManglerKind,
73    /// Enable mangling of declarations in the main module.
74    ///
75    /// By default, WESL does not mangle main module declarations.
76    pub mangle_main: bool,
77    /// If `Some`, specify a list of main module declarations to keep.
78    /// If `None`, only the entrypoint functions (and their dependencies) are kept.
79    ///
80    /// This option has no effect if [`Self::keep_main`] is enabled or  [`Self::strip`] is
81    /// disabled.
82    pub keep: Option<Vec<String>>,
83    /// If `true`, all main module declarations are preserved when stripping is enabled.
84    ///
85    /// This option takes precedence over [`Self::keep`], and has no effect if
86    /// [`Self::strip`] is disabled.
87    pub keep_main: bool,
88    /// Conditional Translation feature flags.
89    ///
90    /// This option has no effect if [`Self::condcomp`] is disabled.
91    ///
92    /// See [`Features`].
93    pub features: Features,
94    /// Literal constants in the `constants` virtual module.
95    ///
96    /// See [`Constants`].
97    pub constants: Constants,
98    /// Importable packages dependencies.
99    pub dependencies: Vec<&'static StaticPackage>,
100}
101
102/// Declaration name mangling scheme. Used in [`CompileOptions::mangler`].
103#[derive(Default, Clone, Copy, Debug, PartialEq, Eq)]
104pub enum ManglerKind {
105    /// Escaped path mangler.
106    /// `foo_bar::item -> _1foo_bar_item`
107    #[default]
108    Escape,
109    /// Hash mangler.
110    /// `foo::bar::item -> item_1985638328947`
111    Hash,
112    /// Make valid identifiers with unicode "confusables" characters.
113    /// `foo::bar<baz, moo> -> foo::barᐸbazˏmooᐳ`
114    Unicode,
115    /// Disable mangling. (warning: will break shaders if case of name conflicts!)
116    None,
117}
118
119impl From<ManglerKind> for Box<dyn Mangler> {
120    fn from(kind: ManglerKind) -> Self {
121        match kind {
122            ManglerKind::Escape => Box::new(mangler::EscapeMangler),
123            ManglerKind::Hash => Box::new(mangler::HashMangler),
124            ManglerKind::Unicode => Box::new(mangler::UnicodeMangler),
125            ManglerKind::None => Box::new(mangler::NoMangler),
126        }
127    }
128}
129
130impl Default for CompileOptions {
131    fn default() -> Self {
132        Self {
133            imports: true,
134            condcomp: true,
135            visibility: true,
136            generics: false,
137            strip: true,
138            lower: false,
139            validate: true,
140            sourcemap: true,
141            mangler: Default::default(),
142            mangle_main: false,
143            keep: Default::default(),
144            keep_main: false,
145            features: Default::default(),
146            constants: Default::default(),
147            dependencies: Default::default(),
148        }
149    }
150}
151
152/// The WESL compiler.
153///
154/// # Basic Usage
155///
156/// ```rust
157/// # use wesl::{Compiler, resolver::VirtualResolver};
158/// #
159/// let compiler = Compiler::default();
160/// #
161/// # // just adding a virtual file here so the doctest runs without a filesystem
162/// # let mut resolver = VirtualResolver::new();
163/// # let shader_string = "fn my_fn() {\n\n}\n";
164/// # resolver.add_module("package::path::to::shader".parse().unwrap(), shader_string.into());
165/// # let mut compiler = compiler.with_resolver(resolver);
166/// # compiler.options.keep_main = true; // prevent dead code elimination
167/// #
168/// let wgsl_string = compiler
169///     .compile("path/to/shader.wgsl")
170///     .inspect_err(|e| eprintln!("{e}")) // pretty-print errors
171///     .unwrap()
172///     .to_string();
173/// #
174/// # assert!(wgsl_string == shader_string);
175/// ```
176#[derive(Clone, Debug)]
177pub struct Compiler<R = ()> {
178    pub options: CompileOptions,
179    pub resolver: R,
180}
181
182impl Default for Compiler<()> {
183    fn default() -> Self {
184        Self {
185            options: Default::default(),
186            resolver: Default::default(),
187        }
188    }
189}
190
191impl<R1> Compiler<R1> {
192    /// Set the compilation [`crate::Resolver`] or [`crate::AsyncResolver`].
193    pub fn with_resolver<R2>(self, resolver: R2) -> Compiler<R2> {
194        Compiler::<R2> {
195            options: self.options,
196            resolver,
197        }
198    }
199}
200
201impl<R> Compiler<R> {
202    /// Shorthand for `Compiler::new(options).with_resolver(resolver)`.
203    pub fn new_with_resolver(options: CompileOptions, resolver: R) -> Self {
204        Self { options, resolver }
205    }
206}
207
208impl Compiler<()> {
209    /// Create a new compiler.
210    ///
211    /// By default, the compiler will use a [`StandardResolver`] when compiling.
212    pub fn new(options: CompileOptions) -> Self {
213        let resolver = ();
214        Self { options, resolver }
215    }
216}
217
218/// Standalone version of [`Compiler::compile_module`].
219pub fn compile(
220    main_path: &ModulePath,
221    options: &CompileOptions,
222    resolver: &impl Resolver,
223) -> Result<CompileResult, Error> {
224    let mangler = Box::<dyn Mangler>::from(options.mangler);
225
226    if options.sourcemap {
227        let sourcemapper = SourceMapper::new(main_path.clone(), &resolver, &mangler);
228        let mut pass = CompilationPass::new(main_path, options, &sourcemapper, &sourcemapper);
229        let res = CompilerDriver::compile(&mut pass);
230        let sourcemap = sourcemapper.finish();
231        let res = res.map_err(|e| Diagnostic::from(e).with_sourcemap(&sourcemap))?;
232
233        Ok(CompileResult {
234            syntax: res.syntax,
235            modules: res.modules,
236            sourcemap: Some(sourcemap),
237            used_items: res.used_items,
238        })
239    } else {
240        let mut pass = CompilationPass::new(main_path, options, &resolver, &mangler);
241        let res = CompilerDriver::compile(&mut pass)?;
242        Ok(CompileResult {
243            syntax: res.syntax,
244            modules: res.modules,
245            sourcemap: None,
246            used_items: res.used_items,
247        })
248    }
249}
250
251/// Async version of [`compile`].
252pub async fn compile_async(
253    main_path: &ModulePath,
254    options: &CompileOptions,
255    resolver: &impl Resolver,
256) -> Result<CompileResult, Error> {
257    let mangler = Box::<dyn Mangler>::from(options.mangler);
258
259    if options.sourcemap {
260        let sourcemapper = SourceMapper::new(main_path.clone(), &resolver, &mangler);
261        let mut pass = CompilationPass::new(main_path, options, &sourcemapper, &sourcemapper);
262        let res = CompilerDriver::compile_async(&mut pass).await;
263        let sourcemap = sourcemapper.finish();
264        let res = res.map_err(|e| Diagnostic::from(e).with_sourcemap(&sourcemap))?;
265
266        Ok(CompileResult {
267            syntax: res.syntax,
268            modules: res.modules,
269            sourcemap: Some(sourcemap),
270            used_items: res.used_items,
271        })
272    } else {
273        let mut pass = CompilationPass::new(main_path, options, &resolver, &mangler);
274        let res = CompilerDriver::compile_async(&mut pass).await?;
275        Ok(CompileResult {
276            syntax: res.syntax,
277            modules: res.modules,
278            sourcemap: None,
279            used_items: res.used_items,
280        })
281    }
282}
283
284/// Get the module path for a filesystem path relative to the package root directory,
285/// with a fallback to the current working directory if the resolver doesn't support
286/// filesystem mappings.
287///
288/// # Panics
289///
290/// [`ModulePath::from_path`] can panic.
291fn main_module_path(path: &Path, resolver: &impl Resolver) -> Result<ModulePath, Error> {
292    let mut main_path = resolver.module_path(path).or_else(|e| {
293        if matches!(e, ResolveError::FilesystemNotSupported) {
294            Ok(ModulePath::from_path(Path::new(".").join(path)))
295        } else {
296            Err(e)
297        }
298    })?;
299
300    // we force the origin to be absolute if it was relative.
301    main_path.origin = PathOrigin::Absolute;
302    Ok(main_path)
303}
304
305impl Compiler<()> {
306    // TODO: implement and validate wesl-toml semantics.
307    /// Compile a WESL shader to WGSL.
308    ///
309    /// `path` defines where to look for shader files.
310    /// It can point to a `wesl.toml` file, a directory or a shader file.
311    ///
312    /// The main module (which exposes entry points, bindings and overrides) is determined from `path`:
313    /// It is the package root module if `path` points to a toml file or a directory,
314    /// otherwise it is the file that `path` points to.
315    /// See [`Self::compile_module`] to compile a different main module.
316    ///
317    /// | Path         | Package root directory      | Main module                |
318    /// | ------------ | --------------------------- | -------------------------- |
319    /// | `wesl.toml`  | `root` field in `wesl.toml` | `package.wesl` in root dir |
320    /// | directory    | the directory               | `package.wesl` in root dir |
321    /// | `.wesl` file | the parent directory        | the file specified         |
322    ///
323    /// Note: `.wgsl` extensions are also supported, but `.wesl` takes priority.
324    pub fn compile(&self, path: impl AsRef<Path>) -> Result<CompileResult, Error> {
325        let (pkg_root_dir, main_path) = self.root_and_main(path.as_ref())?;
326        self.compile_module(pkg_root_dir, &main_path)
327    }
328
329    /// Variant of [`Self::compile`] with a custom main module path.
330    pub fn compile_module(
331        &self,
332        pkg_root_dir: impl AsRef<Path>,
333        main_path: &ModulePath,
334    ) -> Result<CompileResult, Error> {
335        let resolver = self.create_resolver(pkg_root_dir.as_ref());
336        compile(main_path, &self.options, &resolver)
337    }
338
339    /// Async version of [`Self::compile`].
340    pub async fn compile_async(&self, path: &Path) -> Result<CompileResult, Error> {
341        let (pkg_root_dir, main_path) = self.root_and_main(path.as_ref())?;
342        self.compile_module_async(&pkg_root_dir, &main_path).await
343    }
344
345    /// Async version of [`Self::compile_module`].
346    pub async fn compile_module_async(
347        &self,
348        pkg_root_dir: impl AsRef<Path>,
349        main_path: &ModulePath,
350    ) -> Result<CompileResult, Error> {
351        let resolver = self.create_resolver(pkg_root_dir.as_ref());
352        compile_async(main_path, &self.options, &resolver).await
353    }
354
355    fn create_resolver(&self, pkg_root_dir: &Path) -> StandardResolver {
356        let mut resolver = StandardResolver::new(pkg_root_dir);
357
358        for (name, value) in self.options.constants.iter() {
359            resolver.add_constant(name.clone(), *value);
360        }
361
362        for package in self.options.dependencies.iter() {
363            resolver.add_package(package);
364        }
365
366        resolver
367    }
368
369    fn root_and_main(&self, path: &Path) -> Result<(PathBuf, ModulePath), Error> {
370        let (pkg_root_dir, main_path) = if let Some(file_name) = path.file_name()
371            && file_name == "wesl.toml"
372        {
373            let cfg = crate::toml_cfg::WeslToml::from_file(path)?;
374            let root = path
375                .parent()
376                .unwrap(/* SAFETY: cannot fail if `file_name` succeeds */)
377                .join(&cfg.package.root);
378            let main = ModulePath::new_root();
379            (root, main)
380        } else if let Some(name) = path.file_stem()
381            && !path.is_dir()
382        {
383            let root = path.parent().unwrap(/* SAFETY: cannot fail if `file_name` succeeds */).to_path_buf();
384            let main = ModulePath::new(
385                PathOrigin::Absolute,
386                vec![name.to_string_lossy().to_string()],
387            );
388            (root, main)
389        } else {
390            let root = path.to_path_buf();
391            let main = ModulePath::new_root();
392            (root, main)
393        };
394
395        Ok((pkg_root_dir, main_path))
396    }
397}
398
399impl<R: Resolver> Compiler<R> {
400    /// Compile a WESL shader to WGSL.
401    ///
402    /// The main module defaults to the package root module.
403    /// See [`Self::compile_module`] to compile a different main module.
404    pub fn compile_root(&self) -> Result<CompileResult, Error> {
405        compile(&ModulePath::new_root(), &self.options, &self.resolver)
406    }
407
408    /// Variant of [`Self::compile`] with a custom main module path.
409    pub fn compile_module(&self, main_path: &ModulePath) -> Result<CompileResult, Error> {
410        compile(main_path, &self.options, &self.resolver)
411    }
412
413    /// Variant of [`Self::compile_root`] with a custom main module path.
414    ///
415    /// `fs_main_path` defines the main module path according to the resolver's file system mapping implemented in [`Resolver::module_path`].
416    ///
417    /// # Warning
418    ///
419    /// This function works best with filesystem resolvers which implement [`Resolver::module_path`].
420    /// If not, this function assumes that the package root directory is the current working directory.
421    ///
422    /// # Panics
423    ///
424    /// Can panic if [`ModulePath::from_path`] fails.
425    // TODO: we don't want that panic.
426    pub fn compile(&self, fs_main_path: impl AsRef<Path>) -> Result<CompileResult, Error> {
427        let main_path = main_module_path(fs_main_path.as_ref(), &self.resolver)?;
428        compile(&main_path, &self.options, &self.resolver)
429    }
430
431    /// Async version of [`Self::compile_root`].
432    pub async fn compile_root_async(&self) -> Result<CompileResult, Error> {
433        compile_async(&ModulePath::new_root(), &self.options, &self.resolver).await
434    }
435
436    /// Async version of [`Self::compile_module`].
437    pub async fn compile_module_async(
438        &self,
439        main_path: &ModulePath,
440    ) -> Result<CompileResult, Error> {
441        compile_async(main_path, &self.options, &self.resolver).await
442    }
443
444    /// Async version of [`Self::compile`].
445    pub async fn compile_async(
446        &self,
447        fs_main_path: &impl AsRef<Path>,
448    ) -> Result<CompileResult, Error> {
449        let main_path = main_module_path(fs_main_path.as_ref(), &self.resolver)?;
450        compile_async(&main_path, &self.options, &self.resolver).await
451    }
452}
453
454/// Result of [`Compiler::compile`].
455///
456/// This type contains the resulting WGSL syntax tree, the sourcemap (if enabled),
457/// and the list of used modules/declarations.
458///
459/// It implements [`std::fmt::Display`], call `to_string()` to get the compiled WGSL.
460#[derive(Default, Clone)]
461pub struct CompileResult {
462    /// The syntax tree of the resulting
463    pub syntax: TranslationUnit,
464    pub modules: Vec<Module>,
465    pub sourcemap: Option<BasicSourceMap>,
466    pub used_items: UsedItems,
467}
468
469impl CompileResult {
470    /// Write the compiled result to a file.
471    pub fn write_to_file(&self, path: impl AsRef<Path>) -> std::io::Result<()> {
472        std::fs::write(path, self.to_string())
473    }
474
475    /// Emit `rerun-if-changed` instructions so the build script reruns only if the
476    /// shader files are modified.
477    pub fn emit_rerun_if_changed(&self) {
478        let Some(sourcemap) = &self.sourcemap else {
479            println!("cargo::warning=cannot emit rerun-if-changed directive without a sourcemap");
480            return;
481        };
482
483        for (module_path, _) in self.used_items.iter() {
484            if module_path.origin.is_package() {
485                continue;
486            }
487            assert!(
488                !module_path.origin.is_relative(),
489                "the modules passed to emit_rerun_if_changed must be absolute"
490            );
491            if let Some(source) = sourcemap.file(module_path)
492                && let Some(fs_path) = &source.path
493            {
494                // Path::display is safe here because of the ModulePath naming restrictions
495                println!("cargo::rerun-if-changed={}", fs_path.display());
496
497                // If it's a fallback path, we need to react to the higher priority path as well
498                if fs_path.extension().unwrap() == "wgsl" {
499                    let fs_path = fs_path.with_extension("wesl");
500                    println!("cargo::rerun-if-changed={}", fs_path.display());
501                }
502            }
503        }
504    }
505
506    /// Write the result in rust's `OUT_DIR`.
507    ///
508    /// This function is meant to be used in a `build.rs` workflow. The output WGSL will
509    /// be accessed with the [`wesl_core::include_wesl`] macro. See the crate documentation for a
510    /// usage example.
511    ///
512    /// # Panics
513    ///
514    /// Panics when the output file cannot be written.
515    pub fn write_artifact(&self, artifact_name: &str) {
516        let dirname = std::env::var("OUT_DIR").unwrap();
517        let out_name = Path::new(artifact_name);
518        if out_name.iter().count() != 1 || out_name.extension().is_some() {
519            eprintln!("`out_name` cannot contain path separators or file extension");
520            panic!()
521        }
522        let mut output = Path::new(&dirname).join(out_name);
523        output.set_extension("wgsl");
524        self.write_to_file(output)
525            .expect("failed to write output shader");
526    }
527}
528
529impl std::fmt::Display for CompileResult {
530    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
531        self.syntax.fmt(f)
532    }
533}
534
535/// Ephemeral type that implements [`CompilerDriver`] for a single compilation pass in [`Compiler::compile`]
536struct CompilationPass<'a> {
537    main_path: &'a ModulePath,
538    options: &'a CompileOptions,
539    resolver: &'a dyn Resolver,
540    mangler: &'a dyn Mangler,
541}
542
543impl<'a> CompilationPass<'a> {
544    fn new(
545        main_path: &'a ModulePath,
546        options: &'a CompileOptions,
547        resolver: &'a dyn Resolver,
548        mangler: &'a dyn Mangler,
549    ) -> Self {
550        Self {
551            main_path,
552            options,
553            resolver,
554            mangler,
555        }
556    }
557}
558
559impl CompilerDriver for CompilationPass<'_> {
560    fn main_path(&self) -> &ModulePath {
561        self.main_path
562    }
563
564    fn canonical_path(&self, path: &ModulePath) -> ModulePath {
565        self.resolver.canonical_path(path)
566    }
567
568    fn main_entry_points(&self, main_module: &TranslationUnit) -> Result<HashSet<Ident>, Error> {
569        // keep all declarations when strip is disabled or keep_main is enabled.
570        if !self.options.strip || self.options.keep_main {
571            Ok(main_module
572                .global_declarations
573                .iter()
574                .filter_map(|decl| decl.ident())
575                .collect())
576        }
577        // user provided an explicit list of entry points to start from.
578        else if let Some(keep) = &self.options.keep {
579            keep.iter()
580                .map(|name| {
581                    main_module.decl_ident(name).ok_or_else(|| {
582                        UsageError::NotFound(self.main_path.clone(), name.to_string()).into()
583                    })
584                })
585                .collect::<Result<HashSet<Ident>, Error>>()
586        }
587        // otherwise, we keep the WGSL entry points. this is the default.
588        else {
589            Ok(main_module.entry_points().collect())
590        }
591    }
592
593    fn module_usage_analysis(
594        &self,
595        module: &Module,
596        already_used: &mut UsedItems,
597        to_analyze: &mut UsedItems,
598    ) -> Result<(), Error> {
599        pass::module_usage_analysis(module, already_used, to_analyze, !self.options.visibility)?;
600
601        // when strip is disabled, all declarations in the module are included so they
602        // must be usage-analyzed.
603        if !self.options.strip {
604            for decl in &module.syntax.global_declarations {
605                if let Some(ident) = decl.ident() {
606                    self.usage_analysis(
607                        module,
608                        &ident.name(),
609                        decl.visibility(),
610                        already_used,
611                        to_analyze,
612                    )?;
613                }
614            }
615        }
616
617        Ok(())
618    }
619
620    fn usage_analysis(
621        &self,
622        module: &Module,
623        decl_name: &str,
624        min_vis: Visibility,
625        already_used: &mut UsedItems,
626        to_analyze: &mut UsedItems,
627    ) -> Result<(), Error> {
628        pass::usage_analysis(
629            module,
630            decl_name,
631            min_vis,
632            already_used,
633            to_analyze,
634            !self.options.visibility,
635        )?;
636        Ok(())
637    }
638
639    fn load_module(&mut self, path: &ModulePath) -> Result<TranslationUnit, Error> {
640        let mut module = pass::load_module(path, &self.resolver)?;
641
642        if self.options.condcomp {
643            pass::condcomp(&mut module, &self.options.features)?;
644        }
645
646        pass::retarget_idents(&mut module);
647
648        if self.options.validate {
649            pass::validate_wesl(&module)?;
650        }
651
652        Ok(module)
653    }
654
655    fn link(
656        &self,
657        modules: &mut Vec<Module>,
658        used_items: &UsedItems,
659    ) -> Result<TranslationUnit, Error> {
660        pass::retarget_modules(modules, used_items, &self.resolver);
661
662        for module in modules.iter_mut() {
663            if !self.options.mangle_main && module.path == *self.main_path {
664                continue;
665            }
666            pass::mangle(&mut module.syntax, &module.path, &self.mangler);
667        }
668
669        let mut module = pass::link(modules, self.options.strip.then_some(used_items));
670
671        if self.options.lower {
672            pass::lower(&mut module)?;
673        }
674
675        if self.options.validate {
676            pass::validate_wgsl(&module)?;
677        }
678
679        Ok(module)
680    }
681}
682
683#[test]
684fn test_send_sync() {
685    fn assert_send_sync<T: Send + Sync>() {}
686    assert_send_sync::<Compiler<()>>();
687}
wesl/frontend.rs

wesl/
frontend.rs
