lagrange_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
691#[test]
692fn test_matrix_size() {
693    let module = crate::Module::default();
694    assert_eq!(
695        crate::TypeInner::Matrix {
696            columns: crate::VectorSize::Tri,
697            rows: crate::VectorSize::Tri,
698            scalar: crate::Scalar::F32,
699        }
700        .size(module.to_ctx()),
701        48,
702    );
703}
704
705impl crate::Module {
706    /// Extracts mesh shader info from a mesh output global variable. Used in frontends
707    /// and by validators. This only validates the output variable itself, and not the
708    /// vertex and primitive output types.
709    ///
710    /// The output contains the extracted mesh stage info, with overrides unset,
711    /// and then the overrides separately. This is because the overrides should be
712    /// treated as expressions elsewhere, but that requires mutably modifying the
713    /// module and the expressions should only be created at parse time, not validation
714    /// time.
715    #[allow(clippy::type_complexity)]
716    pub fn analyze_mesh_shader_info(
717        &self,
718        gv: crate::Handle<crate::GlobalVariable>,
719    ) -> (
720        crate::MeshStageInfo,
721        [Option<crate::Handle<crate::Override>>; 2],
722        Option<crate::WithSpan<crate::valid::EntryPointError>>,
723    ) {
724        use crate::span::AddSpan;
725        use crate::valid::EntryPointError;
726        #[derive(Default)]
727        struct OutError {
728            pub inner: Option<EntryPointError>,
729        }
730        impl OutError {
731            pub fn set(&mut self, err: EntryPointError) {
732                if self.inner.is_none() {
733                    self.inner = Some(err);
734                }
735            }
736        }
737
738        // Used to temporarily initialize stuff
739        let null_type = crate::Handle::new(NonMaxU32::new(0).unwrap());
740        let mut output = crate::MeshStageInfo {
741            topology: crate::MeshOutputTopology::Triangles,
742            max_vertices: 0,
743            max_vertices_override: None,
744            max_primitives: 0,
745            max_primitives_override: None,
746            vertex_output_type: null_type,
747            primitive_output_type: null_type,
748            output_variable: gv,
749        };
750        // Stores the error to output, if any.
751        let mut error = OutError::default();
752        let r#type = &self.types[self.global_variables[gv].ty].inner;
753
754        let mut topology = output.topology;
755        // Max, max override, type
756        let mut vertex_info = (0, None, null_type);
757        let mut primitive_info = (0, None, null_type);
758
759        match r#type {
760            &crate::TypeInner::Struct { ref members, .. } => {
761                let mut builtins = crate::FastHashSet::default();
762                for member in members {
763                    match member.binding {
764                        Some(crate::Binding::BuiltIn(crate::BuiltIn::VertexCount)) => {
765                            // Must have type u32
766                            if self.types[member.ty].inner.scalar() != Some(crate::Scalar::U32) {
767                                error.set(EntryPointError::BadMeshOutputVariableField);
768                            }
769                            // Each builtin should only occur once
770                            if builtins.contains(&crate::BuiltIn::VertexCount) {
771                                error.set(EntryPointError::BadMeshOutputVariableType);
772                            }
773                            builtins.insert(crate::BuiltIn::VertexCount);
774                        }
775                        Some(crate::Binding::BuiltIn(crate::BuiltIn::PrimitiveCount)) => {
776                            // Must have type u32
777                            if self.types[member.ty].inner.scalar() != Some(crate::Scalar::U32) {
778                                error.set(EntryPointError::BadMeshOutputVariableField);
779                            }
780                            // Each builtin should only occur once
781                            if builtins.contains(&crate::BuiltIn::PrimitiveCount) {
782                                error.set(EntryPointError::BadMeshOutputVariableType);
783                            }
784                            builtins.insert(crate::BuiltIn::PrimitiveCount);
785                        }
786                        Some(crate::Binding::BuiltIn(
787                            crate::BuiltIn::Vertices | crate::BuiltIn::Primitives,
788                        )) => {
789                            let ty = &self.types[member.ty].inner;
790                            // Analyze the array type to determine size and vertex/primitive type
791                            let (a, b, c) = match ty {
792                                &crate::TypeInner::Array { base, size, .. } => {
793                                    let ty = base;
794                                    let (max, max_override) = match size {
795                                        crate::ArraySize::Constant(a) => (a.get(), None),
796                                        crate::ArraySize::Pending(o) => (0, Some(o)),
797                                        crate::ArraySize::Dynamic => {
798                                            error.set(EntryPointError::BadMeshOutputVariableField);
799                                            (0, None)
800                                        }
801                                    };
802                                    (max, max_override, ty)
803                                }
804                                _ => {
805                                    error.set(EntryPointError::BadMeshOutputVariableField);
806                                    (0, None, null_type)
807                                }
808                            };
809                            if matches!(
810                                member.binding,
811                                Some(crate::Binding::BuiltIn(crate::BuiltIn::Primitives))
812                            ) {
813                                // Primitives require special analysis to determine topology
814                                primitive_info = (a, b, c);
815                                match self.types[c].inner {
816                                    crate::TypeInner::Struct { ref members, .. } => {
817                                        for member in members {
818                                            match member.binding {
819                                                Some(crate::Binding::BuiltIn(
820                                                    crate::BuiltIn::PointIndex,
821                                                )) => {
822                                                    topology = crate::MeshOutputTopology::Points;
823                                                }
824                                                Some(crate::Binding::BuiltIn(
825                                                    crate::BuiltIn::LineIndices,
826                                                )) => {
827                                                    topology = crate::MeshOutputTopology::Lines;
828                                                }
829                                                Some(crate::Binding::BuiltIn(
830                                                    crate::BuiltIn::TriangleIndices,
831                                                )) => {
832                                                    topology = crate::MeshOutputTopology::Triangles;
833                                                }
834                                                _ => (),
835                                            }
836                                        }
837                                    }
838                                    _ => (),
839                                }
840                                // Each builtin should only occur once
841                                if builtins.contains(&crate::BuiltIn::Primitives) {
842                                    error.set(EntryPointError::BadMeshOutputVariableType);
843                                }
844                                builtins.insert(crate::BuiltIn::Primitives);
845                            } else {
846                                vertex_info = (a, b, c);
847                                // Each builtin should only occur once
848                                if builtins.contains(&crate::BuiltIn::Vertices) {
849                                    error.set(EntryPointError::BadMeshOutputVariableType);
850                                }
851                                builtins.insert(crate::BuiltIn::Vertices);
852                            }
853                        }
854                        _ => error.set(EntryPointError::BadMeshOutputVariableType),
855                    }
856                }
857                output = crate::MeshStageInfo {
858                    topology,
859                    max_vertices: vertex_info.0,
860                    max_vertices_override: None,
861                    vertex_output_type: vertex_info.2,
862                    max_primitives: primitive_info.0,
863                    max_primitives_override: None,
864                    primitive_output_type: primitive_info.2,
865                    ..output
866                }
867            }
868            _ => error.set(EntryPointError::BadMeshOutputVariableType),
869        }
870        (
871            output,
872            [vertex_info.1, primitive_info.1],
873            error
874                .inner
875                .map(|a| a.with_span_handle(self.global_variables[gv].ty, &self.types)),
876        )
877    }
878
879    pub fn uses_mesh_shaders(&self) -> bool {
880        let binding_uses_mesh = |b: &crate::Binding| {
881            matches!(
882                b,
883                crate::Binding::BuiltIn(
884                    crate::BuiltIn::MeshTaskSize
885                        | crate::BuiltIn::CullPrimitive
886                        | crate::BuiltIn::PointIndex
887                        | crate::BuiltIn::LineIndices
888                        | crate::BuiltIn::TriangleIndices
889                        | crate::BuiltIn::VertexCount
890                        | crate::BuiltIn::Vertices
891                        | crate::BuiltIn::PrimitiveCount
892                        | crate::BuiltIn::Primitives,
893                ) | crate::Binding::Location {
894                    per_primitive: true,
895                    ..
896                }
897            )
898        };
899        for (_, ty) in self.types.iter() {
900            match ty.inner {
901                crate::TypeInner::Struct { ref members, .. } => {
902                    for binding in members.iter().filter_map(|m| m.binding.as_ref()) {
903                        if binding_uses_mesh(binding) {
904                            return true;
905                        }
906                    }
907                }
908                _ => (),
909            }
910        }
911        for ep in &self.entry_points {
912            if matches!(
913                ep.stage,
914                crate::ShaderStage::Mesh | crate::ShaderStage::Task
915            ) {
916                return true;
917            }
918            for binding in ep
919                .function
920                .arguments
921                .iter()
922                .filter_map(|arg| arg.binding.as_ref())
923                .chain(
924                    ep.function
925                        .result
926                        .iter()
927                        .filter_map(|res| res.binding.as_ref()),
928                )
929            {
930                if binding_uses_mesh(binding) {
931                    return true;
932                }
933            }
934        }
935        if self
936            .global_variables
937            .iter()
938            .any(|gv| gv.1.space == crate::AddressSpace::TaskPayload)
939        {
940            return true;
941        }
942        false
943    }
944
945    pub fn uses_ray_tracing(&self, ep_index: Option<usize>) -> RayTracingUses {
946        let mut uses = RayTracingUses::default();
947        // Whether this uses ray tracing (unknown whether the usage is pipelines or ray queries).
948        let mut uses_ray_tracing = self.special_types.ray_desc.is_some();
949
950        uses.queries |= self.special_types.ray_intersection.is_some();
951
952        for (_, &crate::Type { ref inner, .. }) in self.types.iter() {
953            // Backends do not know whether these have vertex return - that is done by us
954            match *inner {
955                crate::TypeInner::AccelerationStructure { .. } => {
956                    uses_ray_tracing = true;
957                }
958                crate::TypeInner::RayQuery { .. } => uses.queries = true,
959                _ => {}
960            }
961        }
962
963        for (index, ep) in self.entry_points.iter().enumerate() {
964            if ep_index.is_some() && ep_index != Some(index) {
965                continue;
966            }
967
968            // if we have a ray tracing pipeline shader we are definitely using
969            // pipelines, otherwise, if we have a ray tracing type, we might
970            // be using it in the shader (which would require ray queries),
971            // so we should use queries.
972            if matches!(
973                ep.stage,
974                crate::ShaderStage::RayGeneration
975                    | crate::ShaderStage::AnyHit
976                    | crate::ShaderStage::ClosestHit
977                    | crate::ShaderStage::Miss
978            ) {
979                uses.pipelines = true;
980            } else {
981                uses.queries |= uses_ray_tracing;
982            }
983        }
984
985        uses
986    }
987}
988
989#[derive(Default, Copy, Clone)]
990pub struct RayTracingUses {
991    pub pipelines: bool,
992    pub queries: bool,
993}
994
995impl crate::MeshOutputTopology {
996    pub const fn to_builtin(self) -> crate::BuiltIn {
997        match self {
998            Self::Points => crate::BuiltIn::PointIndex,
999            Self::Lines => crate::BuiltIn::LineIndices,
1000            Self::Triangles => crate::BuiltIn::TriangleIndices,
1001        }
1002    }
1003}
1004
1005impl crate::AddressSpace {
1006    pub const fn is_workgroup_like(self) -> bool {
1007        matches!(self, Self::WorkGroup | Self::TaskPayload)
1008    }
1009}
1010
1011impl TryFrom<crate::ScalarKind> for nt::glsl::GlslScalarKind {
1012    type Error = ();
1013
1014    fn try_from(value: crate::ScalarKind) -> Result<Self, Self::Error> {
1015        Ok(match value {
1016            crate::ScalarKind::Sint => nt::glsl::GlslScalarKind::Sint,
1017            crate::ScalarKind::Uint => nt::glsl::GlslScalarKind::Uint,
1018            crate::ScalarKind::Float => nt::glsl::GlslScalarKind::Float,
1019            _ => return Err(()),
1020        })
1021    }
1022}
1023
1024impl From<crate::VectorSize> for nt::glsl::GlslVectorSize {
1025    fn from(val: crate::VectorSize) -> Self {
1026        match val {
1027            crate::VectorSize::Bi => nt::glsl::GlslVectorSize::Bi,
1028            crate::VectorSize::Tri => nt::glsl::GlslVectorSize::Tri,
1029            crate::VectorSize::Quad => nt::glsl::GlslVectorSize::Quad,
1030        }
1031    }
1032}
1033
1034impl TryFrom<crate::Scalar> for nt::glsl::GlslScalar {
1035    type Error = ();
1036
1037    fn try_from(value: crate::Scalar) -> Result<Self, Self::Error> {
1038        Ok(nt::glsl::GlslScalar {
1039            kind: value.kind.try_into()?,
1040            width: value.width,
1041        })
1042    }
1043}
1044
1045impl TryFrom<&crate::TypeInner> for nt::glsl::GlslUniformType {
1046    type Error = ();
1047    fn try_from(value: &crate::TypeInner) -> Result<Self, Self::Error> {
1048        match *value {
1049            crate::TypeInner::Scalar(scalar) => {
1050                Ok(nt::glsl::GlslUniformType::Scalar(scalar.try_into()?))
1051            }
1052            crate::TypeInner::Vector { size, scalar } => Ok(nt::glsl::GlslUniformType::Vector {
1053                size: size.into(),
1054                scalar: scalar.try_into()?,
1055            }),
1056            crate::TypeInner::Matrix {
1057                columns,
1058                rows,
1059                scalar,
1060            } => Ok(nt::glsl::GlslUniformType::Matrix {
1061                columns: columns.into(),
1062                rows: rows.into(),
1063                scalar: scalar.try_into()?,
1064            }),
1065            _ => Err(()),
1066        }
1067    }
1068}
lagrange_naga/proc/mod.rs

lagrange_naga/proc/
mod.rs
