naga/front/wgsl/parse/
ast.rs

1use alloc::vec::Vec;
2use core::hash::Hash;
3
4use crate::diagnostic_filter::DiagnosticFilterNode;
5use crate::front::wgsl::parse::directive::enable_extension::EnableExtensions;
6use crate::front::wgsl::parse::number::Number;
7use crate::front::wgsl::Scalar;
8use crate::{Arena, FastIndexSet, Handle, Span};
9
10#[derive(Debug, Default)]
11pub struct TranslationUnit<'a> {
12    pub enable_extensions: EnableExtensions,
13    pub decls: Arena<GlobalDecl<'a>>,
14    /// The common expressions arena for the entire translation unit.
15    ///
16    /// All functions, global initializers, array lengths, etc. store their
17    /// expressions here. We apportion these out to individual Naga
18    /// [`Function`]s' expression arenas at lowering time. Keeping them all in a
19    /// single arena simplifies handling of things like array lengths (which are
20    /// effectively global and thus don't clearly belong to any function) and
21    /// initializers (which can appear in both function-local and module-scope
22    /// contexts).
23    ///
24    /// [`Function`]: crate::Function
25    pub expressions: Arena<Expression<'a>>,
26
27    /// Non-user-defined types, like `vec4<f32>` or `array<i32, 10>`.
28    ///
29    /// These are referred to by `Handle<ast::Type<'a>>` values.
30    /// User-defined types are referred to by name until lowering.
31    pub types: Arena<Type<'a>>,
32
33    /// Arena for all diagnostic filter rules parsed in this module, including those in functions.
34    ///
35    /// See [`DiagnosticFilterNode`] for details on how the tree is represented and used in
36    /// validation.
37    pub diagnostic_filters: Arena<DiagnosticFilterNode>,
38    /// The leaf of all `diagnostic(…)` directives in this module.
39    ///
40    /// See [`DiagnosticFilterNode`] for details on how the tree is represented and used in
41    /// validation.
42    pub diagnostic_filter_leaf: Option<Handle<DiagnosticFilterNode>>,
43
44    /// Doc comments appearing first in the file.
45    /// This serves as documentation for the whole TranslationUnit.
46    pub doc_comments: Vec<&'a str>,
47}
48
49#[derive(Debug, Clone, Copy)]
50pub struct Ident<'a> {
51    pub name: &'a str,
52    pub span: Span,
53}
54
55#[derive(Debug)]
56pub enum IdentExpr<'a> {
57    Unresolved(&'a str),
58    Local(Handle<Local>),
59}
60
61/// A reference to a module-scope definition or predeclared object.
62///
63/// Each [`GlobalDecl`] holds a set of these values, to be resolved to
64/// specific definitions later. To support de-duplication, `Eq` and
65/// `Hash` on a `Dependency` value consider only the name, not the
66/// source location at which the reference occurs.
67#[derive(Debug)]
68pub struct Dependency<'a> {
69    /// The name referred to.
70    pub ident: &'a str,
71
72    /// The location at which the reference to that name occurs.
73    pub usage: Span,
74}
75
76impl Hash for Dependency<'_> {
77    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
78        self.ident.hash(state);
79    }
80}
81
82impl PartialEq for Dependency<'_> {
83    fn eq(&self, other: &Self) -> bool {
84        self.ident == other.ident
85    }
86}
87
88impl Eq for Dependency<'_> {}
89
90/// A module-scope declaration.
91#[derive(Debug)]
92pub struct GlobalDecl<'a> {
93    pub kind: GlobalDeclKind<'a>,
94
95    /// Names of all module-scope or predeclared objects this
96    /// declaration uses.
97    pub dependencies: FastIndexSet<Dependency<'a>>,
98}
99
100#[derive(Debug)]
101pub enum GlobalDeclKind<'a> {
102    Fn(Function<'a>),
103    Var(GlobalVariable<'a>),
104    Const(Const<'a>),
105    Override(Override<'a>),
106    Struct(Struct<'a>),
107    Type(TypeAlias<'a>),
108    ConstAssert(Handle<Expression<'a>>),
109}
110
111#[derive(Debug)]
112pub struct FunctionArgument<'a> {
113    pub name: Ident<'a>,
114    pub ty: Handle<Type<'a>>,
115    pub binding: Option<Binding<'a>>,
116    pub handle: Handle<Local>,
117}
118
119#[derive(Debug)]
120pub struct FunctionResult<'a> {
121    pub ty: Handle<Type<'a>>,
122    pub binding: Option<Binding<'a>>,
123    pub must_use: bool,
124}
125
126#[derive(Debug)]
127pub struct EntryPoint<'a> {
128    pub stage: crate::ShaderStage,
129    pub early_depth_test: Option<crate::EarlyDepthTest>,
130    pub workgroup_size: Option<[Option<Handle<Expression<'a>>>; 3]>,
131    pub mesh_output_variable: Option<(&'a str, Span)>,
132    pub task_payload: Option<(&'a str, Span)>,
133}
134
135#[cfg(doc)]
136use crate::front::wgsl::lower::{LocalExpressionContext, StatementContext};
137
138#[derive(Debug)]
139pub struct Function<'a> {
140    pub entry_point: Option<EntryPoint<'a>>,
141    pub name: Ident<'a>,
142    pub arguments: Vec<FunctionArgument<'a>>,
143    pub result: Option<FunctionResult<'a>>,
144    pub body: Block<'a>,
145    pub diagnostic_filter_leaf: Option<Handle<DiagnosticFilterNode>>,
146    pub doc_comments: Vec<&'a str>,
147}
148
149#[derive(Debug)]
150pub enum Binding<'a> {
151    BuiltIn(crate::BuiltIn),
152    Location {
153        location: Handle<Expression<'a>>,
154        interpolation: Option<crate::Interpolation>,
155        sampling: Option<crate::Sampling>,
156        blend_src: Option<Handle<Expression<'a>>>,
157        per_primitive: bool,
158    },
159}
160
161#[derive(Debug)]
162pub struct ResourceBinding<'a> {
163    pub group: Handle<Expression<'a>>,
164    pub binding: Handle<Expression<'a>>,
165}
166
167#[derive(Debug)]
168pub struct GlobalVariable<'a> {
169    pub name: Ident<'a>,
170    pub space: crate::AddressSpace,
171    pub binding: Option<ResourceBinding<'a>>,
172    pub ty: Option<Handle<Type<'a>>>,
173    pub init: Option<Handle<Expression<'a>>>,
174    pub doc_comments: Vec<&'a str>,
175}
176
177#[derive(Debug)]
178pub struct StructMember<'a> {
179    pub name: Ident<'a>,
180    pub ty: Handle<Type<'a>>,
181    pub binding: Option<Binding<'a>>,
182    pub align: Option<Handle<Expression<'a>>>,
183    pub size: Option<Handle<Expression<'a>>>,
184    pub doc_comments: Vec<&'a str>,
185}
186
187#[derive(Debug)]
188pub struct Struct<'a> {
189    pub name: Ident<'a>,
190    pub members: Vec<StructMember<'a>>,
191    pub doc_comments: Vec<&'a str>,
192}
193
194#[derive(Debug)]
195pub struct TypeAlias<'a> {
196    pub name: Ident<'a>,
197    pub ty: Handle<Type<'a>>,
198}
199
200#[derive(Debug)]
201pub struct Const<'a> {
202    pub name: Ident<'a>,
203    pub ty: Option<Handle<Type<'a>>>,
204    pub init: Handle<Expression<'a>>,
205    pub doc_comments: Vec<&'a str>,
206}
207
208#[derive(Debug)]
209pub struct Override<'a> {
210    pub name: Ident<'a>,
211    pub id: Option<Handle<Expression<'a>>>,
212    pub ty: Option<Handle<Type<'a>>>,
213    pub init: Option<Handle<Expression<'a>>>,
214}
215
216/// The size of an [`Array`] or [`BindingArray`].
217///
218/// [`Array`]: Type::Array
219/// [`BindingArray`]: Type::BindingArray
220#[derive(Debug, Copy, Clone)]
221pub enum ArraySize<'a> {
222    /// The length as a constant expression.
223    Constant(Handle<Expression<'a>>),
224    Dynamic,
225}
226
227#[derive(Debug)]
228pub enum Type<'a> {
229    Scalar(Scalar),
230    Vector {
231        size: crate::VectorSize,
232        ty: Handle<Type<'a>>,
233        ty_span: Span,
234    },
235    Matrix {
236        columns: crate::VectorSize,
237        rows: crate::VectorSize,
238        ty: Handle<Type<'a>>,
239        ty_span: Span,
240    },
241    CooperativeMatrix {
242        columns: crate::CooperativeSize,
243        rows: crate::CooperativeSize,
244        ty: Handle<Type<'a>>,
245        ty_span: Span,
246        role: crate::CooperativeRole,
247    },
248    Atomic(Scalar),
249    Pointer {
250        base: Handle<Type<'a>>,
251        space: crate::AddressSpace,
252    },
253    Array {
254        base: Handle<Type<'a>>,
255        size: ArraySize<'a>,
256    },
257    Image {
258        dim: crate::ImageDimension,
259        arrayed: bool,
260        class: crate::ImageClass,
261    },
262    Sampler {
263        comparison: bool,
264    },
265    AccelerationStructure {
266        vertex_return: bool,
267    },
268    RayQuery {
269        vertex_return: bool,
270    },
271    RayDesc,
272    RayIntersection,
273    BindingArray {
274        base: Handle<Type<'a>>,
275        size: ArraySize<'a>,
276    },
277
278    /// A user-defined type, like a struct or a type alias.
279    User(Ident<'a>),
280}
281
282#[derive(Debug, Default)]
283pub struct Block<'a> {
284    pub stmts: Vec<Statement<'a>>,
285}
286
287#[derive(Debug)]
288pub struct Statement<'a> {
289    pub kind: StatementKind<'a>,
290    pub span: Span,
291}
292
293#[derive(Debug)]
294pub enum StatementKind<'a> {
295    LocalDecl(LocalDecl<'a>),
296    Block(Block<'a>),
297    If {
298        condition: Handle<Expression<'a>>,
299        accept: Block<'a>,
300        reject: Block<'a>,
301    },
302    Switch {
303        selector: Handle<Expression<'a>>,
304        cases: Vec<SwitchCase<'a>>,
305    },
306    Loop {
307        body: Block<'a>,
308        continuing: Block<'a>,
309        break_if: Option<Handle<Expression<'a>>>,
310    },
311    Break,
312    Continue,
313    Return {
314        value: Option<Handle<Expression<'a>>>,
315    },
316    Kill,
317    Call {
318        function: Ident<'a>,
319        arguments: Vec<Handle<Expression<'a>>>,
320    },
321    Assign {
322        target: Handle<Expression<'a>>,
323        op: Option<crate::BinaryOperator>,
324        value: Handle<Expression<'a>>,
325    },
326    Increment(Handle<Expression<'a>>),
327    Decrement(Handle<Expression<'a>>),
328    Phony(Handle<Expression<'a>>),
329    ConstAssert(Handle<Expression<'a>>),
330}
331
332#[derive(Debug)]
333pub enum SwitchValue<'a> {
334    Expr(Handle<Expression<'a>>),
335    Default,
336}
337
338#[derive(Debug)]
339pub struct SwitchCase<'a> {
340    pub value: SwitchValue<'a>,
341    pub body: Block<'a>,
342    pub fall_through: bool,
343}
344
345/// A type at the head of a [`Construct`] expression.
346///
347/// WGSL has two types of [`type constructor expressions`]:
348///
349/// - Those that fully specify the type being constructed, like
350///   `vec3<f32>(x,y,z)`, which obviously constructs a `vec3<f32>`.
351///
352/// - Those that leave the component type of the composite being constructed
353///   implicit, to be inferred from the argument types, like `vec3(x,y,z)`,
354///   which constructs a `vec3<T>` where `T` is the type of `x`, `y`, and `z`.
355///
356/// This enum represents the head type of both cases. The `PartialFoo` variants
357/// represent the second case, where the component type is implicit.
358///
359/// This does not cover structs or types referred to by type aliases. See the
360/// documentation for [`Construct`] and [`Call`] expressions for details.
361///
362/// [`Construct`]: Expression::Construct
363/// [`type constructor expressions`]: https://gpuweb.github.io/gpuweb/wgsl/#type-constructor-expr
364/// [`Call`]: Expression::Call
365#[derive(Debug)]
366pub enum ConstructorType<'a> {
367    /// A scalar type or conversion: `f32(1)`.
368    Scalar(Scalar),
369
370    /// A vector construction whose component type is inferred from the
371    /// argument: `vec3(1.0)`.
372    PartialVector { size: crate::VectorSize },
373
374    /// A vector construction whose component type is written out:
375    /// `vec3<f32>(1.0)`.
376    Vector {
377        size: crate::VectorSize,
378        ty: Handle<Type<'a>>,
379        ty_span: Span,
380    },
381
382    /// A matrix construction whose component type is inferred from the
383    /// argument: `mat2x2(1,2,3,4)`.
384    PartialMatrix {
385        columns: crate::VectorSize,
386        rows: crate::VectorSize,
387    },
388
389    /// A matrix construction whose component type is written out:
390    /// `mat2x2<f32>(1,2,3,4)`.
391    Matrix {
392        columns: crate::VectorSize,
393        rows: crate::VectorSize,
394        ty: Handle<Type<'a>>,
395        ty_span: Span,
396    },
397
398    /// A cooperative matrix construction base `coop_mat8x8(...)`.
399    PartialCooperativeMatrix {
400        columns: crate::CooperativeSize,
401        rows: crate::CooperativeSize,
402    },
403
404    /// A full cooperative matrix construction `coop_mat8x8<f32,A>(...)`.
405    CooperativeMatrix {
406        columns: crate::CooperativeSize,
407        rows: crate::CooperativeSize,
408        ty: Handle<Type<'a>>,
409        ty_span: Span,
410        role: crate::CooperativeRole,
411    },
412
413    /// An array whose component type and size are inferred from the arguments:
414    /// `array(3,4,5)`.
415    PartialArray,
416
417    /// An array whose component type and size are written out:
418    /// `array<u32, 4>(3,4,5)`.
419    Array {
420        base: Handle<Type<'a>>,
421        size: ArraySize<'a>,
422    },
423
424    /// Constructing a value of a known Naga IR type.
425    ///
426    /// This variant is produced only during lowering, when we have Naga types
427    /// available, never during parsing.
428    Type(Handle<crate::Type>),
429}
430
431#[derive(Debug, Copy, Clone)]
432pub enum Literal {
433    Bool(bool),
434    Number(Number),
435}
436
437#[cfg(doc)]
438use crate::front::wgsl::lower::Lowerer;
439
440#[derive(Debug)]
441pub enum Expression<'a> {
442    Literal(Literal),
443    Ident(IdentExpr<'a>),
444
445    /// A type constructor expression.
446    ///
447    /// This is only used for expressions like `KEYWORD(EXPR...)` and
448    /// `KEYWORD<PARAM>(EXPR...)`, where `KEYWORD` is a [type-defining keyword] like
449    /// `vec3`. These keywords cannot be shadowed by user definitions, so we can
450    /// tell that such an expression is a construction immediately.
451    ///
452    /// For ordinary identifiers, we can't tell whether an expression like
453    /// `IDENTIFIER(EXPR, ...)` is a construction expression or a function call
454    /// until we know `IDENTIFIER`'s definition, so we represent those as
455    /// [`Call`] expressions.
456    ///
457    /// [type-defining keyword]: https://gpuweb.github.io/gpuweb/wgsl/#type-defining-keywords
458    /// [`Call`]: Expression::Call
459    Construct {
460        ty: ConstructorType<'a>,
461        ty_span: Span,
462        components: Vec<Handle<Expression<'a>>>,
463    },
464    Unary {
465        op: crate::UnaryOperator,
466        expr: Handle<Expression<'a>>,
467    },
468    AddrOf(Handle<Expression<'a>>),
469    Deref(Handle<Expression<'a>>),
470    Binary {
471        op: crate::BinaryOperator,
472        left: Handle<Expression<'a>>,
473        right: Handle<Expression<'a>>,
474    },
475
476    /// A function call or type constructor expression.
477    ///
478    /// We can't tell whether an expression like `IDENTIFIER(EXPR, ...)` is a
479    /// construction expression or a function call until we know `IDENTIFIER`'s
480    /// definition, so we represent everything of that form as one of these
481    /// expressions until lowering. At that point, [`Lowerer::call`] has
482    /// everything's definition in hand, and can decide whether to emit a Naga
483    /// [`Constant`], [`As`], [`Splat`], or [`Compose`] expression.
484    ///
485    /// [`Lowerer::call`]: Lowerer::call
486    /// [`Constant`]: crate::Expression::Constant
487    /// [`As`]: crate::Expression::As
488    /// [`Splat`]: crate::Expression::Splat
489    /// [`Compose`]: crate::Expression::Compose
490    Call {
491        function: Ident<'a>,
492        arguments: Vec<Handle<Expression<'a>>>,
493        result_ty: Option<(Handle<Type<'a>>, Span)>,
494    },
495    Index {
496        base: Handle<Expression<'a>>,
497        index: Handle<Expression<'a>>,
498    },
499    Member {
500        base: Handle<Expression<'a>>,
501        field: Ident<'a>,
502    },
503    Bitcast {
504        expr: Handle<Expression<'a>>,
505        to: Handle<Type<'a>>,
506        ty_span: Span,
507    },
508}
509
510#[derive(Debug)]
511pub struct LocalVariable<'a> {
512    pub name: Ident<'a>,
513    pub ty: Option<Handle<Type<'a>>>,
514    pub init: Option<Handle<Expression<'a>>>,
515    pub handle: Handle<Local>,
516}
517
518#[derive(Debug)]
519pub struct Let<'a> {
520    pub name: Ident<'a>,
521    pub ty: Option<Handle<Type<'a>>>,
522    pub init: Handle<Expression<'a>>,
523    pub handle: Handle<Local>,
524}
525
526#[derive(Debug)]
527pub struct LocalConst<'a> {
528    pub name: Ident<'a>,
529    pub ty: Option<Handle<Type<'a>>>,
530    pub init: Handle<Expression<'a>>,
531    pub handle: Handle<Local>,
532}
533
534#[derive(Debug)]
535pub enum LocalDecl<'a> {
536    Var(LocalVariable<'a>),
537    Let(Let<'a>),
538    Const(LocalConst<'a>),
539}
540
541#[derive(Debug)]
542/// A placeholder for a local variable declaration.
543///
544/// See [`super::ExpressionContext::locals`] for more information.
545pub struct Local;