WIP: constant evaluation for select

Author: ErichDonGubler

naga/src/proc/constant_evaluator.rs

@@ -563,6 +563,24 @@ pub enum ConstantEvaluatorError {
     RuntimeExpr,
     #[error("Unexpected override-expression")]
     OverrideExpr,
+    #[error("Expected scalars of the same type for reject and accept args., got something else")]
+    SelectScalarRejectAcceptTypeMismatch,
+    #[error("Expected boolean expression for condition argument of `select`, got something else")]
+    SelectScalarConditionNotABool,
+    #[error(
+        "Expected vectors of the same component type for reject and accept args., got something else",
+    )]
+    SelectVecRejectAcceptTypeMismatch,
+    #[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 +922,118 @@ impl<'a> ConstantEvaluator<'a> {
                     )),
                 }
             }
-            Expression::Select { .. } => Err(ConstantEvaluatorError::NotImplemented(
-                "select built-in function".into(),
-            )),
+            Expression::Select {
+                reject,
+                accept,
+                condition,
+            } => {
+                let arg = |this: &mut Self, expr| {
+                    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)
+                        }
+                    };
+
+                // NOTE: Frontends are on the hook for any conversions between arguments, if the
+                // language supports it. If types don't match, the buck stops here.
+                match (&self.expressions[reject], &self.expressions[accept]) {
+                    (&Expression::Literal(reject_lit), &Expression::Literal(accept_lit)) => {
+                        let reject_scalar = reject_lit.scalar();
+                        // if reject_scalar != accept_lit.scalar() {
+                        //     return Err(
+                        //         ConstantEvaluatorError::SelectScalarRejectAcceptTypeMismatch,
+                        //     );
+                        // }
+                        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,
+                        },
+                    ) => {
+                        if reject_ty != accept_ty {
+                            return Err(ConstantEvaluatorError::SelectVecRejectAcceptTypeMismatch);
+                        }
+                        let ty = reject_ty;
+
+                        let ty_deets = |ty| {
+                            let (size, scalar) =
+                                self.types[ty].inner.vector_size_and_scalar().unwrap();
+                            (size.unwrap(), scalar)
+                        };
+                        let (expected_vec_size, _scalar) = ty_deets(ty);
+
+                        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,
+                            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),
+                }
+            }
             Expression::Relational { fun, argument } => {
                 let argument = self.check_and_get(argument)?;
                 self.relational(fun, argument, span)