naga/front/
type_gen.rs

1/*!
2Type generators.
3*/
4
5use alloc::{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        if self.special_types.external_texture_params.is_some() {
292            return;
293        }
294
295        let ty_f32 = self.types.insert(
296            crate::Type {
297                name: None,
298                inner: crate::TypeInner::Scalar(crate::Scalar::F32),
299            },
300            Span::UNDEFINED,
301        );
302        let ty_u32 = self.types.insert(
303            crate::Type {
304                name: None,
305                inner: crate::TypeInner::Scalar(crate::Scalar::U32),
306            },
307            Span::UNDEFINED,
308        );
309        let ty_vec2u = self.types.insert(
310            crate::Type {
311                name: None,
312                inner: crate::TypeInner::Vector {
313                    size: crate::VectorSize::Bi,
314                    scalar: crate::Scalar::U32,
315                },
316            },
317            Span::UNDEFINED,
318        );
319        let ty_mat3x2f = self.types.insert(
320            crate::Type {
321                name: None,
322                inner: crate::TypeInner::Matrix {
323                    columns: crate::VectorSize::Tri,
324                    rows: crate::VectorSize::Bi,
325                    scalar: crate::Scalar::F32,
326                },
327            },
328            Span::UNDEFINED,
329        );
330        let ty_mat3x3f = self.types.insert(
331            crate::Type {
332                name: None,
333                inner: crate::TypeInner::Matrix {
334                    columns: crate::VectorSize::Tri,
335                    rows: crate::VectorSize::Tri,
336                    scalar: crate::Scalar::F32,
337                },
338            },
339            Span::UNDEFINED,
340        );
341        let ty_mat4x4f = self.types.insert(
342            crate::Type {
343                name: None,
344                inner: crate::TypeInner::Matrix {
345                    columns: crate::VectorSize::Quad,
346                    rows: crate::VectorSize::Quad,
347                    scalar: crate::Scalar::F32,
348                },
349            },
350            Span::UNDEFINED,
351        );
352
353        let transfer_fn_handle = self.types.insert(
354            crate::Type {
355                name: Some("NagaExternalTextureTransferFn".to_string()),
356                inner: crate::TypeInner::Struct {
357                    members: vec![
358                        crate::StructMember {
359                            name: Some("a".to_string()),
360                            ty: ty_f32,
361                            binding: None,
362                            offset: 0,
363                        },
364                        crate::StructMember {
365                            name: Some("b".to_string()),
366                            ty: ty_f32,
367                            binding: None,
368                            offset: 4,
369                        },
370                        crate::StructMember {
371                            name: Some("g".to_string()),
372                            ty: ty_f32,
373                            binding: None,
374                            offset: 8,
375                        },
376                        crate::StructMember {
377                            name: Some("k".to_string()),
378                            ty: ty_f32,
379                            binding: None,
380                            offset: 12,
381                        },
382                    ],
383                    span: 16,
384                },
385            },
386            Span::UNDEFINED,
387        );
388        self.special_types.external_texture_transfer_function = Some(transfer_fn_handle);
389
390        let params_handle = self.types.insert(
391            crate::Type {
392                name: Some("NagaExternalTextureParams".to_string()),
393                inner: crate::TypeInner::Struct {
394                    members: vec![
395                        crate::StructMember {
396                            name: Some("yuv_conversion_matrix".to_string()),
397                            ty: ty_mat4x4f,
398                            binding: None,
399                            offset: 0,
400                        },
401                        crate::StructMember {
402                            name: Some("gamut_conversion_matrix".to_string()),
403                            ty: ty_mat3x3f,
404                            binding: None,
405                            offset: 64,
406                        },
407                        crate::StructMember {
408                            name: Some("src_tf".to_string()),
409                            ty: transfer_fn_handle,
410                            binding: None,
411                            offset: 112,
412                        },
413                        crate::StructMember {
414                            name: Some("dst_tf".to_string()),
415                            ty: transfer_fn_handle,
416                            binding: None,
417                            offset: 128,
418                        },
419                        crate::StructMember {
420                            name: Some("sample_transform".to_string()),
421                            ty: ty_mat3x2f,
422                            binding: None,
423                            offset: 144,
424                        },
425                        crate::StructMember {
426                            name: Some("load_transform".to_string()),
427                            ty: ty_mat3x2f,
428                            binding: None,
429                            offset: 168,
430                        },
431                        crate::StructMember {
432                            name: Some("size".to_string()),
433                            ty: ty_vec2u,
434                            binding: None,
435                            offset: 192,
436                        },
437                        crate::StructMember {
438                            name: Some("num_planes".to_string()),
439                            ty: ty_u32,
440                            binding: None,
441                            offset: 200,
442                        },
443                    ],
444                    span: 208,
445                },
446            },
447            Span::UNDEFINED,
448        );
449        self.special_types.external_texture_params = Some(params_handle);
450    }
451
452    /// Populate this module's [`SpecialTypes::predeclared_types`] type and return the handle.
453    ///
454    /// [`SpecialTypes::predeclared_types`]: crate::SpecialTypes::predeclared_types
455    pub fn generate_predeclared_type(
456        &mut self,
457        special_type: crate::PredeclaredType,
458    ) -> Handle<crate::Type> {
459        if let Some(value) = self.special_types.predeclared_types.get(&special_type) {
460            return *value;
461        }
462
463        let name = special_type.struct_name();
464        let ty = match special_type {
465            crate::PredeclaredType::AtomicCompareExchangeWeakResult(scalar) => {
466                let bool_ty = self.types.insert(
467                    crate::Type {
468                        name: None,
469                        inner: crate::TypeInner::Scalar(crate::Scalar::BOOL),
470                    },
471                    Span::UNDEFINED,
472                );
473                let scalar_ty = self.types.insert(
474                    crate::Type {
475                        name: None,
476                        inner: crate::TypeInner::Scalar(scalar),
477                    },
478                    Span::UNDEFINED,
479                );
480
481                crate::Type {
482                    name: Some(name),
483                    inner: crate::TypeInner::Struct {
484                        members: vec![
485                            crate::StructMember {
486                                name: Some("old_value".to_string()),
487                                ty: scalar_ty,
488                                binding: None,
489                                offset: 0,
490                            },
491                            crate::StructMember {
492                                name: Some("exchanged".to_string()),
493                                ty: bool_ty,
494                                binding: None,
495                                offset: scalar.width as u32,
496                            },
497                        ],
498                        span: scalar.width as u32 * 2,
499                    },
500                }
501            }
502            crate::PredeclaredType::ModfResult { size, scalar } => {
503                let float_ty = self.types.insert(
504                    crate::Type {
505                        name: None,
506                        inner: crate::TypeInner::Scalar(scalar),
507                    },
508                    Span::UNDEFINED,
509                );
510
511                let (member_ty, second_offset) = if let Some(size) = size {
512                    let vec_ty = self.types.insert(
513                        crate::Type {
514                            name: None,
515                            inner: crate::TypeInner::Vector { size, scalar },
516                        },
517                        Span::UNDEFINED,
518                    );
519                    (vec_ty, size as u32 * scalar.width as u32)
520                } else {
521                    (float_ty, scalar.width as u32)
522                };
523
524                crate::Type {
525                    name: Some(name),
526                    inner: crate::TypeInner::Struct {
527                        members: vec![
528                            crate::StructMember {
529                                name: Some("fract".to_string()),
530                                ty: member_ty,
531                                binding: None,
532                                offset: 0,
533                            },
534                            crate::StructMember {
535                                name: Some("whole".to_string()),
536                                ty: member_ty,
537                                binding: None,
538                                offset: second_offset,
539                            },
540                        ],
541                        span: second_offset * 2,
542                    },
543                }
544            }
545            crate::PredeclaredType::FrexpResult { size, scalar } => {
546                let float_ty = self.types.insert(
547                    crate::Type {
548                        name: None,
549                        inner: crate::TypeInner::Scalar(scalar),
550                    },
551                    Span::UNDEFINED,
552                );
553
554                let int_ty = self.types.insert(
555                    crate::Type {
556                        name: None,
557                        inner: crate::TypeInner::Scalar(crate::Scalar {
558                            kind: crate::ScalarKind::Sint,
559                            width: scalar.width,
560                        }),
561                    },
562                    Span::UNDEFINED,
563                );
564
565                let (fract_member_ty, exp_member_ty, second_offset) = if let Some(size) = size {
566                    let vec_float_ty = self.types.insert(
567                        crate::Type {
568                            name: None,
569                            inner: crate::TypeInner::Vector { size, scalar },
570                        },
571                        Span::UNDEFINED,
572                    );
573                    let vec_int_ty = self.types.insert(
574                        crate::Type {
575                            name: None,
576                            inner: crate::TypeInner::Vector {
577                                size,
578                                scalar: crate::Scalar {
579                                    kind: crate::ScalarKind::Sint,
580                                    width: scalar.width,
581                                },
582                            },
583                        },
584                        Span::UNDEFINED,
585                    );
586                    (vec_float_ty, vec_int_ty, size as u32 * scalar.width as u32)
587                } else {
588                    (float_ty, int_ty, scalar.width as u32)
589                };
590
591                crate::Type {
592                    name: Some(name),
593                    inner: crate::TypeInner::Struct {
594                        members: vec![
595                            crate::StructMember {
596                                name: Some("fract".to_string()),
597                                ty: fract_member_ty,
598                                binding: None,
599                                offset: 0,
600                            },
601                            crate::StructMember {
602                                name: Some("exp".to_string()),
603                                ty: exp_member_ty,
604                                binding: None,
605                                offset: second_offset,
606                            },
607                        ],
608                        span: second_offset * 2,
609                    },
610                }
611            }
612        };
613
614        let handle = self.types.insert(ty, Span::UNDEFINED);
615        self.special_types
616            .predeclared_types
617            .insert(special_type, handle);
618        handle
619    }
620}