[naga] Correct override resolution in array lengths.

When the user provides values for a module's overrides, rather than
replacing override-sized array types with ordinary array types (which
could require adjusting type handles throughout the module), instead
edit all overrides to have initializers that are fully-evaluated
constant expressions. Then, change all backends to handle
override-sized arrays by retrieving their overrides' values.

For arrays whose sizes are override expressions, not simple references
to a specific override's value, let front ends built array types that
refer to anonymous overrides whose initializers are the necessary
expression.

This means that all arrays whose sizes are override expressions are
references to some `Override`. Remove `naga::PendingArraySize`, and
let `ArraySize::Pending` hold a `Handle<Override>` in all cases.

Expand `tests/gpu-tests/shader/array_size_overrides.rs` to include the
test case that motivated this approach.

Author: kentslaney

CHANGELOG.md

@@ -199,6 +199,7 @@ By @Vecvec in [#6905](https://github.com/gfx-rs/wgpu/pull/6905), [#7086](https:/
 - Forward '--keep-coordinate-space' flag to GLSL backend in naga-cli. By @cloone8 in [#7206](https://github.com/gfx-rs/wgpu/pull/7206).
 - Allow template lists to have a trailing comma. By @KentSlaney in [#7142](https://github.com/gfx-rs/wgpu/pull/7142).
 - Allow WGSL const declarations to have abstract types. By @jamienicol in [#7055](https://github.com/gfx-rs/wgpu/pull/7055) and [#7222](https://github.com/gfx-rs/wgpu/pull/7222).
+- Allows override-sized arrays to resolve to the same size without causing the type arena to panic. By @KentSlaney in [#7082](https://github.com/gfx-rs/wgpu/pull/7082).
 
 #### General
 

naga/src/arena/mod.rs

@@ -102,6 +102,18 @@ impl<T> Arena<T> {
             .map(|(i, v)| unsafe { (Handle::from_usize_unchecked(i), v) })
     }
 
+    /// Returns an iterator over the items stored in this arena, returning both
+    /// the item's handle and a reference to it.
+    pub fn iter_mut_span(
+        &mut self,
+    ) -> impl DoubleEndedIterator<Item = (Handle<T>, &mut T, &Span)> + ExactSizeIterator {
+        self.data
+            .iter_mut()
+            .zip(self.span_info.iter())
+            .enumerate()
+            .map(|(i, (v, span))| unsafe { (Handle::from_usize_unchecked(i), v, span) })
+    }
+
     /// Drains the arena, returning an iterator over the items stored.
     pub fn drain(&mut self) -> impl DoubleEndedIterator<Item = (Handle<T>, T, Span)> {
         let arena = core::mem::take(self);

naga/src/back/glsl/mod.rs

@@ -539,6 +539,8 @@ pub enum Error {
     /// [`crate::Sampling::First`] is unsupported.
     #[error("`{:?}` sampling is unsupported", crate::Sampling::First)]
     FirstSamplingNotSupported,
+    #[error(transparent)]
+    ResolveArraySizeError(#[from] proc::ResolveArraySizeError),
 }
 
 /// Binary operation with a different logic on the GLSL side.
@@ -612,10 +614,6 @@ impl<'a, W: Write> Writer<'a, W> {
         pipeline_options: &'a PipelineOptions,
         policies: proc::BoundsCheckPolicies,
     ) -> Result<Self, Error> {
-        if !module.overrides.is_empty() {
-            return Err(Error::Override);
-        }
-
         // Check if the requested version is supported
         if !options.version.is_supported() {
             log::error!("Version {}", options.version);
@@ -1013,13 +1011,12 @@ impl<'a, W: Write> Writer<'a, W> {
         write!(self.out, "[")?;
 
         // Write the array size
-        // Writes nothing if `ArraySize::Dynamic`
-        match size {
-            crate::ArraySize::Constant(size) => {
+        // Writes nothing if `IndexableLength::Dynamic`
+        match size.resolve(self.module.to_ctx())? {
+            proc::IndexableLength::Known(size) => {
                 write!(self.out, "{size}")?;
             }
-            crate::ArraySize::Pending(_) => unreachable!(),
-            crate::ArraySize::Dynamic => (),
+            proc::IndexableLength::Dynamic => (),
         }
 
         write!(self.out, "]")?;
@@ -2759,7 +2756,9 @@ impl<'a, W: Write> Writer<'a, W> {
                 write_expression(self, value)?;
                 write!(self.out, ")")?
             }
-            _ => unreachable!(),
+            _ => {
+                return Err(Error::Override);
+            }
         }
 
         Ok(())
@@ -4574,12 +4573,8 @@ impl<'a, W: Write> Writer<'a, W> {
                 write!(self.out, ")")?;
             }
             TypeInner::Array { base, size, .. } => {
-                let count = match size
-                    .to_indexable_length(self.module)
-                    .expect("Bad array size")
-                {
+                let count = match size.resolve(self.module.to_ctx())? {
                     proc::IndexableLength::Known(count) => count,
-                    proc::IndexableLength::Pending => unreachable!(),
                     proc::IndexableLength::Dynamic => return Ok(()),
                 };
                 self.write_type(base)?;

naga/src/back/hlsl/conv.rs

@@ -54,7 +54,7 @@ impl crate::TypeInner {
         }
     }
 
-    pub(super) fn size_hlsl(&self, gctx: crate::proc::GlobalCtx) -> u32 {
+    pub(super) fn size_hlsl(&self, gctx: crate::proc::GlobalCtx) -> Result<u32, Error> {
         match *self {
             Self::Matrix {
                 columns,
@@ -63,19 +63,18 @@ impl crate::TypeInner {
             } => {
                 let stride = Alignment::from(rows) * scalar.width as u32;
                 let last_row_size = rows as u32 * scalar.width as u32;
-                ((columns as u32 - 1) * stride) + last_row_size
+                Ok(((columns as u32 - 1) * stride) + last_row_size)
             }
             Self::Array { base, size, stride } => {
-                let count = match size {
-                    crate::ArraySize::Constant(size) => size.get(),
+                let count = match size.resolve(gctx)? {
+                    crate::proc::IndexableLength::Known(size) => size,
                     // A dynamically-sized array has to have at least one element
-                    crate::ArraySize::Pending(_) => unreachable!(),
-                    crate::ArraySize::Dynamic => 1,
+                    crate::proc::IndexableLength::Dynamic => 1,
                 };
-                let last_el_size = gctx.types[base].inner.size_hlsl(gctx);
-                ((count - 1) * stride) + last_el_size
+                let last_el_size = gctx.types[base].inner.size_hlsl(gctx)?;
+                Ok(((count - 1) * stride) + last_el_size)
             }
-            _ => self.size(gctx),
+            _ => Ok(self.size(gctx)),
         }
     }
 

naga/src/back/hlsl/mod.rs

@@ -442,6 +442,8 @@ pub enum Error {
     Custom(String),
     #[error("overrides should not be present at this stage")]
     Override,
+    #[error(transparent)]
+    ResolveArraySizeError(#[from] proc::ResolveArraySizeError),
 }
 
 #[derive(PartialEq, Eq, Hash)]

naga/src/back/hlsl/writer.rs

@@ -268,10 +268,6 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
         module_info: &valid::ModuleInfo,
         fragment_entry_point: Option<&FragmentEntryPoint<'_>>,
     ) -> Result<super::ReflectionInfo, Error> {
-        if !module.overrides.is_empty() {
-            return Err(Error::Override);
-        }
-
         self.reset(module);
 
         // Write special constants, if needed
@@ -1129,12 +1125,11 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
     ) -> BackendResult {
         write!(self.out, "[")?;
 
-        match size {
-            crate::ArraySize::Constant(size) => {
+        match size.resolve(module.to_ctx())? {
+            proc::IndexableLength::Known(size) => {
                 write!(self.out, "{size}")?;
             }
-            crate::ArraySize::Pending(_) => unreachable!(),
-            crate::ArraySize::Dynamic => unreachable!(),
+            proc::IndexableLength::Dynamic => unreachable!(),
         }
 
         write!(self.out, "]")?;
@@ -1179,7 +1174,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 }
             }
             let ty_inner = &module.types[member.ty].inner;
-            last_offset = member.offset + ty_inner.size_hlsl(module.to_ctx());
+            last_offset = member.offset + ty_inner.size_hlsl(module.to_ctx())?;
 
             // The indentation is only for readability
             write!(self.out, "{}", back::INDENT)?;
@@ -2701,7 +2696,9 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 write_expression(self, value)?;
                 write!(self.out, ").{number_of_components}")?
             }
-            _ => unreachable!(),
+            _ => {
+                return Err(Error::Override);
+            }
         }
 
         Ok(())
@@ -2971,7 +2968,6 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                             index::IndexableLength::Known(limit) => {
                                 write!(self.out, "{}u", limit - 1)?;
                             }
-                            index::IndexableLength::Pending => unreachable!(),
                             index::IndexableLength::Dynamic => unreachable!(),
                         }
                         write!(self.out, ")")?;

naga/src/back/msl/mod.rs

@@ -182,6 +182,8 @@ pub enum Error {
     Override,
     #[error("bitcasting to {0:?} is not supported")]
     UnsupportedBitCast(crate::TypeInner),
+    #[error(transparent)]
+    ResolveArraySizeError(#[from] crate::proc::ResolveArraySizeError),
 }
 
 #[derive(Clone, Debug, PartialEq, thiserror::Error)]

naga/src/back/msl/writer.rs

@@ -1563,7 +1563,9 @@ impl<W: Write> Writer<W> {
                 put_expression(self, ctx, value)?;
                 write!(self.out, ")")?;
             }
-            _ => unreachable!(),
+            _ => {
+                return Err(Error::Override);
+            }
         }
 
         Ok(())
@@ -2612,7 +2614,6 @@ impl<W: Write> Writer<W> {
                     self.out.write_str(") < ")?;
                     match length {
                         index::IndexableLength::Known(value) => write!(self.out, "{value}")?,
-                        index::IndexableLength::Pending => unreachable!(),
                         index::IndexableLength::Dynamic => {
                             let global =
                                 context.function.originating_global(base).ok_or_else(|| {
@@ -2749,7 +2750,7 @@ impl<W: Write> Writer<W> {
     ) -> BackendResult {
         let accessing_wrapped_array = match *base_ty {
             crate::TypeInner::Array {
-                size: crate::ArraySize::Constant(_),
+                size: crate::ArraySize::Constant(_) | crate::ArraySize::Pending(_),
                 ..
             } => true,
             _ => false,
@@ -2777,7 +2778,6 @@ impl<W: Write> Writer<W> {
                 index::IndexableLength::Known(limit) => {
                     write!(self.out, "{}u", limit - 1)?;
                 }
-                index::IndexableLength::Pending => unreachable!(),
                 index::IndexableLength::Dynamic => {
                     let global = context.function.originating_global(base).ok_or_else(|| {
                         Error::GenericValidation("Could not find originating global".into())
@@ -3795,10 +3795,6 @@ impl<W: Write> Writer<W> {
         options: &Options,
         pipeline_options: &PipelineOptions,
     ) -> Result<TranslationInfo, Error> {
-        if !module.overrides.is_empty() {
-            return Err(Error::Override);
-        }
-
         self.names.clear();
         self.namer.reset(
             module,
@@ -3995,8 +3991,8 @@ impl<W: Write> Writer<W> {
                         first_time: false,
                     };
 
-                    match size {
-                        crate::ArraySize::Constant(size) => {
+                    match size.resolve(module.to_ctx())? {
+                        proc::IndexableLength::Known(size) => {
                             writeln!(self.out, "struct {name} {{")?;
                             writeln!(
                                 self.out,
@@ -4008,10 +4004,7 @@ impl<W: Write> Writer<W> {
                             )?;
                             writeln!(self.out, "}};")?;
                         }
-                        crate::ArraySize::Pending(_) => {
-                            unreachable!()
-                        }
-                        crate::ArraySize::Dynamic => {
+                        proc::IndexableLength::Dynamic => {
                             writeln!(self.out, "typedef {base_name} {name}[1];")?;
                         }
                     }
@@ -6694,10 +6687,8 @@ mod workgroup_mem_init {
                         writeln!(self.out, ", 0, {NAMESPACE}::memory_order_relaxed);")?;
                     }
                     crate::TypeInner::Array { base, size, .. } => {
-                        let count = match size.to_indexable_length(module).expect("Bad array size")
-                        {
+                        let count = match size.resolve(module.to_ctx())? {
                             proc::IndexableLength::Known(count) => count,
-                            proc::IndexableLength::Pending => unreachable!(),
                             proc::IndexableLength::Dynamic => unreachable!(),
                         };