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    /// Populate this module's [`SpecialTypes::predeclared_types`] type and return the handle.
280    ///
281    /// [`SpecialTypes::predeclared_types`]: crate::SpecialTypes::predeclared_types
282    pub fn generate_predeclared_type(
283        &mut self,
284        special_type: crate::PredeclaredType,
285    ) -> Handle<crate::Type> {
286        if let Some(value) = self.special_types.predeclared_types.get(&special_type) {
287            return *value;
288        }
289
290        let name = special_type.struct_name();
291        let ty = match special_type {
292            crate::PredeclaredType::AtomicCompareExchangeWeakResult(scalar) => {
293                let bool_ty = self.types.insert(
294                    crate::Type {
295                        name: None,
296                        inner: crate::TypeInner::Scalar(crate::Scalar::BOOL),
297                    },
298                    Span::UNDEFINED,
299                );
300                let scalar_ty = self.types.insert(
301                    crate::Type {
302                        name: None,
303                        inner: crate::TypeInner::Scalar(scalar),
304                    },
305                    Span::UNDEFINED,
306                );
307
308                crate::Type {
309                    name: Some(name),
310                    inner: crate::TypeInner::Struct {
311                        members: vec![
312                            crate::StructMember {
313                                name: Some("old_value".to_string()),
314                                ty: scalar_ty,
315                                binding: None,
316                                offset: 0,
317                            },
318                            crate::StructMember {
319                                name: Some("exchanged".to_string()),
320                                ty: bool_ty,
321                                binding: None,
322                                offset: scalar.width as u32,
323                            },
324                        ],
325                        span: scalar.width as u32 * 2,
326                    },
327                }
328            }
329            crate::PredeclaredType::ModfResult { size, scalar } => {
330                let float_ty = self.types.insert(
331                    crate::Type {
332                        name: None,
333                        inner: crate::TypeInner::Scalar(scalar),
334                    },
335                    Span::UNDEFINED,
336                );
337
338                let (member_ty, second_offset) = if let Some(size) = size {
339                    let vec_ty = self.types.insert(
340                        crate::Type {
341                            name: None,
342                            inner: crate::TypeInner::Vector { size, scalar },
343                        },
344                        Span::UNDEFINED,
345                    );
346                    (vec_ty, size as u32 * scalar.width as u32)
347                } else {
348                    (float_ty, scalar.width as u32)
349                };
350
351                crate::Type {
352                    name: Some(name),
353                    inner: crate::TypeInner::Struct {
354                        members: vec![
355                            crate::StructMember {
356                                name: Some("fract".to_string()),
357                                ty: member_ty,
358                                binding: None,
359                                offset: 0,
360                            },
361                            crate::StructMember {
362                                name: Some("whole".to_string()),
363                                ty: member_ty,
364                                binding: None,
365                                offset: second_offset,
366                            },
367                        ],
368                        span: second_offset * 2,
369                    },
370                }
371            }
372            crate::PredeclaredType::FrexpResult { size, scalar } => {
373                let float_ty = self.types.insert(
374                    crate::Type {
375                        name: None,
376                        inner: crate::TypeInner::Scalar(scalar),
377                    },
378                    Span::UNDEFINED,
379                );
380
381                let int_ty = self.types.insert(
382                    crate::Type {
383                        name: None,
384                        inner: crate::TypeInner::Scalar(crate::Scalar {
385                            kind: crate::ScalarKind::Sint,
386                            width: scalar.width,
387                        }),
388                    },
389                    Span::UNDEFINED,
390                );
391
392                let (fract_member_ty, exp_member_ty, second_offset) = if let Some(size) = size {
393                    let vec_float_ty = self.types.insert(
394                        crate::Type {
395                            name: None,
396                            inner: crate::TypeInner::Vector { size, scalar },
397                        },
398                        Span::UNDEFINED,
399                    );
400                    let vec_int_ty = self.types.insert(
401                        crate::Type {
402                            name: None,
403                            inner: crate::TypeInner::Vector {
404                                size,
405                                scalar: crate::Scalar {
406                                    kind: crate::ScalarKind::Sint,
407                                    width: scalar.width,
408                                },
409                            },
410                        },
411                        Span::UNDEFINED,
412                    );
413                    (vec_float_ty, vec_int_ty, size as u32 * scalar.width as u32)
414                } else {
415                    (float_ty, int_ty, scalar.width as u32)
416                };
417
418                crate::Type {
419                    name: Some(name),
420                    inner: crate::TypeInner::Struct {
421                        members: vec![
422                            crate::StructMember {
423                                name: Some("fract".to_string()),
424                                ty: fract_member_ty,
425                                binding: None,
426                                offset: 0,
427                            },
428                            crate::StructMember {
429                                name: Some("exp".to_string()),
430                                ty: exp_member_ty,
431                                binding: None,
432                                offset: second_offset,
433                            },
434                        ],
435                        span: second_offset * 2,
436                    },
437                }
438            }
439        };
440
441        let handle = self.types.insert(ty, Span::UNDEFINED);
442        self.special_types
443            .predeclared_types
444            .insert(special_type, handle);
445        handle
446    }
447}