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