Skip to main content

wgsl_parse/
parser_support.rs

1//! support functions to be injected in the lalrpop parser.
2
3use std::str::FromStr;
4
5use itertools::Itertools;
6
7use crate::{
8    error::ParseError,
9    span::{Span, Spanned},
10    syntax::*,
11};
12
13type E = ParseError;
14
15// HACK: parsing entrypoints.
16// see: https://github.com/lalrpop/lalrpop/issues/65#issuecomment-516769995
17pub enum ParseEntryPoint {
18    TryTemplateList(Span),
19    TranslationUnit(TranslationUnit),
20    GlobalDecl(GlobalDeclaration),
21    Literal(LiteralExpression),
22    GlobalDirective(GlobalDirective),
23    Expression(Expression),
24    Statement(Statement),
25    #[cfg(feature = "imports")]
26    ImportStatement(ImportStatement),
27}
28
29pub(crate) enum Component {
30    Named(Ident),
31    Index(ExpressionNode),
32}
33
34pub(crate) fn apply_components(
35    expr: Expression,
36    span: Span,
37    components: Vec<Spanned<Component>>,
38) -> Expression {
39    components.into_iter().fold(expr, |base, comp| {
40        let span = span.extend(comp.span());
41        let base = Spanned::new(base, span);
42        match comp.into_inner() {
43            Component::Named(component) => {
44                Expression::NamedComponent(NamedComponentExpression { base, component })
45            }
46            Component::Index(index) => Expression::Indexing(IndexingExpression { base, index }),
47        }
48    })
49}
50
51impl FromStr for DeclarationKind {
52    type Err = ();
53
54    fn from_str(s: &str) -> Result<Self, Self::Err> {
55        match s {
56            "const" => Ok(Self::Const),
57            "override" => Ok(Self::Override),
58            "let" => Ok(Self::Let),
59            "var" => Ok(Self::Var(None)),
60            _ => Err(()),
61        }
62    }
63}
64
65fn one_arg(arguments: Option<Vec<ExpressionNode>>) -> Option<ExpressionNode> {
66    match arguments {
67        Some(mut args) => (args.len() == 1).then(|| args.pop().unwrap()),
68        None => None,
69    }
70}
71fn two_args(arguments: Option<Vec<ExpressionNode>>) -> Option<(ExpressionNode, ExpressionNode)> {
72    match arguments {
73        Some(args) => (args.len() == 2).then(|| args.into_iter().collect_tuple().unwrap()),
74        None => None,
75    }
76}
77fn zero_args(arguments: Option<Vec<ExpressionNode>>) -> bool {
78    arguments.is_none()
79}
80fn ident(expr: ExpressionNode) -> Option<Ident> {
81    match expr.into_inner() {
82        Expression::TypeOrIdentifier(TypeExpression {
83            #[cfg(feature = "imports")]
84                path: _,
85            ident,
86            template_args: None,
87        }) => Some(ident),
88        _ => None,
89    }
90}
91
92pub(crate) fn parse_attribute(
93    name: String,
94    args: Option<Vec<ExpressionNode>>,
95) -> Result<Attribute, E> {
96    match name.as_str() {
97        "align" => match one_arg(args) {
98            Some(expr) => Ok(Attribute::Align(expr)),
99            _ => Err(E::Attribute("align", "expected 1 argument")),
100        },
101        "binding" => match one_arg(args) {
102            Some(expr) => Ok(Attribute::Binding(expr)),
103            _ => Err(E::Attribute("binding", "expected 1 argument")),
104        },
105        "blend_src" => match one_arg(args) {
106            Some(expr) => Ok(Attribute::BlendSrc(expr)),
107            _ => Err(E::Attribute("blend_src", "expected 1 argument")),
108        },
109        "builtin" => match one_arg(args) {
110            Some(expr) => match ident(expr).and_then(|id| id.name().parse().ok()) {
111                Some(b) => Ok(Attribute::Builtin(b)),
112                _ => Err(E::Attribute(
113                    "builtin",
114                    "the argument is not a valid built-in value name",
115                )),
116            },
117            _ => Err(E::Attribute("builtin", "expected 1 argument")),
118        },
119        "const" => match zero_args(args) {
120            true => Ok(Attribute::Const),
121            false => Err(E::Attribute("const", "expected 0 arguments")),
122        },
123        "diagnostic" => match two_args(args) {
124            Some((e1, e2)) => {
125                let severity = ident(e1).and_then(|id| id.name().parse().ok());
126                let rule = match e2.into_inner() {
127                    Expression::TypeOrIdentifier(TypeExpression {
128                        #[cfg(feature = "imports")]
129                            path: _,
130                        ident,
131                        template_args: None,
132                    }) => Some(ident.name().to_string()),
133                    Expression::NamedComponent(e) => {
134                        ident(e.base).map(|id| format!("{}.{}", id.name(), e.component))
135                    }
136                    _ => None,
137                };
138                match (severity, rule) {
139                    (Some(severity), Some(rule)) => {
140                        Ok(Attribute::Diagnostic(DiagnosticAttribute {
141                            severity,
142                            rule,
143                        }))
144                    }
145                    _ => Err(E::Attribute("diagnostic", "invalid arguments")),
146                }
147            }
148            _ => Err(E::Attribute("diagnostic", "expected 1 argument")),
149        },
150        "group" => match one_arg(args) {
151            Some(expr) => Ok(Attribute::Group(expr)),
152            _ => Err(E::Attribute("group", "expected 1 argument")),
153        },
154        "id" => match one_arg(args) {
155            Some(expr) => Ok(Attribute::Id(expr)),
156            _ => Err(E::Attribute("id", "expected 1 argument")),
157        },
158        "interpolate" => match args {
159            Some(v) if v.len() == 2 => {
160                let (e1, e2) = v.into_iter().collect_tuple().unwrap();
161                let ty = ident(e1).and_then(|id| id.name().parse().ok());
162                let sampling = ident(e2).and_then(|id| id.name().parse().ok());
163                match (ty, sampling) {
164                    (Some(ty), Some(sampling)) => {
165                        Ok(Attribute::Interpolate(InterpolateAttribute {
166                            ty,
167                            sampling: Some(sampling),
168                        }))
169                    }
170                    _ => Err(E::Attribute("interpolate", "invalid arguments")),
171                }
172            }
173            Some(v) if v.len() == 1 => {
174                let e1 = v.into_iter().next().unwrap();
175                let ty = ident(e1).and_then(|id| id.name().parse().ok());
176                match ty {
177                    Some(ty) => Ok(Attribute::Interpolate(InterpolateAttribute {
178                        ty,
179                        sampling: None,
180                    })),
181                    _ => Err(E::Attribute("interpolate", "invalid arguments")),
182                }
183            }
184            _ => Err(E::Attribute("interpolate", "invalid arguments")),
185        },
186
187        "invariant" => match zero_args(args) {
188            true => Ok(Attribute::Invariant),
189            false => Err(E::Attribute("invariant", "expected 0 arguments")),
190        },
191        "location" => match one_arg(args) {
192            Some(expr) => Ok(Attribute::Location(expr)),
193            _ => Err(E::Attribute("location", "expected 1 argument")),
194        },
195        "must_use" => match zero_args(args) {
196            true => Ok(Attribute::MustUse),
197            false => Err(E::Attribute("must_use", "expected 0 arguments")),
198        },
199        "size" => match one_arg(args) {
200            Some(expr) => Ok(Attribute::Size(expr)),
201            _ => Err(E::Attribute("size", "expected 1 argument")),
202        },
203        "workgroup_size" => match args {
204            Some(args) => {
205                let mut it = args.into_iter();
206                match (it.next(), it.next(), it.next(), it.next()) {
207                    (Some(x), y, z, None) => {
208                        Ok(Attribute::WorkgroupSize(WorkgroupSizeAttribute { x, y, z }))
209                    }
210                    _ => Err(E::Attribute("workgroup_size", "expected 1-3 arguments")),
211                }
212            }
213            _ => Err(E::Attribute("workgroup_size", "expected 1-3 arguments")),
214        },
215        "vertex" => match zero_args(args) {
216            true => Ok(Attribute::Vertex),
217            false => Err(E::Attribute("vertex", "expected 0 arguments")),
218        },
219        "fragment" => match zero_args(args) {
220            true => Ok(Attribute::Fragment),
221            false => Err(E::Attribute("fragment", "expected 0 arguments")),
222        },
223        "compute" => match zero_args(args) {
224            true => Ok(Attribute::Compute),
225            false => Err(E::Attribute("compute", "expected 0 arguments")),
226        },
227        #[cfg(feature = "naga-ext")]
228        "task" => match zero_args(args) {
229            true => Ok(Attribute::Task),
230            false => Err(E::Attribute("task", "expected 0 arguments")),
231        },
232        #[cfg(feature = "naga-ext")]
233        "payload" => match one_arg(args) {
234            Some(expr) => Ok(Attribute::Payload(expr)),
235            None => Err(E::Attribute("payload", "expected 1 arguments")),
236        },
237        #[cfg(feature = "naga-ext")]
238        "mesh" => match one_arg(args) {
239            Some(expr) => Ok(Attribute::Mesh(expr)),
240            None => Err(E::Attribute("mesh", "expected 1 arguments")),
241        },
242        #[cfg(feature = "imports")]
243        "publish" => Ok(Attribute::Publish),
244        #[cfg(feature = "condcomp")]
245        "if" => match one_arg(args) {
246            Some(expr) => Ok(Attribute::If(expr)),
247            None => Err(E::Attribute("if", "expected 1 argument")),
248        },
249        #[cfg(feature = "condcomp")]
250        "elif" => match one_arg(args) {
251            Some(expr) => Ok(Attribute::Elif(expr)),
252            None => Err(E::Attribute("elif", "expected 1 argument")),
253        },
254        #[cfg(feature = "condcomp")]
255        "else" => match zero_args(args) {
256            true => Ok(Attribute::Else),
257            false => Err(E::Attribute("else", "expected 0 arguments")),
258        },
259        #[cfg(feature = "generics")]
260        "type" => parse_attr_type(args).map(Attribute::Type),
261        #[cfg(feature = "naga-ext")]
262        "early_depth_test" => match args {
263            Some(args) => {
264                let mut it = args.into_iter();
265                match (it.next(), it.next()) {
266                    (Some(expr), None) => match ident(expr).and_then(|id| id.name().parse().ok()) {
267                        Some(c) => Ok(Attribute::EarlyDepthTest(Some(c))),
268                        _ => Err(E::Attribute(
269                            "early_depth_test",
270                            "the argument must be one of `greater_equal`, `less_equal`, `unchanged`",
271                        )),
272                    },
273                    (None, None) => Ok(Attribute::EarlyDepthTest(None)),
274                    _ => Err(E::Attribute(
275                        "early_depth_test",
276                        "expected 0 or 1 arguments",
277                    )),
278                }
279            }
280            _ => Err(E::Attribute(
281                "early_depth_test",
282                "expected 0 or 1 arguments",
283            )),
284        },
285        _ => Ok(Attribute::Custom(CustomAttribute {
286            name,
287            arguments: args,
288        })),
289    }
290}
291
292// format: @type(T, foo | bar | baz)
293#[cfg(feature = "generics")]
294fn parse_attr_type(arguments: Option<Vec<ExpressionNode>>) -> Result<TypeConstraint, E> {
295    fn parse_rec(expr: Expression) -> Result<Vec<TypeExpression>, E> {
296        match expr {
297            Expression::TypeOrIdentifier(ty) => Ok(vec![ty]),
298            Expression::Binary(BinaryExpression {
299                operator: BinaryOperator::BitwiseOr,
300                left,
301                right,
302            }) => {
303                let ty = match right.into_inner() {
304                    Expression::TypeOrIdentifier(ty) => Ok(ty),
305                    _ => Err(E::Attribute(
306                        "type",
307                        "invalid second argument (type constraint)",
308                    )),
309                }?;
310                let mut v = parse_rec(left.into_inner())?;
311                v.push(ty);
312                Ok(v)
313            }
314            _ => Err(E::Attribute(
315                "type",
316                "invalid second argument (type constraint)",
317            )),
318        }
319    }
320    match two_args(arguments) {
321        Some((e1, e2)) => ident(e1)
322            .map(|ident| {
323                parse_rec(e2.into_inner()).map(|variants| TypeConstraint { ident, variants })
324            })
325            .unwrap_or_else(|| Err(E::Attribute("type", "invalid first argument (type name)"))),
326
327        None => Err(E::Attribute("type", "expected 2 arguments")),
328    }
329}
330
331pub(crate) fn parse_var_template(
332    template_args: TemplateArgs,
333) -> Result<Option<(AddressSpace, Option<AccessMode>)>, E> {
334    match template_args {
335        Some(tplt) => {
336            let mut it = tplt.into_iter();
337            match (it.next(), it.next(), it.next()) {
338                (Some(e1), e2, None) => {
339                    let addr_space = ident(e1.expression)
340                        .and_then(|id| id.name().parse().ok())
341                        .ok_or(E::VarTemplate("invalid address space"))?;
342                    let mut access_mode = None;
343                    if let Some(e2) = e2 {
344                        if addr_space == AddressSpace::Storage {
345                            access_mode = Some(
346                                ident(e2.expression)
347                                    .and_then(|id| id.name().parse().ok())
348                                    .ok_or(E::VarTemplate("invalid access mode"))?,
349                            );
350                        } else {
351                            return Err(E::VarTemplate(
352                                "only variables with `storage` address space can have an access mode",
353                            ));
354                        }
355                    }
356                    Ok(Some((addr_space, access_mode)))
357                }
358                _ => Err(E::VarTemplate("template is empty")),
359            }
360        }
361        None => Ok(None),
362    }
363}