[wgsl-in] add support for override declarations

[ir] split overrides from constants

Author: teoxoy

naga/src/back/dot/mod.rs

@@ -404,6 +404,7 @@ fn write_function_expressions(
         let (label, color_id) = match *expression {
             E::Literal(_) => ("Literal".into(), 2),
             E::Constant(_) => ("Constant".into(), 2),
+            E::Override(_) => ("Override".into(), 2),
             E::ZeroValue(_) => ("ZeroValue".into(), 2),
             E::Compose { ref components, .. } => {
                 payload = Some(Payload::Arguments(components));

naga/src/back/glsl/mod.rs

@@ -2525,6 +2525,7 @@ impl<'a, W: Write> Writer<'a, W> {
                     |writer, expr| writer.write_expr(expr, ctx),
                 )?;
             }
+            Expression::Override(_) => return Err(Error::Custom("overrides are WIP".into())),
             // `Access` is applied to arrays, vectors and matrices and is written as indexing
             Expression::Access { base, index } => {
                 self.write_expr(base, ctx)?;

naga/src/back/hlsl/writer.rs

@@ -2151,6 +2151,9 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                     |writer, expr| writer.write_expr(module, expr, func_ctx),
                 )?;
             }
+            Expression::Override(_) => {
+                return Err(Error::Unimplemented("overrides are WIP".into()))
+            }
             // All of the multiplication can be expressed as `mul`,
             // except vector * vector, which needs to use the "*" operator.
             Expression::Binary {

naga/src/back/msl/writer.rs

@@ -1394,6 +1394,9 @@ impl<W: Write> Writer<W> {
                     |writer, context, expr| writer.put_expression(expr, context, true),
                 )?;
             }
+            crate::Expression::Override(_) => {
+                return Err(Error::FeatureNotImplemented("overrides are WIP".into()))
+            }
             crate::Expression::Access { base, .. }
             | crate::Expression::AccessIndex { base, .. } => {
                 // This is an acceptable place to generate a `ReadZeroSkipWrite` check.

naga/src/back/spv/block.rs

@@ -239,6 +239,9 @@ impl<'w> BlockContext<'w> {
                 let init = self.ir_module.constants[handle].init;
                 self.writer.constant_ids[init.index()]
             }
+            crate::Expression::Override(_) => {
+                return Err(Error::FeatureNotImplemented("overrides are WIP"))
+            }
             crate::Expression::ZeroValue(_) => self.writer.get_constant_null(result_type_id),
             crate::Expression::Compose { ty, ref components } => {
                 self.temp_list.clear();

naga/src/back/wgsl/writer.rs

@@ -1180,6 +1180,9 @@ impl<W: Write> Writer<W> {
                     |writer, expr| writer.write_expr(module, expr, func_ctx),
                 )?;
             }
+            Expression::Override(_) => {
+                return Err(Error::Unimplemented("overrides are WIP".into()))
+            }
             Expression::FunctionArgument(pos) => {
                 let name_key = func_ctx.argument_key(pos);
                 let name = &self.names[&name_key];

naga/src/compact/expressions.rs

@@ -3,6 +3,7 @@ use crate::arena::{Arena, Handle};
 
 pub struct ExpressionTracer<'tracer> {
     pub constants: &'tracer Arena<crate::Constant>,
+    pub overrides: &'tracer Arena<crate::Override>,
 
     /// The arena in which we are currently tracing expressions.
     pub expressions: &'tracer Arena<crate::Expression>,
@@ -13,6 +14,9 @@ pub struct ExpressionTracer<'tracer> {
     /// The used map for `constants`.
     pub constants_used: &'tracer mut HandleSet<crate::Constant>,
 
+    /// The used map for `overrides`.
+    pub overrides_used: &'tracer mut HandleSet<crate::Override>,
+
     /// The used set for `arena`.
     ///
     /// This points to whatever arena holds the expressions we are
@@ -88,6 +92,22 @@ impl<'tracer> ExpressionTracer<'tracer> {
                         None => self.expressions_used.insert(init),
                     }
                 }
+                Ex::Override(handle) => {
+                    self.overrides_used.insert(handle);
+                    // Overrides and expressions are mutually recursive, which
+                    // complicates our nice one-pass algorithm. However, since
+                    // overrides don't refer to each other, we can get around
+                    // this by looking *through* each override and marking its
+                    // initializer as used. Since `expr` refers to the override,
+                    // and the override refers to the initializer, it must
+                    // precede `expr` in the arena.
+                    if let Some(init) = self.overrides[handle].init {
+                        match self.const_expressions_used {
+                            Some(ref mut used) => used.insert(init),
+                            None => self.expressions_used.insert(init),
+                        }
+                    }
+                }
                 Ex::ZeroValue(ty) => self.types_used.insert(ty),
                 Ex::Compose { ty, ref components } => {
                     self.types_used.insert(ty);
@@ -221,6 +241,7 @@ impl ModuleMap {
 
             // Expressions that contain handles that need to be adjusted.
             Ex::Constant(ref mut constant) => self.constants.adjust(constant),
+            Ex::Override(ref mut override_) => self.overrides.adjust(override_),
             Ex::ZeroValue(ref mut ty) => self.types.adjust(ty),
             Ex::Compose {
                 ref mut ty,

naga/src/compact/functions.rs

@@ -4,9 +4,11 @@ use super::{FunctionMap, ModuleMap};
 pub struct FunctionTracer<'a> {
     pub function: &'a crate::Function,
     pub constants: &'a crate::Arena<crate::Constant>,
+    pub overrides: &'a crate::Arena<crate::Override>,
 
     pub types_used: &'a mut HandleSet<crate::Type>,
     pub constants_used: &'a mut HandleSet<crate::Constant>,
+    pub overrides_used: &'a mut HandleSet<crate::Override>,
     pub const_expressions_used: &'a mut HandleSet<crate::Expression>,
 
     /// Function-local expressions used.
@@ -47,10 +49,12 @@ impl<'a> FunctionTracer<'a> {
     fn as_expression(&mut self) -> super::expressions::ExpressionTracer {
         super::expressions::ExpressionTracer {
             constants: self.constants,
+            overrides: self.overrides,
             expressions: &self.function.expressions,
 
             types_used: self.types_used,
             constants_used: self.constants_used,
+            overrides_used: self.overrides_used,
             expressions_used: &mut self.expressions_used,
             const_expressions_used: Some(&mut self.const_expressions_used),
         }

naga/src/compact/mod.rs

@@ -54,6 +54,14 @@ pub fn compact(module: &mut crate::Module) {
         }
     }
 
+    // We treat all overrides as used by definition.
+    for (handle, override_) in module.overrides.iter() {
+        module_tracer.overrides_used.insert(handle);
+        if let Some(init) = override_.init {
+            module_tracer.const_expressions_used.insert(init);
+        }
+    }
+
     // We assume that all functions are used.
     //
     // Observe which types, constant expressions, constants, and
@@ -99,6 +107,11 @@ pub fn compact(module: &mut crate::Module) {
             module_tracer.types_used.insert(constant.ty);
         }
     }
+    for (handle, override_) in module.overrides.iter() {
+        if module_tracer.overrides_used.contains(handle) {
+            module_tracer.types_used.insert(override_.ty);
+        }
+    }
 
     // Treat all named types as used.
     for (handle, ty) in module.types.iter() {
@@ -158,6 +171,20 @@ pub fn compact(module: &mut crate::Module) {
         }
     });
 
+    // Drop unused overrides in place, reusing existing storage.
+    log::trace!("adjusting overrides");
+    module.overrides.retain_mut(|handle, override_| {
+        if module_map.overrides.used(handle) {
+            module_map.types.adjust(&mut override_.ty);
+            if let Some(init) = override_.init.as_mut() {
+                module_map.const_expressions.adjust(init);
+            }
+            true
+        } else {
+            false
+        }
+    });
+
     // Adjust global variables' types and initializers.
     log::trace!("adjusting global variables");
     for (_, global) in module.global_variables.iter_mut() {
@@ -193,6 +220,7 @@ struct ModuleTracer<'module> {
     module: &'module crate::Module,
     types_used: HandleSet<crate::Type>,
     constants_used: HandleSet<crate::Constant>,
+    overrides_used: HandleSet<crate::Override>,
     const_expressions_used: HandleSet<crate::Expression>,
 }
 
@@ -202,6 +230,7 @@ impl<'module> ModuleTracer<'module> {
             module,
             types_used: HandleSet::for_arena(&module.types),
             constants_used: HandleSet::for_arena(&module.constants),
+            overrides_used: HandleSet::for_arena(&module.overrides),
             const_expressions_used: HandleSet::for_arena(&module.const_expressions),
         }
     }
@@ -235,8 +264,10 @@ impl<'module> ModuleTracer<'module> {
         expressions::ExpressionTracer {
             expressions: &self.module.const_expressions,
             constants: &self.module.constants,
+            overrides: &self.module.overrides,
             types_used: &mut self.types_used,
             constants_used: &mut self.constants_used,
+            overrides_used: &mut self.overrides_used,
             expressions_used: &mut self.const_expressions_used,
             const_expressions_used: None,
         }
@@ -249,8 +280,10 @@ impl<'module> ModuleTracer<'module> {
         FunctionTracer {
             function,
             constants: &self.module.constants,
+            overrides: &self.module.overrides,
             types_used: &mut self.types_used,
             constants_used: &mut self.constants_used,
+            overrides_used: &mut self.overrides_used,
             const_expressions_used: &mut self.const_expressions_used,
             expressions_used: HandleSet::for_arena(&function.expressions),
         }
@@ -260,6 +293,7 @@ impl<'module> ModuleTracer<'module> {
 struct ModuleMap {
     types: HandleMap<crate::Type>,
     constants: HandleMap<crate::Constant>,
+    overrides: HandleMap<crate::Override>,
     const_expressions: HandleMap<crate::Expression>,
 }
 
@@ -268,6 +302,7 @@ impl From<ModuleTracer<'_>> for ModuleMap {
         ModuleMap {
             types: HandleMap::from_set(used.types_used),
             constants: HandleMap::from_set(used.constants_used),
+            overrides: HandleMap::from_set(used.overrides_used),
             const_expressions: HandleMap::from_set(used.const_expressions_used),
         }
     }

naga/src/front/spv/function.rs

@@ -128,6 +128,7 @@ impl<I: Iterator<Item = u32>> super::Frontend<I> {
             expressions: &mut fun.expressions,
             local_arena: &mut fun.local_variables,
             const_arena: &mut module.constants,
+            overrides: &mut module.overrides,
             const_expressions: &mut module.const_expressions,
             type_arena: &module.types,
             global_arena: &module.global_variables,
@@ -573,6 +574,7 @@ impl<'function> BlockContext<'function> {
         crate::proc::GlobalCtx {
             types: self.type_arena,
             constants: self.const_arena,
+            overrides: self.overrides,
             const_expressions: self.const_expressions,
         }
     }

naga/src/front/spv/mod.rs

@@ -532,6 +532,7 @@ struct BlockContext<'function> {
     local_arena: &'function mut Arena<crate::LocalVariable>,
     /// Constants arena of the module being processed
     const_arena: &'function mut Arena<crate::Constant>,
+    overrides: &'function mut Arena<crate::Override>,
     const_expressions: &'function mut Arena<crate::Expression>,
     /// Type arena of the module being processed
     type_arena: &'function UniqueArena<crate::Type>,
@@ -3933,7 +3934,7 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
                 Op::TypeImage => self.parse_type_image(inst, &mut module),
                 Op::TypeSampledImage => self.parse_type_sampled_image(inst),
                 Op::TypeSampler => self.parse_type_sampler(inst, &mut module),
-                Op::Constant | Op::SpecConstant => self.parse_constant(inst, &mut module),
+                Op::Constant => self.parse_constant(inst, &mut module),
                 Op::ConstantComposite => self.parse_composite_constant(inst, &mut module),
                 Op::ConstantNull | Op::Undef => self.parse_null_constant(inst, &mut module),
                 Op::ConstantTrue => self.parse_bool_constant(inst, true, &mut module),

naga/src/front/wgsl/error.rs

@@ -190,7 +190,7 @@ pub enum Error<'a> {
         expected: String,
         got: String,
     },
-    MissingType(Span),
+    DeclMissingTypeAndInit(Span),
     MissingAttribute(&'static str, Span),
     InvalidAtomicPointer(Span),
     InvalidAtomicOperandType(Span),
@@ -251,6 +251,7 @@ pub enum Error<'a> {
     ExpectedPositiveArrayLength(Span),
     MissingWorkgroupSize(Span),
     ConstantEvaluatorError(ConstantEvaluatorError, Span),
+    PipelineConstantIDValue(Span),
 }
 
 impl<'a> Error<'a> {
@@ -500,11 +501,11 @@ impl<'a> Error<'a> {
                     notes: vec![],
                 }
             }
-            Error::MissingType(name_span) => ParseError {
-                message: format!("variable `{}` needs a type", &source[name_span]),
+            Error::DeclMissingTypeAndInit(name_span) => ParseError {
+                message: format!("declaration of `{}` needs a type specifier or initializer", &source[name_span]),
                 labels: vec![(
                     name_span,
-                    format!("definition of `{}`", &source[name_span]).into(),
+                    "needs a type specifier or initializer".into(),
                 )],
                 notes: vec![],
             },
@@ -712,6 +713,14 @@ impl<'a> Error<'a> {
                 )],
                 notes: vec![],
             },
+            Error::PipelineConstantIDValue(span) => ParseError {
+                message: "pipeline constant ID must be between 0 and 65535 inclusive".to_string(),
+                labels: vec![(
+                    span,
+                    "must be between 0 and 65535 inclusive".into(),
+                )],
+                notes: vec![],
+            },
         }
     }
 }

naga/src/front/wgsl/index.rs

@@ -187,6 +187,7 @@ const fn decl_ident<'a>(decl: &ast::GlobalDecl<'a>) -> ast::Ident<'a> {
         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,
     }