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

naga/proc/
mod.rs
