Skip to main content

naga/front/
type_gen.rs

1/*!
2Type generators.
3*/
4
5use alloc::{boxed::Box, string::ToString, vec};
6
7use crate::{arena::Handle, span::Span};
8
9impl crate::Module {
10    /// Populate this module's [`SpecialTypes::ray_desc`] type.
11    ///
12    /// [`SpecialTypes::ray_desc`] is the type of the [`descriptor`] operand of
13    /// an [`Initialize`] [`RayQuery`] statement. In WGSL, it is a struct type
14    /// referred to as `RayDesc`.
15    ///
16    /// Backends consume values of this type to drive platform APIs, so if you
17    /// change any its fields, you must update the backends to match. Look for
18    /// backend code dealing with [`RayQueryFunction::Initialize`].
19    ///
20    /// [`SpecialTypes::ray_desc`]: crate::SpecialTypes::ray_desc
21    /// [`descriptor`]: crate::RayQueryFunction::Initialize::descriptor
22    /// [`Initialize`]: crate::RayQueryFunction::Initialize
23    /// [`RayQuery`]: crate::Statement::RayQuery
24    /// [`RayQueryFunction::Initialize`]: crate::RayQueryFunction::Initialize
25    pub fn generate_ray_desc_type(&mut self) -> Handle<crate::Type> {
26        if let Some(handle) = self.special_types.ray_desc {
27            return handle;
28        }
29
30        let ty_flag = self.types.insert(
31            crate::Type {
32                name: None,
33                inner: crate::TypeInner::Scalar(crate::Scalar::U32),
34            },
35            Span::UNDEFINED,
36        );
37        let ty_scalar = self.types.insert(
38            crate::Type {
39                name: None,
40                inner: crate::TypeInner::Scalar(crate::Scalar::F32),
41            },
42            Span::UNDEFINED,
43        );
44        let ty_vector = self.types.insert(
45            crate::Type {
46                name: None,
47                inner: crate::TypeInner::Vector {
48                    size: crate::VectorSize::Tri,
49                    scalar: crate::Scalar::F32,
50                },
51            },
52            Span::UNDEFINED,
53        );
54
55        let handle = self.types.insert(
56            crate::Type {
57                name: Some("RayDesc".to_string()),
58                inner: crate::TypeInner::Struct {
59                    members: vec![
60                        crate::StructMember {
61                            name: Some("flags".to_string()),
62                            ty: ty_flag,
63                            binding: None,
64                            offset: 0,
65                        },
66                        crate::StructMember {
67                            name: Some("cull_mask".to_string()),
68                            ty: ty_flag,
69                            binding: None,
70                            offset: 4,
71                        },
72                        crate::StructMember {
73                            name: Some("tmin".to_string()),
74                            ty: ty_scalar,
75                            binding: None,
76                            offset: 8,
77                        },
78                        crate::StructMember {
79                            name: Some("tmax".to_string()),
80                            ty: ty_scalar,
81                            binding: None,
82                            offset: 12,
83                        },
84                        crate::StructMember {
85                            name: Some("origin".to_string()),
86                            ty: ty_vector,
87                            binding: None,
88                            offset: 16,
89                        },
90                        crate::StructMember {
91                            name: Some("dir".to_string()),
92                            ty: ty_vector,
93                            binding: None,
94                            offset: 32,
95                        },
96                    ],
97                    span: 48,
98                },
99            },
100            Span::UNDEFINED,
101        );
102
103        self.special_types.ray_desc = Some(handle);
104        handle
105    }
106
107    /// Make sure the types for the vertex return are in the module's type
108    pub fn generate_vertex_return_type(&mut self) -> Handle<crate::Type> {
109        if let Some(handle) = self.special_types.ray_vertex_return {
110            return handle;
111        }
112        let ty_vec3f = self.types.insert(
113            crate::Type {
114                name: None,
115                inner: crate::TypeInner::Vector {
116                    size: crate::VectorSize::Tri,
117                    scalar: crate::Scalar::F32,
118                },
119            },
120            Span::UNDEFINED,
121        );
122        let array = self.types.insert(
123            crate::Type {
124                name: None,
125                inner: crate::TypeInner::Array {
126                    base: ty_vec3f,
127                    size: crate::ArraySize::Constant(core::num::NonZeroU32::new(3).unwrap()),
128                    stride: 16,
129                },
130            },
131            Span::UNDEFINED,
132        );
133        self.special_types.ray_vertex_return = Some(array);
134        array
135    }
136
137    /// Populate this module's [`SpecialTypes::ray_intersection`] type.
138    ///
139    /// [`SpecialTypes::ray_intersection`] is the type of a
140    /// `RayQueryGetIntersection` expression. In WGSL, it is a struct type
141    /// referred to as `RayIntersection`.
142    ///
143    /// Backends construct values of this type based on platform APIs, so if you
144    /// change any its fields, you must update the backends to match. Look for
145    /// the backend's handling for [`Expression::RayQueryGetIntersection`].
146    ///
147    /// [`SpecialTypes::ray_intersection`]: crate::SpecialTypes::ray_intersection
148    /// [`Expression::RayQueryGetIntersection`]: crate::Expression::RayQueryGetIntersection
149    pub fn generate_ray_intersection_type(&mut self) -> Handle<crate::Type> {
150        if let Some(handle) = self.special_types.ray_intersection {
151            return handle;
152        }
153
154        let ty_flag = self.types.insert(
155            crate::Type {
156                name: None,
157                inner: crate::TypeInner::Scalar(crate::Scalar::U32),
158            },
159            Span::UNDEFINED,
160        );
161        let ty_scalar = self.types.insert(
162            crate::Type {
163                name: None,
164                inner: crate::TypeInner::Scalar(crate::Scalar::F32),
165            },
166            Span::UNDEFINED,
167        );
168        let ty_barycentrics = self.types.insert(
169            crate::Type {
170                name: None,
171                inner: crate::TypeInner::Vector {
172                    size: crate::VectorSize::Bi,
173                    scalar: crate::Scalar::F32,
174                },
175            },
176            Span::UNDEFINED,
177        );
178        let ty_bool = self.types.insert(
179            crate::Type {
180                name: None,
181                inner: crate::TypeInner::Scalar(crate::Scalar::BOOL),
182            },
183            Span::UNDEFINED,
184        );
185        let ty_transform = self.types.insert(
186            crate::Type {
187                name: None,
188                inner: crate::TypeInner::Matrix {
189                    columns: crate::VectorSize::Quad,
190                    rows: crate::VectorSize::Tri,
191                    scalar: crate::Scalar::F32,
192                },
193            },
194            Span::UNDEFINED,
195        );
196
197        let handle = self.types.insert(
198            crate::Type {
199                name: Some("RayIntersection".to_string()),
200                inner: crate::TypeInner::Struct {
201                    members: vec![
202                        crate::StructMember {
203                            name: Some("kind".to_string()),
204                            ty: ty_flag,
205                            binding: None,
206                            offset: 0,
207                        },
208                        crate::StructMember {
209                            name: Some("t".to_string()),
210                            ty: ty_scalar,
211                            binding: None,
212                            offset: 4,
213                        },
214                        crate::StructMember {
215                            name: Some("instance_custom_data".to_string()),
216                            ty: ty_flag,
217                            binding: None,
218                            offset: 8,
219                        },
220                        crate::StructMember {
221                            name: Some("instance_index".to_string()),
222                            ty: ty_flag,
223                            binding: None,
224                            offset: 12,
225                        },
226                        crate::StructMember {
227                            name: Some("sbt_record_offset".to_string()),
228                            ty: ty_flag,
229                            binding: None,
230                            offset: 16,
231                        },
232                        crate::StructMember {
233                            name: Some("geometry_index".to_string()),
234                            ty: ty_flag,
235                            binding: None,
236                            offset: 20,
237                        },
238                        crate::StructMember {
239                            name: Some("primitive_index".to_string()),
240                            ty: ty_flag,
241                            binding: None,
242                            offset: 24,
243                        },
244                        crate::StructMember {
245                            name: Some("barycentrics".to_string()),
246                            ty: ty_barycentrics,
247                            binding: None,
248                            offset: 28,
249                        },
250                        crate::StructMember {
251                            name: Some("front_face".to_string()),
252                            ty: ty_bool,
253                            binding: None,
254                            offset: 36,
255                        },
256                        crate::StructMember {
257                            name: Some("object_to_world".to_string()),
258                            ty: ty_transform,
259                            binding: None,
260                            offset: 48,
261                        },
262                        crate::StructMember {
263                            name: Some("world_to_object".to_string()),
264                            ty: ty_transform,
265                            binding: None,
266                            offset: 112,
267                        },
268                    ],
269                    span: 176,
270                },
271            },
272            Span::UNDEFINED,
273        );
274
275        self.special_types.ray_intersection = Some(handle);
276        handle
277    }
278
279    /// Generate [`SpecialTypes::external_texture_params`] and
280    /// [`SpecialTypes::external_texture_transfer_function`].
281    ///
282    /// Other than the WGSL backend, every backend that supports external
283    /// textures does so by lowering them to a set of ordinary textures and
284    /// some parameters saying how to sample from them. These types are used
285    /// for said parameters. Note that they are not used by the IR, but
286    /// generated purely as a convenience for the backends.
287    ///
288    /// [`SpecialTypes::external_texture_params`]: crate::ir::SpecialTypes::external_texture_params
289    /// [`SpecialTypes::external_texture_transfer_function`]: crate::ir::SpecialTypes::external_texture_transfer_function
290    pub fn generate_external_texture_types(&mut self) {
291        use crate::common::ExternalTextureParameter as P;
292
293        if self.special_types.external_texture_params.is_some() {
294            return;
295        }
296
297        let ty_f32 = self.types.insert(
298            crate::Type {
299                name: None,
300                inner: crate::TypeInner::Scalar(crate::Scalar::F32),
301            },
302            Span::UNDEFINED,
303        );
304        let ty_u32 = self.types.insert(
305            crate::Type {
306                name: None,
307                inner: crate::TypeInner::Scalar(crate::Scalar::U32),
308            },
309            Span::UNDEFINED,
310        );
311        let ty_vec2u = self.types.insert(
312            crate::Type {
313                name: None,
314                inner: crate::TypeInner::Vector {
315                    size: crate::VectorSize::Bi,
316                    scalar: crate::Scalar::U32,
317                },
318            },
319            Span::UNDEFINED,
320        );
321        let ty_mat3x2f = self.types.insert(
322            crate::Type {
323                name: None,
324                inner: crate::TypeInner::Matrix {
325                    columns: crate::VectorSize::Tri,
326                    rows: crate::VectorSize::Bi,
327                    scalar: crate::Scalar::F32,
328                },
329            },
330            Span::UNDEFINED,
331        );
332        let ty_mat3x3f = self.types.insert(
333            crate::Type {
334                name: None,
335                inner: crate::TypeInner::Matrix {
336                    columns: crate::VectorSize::Tri,
337                    rows: crate::VectorSize::Tri,
338                    scalar: crate::Scalar::F32,
339                },
340            },
341            Span::UNDEFINED,
342        );
343        let ty_mat4x4f = self.types.insert(
344            crate::Type {
345                name: None,
346                inner: crate::TypeInner::Matrix {
347                    columns: crate::VectorSize::Quad,
348                    rows: crate::VectorSize::Quad,
349                    scalar: crate::Scalar::F32,
350                },
351            },
352            Span::UNDEFINED,
353        );
354
355        let transfer_fn_handle = self.types.insert(
356            crate::Type {
357                name: Some("NagaExternalTextureTransferFn".to_string()),
358                inner: crate::TypeInner::Struct {
359                    members: vec![
360                        crate::StructMember {
361                            name: Some("a".to_string()),
362                            ty: ty_f32,
363                            binding: None,
364                            offset: 0,
365                        },
366                        crate::StructMember {
367                            name: Some("b".to_string()),
368                            ty: ty_f32,
369                            binding: None,
370                            offset: 4,
371                        },
372                        crate::StructMember {
373                            name: Some("g".to_string()),
374                            ty: ty_f32,
375                            binding: None,
376                            offset: 8,
377                        },
378                        crate::StructMember {
379                            name: Some("k".to_string()),
380                            ty: ty_f32,
381                            binding: None,
382                            offset: 12,
383                        },
384                    ],
385                    span: 16,
386                },
387            },
388            Span::UNDEFINED,
389        );
390        self.special_types.external_texture_transfer_function = Some(transfer_fn_handle);
391
392        // This definition must match [`wgpu_core::device::resource::ExternalTextureParams`].
393        let members: Box<[_; P::COUNT]> = Box::new([
394            crate::StructMember {
395                name: Some(P::YuvConversionMatrix.name().to_string()),
396                ty: ty_mat4x4f,
397                binding: None,
398                offset: 0,
399            },
400            crate::StructMember {
401                name: Some(P::GamutConversionMatrix.name().to_string()),
402                ty: ty_mat3x3f,
403                binding: None,
404                offset: 64,
405            },
406            crate::StructMember {
407                name: Some(P::SrcTf.name().to_string()),
408                ty: transfer_fn_handle,
409                binding: None,
410                offset: 112,
411            },
412            crate::StructMember {
413                name: Some(P::DstTf.name().to_string()),
414                ty: transfer_fn_handle,
415                binding: None,
416                offset: 128,
417            },
418            crate::StructMember {
419                name: Some(P::SampleTransform.name().to_string()),
420                ty: ty_mat3x2f,
421                binding: None,
422                offset: 144,
423            },
424            crate::StructMember {
425                name: Some(P::LoadTransform.name().to_string()),
426                ty: ty_mat3x2f,
427                binding: None,
428                offset: 168,
429            },
430            crate::StructMember {
431                name: Some(P::Size.name().to_string()),
432                ty: ty_vec2u,
433                binding: None,
434                offset: 192,
435            },
436            crate::StructMember {
437                name: Some(P::NumPlanes.name().to_string()),
438                ty: ty_u32,
439                binding: None,
440                offset: 200,
441            },
442        ]);
443
444        let ty = crate::Type {
445            name: Some("NagaExternalTextureParams".to_string()),
446            inner: crate::TypeInner::Struct {
447                members: <[_]>::into_vec(members),
448                span: 208,
449            },
450        };
451
452        let params_handle = self.types.insert(ty, Span::UNDEFINED);
453        self.special_types.external_texture_params = Some(params_handle);
454    }
455
456    /// Populate this module's [`SpecialTypes::predeclared_types`] type and return the handle.
457    ///
458    /// [`SpecialTypes::predeclared_types`]: crate::SpecialTypes::predeclared_types
459    pub fn generate_predeclared_type(
460        &mut self,
461        special_type: crate::PredeclaredType,
462    ) -> Handle<crate::Type> {
463        if let Some(value) = self.special_types.predeclared_types.get(&special_type) {
464            return *value;
465        }
466
467        let name = special_type.struct_name();
468        let ty = match special_type {
469            crate::PredeclaredType::AtomicCompareExchangeWeakResult(scalar) => {
470                let bool_ty = self.types.insert(
471                    crate::Type {
472                        name: None,
473                        inner: crate::TypeInner::Scalar(crate::Scalar::BOOL),
474                    },
475                    Span::UNDEFINED,
476                );
477                let scalar_ty = self.types.insert(
478                    crate::Type {
479                        name: None,
480                        inner: crate::TypeInner::Scalar(scalar),
481                    },
482                    Span::UNDEFINED,
483                );
484
485                crate::Type {
486                    name: Some(name),
487                    inner: crate::TypeInner::Struct {
488                        members: vec![
489                            crate::StructMember {
490                                name: Some("old_value".to_string()),
491                                ty: scalar_ty,
492                                binding: None,
493                                offset: 0,
494                            },
495                            crate::StructMember {
496                                name: Some("exchanged".to_string()),
497                                ty: bool_ty,
498                                binding: None,
499                                offset: scalar.width as u32,
500                            },
501                        ],
502                        span: scalar.width as u32 * 2,
503                    },
504                }
505            }
506            crate::PredeclaredType::ModfResult { size, scalar } => {
507                let float_ty = self.types.insert(
508                    crate::Type {
509                        name: None,
510                        inner: crate::TypeInner::Scalar(scalar),
511                    },
512                    Span::UNDEFINED,
513                );
514
515                let (member_ty, second_offset) = if let Some(size) = size {
516                    let vec_ty = self.types.insert(
517                        crate::Type {
518                            name: None,
519                            inner: crate::TypeInner::Vector { size, scalar },
520                        },
521                        Span::UNDEFINED,
522                    );
523                    (vec_ty, size as u32 * scalar.width as u32)
524                } else {
525                    (float_ty, scalar.width as u32)
526                };
527
528                crate::Type {
529                    name: Some(name),
530                    inner: crate::TypeInner::Struct {
531                        members: vec![
532                            crate::StructMember {
533                                name: Some("fract".to_string()),
534                                ty: member_ty,
535                                binding: None,
536                                offset: 0,
537                            },
538                            crate::StructMember {
539                                name: Some("whole".to_string()),
540                                ty: member_ty,
541                                binding: None,
542                                offset: second_offset,
543                            },
544                        ],
545                        span: second_offset * 2,
546                    },
547                }
548            }
549            crate::PredeclaredType::FrexpResult { size, scalar } => {
550                let float_ty = self.types.insert(
551                    crate::Type {
552                        name: None,
553                        inner: crate::TypeInner::Scalar(scalar),
554                    },
555                    Span::UNDEFINED,
556                );
557
558                let int_ty = self.types.insert(
559                    crate::Type {
560                        name: None,
561                        inner: crate::TypeInner::Scalar(crate::Scalar {
562                            kind: crate::ScalarKind::Sint,
563                            width: scalar.width,
564                        }),
565                    },
566                    Span::UNDEFINED,
567                );
568
569                let (fract_member_ty, exp_member_ty, second_offset) = if let Some(size) = size {
570                    let vec_float_ty = self.types.insert(
571                        crate::Type {
572                            name: None,
573                            inner: crate::TypeInner::Vector { size, scalar },
574                        },
575                        Span::UNDEFINED,
576                    );
577                    let vec_int_ty = self.types.insert(
578                        crate::Type {
579                            name: None,
580                            inner: crate::TypeInner::Vector {
581                                size,
582                                scalar: crate::Scalar {
583                                    kind: crate::ScalarKind::Sint,
584                                    width: scalar.width,
585                                },
586                            },
587                        },
588                        Span::UNDEFINED,
589                    );
590                    (vec_float_ty, vec_int_ty, size as u32 * scalar.width as u32)
591                } else {
592                    (float_ty, int_ty, scalar.width as u32)
593                };
594
595                crate::Type {
596                    name: Some(name),
597                    inner: crate::TypeInner::Struct {
598                        members: vec![
599                            crate::StructMember {
600                                name: Some("fract".to_string()),
601                                ty: fract_member_ty,
602                                binding: None,
603                                offset: 0,
604                            },
605                            crate::StructMember {
606                                name: Some("exp".to_string()),
607                                ty: exp_member_ty,
608                                binding: None,
609                                offset: second_offset,
610                            },
611                        ],
612                        span: second_offset * 2,
613                    },
614                }
615            }
616        };
617
618        let handle = self.types.insert(ty, Span::UNDEFINED);
619        self.special_types
620            .predeclared_types
621            .insert(special_type, handle);
622        handle
623    }
624}