Compile with context

CHANGELOG.md

@@ -95,6 +95,27 @@ By @stefnotch in [#5410](https://github.com/gfx-rs/wgpu/pull/5410)
 
 - Implement `WGSL`'s `unpack4xI8`,`unpack4xU8`,`pack4xI8` and `pack4xU8`. By @VlaDexa in [#5424](https://github.com/gfx-rs/wgpu/pull/5424)
 - Began work adding support for atomics to the SPIR-V frontend. Tracking issue is [here](https://github.com/gfx-rs/wgpu/issues/4489). By @schell in [#5702](https://github.com/gfx-rs/wgpu/pull/5702).
+- Compile shaders with a "base module". This allows for cleanly implementing imports, or shadertoy-like environments. By @stefnotch in [#5791](https://github.com/gfx-rs/wgpu/pull/5791).
+  ```rust
+  use naga::front::wgsl::Frontend;
+
+  let base_module = Frontend::new()
+      .parse("
+      fn main_image(frag_coord: vec2f) -> vec4f {
+          return vec4f(sin(frag_coord.x), cos(frag_coord.y), 0.0, 1.0);
+      }").unwrap();
+  // A full shadertoy implementation would rename all globals from base_module, except for `main_image`.
+  let result = Frontend::new()
+      .parse_to_ast("
+      @fragment
+      fn fs_main(@builtin(position) pos : vec4f, @location(0) uv: vec2f) -> vec4f {
+          return main_image(uv);
+      }")
+      .unwrap()
+      .to_module(Some(&base_module))
+      .unwrap();
+  ```
+
 
 ### Changes
 

naga/src/front/wgsl/lower/mod.rs

@@ -913,8 +913,9 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
     pub fn lower(
         &mut self,
         tu: &'temp ast::TranslationUnit<'source>,
+        base_module: Option<&'source crate::Module>,
     ) -> Result<crate::Module, Error<'source>> {
-        let mut module = crate::Module::default();
+        let mut module = base_module.map(|v| v.to_owned()).unwrap_or_default();
 
         let mut ctx = GlobalContext {
             ast_expressions: &tu.expressions,
@@ -925,6 +926,70 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
             global_expression_kind_tracker: &mut crate::proc::ExpressionKindTracker::new(),
         };
 
+        if let Some(base_module) = base_module {
+            // The handles for base_module are equal to the handles for ctx.module, because we just cloned the arenas.
+            for (handle, f) in base_module.functions.iter() {
+                if let Some(name) = f.name.as_ref() {
+                    ctx.globals
+                        .insert(name, LoweredGlobalDecl::Function(handle));
+                }
+            }
+            for (handle, v) in base_module.global_variables.iter() {
+                if let Some(name) = v.name.as_ref() {
+                    ctx.globals.insert(name, LoweredGlobalDecl::Var(handle));
+                }
+            }
+            for (handle, c) in base_module.constants.iter() {
+                if let Some(name) = c.name.as_ref() {
+                    ctx.globals.insert(name, LoweredGlobalDecl::Const(handle));
+                }
+            }
+            for (handle, o) in base_module.overrides.iter() {
+                if let Some(name) = o.name.as_ref() {
+                    ctx.globals
+                        .insert(name, LoweredGlobalDecl::Override(handle));
+                }
+            }
+            for (handle, t) in base_module.types.iter() {
+                if let Some(name) = t.name.as_ref() {
+                    ctx.globals.insert(name, LoweredGlobalDecl::Type(handle));
+                }
+            }
+            for entry_point in base_module.entry_points.iter() {
+                ctx.globals
+                    .insert(entry_point.name.as_str(), LoweredGlobalDecl::EntryPoint);
+            }
+            *ctx.global_expression_kind_tracker =
+                crate::proc::ExpressionKindTracker::from_arena(&ctx.module.global_expressions);
+        }
+
+        // check for redefinitions
+        for (_, decl) in tu.decls.iter() {
+            let ident = match decl.kind {
+                ast::GlobalDeclKind::Fn(ref f) => f.name,
+                ast::GlobalDeclKind::Var(ref v) => v.name,
+                ast::GlobalDeclKind::Const(ref c) => c.name,
+                ast::GlobalDeclKind::Override(ref o) => o.name,
+                ast::GlobalDeclKind::Struct(ref s) => s.name,
+                ast::GlobalDeclKind::Type(ref t) => t.name,
+            };
+            if let Some(old) = ctx.globals.get(ident.name) {
+                let span = match *old {
+                    LoweredGlobalDecl::Function(handle) => ctx.module.functions.get_span(handle),
+                    LoweredGlobalDecl::Var(handle) => ctx.module.global_variables.get_span(handle),
+                    LoweredGlobalDecl::Const(handle) => ctx.module.constants.get_span(handle),
+                    LoweredGlobalDecl::Override(handle) => ctx.module.overrides.get_span(handle),
+                    LoweredGlobalDecl::Type(handle) => ctx.module.types.get_span(handle),
+                    // We don't have good spans for entry points
+                    LoweredGlobalDecl::EntryPoint => Default::default(),
+                };
+                return Err(Error::Redefinition {
+                    previous: span,
+                    current: ident.span,
+                });
+            }
+        }
+
         for decl_handle in self.index.visit_ordered() {
             let span = tu.decls.get_span(decl_handle);
             let decl = &tu.decls[decl_handle];

naga/src/front/wgsl/mod.rs

@@ -20,6 +20,8 @@ pub use crate::front::wgsl::error::ParseError;
 use crate::front::wgsl::lower::Lowerer;
 use crate::Scalar;
 
+use self::parse::ast::TranslationUnit;
+
 pub struct Frontend {
     parser: Parser,
 }
@@ -32,18 +34,46 @@ impl Frontend {
     }
 
     pub fn parse(&mut self, source: &str) -> Result<crate::Module, ParseError> {
-        self.inner(source).map_err(|x| x.as_parse_error(source))
+        self.parse_to_ast(source)?.to_module(None)
     }
 
-    fn inner<'a>(&mut self, source: &'a str) -> Result<crate::Module, Error<'a>> {
-        let tu = self.parser.parse(source)?;
-        let index = index::Index::generate(&tu)?;
-        let module = Lowerer::new(&index).lower(&tu)?;
+    /// Two-step module conversion, can be used to compile with a "base module".
+    pub fn parse_to_ast<'a>(&mut self, source: &'a str) -> Result<ParsedWgsl<'a>, ParseError> {
+        self.inner_to_ast(source)
+            .map_err(|x| x.as_parse_error(source))
+    }
 
-        Ok(module)
+    fn inner_to_ast<'a>(&mut self, source: &'a str) -> Result<ParsedWgsl<'a>, Error<'a>> {
+        let translation_unit = self.parser.parse(source)?;
+        let index = index::Index::generate(&translation_unit)?;
+        Ok(ParsedWgsl {
+            source,
+            translation_unit,
+            index,
+        })
     }
 }
 
+pub struct ParsedWgsl<'a> {
+    source: &'a str,
+    translation_unit: TranslationUnit<'a>,
+    index: index::Index<'a>,
+}
+impl<'a> ParsedWgsl<'a> {
+    pub fn to_module(
+        &self,
+        base_module: Option<&crate::Module>,
+    ) -> Result<crate::Module, ParseError> {
+        self.inner_to_module(base_module)
+            .map_err(|x| x.as_parse_error(self.source))
+    }
+    fn inner_to_module(
+        &self,
+        base_module: Option<&'a crate::Module>,
+    ) -> Result<crate::Module, Error<'a>> {
+        Lowerer::new(&self.index).lower(&self.translation_unit, base_module)
+    }
+}
 /// <div class="warning">
 // NOTE: Keep this in sync with `wgpu::Device::create_shader_module`!
 // NOTE: Keep this in sync with `wgpu_core::Global::device_create_shader_module`!