1 files changed, 29 insertions, 10 deletions

diff --git a/third_party/rust/naga/src/valid/expression.rs b/third_party/rust/naga/src/valid/expression.rs
index 838ecc4e27..525bd28c17 100644
--- a/third_party/rust/naga/src/valid/expression.rs
+++ b/third_party/rust/naga/src/valid/expression.rs
@@ -90,6 +90,8 @@ pub enum ExpressionError {
         sampler: bool,
         has_ref: bool,
     },
+    #[error("Sample offset must be a const-expression")]
+    InvalidSampleOffsetExprType,
     #[error("Sample offset constant {1:?} doesn't match the image dimension {0:?}")]
     InvalidSampleOffset(crate::ImageDimension, Handle<crate::Expression>),
     #[error("Depth reference {0:?} is not a scalar float")]
@@ -129,9 +131,12 @@ pub enum ExpressionError {
 }
 
 #[derive(Clone, Debug, thiserror::Error)]
+#[cfg_attr(test, derive(PartialEq))]
 pub enum ConstExpressionError {
-    #[error("The expression is not a constant expression")]
-    NonConst,
+    #[error("The expression is not a constant or override expression")]
+    NonConstOrOverride,
+    #[error("The expression is not a fully evaluated constant expression")]
+    NonFullyEvaluatedConst,
     #[error(transparent)]
     Compose(#[from] super::ComposeError),
     #[error("Splatting {0:?} can't be done")]
@@ -184,10 +189,15 @@ impl super::Validator {
         handle: Handle<crate::Expression>,
         gctx: crate::proc::GlobalCtx,
         mod_info: &ModuleInfo,
+        global_expr_kind: &crate::proc::ExpressionKindTracker,
     ) -> Result<(), ConstExpressionError> {
         use crate::Expression as E;
 
-        match gctx.const_expressions[handle] {
+        if !global_expr_kind.is_const_or_override(handle) {
+            return Err(ConstExpressionError::NonConstOrOverride);
+        }
+
+        match gctx.global_expressions[handle] {
             E::Literal(literal) => {
                 self.validate_literal(literal)?;
             }
@@ -201,14 +211,19 @@ impl super::Validator {
             }
             E::Splat { value, .. } => match *mod_info[value].inner_with(gctx.types) {
                 crate::TypeInner::Scalar { .. } => {}
-                _ => return Err(super::ConstExpressionError::InvalidSplatType(value)),
+                _ => return Err(ConstExpressionError::InvalidSplatType(value)),
             },
-            _ => return Err(super::ConstExpressionError::NonConst),
+            _ if global_expr_kind.is_const(handle) || !self.allow_overrides => {
+                return Err(ConstExpressionError::NonFullyEvaluatedConst)
+            }
+            // the constant evaluator will report errors about override-expressions
+            _ => {}
         }
 
         Ok(())
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub(super) fn validate_expression(
         &self,
         root: Handle<crate::Expression>,
@@ -217,6 +232,7 @@ impl super::Validator {
         module: &crate::Module,
         info: &FunctionInfo,
         mod_info: &ModuleInfo,
+        global_expr_kind: &crate::proc::ExpressionKindTracker,
     ) -> Result<ShaderStages, ExpressionError> {
         use crate::{Expression as E, Scalar as Sc, ScalarKind as Sk, TypeInner as Ti};
 
@@ -252,9 +268,7 @@ impl super::Validator {
                         return Err(ExpressionError::InvalidIndexType(index));
                     }
                 }
-                if dynamic_indexing_restricted
-                    && function.expressions[index].is_dynamic_index(module)
-                {
+                if dynamic_indexing_restricted && function.expressions[index].is_dynamic_index() {
                     return Err(ExpressionError::IndexMustBeConstant(base));
                 }
 
@@ -347,7 +361,7 @@ impl super::Validator {
                 self.validate_literal(literal)?;
                 ShaderStages::all()
             }
-            E::Constant(_) | E::ZeroValue(_) => ShaderStages::all(),
+            E::Constant(_) | E::Override(_) | E::ZeroValue(_) => ShaderStages::all(),
             E::Compose { ref components, ty } => {
                 validate_compose(
                     ty,
@@ -464,6 +478,10 @@ impl super::Validator {
 
                 // check constant offset
                 if let Some(const_expr) = offset {
+                    if !global_expr_kind.is_const(const_expr) {
+                        return Err(ExpressionError::InvalidSampleOffsetExprType);
+                    }
+
                     match *mod_info[const_expr].inner_with(&module.types) {
                         Ti::Scalar(Sc { kind: Sk::Sint, .. }) if num_components == 1 => {}
                         Ti::Vector {
@@ -1623,6 +1641,7 @@ impl super::Validator {
                     return Err(ExpressionError::InvalidRayQueryType(query));
                 }
             },
+            E::SubgroupBallotResult | E::SubgroupOperationResult { .. } => self.subgroup_stages,
         };
         Ok(stages)
     }
@@ -1716,7 +1735,7 @@ fn validate_with_const_expression(
     use crate::span::Span;
 
     let mut module = crate::Module::default();
-    module.const_expressions.append(expr, Span::default());
+    module.global_expressions.append(expr, Span::default());
 
     let mut validator = super::Validator::new(super::ValidationFlags::CONSTANTS, caps);