[naga wgsl] Implement local const declarations (#6156)

Author: sagudev

naga/src/back/glsl/mod.rs

@@ -47,7 +47,7 @@ pub use features::Features;
 
 use crate::{
     back::{self, Baked},
-    proc::{self, NameKey},
+    proc::{self, ExpressionKindTracker, NameKey},
     valid, Handle, ShaderStage, TypeInner,
 };
 use features::FeaturesManager;
@@ -1584,6 +1584,7 @@ impl<'a, W: Write> Writer<'a, W> {
             info,
             expressions: &func.expressions,
             named_expressions: &func.named_expressions,
+            expr_kind_tracker: ExpressionKindTracker::from_arena(&func.expressions),
         };
 
         self.named_expressions.clear();

naga/src/back/hlsl/writer.rs

@@ -8,7 +8,7 @@ use super::{
 };
 use crate::{
     back::{self, Baked},
-    proc::{self, NameKey},
+    proc::{self, ExpressionKindTracker, NameKey},
     valid, Handle, Module, Scalar, ScalarKind, ShaderStage, TypeInner,
 };
 use std::{fmt, mem};
@@ -346,6 +346,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 info,
                 expressions: &function.expressions,
                 named_expressions: &function.named_expressions,
+                expr_kind_tracker: ExpressionKindTracker::from_arena(&function.expressions),
             };
             let name = self.names[&NameKey::Function(handle)].clone();
 
@@ -386,6 +387,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 info,
                 expressions: &ep.function.expressions,
                 named_expressions: &ep.function.named_expressions,
+                expr_kind_tracker: ExpressionKindTracker::from_arena(&ep.function.expressions),
             };
 
             self.write_wrapped_functions(module, &ctx)?;

naga/src/back/mod.rs

@@ -3,6 +3,8 @@ Backend functions that export shader [`Module`](super::Module)s into binary and
 */
 #![allow(dead_code)] // can be dead if none of the enabled backends need it
 
+use crate::proc::ExpressionKindTracker;
+
 #[cfg(dot_out)]
 pub mod dot;
 #[cfg(glsl_out)]
@@ -118,6 +120,8 @@ pub struct FunctionCtx<'a> {
     pub expressions: &'a crate::Arena<crate::Expression>,
     /// Map of expressions that have associated variable names
     pub named_expressions: &'a crate::NamedExpressions,
+    /// For constness checks
+    pub expr_kind_tracker: ExpressionKindTracker,
 }
 
 impl FunctionCtx<'_> {

naga/src/back/pipeline_constants.rs

@@ -289,6 +289,7 @@ fn process_function(
         &mut local_expression_kind_tracker,
         &mut emitter,
         &mut block,
+        false,
     );
 
     for (old_h, mut expr, span) in expressions.drain() {

naga/src/back/wgsl/writer.rs

@@ -1,7 +1,7 @@
 use super::Error;
 use crate::{
     back::{self, Baked},
-    proc::{self, NameKey},
+    proc::{self, ExpressionKindTracker, NameKey},
     valid, Handle, Module, ShaderStage, TypeInner,
 };
 use std::fmt::Write;
@@ -166,6 +166,7 @@ impl<W: Write> Writer<W> {
                 info: fun_info,
                 expressions: &function.expressions,
                 named_expressions: &function.named_expressions,
+                expr_kind_tracker: ExpressionKindTracker::from_arena(&function.expressions),
             };
 
             // Write the function
@@ -193,6 +194,7 @@ impl<W: Write> Writer<W> {
                 info: info.get_entry_point(index),
                 expressions: &ep.function.expressions,
                 named_expressions: &ep.function.named_expressions,
+                expr_kind_tracker: ExpressionKindTracker::from_arena(&ep.function.expressions),
             };
             self.write_function(module, &ep.function, &func_ctx)?;
 
@@ -1115,8 +1117,14 @@ impl<W: Write> Writer<W> {
         func_ctx: &back::FunctionCtx,
         name: &str,
     ) -> BackendResult {
+        // Some functions are marked as const, but are not yet implemented as constant expression
+        let quantifier = if func_ctx.expr_kind_tracker.is_impl_const(handle) {
+            "const"
+        } else {
+            "let"
+        };
         // Write variable name
-        write!(self.out, "let {name}")?;
+        write!(self.out, "{quantifier} {name}")?;
         if self.flags.contains(WriterFlags::EXPLICIT_TYPES) {
             write!(self.out, ": ")?;
             let ty = &func_ctx.info[handle].ty;