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