fix(naga): properly impl. auto. type conv. for select

Author: ErichDonGubler

naga/src/front/wgsl/error.rs

@@ -399,6 +399,16 @@ pub(crate) enum Error<'a> {
         on_what: DiagnosticAttributeNotSupportedPosition,
         spans: Vec<Span>,
     },
+    SelectUnexpectedArgumentType {
+        arg_span: Span,
+        arg_type: String,
+    },
+    SelectRejectAndAcceptHaveNoCommonType {
+        reject_span: Span,
+        reject_type: String,
+        accept_span: Span,
+        accept_type: String,
+    },
 }
 
 impl From<ConflictingDiagnosticRuleError> for Error<'_> {
@@ -1340,6 +1350,24 @@ impl<'a> Error<'a> {
                     ],
                 }
             }
+            Error::SelectUnexpectedArgumentType { arg_span, ref arg_type } => ParseError {
+                message: "unexpected argument type for `select` call".into(),
+                labels: vec![(arg_span, format!("this value of type {arg_type}").into())],
+                notes: vec!["expected a scalar or a `vecN` of scalars".into()],
+            },
+            Error::SelectRejectAndAcceptHaveNoCommonType {
+                reject_span,
+                ref reject_type,
+                accept_span,
+                ref accept_type,
+            } => ParseError {
+                message: "type mismatch for reject and accept values in `select` call".into(),
+                labels: vec![
+                    (reject_span, format!("reject value of type {reject_type}").into()),
+                    (accept_span, format!("accept value of type {accept_type}").into()),
+                ],
+                notes: vec![],
+            },
         }
     }
 }

naga/src/front/wgsl/lower/mod.rs

@@ -1,6 +1,7 @@
 use alloc::{
     borrow::ToOwned,
     boxed::Box,
+    format,
     string::{String, ToString},
     vec::Vec,
 };
@@ -2483,12 +2484,70 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                         "select" => {
                             let mut args = ctx.prepare_args(arguments, 3, span);
 
-                            let reject = self.expression(args.next()?, ctx)?;
-                            let accept = self.expression(args.next()?, ctx)?;
+                            let reject_orig = args.next()?;
+                            let accept_orig = args.next()?;
+                            let mut values = [
+                                self.expression_for_abstract(reject_orig, ctx)?,
+                                self.expression_for_abstract(accept_orig, ctx)?,
+                            ];
                             let condition = self.expression(args.next()?, ctx)?;
 
                             args.finish()?;
 
+                            fn expr_ty<'a>(
+                                ctx: &'a ExpressionContext<'_, '_, '_>,
+                                expr: Handle<ir::Expression>,
+                            ) -> &'a ir::TypeInner {
+                                ctx.typifier()[expr].inner_with(&ctx.module.types)
+                            }
+                            let diag_deets =
+                                |module: &ir::Module, ty: &ir::TypeInner, orig_expr| {
+                                    (
+                                        ctx.ast_expressions.get_span(orig_expr),
+                                        format!("`{:?}`", ty.for_debug(&module.types)),
+                                    )
+                                };
+                            for (&value, orig_value) in
+                                values.iter().zip([reject_orig, accept_orig])
+                            {
+                                ctx.grow_types(value)?;
+
+                                let value_type = expr_ty(ctx, value);
+                                if value_type.vector_size_and_scalar().is_none() {
+                                    let (arg_span, arg_type) =
+                                        diag_deets(ctx.module, value_type, orig_value);
+                                    return Err(Box::new(Error::SelectUnexpectedArgumentType {
+                                        arg_span,
+                                        arg_type,
+                                    }));
+                                }
+                            }
+                            let mut consensus_scalar = ctx
+                                .automatic_conversion_consensus(&values)
+                                .map_err(|_idx| {
+                                    let [reject, accept] = values;
+                                    let [(reject_span, reject_type), (accept_span, accept_type)] =
+                                        [(reject_orig, reject), (accept_orig, accept)].map(
+                                            |(orig_expr, expr)| {
+                                                let ty = expr_ty(ctx, expr);
+                                                diag_deets(ctx.module, ty, orig_expr)
+                                            },
+                                        );
+                                    Error::SelectRejectAndAcceptHaveNoCommonType {
+                                        reject_span,
+                                        reject_type,
+                                        accept_span,
+                                        accept_type,
+                                    }
+                                })?;
+                            if !ctx.is_const(condition) {
+                                consensus_scalar = consensus_scalar.concretize();
+                            }
+
+                            ctx.convert_slice_to_common_leaf_scalar(&mut values, consensus_scalar)?;
+
+                            let [reject, accept] = values;
+
                             ir::Expression::Select {
                                 reject,
                                 accept,