wesl/pass/
condcomp.rs

1use crate::error::{CondCompError, Diagnostic, Error};
2use std::collections::HashMap;
3use wgsl_parse::{SyntaxNode, span::Spanned, syntax::*};
4
5type E = Error;
6
7/// Set the behavior for a feature flag during conditional translation.
8///
9/// * `Keep` means that the feature flag will be left as-is. This is useful for
10///   incremental compilation, e.g. for generating shader variants
11/// * `Error` means that unspecified feature flags will trigger a
12///   [`CondCompError::UnexpectedFeatureFlag`].
13///
14/// Default is `Disable`.
15#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
16pub enum Feature {
17    Enable,
18    #[default]
19    Disable,
20    Keep,
21    Error,
22}
23
24/// Toggle conditional compilation feature flags.
25///
26/// Conditional translation is incremental if [`field@Self::default`] is set to [`Feature::Keep`].
27/// If not all feature flags are handled, the output will contain unevaluated `@if`
28/// attributes and will therefore *not* be valid WGSL.
29///
30/// Feature flags set to `true` are enabled, and `false` are disabled. Feature flags not
31/// present in `flags` are treated according to `default`.
32#[derive(Clone, Debug, Default, PartialEq, Eq)]
33pub struct Features {
34    pub default: Feature,
35    pub flags: HashMap<String, Feature>,
36}
37
38impl From<bool> for Feature {
39    fn from(value: bool) -> Self {
40        if value {
41            Feature::Enable
42        } else {
43            Feature::Disable
44        }
45    }
46}
47
48impl Features {
49    pub fn new() -> Self {
50        Default::default()
51    }
52
53    pub fn set(&mut self, name: impl ToString, value: impl Into<Feature>) {
54        self.flags.insert(name.to_string(), value.into());
55    }
56}
57
58const EXPR_TRUE: Expression = Expression::Literal(LiteralExpression::Bool(true));
59const EXPR_FALSE: Expression = Expression::Literal(LiteralExpression::Bool(false));
60
61fn eval_attr(expr: &ExpressionNode, features: &Features) -> Result<Expression, E> {
62    eval_attr_impl(expr, features).map_err(|e| Diagnostic::from(e).with_span(expr.span()).into())
63}
64
65fn eval_attr_impl(expr: &Expression, features: &Features) -> Result<Expression, E> {
66    fn eval_rec(expr: &ExpressionNode, features: &Features) -> Result<Expression, E> {
67        eval_attr(expr, features).map_err(|e| Diagnostic::from(e).with_span(expr.span()).into())
68    }
69
70    match expr {
71        Expression::Literal(LiteralExpression::Bool(_)) => Ok(expr.clone()),
72        Expression::Parenthesized(paren) => {
73            let expr = eval_rec(&paren.expression, features)?;
74            Ok(match expr {
75                Expression::Binary(_) => ParenthesizedExpression {
76                    expression: Spanned::new(expr, paren.expression.span()),
77                }
78                .into(),
79                _ => expr,
80            })
81        }
82        Expression::Unary(unary) => {
83            let operand = eval_rec(&unary.operand, features)?;
84            match &unary.operator {
85                UnaryOperator::LogicalNegation => {
86                    let expr = if operand == EXPR_TRUE {
87                        EXPR_FALSE.clone()
88                    } else if operand == EXPR_FALSE {
89                        EXPR_TRUE.clone()
90                    } else {
91                        expr.clone()
92                    };
93                    Ok(expr)
94                }
95                _ => Err(CondCompError::InvalidExpression(expr.clone()).into()),
96            }
97        }
98        Expression::Binary(binary) => {
99            let left = eval_rec(&binary.left, features)?;
100            let right = eval_rec(&binary.right, features)?;
101            match &binary.operator {
102                BinaryOperator::ShortCircuitOr => {
103                    let expr = if left == EXPR_TRUE || right == EXPR_TRUE {
104                        EXPR_TRUE.clone()
105                    } else if left == EXPR_FALSE && right == EXPR_FALSE {
106                        left // false
107                    } else if left == EXPR_FALSE {
108                        right
109                    } else if right == EXPR_FALSE {
110                        left
111                    } else {
112                        BinaryExpression {
113                            operator: binary.operator,
114                            left: Spanned::new(left, binary.left.span()),
115                            right: Spanned::new(right, binary.right.span()),
116                        }
117                        .into()
118                    };
119                    Ok(expr)
120                }
121                BinaryOperator::ShortCircuitAnd => {
122                    let expr = if left == EXPR_TRUE && right == EXPR_TRUE {
123                        left // true
124                    } else if left == EXPR_FALSE || right == EXPR_FALSE {
125                        EXPR_FALSE.clone()
126                    } else if left == EXPR_TRUE {
127                        right
128                    } else if right == EXPR_TRUE {
129                        left
130                    } else {
131                        BinaryExpression {
132                            operator: binary.operator,
133                            left: Spanned::new(left, binary.left.span()),
134                            right: Spanned::new(right, binary.right.span()),
135                        }
136                        .into()
137                    };
138                    Ok(expr)
139                }
140                _ => Err(CondCompError::InvalidExpression(expr.clone()).into()),
141            }
142        }
143        Expression::TypeOrIdentifier(ty) => {
144            if ty.template_args.is_some() {
145                return Err(CondCompError::InvalidFeatureFlag(ty.to_string()).into());
146            }
147            let feat = features
148                .flags
149                .get(&*ty.ident.name())
150                .unwrap_or(&features.default);
151            let expr = match feat {
152                Feature::Enable => EXPR_TRUE.clone(),
153                Feature::Disable => EXPR_FALSE.clone(),
154                Feature::Keep => expr.clone(),
155                Feature::Error => {
156                    return Err(
157                        CondCompError::UnexpectedFeatureFlag(ty.ident.name().to_string()).into(),
158                    );
159                }
160            };
161            Ok(expr)
162        }
163        _ => Err(CondCompError::InvalidExpression(expr.clone()).into()),
164    }
165}
166
167fn get_single_attr(attrs: &mut [AttributeNode]) -> Result<Option<&mut AttributeNode>, E> {
168    let mut it = attrs.iter_mut().filter(|attr| {
169        matches!(
170            attr.node(),
171            Attribute::If(_) | Attribute::Elif(_) | Attribute::Else
172        )
173    });
174    let attr = it.next();
175
176    if it.next().is_some() {
177        Err(CondCompError::DuplicateIf.into())
178    } else {
179        Ok(attr)
180    }
181}
182
183/// Conditional state of a syntax node.
184#[derive(Debug, Clone, Copy, PartialEq, Eq)]
185struct NodeEval {
186    /// Whether the syntax node has an if, elif or else attribute.
187    has_condcomp: bool,
188    /// Whether the node evaluates to false and is removed.
189    /// e.g. `@if(false)` or `@if(true) ... @else`
190    is_false: bool,
191    /// Whether this node, or a previous node in the if/elif chain, evaluates to true.
192    /// All following elif/else nodes in the chain therefore evaluate to false.
193    chain_has_true: bool,
194}
195
196/// * ensure there is at most one if/elif/else node.
197/// * ensure elif/else nodes are preceded by if/elif.
198/// * remove the attributes which evaluate to true.
199/// * turn elifs into ifs when previous node was deleted.
200/// * turn elifs into elses when it evaluates to true.
201fn eval_if_attr(
202    node: &mut impl SyntaxNode,
203    prev: &mut NodeEval,
204    features: &Features,
205) -> Result<(), E> {
206    let span = node.span();
207    eval_if_attr_impl(node, prev, features).map_err(|e| {
208        if let Some(span) = span {
209            Diagnostic::from(e).with_span(span).into()
210        } else {
211            e
212        }
213    })
214}
215
216fn eval_if_attr_impl(
217    node: &mut impl SyntaxNode,
218    prev: &mut NodeEval,
219    features: &Features,
220) -> Result<(), E> {
221    let attr = get_single_attr(node.attributes_mut())?;
222    if let Some(attr) = attr {
223        prev.has_condcomp = attr.is_condcomp();
224        if let Attribute::If(expr) = attr.node_mut() {
225            **expr = eval_attr(expr, features)?;
226            // a new `if` starts a new if/elif/else chain
227            prev.chain_has_true = false;
228        } else if let Attribute::Elif(expr) = attr.node_mut() {
229            if !prev.has_condcomp {
230                return Err(CondCompError::NoPrecedingIf.into());
231            } else {
232                **expr = eval_attr(expr, features)?;
233            }
234        } else if let Attribute::Else = attr.node()
235            && !prev.has_condcomp
236        {
237            return Err(CondCompError::NoPrecedingIf.into());
238        }
239    } else {
240        prev.has_condcomp = false;
241    }
242
243    let mut is_false = false;
244
245    node.retain_attributes_mut(|attr| {
246        let mut remove_attr = false;
247        if let Attribute::If(expr) = attr {
248            if **expr == EXPR_TRUE {
249                remove_attr = true;
250                prev.chain_has_true = true;
251            } else if **expr == EXPR_FALSE {
252                is_false = true;
253            }
254        } else if let Attribute::Elif(expr) = attr {
255            if prev.chain_has_true || **expr == EXPR_FALSE {
256                is_false = true; // a previous node was chosen, delete the whole node
257            } else if **expr == EXPR_TRUE {
258                if prev.is_false {
259                    remove_attr = true; // the previous node is false and is deleted, so elif(true) can be removed
260                } else {
261                    *attr = Attribute::Else; // the previous node is undecided, but elif(true) can become else.
262                }
263                prev.chain_has_true = true;
264            } else if prev.is_false {
265                *attr = Attribute::If(expr.clone()); // the previous node is false and is deleted, elif becomes if
266            }
267        } else if let Attribute::Else = attr {
268            if prev.chain_has_true {
269                is_false = true; // a previous node was chosen, delete the whole node
270            } else if prev.is_false {
271                remove_attr = true; // the previous node was deleted, delete this attribute
272                prev.chain_has_true = true;
273            }
274        } else {
275            // we keep non-condcomp attributes
276            return true;
277        }
278
279        !remove_attr
280    });
281
282    prev.is_false = is_false;
283    Ok(())
284}
285
286fn eval_opt_attr(
287    opt_node: &mut Option<impl SyntaxNode>,
288    prev: &mut NodeEval,
289    features: &Features,
290) -> Result<(), E> {
291    if let Some(node) = opt_node {
292        eval_if_attr(node, prev, features)?;
293        if prev.is_false {
294            *opt_node = None;
295        }
296    }
297    Ok(())
298}
299
300fn eval_if_attrs(nodes: &mut Vec<impl SyntaxNode>, features: &Features) -> Result<NodeEval, E> {
301    let mut prev = NodeEval {
302        has_condcomp: false,
303        chain_has_true: false,
304        is_false: false,
305    };
306    let mut err = None;
307
308    // remove the nodes for which the attr evaluate to false.
309    nodes.retain_mut(|node| {
310        let res = eval_if_attr(node, &mut prev, features);
311        if let (Err(e), None) = (res, &err) {
312            err = Some(e);
313        }
314        !prev.is_false
315    });
316
317    if let Some(e) = err {
318        Err(e)
319    } else {
320        Ok(prev)
321    }
322}
323
324fn stmt_eval_if_attrs(statements: &mut Vec<StatementNode>, features: &Features) -> Result<(), E> {
325    fn rec_eval_inside_stmt(stmt: &mut StatementNode, feats: &Features) -> Result<(), E> {
326        match stmt.node_mut() {
327            Statement::Compound(stmt) => {
328                rec(&mut stmt.statements, feats)?;
329            }
330            Statement::If(stmt) => {
331                rec(&mut stmt.if_clause.body.statements, feats)?;
332                for elif in &mut stmt.else_if_clauses {
333                    rec(&mut elif.body.statements, feats)?;
334                }
335                if let Some(el) = &mut stmt.else_clause {
336                    rec(&mut el.body.statements, feats)?;
337                }
338            }
339            Statement::Switch(stmt) => {
340                eval_if_attrs(&mut stmt.clauses, feats)?;
341                for clause in &mut stmt.clauses {
342                    rec(&mut clause.body.statements, feats)?;
343                }
344            }
345            Statement::Loop(stmt) => {
346                let mut prev = rec(&mut stmt.body.statements, feats)?;
347                eval_opt_attr(&mut stmt.continuing, &mut prev, feats)?;
348                if let Some(cont) = &mut stmt.continuing {
349                    rec(&mut cont.body.statements, feats)?;
350                    eval_opt_attr(&mut cont.break_if, &mut prev, feats)?;
351                }
352                rec(&mut stmt.body.statements, feats)?;
353            }
354            Statement::For(stmt) => {
355                if let Some(init) = &mut stmt.initializer {
356                    rec_eval_inside_stmt(&mut *init, feats)?
357                }
358                if let Some(updt) = &mut stmt.update {
359                    rec_eval_inside_stmt(&mut *updt, feats)?
360                }
361                rec(&mut stmt.body.statements, feats)?;
362            }
363            Statement::While(stmt) => {
364                rec(&mut stmt.body.statements, feats)?;
365            }
366            _ => (),
367        };
368        Ok(())
369    }
370
371    fn rec(stmts: &mut Vec<StatementNode>, feats: &Features) -> Result<NodeEval, E> {
372        let mut prev = NodeEval {
373            has_condcomp: false,
374            chain_has_true: false,
375            is_false: false,
376        };
377
378        // If an `@if` decorates a compound statement, the statement gets flattened.
379        // This is the same code as in eval_if_attrs, except it flattens the compound when it evaluates to true.
380        {
381            let mut i = 0;
382
383            // remove the nodes for which the attr evaluate to false.
384            while let Some(node) = stmts.get_mut(i) {
385                eval_if_attr(node, &mut prev, feats)?;
386                if let Statement::Compound(stmt) = &**node
387                    && prev.has_condcomp
388                    && !prev.is_false
389                {
390                    // replace the compound statements with its contents
391                    // TODO: other compound statement attributes are lost. validation has no opportunity to check them.
392                    // COMBAK: this clone is unnecessary and probably inefficient.
393                    let mut body = stmt.statements.clone();
394                    rec(&mut body, feats)?;
395                    let n = body.len();
396                    stmts.splice(i..i + 1, body);
397                    i += n;
398                } else if prev.is_false {
399                    stmts.remove(i);
400                } else {
401                    rec_eval_inside_stmt(node, feats)?;
402                    i += 1;
403                }
404            }
405        }
406
407        Ok(prev)
408    }
409
410    rec(statements, features).map(|_| ())
411}
412
413/// Run the conditional translation phase, eliminating unused code branches.
414pub fn condcomp(module: &mut TranslationUnit, features: &Features) -> Result<(), E> {
415    module.remove_voids();
416    eval_if_attrs(&mut module.imports, features)?;
417    eval_if_attrs(&mut module.global_directives, features)?;
418
419    // If an `@if` decorates a compound global declaration, it gets flattened.
420    // This is the same code as in eval_if_attrs, except it flattens the compound when it evaluates to true.
421    fn eval_flatten_compound(
422        decls: &mut Vec<GlobalDeclarationNode>,
423        features: &Features,
424    ) -> Result<(), E> {
425        let mut prev = NodeEval {
426            has_condcomp: false,
427            chain_has_true: false,
428            is_false: false,
429        };
430
431        let mut i = 0;
432
433        // remove the nodes for which the attr evaluate to false.
434        while let Some(node) = decls.get_mut(i) {
435            eval_if_attr(node, &mut prev, features)?;
436            if let GlobalDeclaration::Compound(stmt) = &**node
437                && prev.has_condcomp
438                && !prev.is_false
439            {
440                // replace the compound statements with its contents
441                // TODO: other compound statement attributes are lost. validation has no opportunity to check them.
442                // COMBAK: this clone is unnecessary and probably inefficient.
443                let mut body = stmt.body.clone();
444                eval_flatten_compound(&mut body, features)?;
445                let n = body.len();
446                decls.splice(i..i + 1, body);
447                i += n;
448            } else if prev.is_false {
449                decls.remove(i);
450            } else {
451                i += 1;
452            }
453        }
454
455        Ok(())
456    }
457
458    eval_flatten_compound(&mut module.global_declarations, features)?;
459
460    for decl in &mut module.global_declarations {
461        if let GlobalDeclaration::Struct(decl) = decl.node_mut() {
462            eval_if_attrs(&mut decl.members, features)
463                .map_err(|e| Diagnostic::from(e).with_declaration(decl.ident.to_string()))?;
464        } else if let GlobalDeclaration::Function(decl) = decl.node_mut() {
465            eval_if_attrs(&mut decl.parameters, features)
466                .map_err(|e| Diagnostic::from(e).with_declaration(decl.ident.to_string()))?;
467            stmt_eval_if_attrs(&mut decl.body.statements, features)
468                .map_err(|e| Diagnostic::from(e).with_declaration(decl.ident.to_string()))?;
469        }
470    }
471
472    Ok(())
473}
wesl/pass/condcomp.rs

wesl/pass/
condcomp.rs
