[naga] Enforce a maximum size for any type (gfx-rs#7950)

* [naga] Enforce a maximum size for any type

Fixes gfx-rs#4339
Fixes gfx-rs#7383

* Doc fix and changelog entry

Author: andyleiserson

CHANGELOG.md

@@ -53,6 +53,12 @@ We have merged the acceleration structure feature into the `RayQuery` feature. T
 
 By @Vecvec in [#7913](https://github.com/gfx-rs/wgpu/pull/7913).
 
+### Changes
+
+#### Naga
+
+Naga now requires that no type be larger than 1 GB. This limit may be lowered in the future; feedback on an appropriate value for the limit is welcome. By @andyleiserson in [#7950](https://github.com/gfx-rs/wgpu/pull/7950).
+
 ## v26.0.1 (2025-07-10)
 
 ### Bug Fixes

naga/src/front/wgsl/error.rs

@@ -409,6 +409,12 @@ pub(crate) enum Error<'a> {
         accept_span: Span,
         accept_type: String,
     },
+    StructMemberTooLarge {
+        member_name_span: Span,
+    },
+    TypeTooLarge {
+        span: Span,
+    },
 }
 
 impl From<ConflictingDiagnosticRuleError> for Error<'_> {
@@ -1367,6 +1373,22 @@ impl<'a> Error<'a> {
                 ],
                 notes: vec![],
             },
+            Error::StructMemberTooLarge { member_name_span } => ParseError {
+                message: "struct member is too large".into(),
+                labels: vec![(member_name_span, "this member exceeds the maximum size".into())],
+                notes: vec![format!(
+                    "the maximum size is {} bytes",
+                    crate::valid::MAX_TYPE_SIZE
+                )],
+            },
+            Error::TypeTooLarge { span } => ParseError {
+                message: "type is too large".into(),
+                labels: vec![(span, "this type exceeds the maximum size".into())],
+                notes: vec![format!(
+                    "the maximum size is {} bytes",
+                    crate::valid::MAX_TYPE_SIZE
+                )],
+            },
         }
     }
 }

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

@@ -7,7 +7,6 @@ use alloc::{
 };
 use core::num::NonZeroU32;
 
-use crate::common::ForDebugWithTypes;
 use crate::front::wgsl::error::{Error, ExpectedToken, InvalidAssignmentType};
 use crate::front::wgsl::index::Index;
 use crate::front::wgsl::parse::number::Number;
@@ -18,6 +17,7 @@ use crate::{
     common::wgsl::{TryToWgsl, TypeContext},
     compact::KeepUnused,
 };
+use crate::{common::ForDebugWithTypes, proc::LayoutErrorInner};
 use crate::{ir, proc};
 use crate::{Arena, FastHashMap, FastIndexMap, Handle, Span};
 
@@ -3709,7 +3709,27 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
         for member in s.members.iter() {
             let ty = self.resolve_ast_type(member.ty, &mut ctx.as_const())?;
 
-            ctx.layouter.update(ctx.module.to_ctx()).unwrap();
+            ctx.layouter.update(ctx.module.to_ctx()).map_err(|err| {
+                let LayoutErrorInner::TooLarge = err.inner else {
+                    unreachable!("unexpected layout error: {err:?}");
+                };
+                // Since anonymous types of struct members don't get a span,
+                // associate the error with the member. The layouter could have
+                // failed on any type that was pending layout, but if it wasn't
+                // the current struct member, it wasn't a struct member at all,
+                // because we resolve struct members one-by-one.
+                if ty == err.ty {
+                    Box::new(Error::StructMemberTooLarge {
+                        member_name_span: member.name.span,
+                    })
+                } else {
+                    // Lots of type definitions don't get spans, so this error
+                    // message may not be very useful.
+                    Box::new(Error::TypeTooLarge {
+                        span: ctx.module.types.get_span(err.ty),
+                    })
+                }
+            })?;
 
             let member_min_size = ctx.layouter[ty].size;
             let member_min_alignment = ctx.layouter[ty].alignment;
@@ -3761,6 +3781,9 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
             });
 
             offset += member_size;
+            if offset > crate::valid::MAX_TYPE_SIZE {
+                return Err(Box::new(Error::TypeTooLarge { span }));
+            }
         }
 
         let size = struct_alignment.round_up(offset);
@@ -3938,7 +3961,17 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                 let base = self.resolve_ast_type(base, &mut ctx.as_const())?;
                 let size = self.array_size(size, ctx)?;
 
-                ctx.layouter.update(ctx.module.to_ctx()).unwrap();
+                // Determine the size of the base type, if needed.
+                ctx.layouter.update(ctx.module.to_ctx()).map_err(|err| {
+                    let LayoutErrorInner::TooLarge = err.inner else {
+                        unreachable!("unexpected layout error: {err:?}");
+                    };
+                    // Lots of type definitions don't get spans, so this error
+                    // message may not be very useful.
+                    Box::new(Error::TypeTooLarge {
+                        span: ctx.module.types.get_span(err.ty),
+                    })
+                })?;
                 let stride = ctx.layouter[base].to_stride();
 
                 ir::TypeInner::Array { base, size, stride }

naga/src/ir/mod.rs

@@ -411,6 +411,18 @@ impl VectorSize {
     pub const MAX: usize = Self::Quad as usize;
 }
 
+impl From<VectorSize> for u8 {
+    fn from(size: VectorSize) -> u8 {
+        size as u8
+    }
+}
+
+impl From<VectorSize> for u32 {
+    fn from(size: VectorSize) -> u32 {
+        size as u32
+    }
+}
+
 /// Primitive type for a scalar.
 #[repr(u8)]
 #[derive(Clone, Copy, Debug, Hash, Eq, Ord, PartialEq, PartialOrd)]

naga/src/proc/layouter.rs

@@ -1,6 +1,9 @@
 use core::{fmt::Display, num::NonZeroU32, ops};
 
-use crate::arena::{Handle, HandleVec};
+use crate::{
+    arena::{Handle, HandleVec},
+    valid::MAX_TYPE_SIZE,
+};
 
 /// A newtype struct where its only valid values are powers of 2
 #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
@@ -121,6 +124,12 @@ impl ops::Index<Handle<crate::Type>> for Layouter {
     }
 }
 
+/// Errors generated by the `Layouter`.
+///
+/// All of these errors can be produced when validating an arbitrary module.
+/// When processing WGSL source, only the `TooLarge` error should be
+/// produced by the `Layouter`, as the front-end should not produce IR
+/// that would result in the other errors.
 #[derive(Clone, Copy, Debug, PartialEq, thiserror::Error)]
 pub enum LayoutErrorInner {
     #[error("Array element type {0:?} doesn't exist")]
@@ -129,6 +138,8 @@ pub enum LayoutErrorInner {
     InvalidStructMemberType(u32, Handle<crate::Type>),
     #[error("Type width must be a power of two")]
     NonPowerOfTwoWidth,
+    #[error("Size exceeds limit of {MAX_TYPE_SIZE} bytes")]
+    TooLarge,
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, thiserror::Error)]
@@ -168,7 +179,10 @@ impl Layouter {
         use crate::TypeInner as Ti;
 
         for (ty_handle, ty) in gctx.types.iter().skip(self.layouts.len()) {
-            let size = ty.inner.size(gctx);
+            let size = ty
+                .inner
+                .try_size(gctx)
+                .ok_or_else(|| LayoutErrorInner::TooLarge.with(ty_handle))?;
             let layout = match ty.inner {
                 Ti::Scalar(scalar) | Ti::Atomic(scalar) => {
                     let alignment = Alignment::new(scalar.width as u32)

naga/src/proc/type_methods.rs

@@ -4,7 +4,7 @@
 //! [`Scalar`]: crate::Scalar
 //! [`ScalarKind`]: crate::ScalarKind
 
-use crate::ir;
+use crate::{ir, valid::MAX_TYPE_SIZE};
 
 use super::TypeResolution;
 
@@ -190,18 +190,19 @@ impl crate::TypeInner {
         }
     }
 
-    /// Get the size of this type.
-    pub fn size(&self, gctx: super::GlobalCtx) -> u32 {
+    /// Attempt to calculate the size of this type. Returns `None` if the size
+    /// exceeds the limit of [`crate::valid::MAX_TYPE_SIZE`].
+    pub fn try_size(&self, gctx: super::GlobalCtx) -> Option<u32> {
         match *self {
-            Self::Scalar(scalar) | Self::Atomic(scalar) => scalar.width as u32,
-            Self::Vector { size, scalar } => size as u32 * scalar.width as u32,
+            Self::Scalar(scalar) | Self::Atomic(scalar) => Some(scalar.width as u32),
+            Self::Vector { size, scalar } => Some(size as u32 * scalar.width as u32),
             // matrices are treated as arrays of aligned columns
             Self::Matrix {
                 columns,
                 rows,
                 scalar,
-            } => super::Alignment::from(rows) * scalar.width as u32 * columns as u32,
-            Self::Pointer { .. } | Self::ValuePointer { .. } => POINTER_SPAN,
+            } => Some(super::Alignment::from(rows) * scalar.width as u32 * columns as u32),
+            Self::Pointer { .. } | Self::ValuePointer { .. } => Some(POINTER_SPAN),
             Self::Array {
                 base: _,
                 size,
@@ -215,17 +216,35 @@ impl crate::TypeInner {
                     // A dynamically-sized array has to have at least one element
                     Ok(crate::proc::IndexableLength::Dynamic) => 1,
                 };
-                count * stride
+                if count > MAX_TYPE_SIZE {
+                    // It shouldn't be possible to have an array of a zero-sized type, but
+                    // let's check just in case.
+                    None
+                } else {
+                    count
+                        .checked_mul(stride)
+                        .filter(|size| *size <= MAX_TYPE_SIZE)
+                }
             }
-            Self::Struct { span, .. } => span,
+            Self::Struct { span, .. } => Some(span),
             Self::Image { .. }
             | Self::Sampler { .. }
             | Self::AccelerationStructure { .. }
             | Self::RayQuery { .. }
-            | Self::BindingArray { .. } => 0,
+            | Self::BindingArray { .. } => Some(0),
         }
     }
 
+    /// Get the size of this type.
+    ///
+    /// Panics if the size exceeds the limit of [`crate::valid::MAX_TYPE_SIZE`].
+    /// Validated modules should not contain such types. Code working with
+    /// modules prior to validation should use [`Self::try_size`] and handle the
+    /// error appropriately.
+    pub fn size(&self, gctx: super::GlobalCtx) -> u32 {
+        self.try_size(gctx).expect("type is too large")
+    }
+
     /// Return the canonical form of `self`, or `None` if it's already in
     /// canonical form.
     ///

naga/src/valid/mod.rs

@@ -35,6 +35,9 @@ pub use r#type::{Disalignment, PushConstantError, TypeError, TypeFlags, WidthErr
 
 use self::handles::InvalidHandleError;
 
+/// Maximum size of a type, in bytes.
+pub const MAX_TYPE_SIZE: u32 = 0x4000_0000; // 1GB
+
 bitflags::bitflags! {
     /// Validation flags.
     ///