Skip to main content

naga/proc/
mod.rs

1/*!
2[`Module`](super::Module) processing functionality.
3*/
4
5mod constant_evaluator;
6mod emitter;
7pub mod index;
8mod keyword_set;
9mod layouter;
10mod namer;
11mod overloads;
12mod terminator;
13mod type_methods;
14mod typifier;
15
16pub use constant_evaluator::{
17    check_bitfield_range, ConstantEvaluator, ConstantEvaluatorError, ExpressionKind,
18    ExpressionKindTracker,
19};
20pub use emitter::Emitter;
21pub use index::{BoundsCheckPolicies, BoundsCheckPolicy, IndexableLength, IndexableLengthError};
22pub use keyword_set::{CaseInsensitiveKeywordSet, KeywordSet};
23pub use layouter::{Alignment, LayoutError, LayoutErrorInner, Layouter, TypeLayout};
24pub use namer::{EntryPointIndex, ExternalTextureNameKey, NameKey, Namer};
25pub use overloads::{Conclusion, MissingSpecialType, OverloadSet, Rule};
26pub use terminator::ensure_block_returns;
27use thiserror::Error;
28pub use type_methods::{
29    concrete_int_scalars, min_max_float_representable_by, vector_size_str, vector_sizes,
30};
31pub use typifier::{compare_types, ResolveContext, ResolveError, TypeResolution};
32
33use crate::non_max_u32::NonMaxU32;
34
35/// Returns `true` if `predicate` returns `true` for any statement in `block`,
36/// including statements nested inside `Block`, `If`, `Loop` and `Switch`
37/// statements.
38///
39/// Does not traverse into the bodies of the functions called by `block`, so
40/// callers that care about a whole module must visit every function and entry
41/// point themselves.
42pub fn any_statement(
43    block: &crate::Block,
44    predicate: &mut impl FnMut(&crate::Statement) -> bool,
45) -> bool {
46    use crate::Statement as S;
47
48    for stmt in block.iter() {
49        if predicate(stmt) {
50            return true;
51        }
52        let nested = match *stmt {
53            S::Block(ref body) => any_statement(body, predicate),
54            S::If {
55                ref accept,
56                ref reject,
57                ..
58            } => any_statement(accept, predicate) || any_statement(reject, predicate),
59            S::Loop {
60                ref body,
61                ref continuing,
62                ..
63            } => any_statement(body, predicate) || any_statement(continuing, predicate),
64            S::Switch { ref cases, .. } => cases
65                .iter()
66                .any(|case| any_statement(&case.body, predicate)),
67            _ => false,
68        };
69        if nested {
70            return true;
71        }
72    }
73
74    false
75}
76
77impl From<super::StorageFormat> for super::Scalar {
78    fn from(format: super::StorageFormat) -> Self {
79        use super::{ScalarKind as Sk, StorageFormat as Sf};
80        let kind = match format {
81            Sf::R8Unorm => Sk::Float,
82            Sf::R8Snorm => Sk::Float,
83            Sf::R8Uint => Sk::Uint,
84            Sf::R8Sint => Sk::Sint,
85            Sf::R16Uint => Sk::Uint,
86            Sf::R16Sint => Sk::Sint,
87            Sf::R16Float => Sk::Float,
88            Sf::Rg8Unorm => Sk::Float,
89            Sf::Rg8Snorm => Sk::Float,
90            Sf::Rg8Uint => Sk::Uint,
91            Sf::Rg8Sint => Sk::Sint,
92            Sf::R32Uint => Sk::Uint,
93            Sf::R32Sint => Sk::Sint,
94            Sf::R32Float => Sk::Float,
95            Sf::Rg16Uint => Sk::Uint,
96            Sf::Rg16Sint => Sk::Sint,
97            Sf::Rg16Float => Sk::Float,
98            Sf::Rgba8Unorm => Sk::Float,
99            Sf::Rgba8Snorm => Sk::Float,
100            Sf::Rgba8Uint => Sk::Uint,
101            Sf::Rgba8Sint => Sk::Sint,
102            Sf::Bgra8Unorm => Sk::Float,
103            Sf::Rgb10a2Uint => Sk::Uint,
104            Sf::Rgb10a2Unorm => Sk::Float,
105            Sf::Rg11b10Ufloat => Sk::Float,
106            Sf::R64Uint => Sk::Uint,
107            Sf::Rg32Uint => Sk::Uint,
108            Sf::Rg32Sint => Sk::Sint,
109            Sf::Rg32Float => Sk::Float,
110            Sf::Rgba16Uint => Sk::Uint,
111            Sf::Rgba16Sint => Sk::Sint,
112            Sf::Rgba16Float => Sk::Float,
113            Sf::Rgba32Uint => Sk::Uint,
114            Sf::Rgba32Sint => Sk::Sint,
115            Sf::Rgba32Float => Sk::Float,
116            Sf::R16Unorm => Sk::Float,
117            Sf::R16Snorm => Sk::Float,
118            Sf::Rg16Unorm => Sk::Float,
119            Sf::Rg16Snorm => Sk::Float,
120            Sf::Rgba16Unorm => Sk::Float,
121            Sf::Rgba16Snorm => Sk::Float,
122        };
123        let width = match format {
124            Sf::R64Uint => 8,
125            _ => 4,
126        };
127        super::Scalar { kind, width }
128    }
129}
130
131#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
132pub enum HashableLiteral {
133    F64(u64),
134    F32(u32),
135    F16(u16),
136    U16(u16),
137    I16(i16),
138    U32(u32),
139    I32(i32),
140    U64(u64),
141    I64(i64),
142    Bool(bool),
143    AbstractInt(i64),
144    AbstractFloat(u64),
145}
146
147impl From<crate::Literal> for HashableLiteral {
148    fn from(l: crate::Literal) -> Self {
149        match l {
150            crate::Literal::F64(v) => Self::F64(v.to_bits()),
151            crate::Literal::F32(v) => Self::F32(v.to_bits()),
152            crate::Literal::F16(v) => Self::F16(v.to_bits()),
153            crate::Literal::U16(v) => Self::U16(v),
154            crate::Literal::I16(v) => Self::I16(v),
155            crate::Literal::U32(v) => Self::U32(v),
156            crate::Literal::I32(v) => Self::I32(v),
157            crate::Literal::U64(v) => Self::U64(v),
158            crate::Literal::I64(v) => Self::I64(v),
159            crate::Literal::Bool(v) => Self::Bool(v),
160            crate::Literal::AbstractInt(v) => Self::AbstractInt(v),
161            crate::Literal::AbstractFloat(v) => Self::AbstractFloat(v.to_bits()),
162        }
163    }
164}
165
166impl crate::Literal {
167    pub const fn new(value: u8, scalar: crate::Scalar) -> Option<Self> {
168        match (value, scalar.kind, scalar.width) {
169            (value, crate::ScalarKind::Float, 8) => Some(Self::F64(value as _)),
170            (value, crate::ScalarKind::Float, 4) => Some(Self::F32(value as _)),
171            (value, crate::ScalarKind::Float, 2) => {
172                Some(Self::F16(half::f16::from_f32_const(value as _)))
173            }
174            (value, crate::ScalarKind::Uint, 2) => Some(Self::U16(value as _)),
175            (value, crate::ScalarKind::Sint, 2) => Some(Self::I16(value as _)),
176            (value, crate::ScalarKind::Uint, 4) => Some(Self::U32(value as _)),
177            (value, crate::ScalarKind::Sint, 4) => Some(Self::I32(value as _)),
178            (value, crate::ScalarKind::Uint, 8) => Some(Self::U64(value as _)),
179            (value, crate::ScalarKind::Sint, 8) => Some(Self::I64(value as _)),
180            (1, crate::ScalarKind::Bool, crate::BOOL_WIDTH) => Some(Self::Bool(true)),
181            (0, crate::ScalarKind::Bool, crate::BOOL_WIDTH) => Some(Self::Bool(false)),
182            (value, crate::ScalarKind::AbstractInt, 8) => Some(Self::AbstractInt(value as _)),
183            (value, crate::ScalarKind::AbstractFloat, 8) => Some(Self::AbstractFloat(value as _)),
184            _ => None,
185        }
186    }
187
188    pub const fn zero(scalar: crate::Scalar) -> Option<Self> {
189        Self::new(0, scalar)
190    }
191
192    pub const fn one(scalar: crate::Scalar) -> Option<Self> {
193        Self::new(1, scalar)
194    }
195
196    pub const fn minus_one(scalar: crate::Scalar) -> Option<Self> {
197        match (scalar.kind, scalar.width) {
198            (crate::ScalarKind::Float, 8) => Some(Self::F64(-1.0)),
199            (crate::ScalarKind::Float, 4) => Some(Self::F32(-1.0)),
200            (crate::ScalarKind::Float, 2) => Some(Self::F16(half::f16::from_f32_const(-1.0))),
201            (crate::ScalarKind::Sint, 8) => Some(Self::I64(-1)),
202            (crate::ScalarKind::Sint, 4) => Some(Self::I32(-1)),
203            (crate::ScalarKind::Sint, 2) => Some(Self::I16(-1)),
204            (crate::ScalarKind::AbstractInt, 8) => Some(Self::AbstractInt(-1)),
205            _ => None,
206        }
207    }
208
209    pub const fn width(&self) -> crate::Bytes {
210        match *self {
211            Self::F64(_) | Self::I64(_) | Self::U64(_) => 8,
212            Self::F32(_) | Self::U32(_) | Self::I32(_) => 4,
213            Self::F16(_) | Self::U16(_) | Self::I16(_) => 2,
214            Self::Bool(_) => crate::BOOL_WIDTH,
215            Self::AbstractInt(_) | Self::AbstractFloat(_) => crate::ABSTRACT_WIDTH,
216        }
217    }
218    pub const fn scalar(&self) -> crate::Scalar {
219        match *self {
220            Self::F64(_) => crate::Scalar::F64,
221            Self::F32(_) => crate::Scalar::F32,
222            Self::F16(_) => crate::Scalar::F16,
223            Self::U16(_) => crate::Scalar::U16,
224            Self::I16(_) => crate::Scalar::I16,
225            Self::U32(_) => crate::Scalar::U32,
226            Self::I32(_) => crate::Scalar::I32,
227            Self::U64(_) => crate::Scalar::U64,
228            Self::I64(_) => crate::Scalar::I64,
229            Self::Bool(_) => crate::Scalar::BOOL,
230            Self::AbstractInt(_) => crate::Scalar::ABSTRACT_INT,
231            Self::AbstractFloat(_) => crate::Scalar::ABSTRACT_FLOAT,
232        }
233    }
234    pub const fn scalar_kind(&self) -> crate::ScalarKind {
235        self.scalar().kind
236    }
237    pub const fn ty_inner(&self) -> crate::TypeInner {
238        crate::TypeInner::Scalar(self.scalar())
239    }
240}
241
242impl TryFrom<crate::Literal> for u32 {
243    type Error = ConstValueError;
244
245    fn try_from(value: crate::Literal) -> Result<Self, Self::Error> {
246        match value {
247            crate::Literal::U16(value) => Ok(value as u32),
248            crate::Literal::I16(value) => value.try_into().map_err(|_| ConstValueError::Negative),
249            crate::Literal::U32(value) => Ok(value),
250            crate::Literal::I32(value) => value.try_into().map_err(|_| ConstValueError::Negative),
251            _ => Err(ConstValueError::InvalidType),
252        }
253    }
254}
255
256impl TryFrom<crate::Literal> for bool {
257    type Error = ConstValueError;
258
259    fn try_from(value: crate::Literal) -> Result<Self, Self::Error> {
260        match value {
261            crate::Literal::Bool(value) => Ok(value),
262            _ => Err(ConstValueError::InvalidType),
263        }
264    }
265}
266
267impl super::AddressSpace {
268    pub fn access(self) -> crate::StorageAccess {
269        use crate::StorageAccess as Sa;
270        match self {
271            crate::AddressSpace::Function
272            | crate::AddressSpace::Private
273            | crate::AddressSpace::WorkGroup => Sa::LOAD | Sa::STORE,
274            crate::AddressSpace::Uniform => Sa::LOAD,
275            crate::AddressSpace::Storage { access } => access,
276            crate::AddressSpace::Handle => Sa::LOAD,
277            crate::AddressSpace::Immediate => Sa::LOAD,
278            // TaskPayload isn't always writable, but this is checked for elsewhere,
279            // when not using multiple payloads and matching the entry payload is checked.
280            crate::AddressSpace::TaskPayload => Sa::LOAD | Sa::STORE,
281            crate::AddressSpace::RayPayload | crate::AddressSpace::IncomingRayPayload => {
282                Sa::LOAD | Sa::STORE
283            }
284        }
285    }
286}
287
288impl super::MathFunction {
289    pub const fn argument_count(&self) -> usize {
290        match *self {
291            // comparison
292            Self::Abs => 1,
293            Self::Min => 2,
294            Self::Max => 2,
295            Self::Clamp => 3,
296            Self::Saturate => 1,
297            // trigonometry
298            Self::Cos => 1,
299            Self::Cosh => 1,
300            Self::Sin => 1,
301            Self::Sinh => 1,
302            Self::Tan => 1,
303            Self::Tanh => 1,
304            Self::Acos => 1,
305            Self::Asin => 1,
306            Self::Atan => 1,
307            Self::Atan2 => 2,
308            Self::Asinh => 1,
309            Self::Acosh => 1,
310            Self::Atanh => 1,
311            Self::Radians => 1,
312            Self::Degrees => 1,
313            // decomposition
314            Self::Ceil => 1,
315            Self::Floor => 1,
316            Self::Round => 1,
317            Self::Fract => 1,
318            Self::Trunc => 1,
319            Self::Modf => 1,
320            Self::Frexp => 1,
321            Self::Ldexp => 2,
322            // exponent
323            Self::Exp => 1,
324            Self::Exp2 => 1,
325            Self::Log => 1,
326            Self::Log2 => 1,
327            Self::Pow => 2,
328            // geometry
329            Self::Dot => 2,
330            Self::Dot4I8Packed => 2,
331            Self::Dot4U8Packed => 2,
332            Self::Outer => 2,
333            Self::Cross => 2,
334            Self::Distance => 2,
335            Self::Length => 1,
336            Self::Normalize => 1,
337            Self::FaceForward => 3,
338            Self::Reflect => 2,
339            Self::Refract => 3,
340            // computational
341            Self::Sign => 1,
342            Self::Fma => 3,
343            Self::Mix => 3,
344            Self::Step => 2,
345            Self::SmoothStep => 3,
346            Self::Sqrt => 1,
347            Self::InverseSqrt => 1,
348            Self::Inverse => 1,
349            Self::Transpose => 1,
350            Self::Determinant => 1,
351            Self::QuantizeToF16 => 1,
352            // bits
353            Self::CountTrailingZeros => 1,
354            Self::CountLeadingZeros => 1,
355            Self::CountOneBits => 1,
356            Self::ReverseBits => 1,
357            Self::ExtractBits => 3,
358            Self::InsertBits => 4,
359            Self::FirstTrailingBit => 1,
360            Self::FirstLeadingBit => 1,
361            // data packing
362            Self::Pack4x8snorm => 1,
363            Self::Pack4x8unorm => 1,
364            Self::Pack2x16snorm => 1,
365            Self::Pack2x16unorm => 1,
366            Self::Pack2x16float => 1,
367            Self::Pack4xI8 => 1,
368            Self::Pack4xU8 => 1,
369            Self::Pack4xI8Clamp => 1,
370            Self::Pack4xU8Clamp => 1,
371            // data unpacking
372            Self::Unpack4x8snorm => 1,
373            Self::Unpack4x8unorm => 1,
374            Self::Unpack2x16snorm => 1,
375            Self::Unpack2x16unorm => 1,
376            Self::Unpack2x16float => 1,
377            Self::Unpack4xI8 => 1,
378            Self::Unpack4xU8 => 1,
379        }
380    }
381}
382
383impl crate::Expression {
384    /// Returns true if the expression is considered emitted at the start of a function.
385    pub const fn needs_pre_emit(&self) -> bool {
386        match *self {
387            Self::Literal(_)
388            | Self::Constant(_)
389            | Self::Override(_)
390            | Self::ZeroValue(_)
391            | Self::FunctionArgument(_)
392            | Self::GlobalVariable(_)
393            | Self::LocalVariable(_) => true,
394            _ => false,
395        }
396    }
397
398    /// Return true if this expression is a dynamic array/vector/matrix index,
399    /// for [`Access`].
400    ///
401    /// This method returns true if this expression is a dynamically computed
402    /// index, and as such can only be used to index matrices when they appear
403    /// behind a pointer. See the documentation for [`Access`] for details.
404    ///
405    /// Note, this does not check the _type_ of the given expression. It's up to
406    /// the caller to establish that the `Access` expression is well-typed
407    /// through other means, like [`ResolveContext`].
408    ///
409    /// [`Access`]: crate::Expression::Access
410    /// [`ResolveContext`]: crate::proc::ResolveContext
411    pub const fn is_dynamic_index(&self) -> bool {
412        match *self {
413            Self::Literal(_) | Self::ZeroValue(_) | Self::Constant(_) => false,
414            _ => true,
415        }
416    }
417}
418
419impl crate::Function {
420    /// Return the global variable being accessed by the expression `pointer`.
421    ///
422    /// Assuming that `pointer` is a series of `Access` and `AccessIndex`
423    /// expressions that ultimately access some part of a `GlobalVariable`,
424    /// return a handle for that global.
425    ///
426    /// If the expression does not ultimately access a global variable, return
427    /// `None`.
428    pub fn originating_global(
429        &self,
430        mut pointer: crate::Handle<crate::Expression>,
431    ) -> Option<crate::Handle<crate::GlobalVariable>> {
432        loop {
433            pointer = match self.expressions[pointer] {
434                crate::Expression::Access { base, .. } => base,
435                crate::Expression::AccessIndex { base, .. } => base,
436                crate::Expression::GlobalVariable(handle) => return Some(handle),
437                // Other expressions are not on this path to a global.
438                _ => return None,
439            }
440        }
441    }
442}
443
444impl crate::SampleLevel {
445    pub const fn implicit_derivatives(&self) -> bool {
446        match *self {
447            Self::Auto | Self::Bias(_) => true,
448            Self::Zero | Self::Exact(_) | Self::Gradient { .. } => false,
449        }
450    }
451}
452
453impl crate::Binding {
454    pub const fn to_built_in(&self) -> Option<crate::BuiltIn> {
455        match *self {
456            crate::Binding::BuiltIn(built_in) => Some(built_in),
457            Self::Location { .. } => None,
458        }
459    }
460}
461
462impl super::SwizzleComponent {
463    pub const XYZW: [Self; 4] = [Self::X, Self::Y, Self::Z, Self::W];
464
465    pub const fn index(&self) -> u32 {
466        match *self {
467            Self::X => 0,
468            Self::Y => 1,
469            Self::Z => 2,
470            Self::W => 3,
471        }
472    }
473    pub const fn from_index(idx: u32) -> Self {
474        match idx {
475            0 => Self::X,
476            1 => Self::Y,
477            2 => Self::Z,
478            _ => Self::W,
479        }
480    }
481}
482
483impl super::ImageClass {
484    pub const fn is_multisampled(self) -> bool {
485        match self {
486            crate::ImageClass::Sampled { multi, .. } | crate::ImageClass::Depth { multi } => multi,
487            crate::ImageClass::Storage { .. } => false,
488            crate::ImageClass::External => false,
489        }
490    }
491
492    pub const fn is_mipmapped(self) -> bool {
493        match self {
494            crate::ImageClass::Sampled { multi, .. } | crate::ImageClass::Depth { multi } => !multi,
495            crate::ImageClass::Storage { .. } => false,
496            crate::ImageClass::External => false,
497        }
498    }
499
500    pub const fn is_depth(self) -> bool {
501        matches!(self, crate::ImageClass::Depth { .. })
502    }
503}
504
505impl crate::Module {
506    pub const fn to_ctx(&self) -> GlobalCtx<'_> {
507        GlobalCtx {
508            types: &self.types,
509            constants: &self.constants,
510            overrides: &self.overrides,
511            global_expressions: &self.global_expressions,
512        }
513    }
514
515    pub fn compare_types(&self, lhs: &TypeResolution, rhs: &TypeResolution) -> bool {
516        compare_types(lhs, rhs, &self.types)
517    }
518}
519
520#[derive(Debug)]
521pub enum ConstValueError {
522    NonConst,
523    Negative,
524    InvalidType,
525}
526
527impl From<core::convert::Infallible> for ConstValueError {
528    fn from(_: core::convert::Infallible) -> Self {
529        unreachable!()
530    }
531}
532
533#[derive(Clone, Copy, Debug)]
534pub struct GlobalCtx<'a> {
535    pub types: &'a crate::UniqueArena<crate::Type>,
536    pub constants: &'a crate::Arena<crate::Constant>,
537    pub overrides: &'a crate::Arena<crate::Override>,
538    pub global_expressions: &'a crate::Arena<crate::Expression>,
539}
540
541impl GlobalCtx<'_> {
542    /// Try to evaluate the expression in `self.global_expressions` using its `handle`
543    /// and return it as a `T: TryFrom<ir::Literal>`.
544    ///
545    /// This currently only evaluates scalar expressions. If adding support for vectors,
546    /// consider changing `valid::expression::validate_constant_shift_amounts` to use that
547    /// support.
548    #[cfg_attr(
549        not(any(
550            feature = "glsl-in",
551            feature = "spv-in",
552            feature = "wgsl-in",
553            glsl_out,
554            hlsl_out,
555            msl_out,
556            wgsl_out
557        )),
558        allow(dead_code)
559    )]
560    pub(super) fn get_const_val<T, E>(
561        &self,
562        handle: crate::Handle<crate::Expression>,
563    ) -> Result<T, ConstValueError>
564    where
565        T: TryFrom<crate::Literal, Error = E>,
566        E: Into<ConstValueError>,
567    {
568        self.get_const_val_from(handle, self.global_expressions)
569    }
570
571    pub(super) fn get_const_val_from<T, E>(
572        &self,
573        handle: crate::Handle<crate::Expression>,
574        arena: &crate::Arena<crate::Expression>,
575    ) -> Result<T, ConstValueError>
576    where
577        T: TryFrom<crate::Literal, Error = E>,
578        E: Into<ConstValueError>,
579    {
580        fn get(
581            gctx: GlobalCtx,
582            handle: crate::Handle<crate::Expression>,
583            arena: &crate::Arena<crate::Expression>,
584        ) -> Option<crate::Literal> {
585            match arena[handle] {
586                crate::Expression::Literal(literal) => Some(literal),
587                crate::Expression::ZeroValue(ty) => match gctx.types[ty].inner {
588                    crate::TypeInner::Scalar(scalar) => crate::Literal::zero(scalar),
589                    _ => None,
590                },
591                _ => None,
592            }
593        }
594        let value = match arena[handle] {
595            crate::Expression::Constant(c) => {
596                get(*self, self.constants[c].init, self.global_expressions)
597            }
598            _ => get(*self, handle, arena),
599        };
600        match value {
601            Some(v) => v.try_into().map_err(Into::into),
602            None => Err(ConstValueError::NonConst),
603        }
604    }
605
606    pub fn compare_types(&self, lhs: &TypeResolution, rhs: &TypeResolution) -> bool {
607        compare_types(lhs, rhs, self.types)
608    }
609}
610
611#[derive(Error, Debug, Clone, Copy, PartialEq)]
612pub enum ResolveArraySizeError {
613    #[error("array element count must be positive (> 0)")]
614    ExpectedPositiveArrayLength,
615    #[error("internal: array size override has not been resolved")]
616    NonConstArrayLength,
617}
618
619impl crate::ArraySize {
620    /// Return the number of elements that `size` represents, if known at code generation time.
621    ///
622    /// If `size` is override-based, return an error unless the override's
623    /// initializer is a fully evaluated constant expression. You can call
624    /// [`pipeline_constants::process_overrides`] to supply values for a
625    /// module's overrides and ensure their initializers are fully evaluated, as
626    /// this function expects.
627    ///
628    /// [`pipeline_constants::process_overrides`]: crate::back::pipeline_constants::process_overrides
629    pub fn resolve(&self, gctx: GlobalCtx) -> Result<IndexableLength, ResolveArraySizeError> {
630        match *self {
631            crate::ArraySize::Constant(length) => Ok(IndexableLength::Known(length.get())),
632            crate::ArraySize::Pending(handle) => {
633                let Some(expr) = gctx.overrides[handle].init else {
634                    return Err(ResolveArraySizeError::NonConstArrayLength);
635                };
636                let length = gctx.get_const_val(expr).map_err(|err| match err {
637                    ConstValueError::NonConst => ResolveArraySizeError::NonConstArrayLength,
638                    ConstValueError::Negative | ConstValueError::InvalidType => {
639                        ResolveArraySizeError::ExpectedPositiveArrayLength
640                    }
641                })?;
642
643                if length == 0 {
644                    return Err(ResolveArraySizeError::ExpectedPositiveArrayLength);
645                }
646
647                Ok(IndexableLength::Known(length))
648            }
649            crate::ArraySize::Dynamic => Ok(IndexableLength::Dynamic),
650        }
651    }
652}
653
654/// Return an iterator over the individual components assembled by a
655/// `Compose` expression.
656///
657/// Given `ty` and `components` from an `Expression::Compose`, return an
658/// iterator over the components of the resulting value.
659///
660/// Normally, this would just be an iterator over `components`. However,
661/// `Compose` expressions can concatenate vectors, in which case the i'th
662/// value being composed is not generally the i'th element of `components`.
663/// This function consults `ty` to decide if this concatenation is occurring,
664/// and returns an iterator that produces the components of the result of
665/// the `Compose` expression in either case.
666pub fn flatten_compose<'arenas>(
667    ty: crate::Handle<crate::Type>,
668    components: &'arenas [crate::Handle<crate::Expression>],
669    expressions: &'arenas crate::Arena<crate::Expression>,
670    types: &'arenas crate::UniqueArena<crate::Type>,
671) -> impl Iterator<Item = crate::Handle<crate::Expression>> + 'arenas {
672    // Returning `impl Iterator` is a bit tricky. We may or may not
673    // want to flatten the components, but we have to settle on a
674    // single concrete type to return. This function returns a single
675    // iterator chain that handles both the flattening and
676    // non-flattening cases.
677    let (size, is_vector) = if let crate::TypeInner::Vector { size, .. } = types[ty].inner {
678        (size as usize, true)
679    } else {
680        (components.len(), false)
681    };
682
683    /// Flatten `Compose` expressions if `is_vector` is true.
684    fn flatten_compose<'c>(
685        component: &'c crate::Handle<crate::Expression>,
686        is_vector: bool,
687        expressions: &'c crate::Arena<crate::Expression>,
688    ) -> &'c [crate::Handle<crate::Expression>] {
689        if is_vector {
690            if let crate::Expression::Compose {
691                ty: _,
692                components: ref subcomponents,
693            } = expressions[*component]
694            {
695                return subcomponents;
696            }
697        }
698        core::slice::from_ref(component)
699    }
700
701    /// Flatten `Splat` expressions if `is_vector` is true.
702    fn flatten_splat<'c>(
703        component: &'c crate::Handle<crate::Expression>,
704        is_vector: bool,
705        expressions: &'c crate::Arena<crate::Expression>,
706    ) -> impl Iterator<Item = crate::Handle<crate::Expression>> {
707        let mut expr = *component;
708        let mut count = 1;
709        if is_vector {
710            if let crate::Expression::Splat { size, value } = expressions[expr] {
711                expr = value;
712                count = size as usize;
713            }
714        }
715        core::iter::repeat_n(expr, count)
716    }
717
718    // Expressions like `vec4(vec3(vec2(6, 7), 8), 9)` require us to
719    // flatten up to two levels of `Compose` expressions.
720    //
721    // Expressions like `vec4(vec3(1.0), 1.0)` require us to flatten
722    // `Splat` expressions. Fortunately, the operand of a `Splat` must
723    // be a scalar, so we can stop there.
724    components
725        .iter()
726        .flat_map(move |component| flatten_compose(component, is_vector, expressions))
727        .flat_map(move |component| flatten_compose(component, is_vector, expressions))
728        .flat_map(move |component| flatten_splat(component, is_vector, expressions))
729        .take(size)
730}
731
732#[test]
733fn test_matrix_size() {
734    let module = crate::Module::default();
735    assert_eq!(
736        crate::TypeInner::Matrix {
737            columns: crate::VectorSize::Tri,
738            rows: crate::VectorSize::Tri,
739            scalar: crate::Scalar::F32,
740        }
741        .size(module.to_ctx()),
742        48,
743    );
744}
745
746impl crate::Module {
747    /// Extracts mesh shader info from a mesh output global variable. Used in frontends
748    /// and by validators. This only validates the output variable itself, and not the
749    /// vertex and primitive output types.
750    ///
751    /// The output contains the extracted mesh stage info, with overrides unset,
752    /// and then the overrides separately. This is because the overrides should be
753    /// treated as expressions elsewhere, but that requires mutably modifying the
754    /// module and the expressions should only be created at parse time, not validation
755    /// time.
756    #[allow(clippy::type_complexity)]
757    pub fn analyze_mesh_shader_info(
758        &self,
759        gv: crate::Handle<crate::GlobalVariable>,
760    ) -> (
761        crate::MeshStageInfo,
762        [Option<crate::Handle<crate::Override>>; 2],
763        Option<crate::WithSpan<crate::valid::EntryPointError>>,
764    ) {
765        use crate::span::AddSpan;
766        use crate::valid::EntryPointError;
767        #[derive(Default)]
768        struct OutError {
769            pub inner: Option<EntryPointError>,
770        }
771        impl OutError {
772            pub fn set(&mut self, err: EntryPointError) {
773                if self.inner.is_none() {
774                    self.inner = Some(err);
775                }
776            }
777        }
778
779        // Used to temporarily initialize stuff
780        let null_type = crate::Handle::new(NonMaxU32::new(0).unwrap());
781        let mut output = crate::MeshStageInfo {
782            topology: crate::MeshOutputTopology::Triangles,
783            max_vertices: 0,
784            max_vertices_override: None,
785            max_primitives: 0,
786            max_primitives_override: None,
787            vertex_output_type: null_type,
788            primitive_output_type: null_type,
789            output_variable: gv,
790        };
791        // Stores the error to output, if any.
792        let mut error = OutError::default();
793        let r#type = &self.types[self.global_variables[gv].ty].inner;
794
795        let mut topology = output.topology;
796        // Max, max override, type
797        let mut vertex_info = (0, None, null_type);
798        let mut primitive_info = (0, None, null_type);
799
800        match r#type {
801            &crate::TypeInner::Struct { ref members, .. } => {
802                let mut builtins = crate::FastHashSet::default();
803                for member in members {
804                    match member.binding {
805                        Some(crate::Binding::BuiltIn(crate::BuiltIn::VertexCount)) => {
806                            // Must have type u32
807                            if self.types[member.ty].inner.scalar() != Some(crate::Scalar::U32) {
808                                error.set(EntryPointError::BadMeshOutputVariableField);
809                            }
810                            // Each builtin should only occur once
811                            if builtins.contains(&crate::BuiltIn::VertexCount) {
812                                error.set(EntryPointError::BadMeshOutputVariableType);
813                            }
814                            builtins.insert(crate::BuiltIn::VertexCount);
815                        }
816                        Some(crate::Binding::BuiltIn(crate::BuiltIn::PrimitiveCount)) => {
817                            // Must have type u32
818                            if self.types[member.ty].inner.scalar() != Some(crate::Scalar::U32) {
819                                error.set(EntryPointError::BadMeshOutputVariableField);
820                            }
821                            // Each builtin should only occur once
822                            if builtins.contains(&crate::BuiltIn::PrimitiveCount) {
823                                error.set(EntryPointError::BadMeshOutputVariableType);
824                            }
825                            builtins.insert(crate::BuiltIn::PrimitiveCount);
826                        }
827                        Some(crate::Binding::BuiltIn(
828                            crate::BuiltIn::Vertices | crate::BuiltIn::Primitives,
829                        )) => {
830                            let ty = &self.types[member.ty].inner;
831                            // Analyze the array type to determine size and vertex/primitive type
832                            let (a, b, c) = match ty {
833                                &crate::TypeInner::Array { base, size, .. } => {
834                                    let ty = base;
835                                    let (max, max_override) = match size {
836                                        crate::ArraySize::Constant(a) => (a.get(), None),
837                                        crate::ArraySize::Pending(o) => (0, Some(o)),
838                                        crate::ArraySize::Dynamic => {
839                                            error.set(EntryPointError::BadMeshOutputVariableField);
840                                            (0, None)
841                                        }
842                                    };
843                                    (max, max_override, ty)
844                                }
845                                _ => {
846                                    error.set(EntryPointError::BadMeshOutputVariableField);
847                                    (0, None, null_type)
848                                }
849                            };
850                            if matches!(
851                                member.binding,
852                                Some(crate::Binding::BuiltIn(crate::BuiltIn::Primitives))
853                            ) {
854                                // Primitives require special analysis to determine topology
855                                primitive_info = (a, b, c);
856                                match self.types[c].inner {
857                                    crate::TypeInner::Struct { ref members, .. } => {
858                                        for member in members {
859                                            match member.binding {
860                                                Some(crate::Binding::BuiltIn(
861                                                    crate::BuiltIn::PointIndex,
862                                                )) => {
863                                                    topology = crate::MeshOutputTopology::Points;
864                                                }
865                                                Some(crate::Binding::BuiltIn(
866                                                    crate::BuiltIn::LineIndices,
867                                                )) => {
868                                                    topology = crate::MeshOutputTopology::Lines;
869                                                }
870                                                Some(crate::Binding::BuiltIn(
871                                                    crate::BuiltIn::TriangleIndices,
872                                                )) => {
873                                                    topology = crate::MeshOutputTopology::Triangles;
874                                                }
875                                                _ => (),
876                                            }
877                                        }
878                                    }
879                                    _ => (),
880                                }
881                                // Each builtin should only occur once
882                                if builtins.contains(&crate::BuiltIn::Primitives) {
883                                    error.set(EntryPointError::BadMeshOutputVariableType);
884                                }
885                                builtins.insert(crate::BuiltIn::Primitives);
886                            } else {
887                                vertex_info = (a, b, c);
888                                // Each builtin should only occur once
889                                if builtins.contains(&crate::BuiltIn::Vertices) {
890                                    error.set(EntryPointError::BadMeshOutputVariableType);
891                                }
892                                builtins.insert(crate::BuiltIn::Vertices);
893                            }
894                        }
895                        _ => error.set(EntryPointError::BadMeshOutputVariableType),
896                    }
897                }
898                output = crate::MeshStageInfo {
899                    topology,
900                    max_vertices: vertex_info.0,
901                    max_vertices_override: None,
902                    vertex_output_type: vertex_info.2,
903                    max_primitives: primitive_info.0,
904                    max_primitives_override: None,
905                    primitive_output_type: primitive_info.2,
906                    ..output
907                }
908            }
909            _ => error.set(EntryPointError::BadMeshOutputVariableType),
910        }
911        (
912            output,
913            [vertex_info.1, primitive_info.1],
914            error
915                .inner
916                .map(|a| a.with_span_handle(self.global_variables[gv].ty, &self.types)),
917        )
918    }
919
920    pub fn uses_mesh_shaders(&self) -> bool {
921        let binding_uses_mesh = |b: &crate::Binding| {
922            matches!(
923                b,
924                crate::Binding::BuiltIn(
925                    crate::BuiltIn::MeshTaskSize
926                        | crate::BuiltIn::CullPrimitive
927                        | crate::BuiltIn::PointIndex
928                        | crate::BuiltIn::LineIndices
929                        | crate::BuiltIn::TriangleIndices
930                        | crate::BuiltIn::VertexCount
931                        | crate::BuiltIn::Vertices
932                        | crate::BuiltIn::PrimitiveCount
933                        | crate::BuiltIn::Primitives,
934                ) | crate::Binding::Location {
935                    per_primitive: true,
936                    ..
937                }
938            )
939        };
940        for (_, ty) in self.types.iter() {
941            match ty.inner {
942                crate::TypeInner::Struct { ref members, .. } => {
943                    for binding in members.iter().filter_map(|m| m.binding.as_ref()) {
944                        if binding_uses_mesh(binding) {
945                            return true;
946                        }
947                    }
948                }
949                _ => (),
950            }
951        }
952        for ep in &self.entry_points {
953            if matches!(
954                ep.stage,
955                crate::ShaderStage::Mesh | crate::ShaderStage::Task
956            ) {
957                return true;
958            }
959            for binding in ep
960                .function
961                .arguments
962                .iter()
963                .filter_map(|arg| arg.binding.as_ref())
964                .chain(
965                    ep.function
966                        .result
967                        .iter()
968                        .filter_map(|res| res.binding.as_ref()),
969                )
970            {
971                if binding_uses_mesh(binding) {
972                    return true;
973                }
974            }
975        }
976        if self
977            .global_variables
978            .iter()
979            .any(|gv| gv.1.space == crate::AddressSpace::TaskPayload)
980        {
981            return true;
982        }
983        false
984    }
985
986    /// Returns `true` if any function or entry point in the module uses
987    /// [`Statement::DebugPrintf`].
988    ///
989    /// [`Statement::DebugPrintf`]: crate::Statement::DebugPrintf
990    pub fn uses_debug_printf(&self) -> bool {
991        let functions = self.functions.iter().map(|(_, f)| f);
992        let entry_points = self.entry_points.iter().map(|ep| &ep.function);
993        functions.chain(entry_points).any(|func| {
994            any_statement(&func.body, &mut |stmt| {
995                matches!(*stmt, crate::Statement::DebugPrintf { .. })
996            })
997        })
998    }
999
1000    pub fn uses_ray_tracing(&self, ep_index: Option<usize>) -> RayTracingUses {
1001        let mut uses = RayTracingUses::default();
1002        // Whether this uses ray tracing (unknown whether the usage is pipelines or ray queries).
1003        let mut uses_ray_tracing = self.special_types.ray_desc.is_some();
1004
1005        uses.queries |= self.special_types.ray_intersection.is_some();
1006
1007        for (_, &crate::Type { ref inner, .. }) in self.types.iter() {
1008            // Backends do not know whether these have vertex return - that is done by us
1009            match *inner {
1010                crate::TypeInner::AccelerationStructure { .. } => {
1011                    uses_ray_tracing = true;
1012                }
1013                crate::TypeInner::RayQuery { .. } => uses.queries = true,
1014                _ => {}
1015            }
1016        }
1017
1018        for (index, ep) in self.entry_points.iter().enumerate() {
1019            if ep_index.is_some() && ep_index != Some(index) {
1020                continue;
1021            }
1022
1023            // if we have a ray tracing pipeline shader we are definitely using
1024            // pipelines, otherwise, if we have a ray tracing type, we might
1025            // be using it in the shader (which would require ray queries),
1026            // so we should use queries.
1027            if matches!(
1028                ep.stage,
1029                crate::ShaderStage::RayGeneration
1030                    | crate::ShaderStage::AnyHit
1031                    | crate::ShaderStage::ClosestHit
1032                    | crate::ShaderStage::Miss
1033            ) {
1034                uses.pipelines = true;
1035            } else {
1036                uses.queries |= uses_ray_tracing;
1037            }
1038        }
1039
1040        uses
1041    }
1042}
1043
1044#[derive(Copy, Clone, Debug, Default)]
1045pub struct RayTracingUses {
1046    pub pipelines: bool,
1047    pub queries: bool,
1048}
1049
1050impl crate::MeshOutputTopology {
1051    pub const fn to_builtin(self) -> crate::BuiltIn {
1052        match self {
1053            Self::Points => crate::BuiltIn::PointIndex,
1054            Self::Lines => crate::BuiltIn::LineIndices,
1055            Self::Triangles => crate::BuiltIn::TriangleIndices,
1056        }
1057    }
1058}
1059
1060impl crate::AddressSpace {
1061    pub const fn is_workgroup_like(self) -> bool {
1062        matches!(self, Self::WorkGroup | Self::TaskPayload)
1063    }
1064}
1065
1066impl TryFrom<crate::ScalarKind> for nt::glsl::GlslScalarKind {
1067    type Error = ();
1068
1069    fn try_from(value: crate::ScalarKind) -> Result<Self, Self::Error> {
1070        Ok(match value {
1071            crate::ScalarKind::Sint => nt::glsl::GlslScalarKind::Sint,
1072            crate::ScalarKind::Uint => nt::glsl::GlslScalarKind::Uint,
1073            crate::ScalarKind::Float => nt::glsl::GlslScalarKind::Float,
1074            _ => return Err(()),
1075        })
1076    }
1077}
1078
1079impl From<crate::VectorSize> for nt::glsl::GlslVectorSize {
1080    fn from(val: crate::VectorSize) -> Self {
1081        match val {
1082            crate::VectorSize::Bi => nt::glsl::GlslVectorSize::Bi,
1083            crate::VectorSize::Tri => nt::glsl::GlslVectorSize::Tri,
1084            crate::VectorSize::Quad => nt::glsl::GlslVectorSize::Quad,
1085        }
1086    }
1087}
1088
1089impl TryFrom<crate::Scalar> for nt::glsl::GlslScalar {
1090    type Error = ();
1091
1092    fn try_from(value: crate::Scalar) -> Result<Self, Self::Error> {
1093        Ok(nt::glsl::GlslScalar {
1094            kind: value.kind.try_into()?,
1095            width: value.width,
1096        })
1097    }
1098}
1099
1100impl TryFrom<&crate::TypeInner> for nt::glsl::GlslUniformType {
1101    type Error = ();
1102    fn try_from(value: &crate::TypeInner) -> Result<Self, Self::Error> {
1103        match *value {
1104            crate::TypeInner::Scalar(scalar) => {
1105                Ok(nt::glsl::GlslUniformType::Scalar(scalar.try_into()?))
1106            }
1107            crate::TypeInner::Vector { size, scalar } => Ok(nt::glsl::GlslUniformType::Vector {
1108                size: size.into(),
1109                scalar: scalar.try_into()?,
1110            }),
1111            crate::TypeInner::Matrix {
1112                columns,
1113                rows,
1114                scalar,
1115            } => Ok(nt::glsl::GlslUniformType::Matrix {
1116                columns: columns.into(),
1117                rows: rows.into(),
1118                scalar: scalar.try_into()?,
1119            }),
1120            _ => Err(()),
1121        }
1122    }
1123}