feat(naga): constant evaluation for select

Author: ErichDonGubler

naga/src/proc/constant_evaluator.rs

@@ -563,6 +563,27 @@ pub enum ConstantEvaluatorError {
     RuntimeExpr,
     #[error("Unexpected override-expression")]
     OverrideExpr,
+    #[error("Expected boolean expression for condition argument of `select`, got something else")]
+    SelectScalarConditionNotABool,
+    #[error(
+        "Expected vectors of the same size for reject and accept args., got {:?} and {:?}",
+        reject,
+        accept
+    )]
+    SelectVecRejectAcceptSizeMismatch {
+        reject: crate::VectorSize,
+        accept: crate::VectorSize,
+    },
+    #[error("Expected boolean vector for condition arg., got something else")]
+    SelectConditionNotAVecBool,
+    #[error(
+        "Expected same number of vector components between condition, accept, and reject args., got something else",
+    )]
+    SelectConditionVecSizeMismatch,
+    #[error(
+        "Expected reject and accept args. to be scalars of vectors of the same type, got something else",
+    )]
+    SelectAcceptRejectTypeMismatch,
 }
 
 impl<'a> ConstantEvaluator<'a> {
@@ -904,9 +925,11 @@ impl<'a> ConstantEvaluator<'a> {
                     )),
                 }
             }
-            Expression::Select { .. } => Err(ConstantEvaluatorError::NotImplemented(
-                "select built-in function".into(),
-            )),
+            Expression::Select {
+                reject,
+                accept,
+                condition,
+            } => self.select(reject, accept, condition, span),
             Expression::Relational { fun, argument } => {
                 let argument = self.check_and_get(argument)?;
                 self.relational(fun, argument, span)
@@ -2497,6 +2520,114 @@ impl<'a> ConstantEvaluator<'a> {
 
         Ok(resolution)
     }
+
+    fn select(
+        &mut self,
+        reject: Handle<Expression>,
+        accept: Handle<Expression>,
+        condition: Handle<Expression>,
+        span: Span,
+    ) -> Result<Handle<Expression>, ConstantEvaluatorError> {
+        let arg = |this: &mut Self, expr| {
+            let expr = this.eval_zero_value_and_splat(expr, span)?;
+            let expr: &Expression = &this.expressions[expr];
+            this.try_eval_and_append(expr.clone(), span)
+        };
+
+        let reject = arg(self, reject)?;
+        let accept = arg(self, accept)?;
+        let condition = arg(self, condition)?;
+
+        let select_single_component =
+            |this: &mut Self, reject_scalar, reject, accept, condition| {
+                let accept = this.cast(accept, reject_scalar, span)?;
+                if condition {
+                    Ok(accept)
+                } else {
+                    Ok(reject)
+                }
+            };
+
+        match (&self.expressions[reject], &self.expressions[accept]) {
+            (&Expression::Literal(reject_lit), &Expression::Literal(_accept_lit)) => {
+                let reject_scalar = reject_lit.scalar();
+                let &Expression::Literal(Literal::Bool(condition)) = &self.expressions[condition]
+                else {
+                    return Err(ConstantEvaluatorError::SelectScalarConditionNotABool);
+                };
+                select_single_component(self, reject_scalar, reject, accept, condition)
+            }
+            (
+                &Expression::Compose {
+                    ty: reject_ty,
+                    components: ref reject_components,
+                },
+                &Expression::Compose {
+                    ty: accept_ty,
+                    components: ref accept_components,
+                },
+            ) => {
+                let ty_deets = |ty| {
+                    let (size, scalar) = self.types[ty].inner.vector_size_and_scalar().unwrap();
+                    (size.unwrap(), scalar)
+                };
+
+                let expected_vec_size = {
+                    let [(reject_vec_size, _), (accept_vec_size, _)] =
+                        [reject_ty, accept_ty].map(ty_deets);
+
+                    if reject_vec_size != accept_vec_size {
+                        return Err(ConstantEvaluatorError::SelectVecRejectAcceptSizeMismatch {
+                            reject: reject_vec_size,
+                            accept: accept_vec_size,
+                        });
+                    }
+                    reject_vec_size
+                };
+
+                let &Expression::Compose {
+                    ty: condition_ty,
+                    components: ref condition_components,
+                } = &self.expressions[condition]
+                else {
+                    return Err(ConstantEvaluatorError::SelectConditionNotAVecBool);
+                };
+
+                {
+                    let (condition_vec_size, condition_scalar) = ty_deets(condition_ty);
+                    if condition_scalar.kind != ScalarKind::Bool {
+                        return Err(ConstantEvaluatorError::SelectConditionNotAVecBool);
+                    }
+                    if condition_vec_size != expected_vec_size {
+                        return Err(ConstantEvaluatorError::SelectConditionVecSizeMismatch);
+                    }
+                }
+
+                let evaluated = Expression::Compose {
+                    ty: reject_ty,
+                    components: reject_components
+                        .clone()
+                        .into_iter()
+                        .zip(accept_components.clone().into_iter())
+                        .zip(condition_components.clone().into_iter())
+                        .map(|((reject, accept), condition)| {
+                            let reject_scalar = match &self.expressions[reject] {
+                                &Expression::Literal(lit) => lit.scalar(),
+                                _ => unreachable!(),
+                            };
+                            let condition = match &self.expressions[condition] {
+                                &Expression::Literal(Literal::Bool(condition)) => condition,
+                                _ => unreachable!(),
+                            };
+                            select_single_component(self, reject_scalar, reject, accept, condition)
+                        })
+                        .collect::<Result<_, _>>()?,
+                };
+                self.register_evaluated_expr(evaluated, span)
+            }
+            _ => Err(ConstantEvaluatorError::SelectAcceptRejectTypeMismatch),
+        }
+    }
 }
 
 fn first_trailing_bit(concrete_int: ConcreteInt<1>) -> ConcreteInt<1> {