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

Author: ErichDonGubler

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

@@ -2496,13 +2496,19 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                 } else {
                     match function.name {
                         "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 condition = self.expression(args.next()?, ctx)?;
+                            const NUM_ARGS: usize = 3;
+                            let mut args_iter = ctx.prepare_args(arguments, NUM_ARGS as u32, span);
+                            let args: [_; NUM_ARGS] =
+                                [args_iter.next()?, args_iter.next()?, args_iter.next()?];
+                            args_iter.finish()?;
 
-                            args.finish()?;
+                            let [reject, accept, condition] = self.function_helper_array(
+                                span,
+                                proc::select::WgslSymbol,
+                                |_args, _ctx| proc::select::overloads(),
+                                args,
+                                ctx,
+                            )?;
 
                             ir::Expression::Select {
                                 reject,
@@ -2999,6 +3005,39 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
         Ok(lowered_arguments)
     }
 
+    /// A convenience around [`Self::function_helper`] that marshals to and from fixed-size arrays.
+    ///
+    /// Useful for resolving function calls with a fixed number of arguments at parse time, like
+    /// [`ir::Expression::Select`].
+    fn function_helper_array<const NUM_ARGS: usize, F, O, R>(
+        &mut self,
+        span: Span,
+        fun: F,
+        resolve_overloads: R,
+        ast_arguments: [Handle<ast::Expression<'source>>; NUM_ARGS],
+        ctx: &mut ExpressionContext<'source, '_, '_>,
+    ) -> Result<'source, [Handle<ir::Expression>; NUM_ARGS]>
+    where
+        F: TryToWgsl + core::fmt::Debug + Copy,
+        O: proc::OverloadSet,
+        R: FnOnce(
+            &[Handle<ir::Expression>; NUM_ARGS],
+            &mut ExpressionContext<'source, '_, '_>,
+        ) -> O,
+    {
+        self.function_helper(
+            span,
+            fun,
+            |args, ctx| {
+                let args = args.try_into().unwrap();
+                resolve_overloads(args, ctx)
+            },
+            ast_arguments.into(),
+            ctx,
+        )
+        .map(|arr| arr.into_inner().unwrap())
+    }
+
     /// Choose the right overload for a function call.
     ///
     /// Return a [`Rule`] representing the most preferred overload in

naga/src/proc/mod.rs

@@ -19,7 +19,7 @@ pub use emitter::Emitter;
 pub use index::{BoundsCheckPolicies, BoundsCheckPolicy, IndexableLength, IndexableLengthError};
 pub use layouter::{Alignment, LayoutError, LayoutErrorInner, Layouter, TypeLayout};
 pub use namer::{EntryPointIndex, NameKey, Namer};
-pub use overloads::{Conclusion, MissingSpecialType, OverloadSet, Rule};
+pub use overloads::{select, Conclusion, MissingSpecialType, OverloadSet, Rule};
 pub use terminator::ensure_block_returns;
 use thiserror::Error;
 pub use type_methods::min_max_float_representable_by;

naga/src/proc/overloads/mod.rs

@@ -21,6 +21,7 @@ mod list;
 mod mathfunction;
 mod one_bits_iter;
 mod rule;
+pub mod select;
 mod utils;
 
 pub use rule::{Conclusion, MissingSpecialType, Rule};

naga/src/proc/overloads/select.rs

@@ -0,0 +1,43 @@
+use crate::common::wgsl::{ToWgsl, TryToWgsl};
+use crate::ir;
+use crate::proc::overloads::utils::{list, rule, scalar_or_vecn, scalars};
+use crate::proc::overloads::OverloadSet;
+
+pub fn overloads() -> impl OverloadSet {
+    list(scalars().flat_map(|scalar| {
+        scalar_or_vecn(scalar).map(|input| {
+            let bool_arg = match input.clone() {
+                ir::TypeInner::Scalar(_) => ir::TypeInner::Scalar(ir::Scalar::BOOL),
+                ir::TypeInner::Vector { size, scalar: _ } => ir::TypeInner::Vector {
+                    size,
+                    scalar: ir::Scalar::BOOL,
+                },
+                _ => unreachable!(),
+            };
+            rule([input.clone(), input.clone(), bool_arg], input)
+        })
+    }))
+}
+
+#[derive(Clone, Copy)]
+pub struct WgslSymbol;
+
+impl ToWgsl for WgslSymbol {
+    fn to_wgsl(self) -> &'static str {
+        "select"
+    }
+}
+
+impl TryToWgsl for WgslSymbol {
+    fn try_to_wgsl(self) -> Option<&'static str> {
+        Some(self.to_wgsl())
+    }
+
+    const DESCRIPTION: &'static str = "`select` built-in";
+}
+
+impl core::fmt::Debug for WgslSymbol {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.write_str(Self::DESCRIPTION)
+    }
+}

naga/src/proc/overloads/utils.rs

@@ -34,6 +34,25 @@ pub fn float_scalars() -> impl Iterator<Item = ir::Scalar> + Clone {
     .into_iter()
 }
 
+/// Produce all [`ir::Scalar`]s.
+///
+/// Note that `*32` and `F16` must appear before other sizes; this is how we
+/// represent conversion rank.
+pub fn scalars() -> impl Iterator<Item = ir::Scalar> + Clone {
+    [
+        ir::Scalar::ABSTRACT_INT,
+        ir::Scalar::ABSTRACT_FLOAT,
+        ir::Scalar::I32,
+        ir::Scalar::U32,
+        ir::Scalar::F32,
+        ir::Scalar::F16,
+        ir::Scalar::I64,
+        ir::Scalar::U64,
+        ir::Scalar::F64,
+    ]
+    .into_iter()
+}
+
 /// Produce all the floating-point [`ir::Scalar`]s, but omit
 /// abstract types, for #7405.
 pub fn float_scalars_unimplemented_abstract() -> impl Iterator<Item = ir::Scalar> + Clone {

naga/src/valid/expression.rs

@@ -984,47 +984,16 @@ impl super::Validator {
                 accept,
                 reject,
             } => {
-                let accept_inner = &resolver[accept];
-                let reject_inner = &resolver[reject];
-                let condition_ty = &resolver[condition];
-                let condition_good = match *condition_ty {
-                    Ti::Scalar(Sc {
-                        kind: Sk::Bool,
-                        width: _,
-                    }) => {
-                        // When `condition` is a single boolean, `accept` and
-                        // `reject` can be vectors or scalars.
-                        match *accept_inner {
-                            Ti::Scalar { .. } | Ti::Vector { .. } => true,
-                            _ => false,
-                        }
-                    }
-                    Ti::Vector {
-                        size,
-                        scalar:
-                            Sc {
-                                kind: Sk::Bool,
-                                width: _,
-                            },
-                    } => match *accept_inner {
-                        Ti::Vector {
-                            size: other_size, ..
-                        } => size == other_size,
-                        _ => false,
-                    },
-                    _ => false,
-                };
-                if accept_inner != reject_inner {
-                    return Err(ExpressionError::SelectValuesTypeMismatch {
-                        accept: accept_inner.clone(),
-                        reject: reject_inner.clone(),
-                    });
-                }
-                if !condition_good {
-                    return Err(ExpressionError::SelectConditionNotABool {
-                        actual: condition_ty.clone(),
-                    });
-                }
+                self.validate_func_call_with_overloads(
+                    module,
+                    proc::select::WgslSymbol,
+                    proc::select::overloads(),
+                    [reject, accept, condition]
+                        .iter()
+                        .copied()
+                        .map(|arg| (arg, &resolver[arg])),
+                )?;
+
                 ShaderStages::all()
             }
             E::Derivative { expr, .. } => {