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

Author: ErichDonGubler

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

@@ -4,6 +4,7 @@ use alloc::{
     string::{String, ToString},
     vec::Vec,
 };
+use arrayvec::ArrayVec;
 use core::num::NonZeroU32;
 
 use crate::common::wgsl::{TryToWgsl, TypeContext};
@@ -2478,13 +2479,58 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                 } else {
                     match function.name {
                         "select" => {
-                            let mut args = ctx.prepare_args(arguments, 3, span);
+                            const NUM_ARGS: usize = 3;
+
+                            // TODO: dedupe with `math_function_helper`
 
-                            let reject = self.expression(args.next()?, ctx)?;
-                            let accept = self.expression(args.next()?, ctx)?;
-                            let condition = self.expression(args.next()?, ctx)?;
+                            let mut lowered_arguments = ArrayVec::<_, NUM_ARGS>::new();
+                            let mut args = ctx.prepare_args(arguments, NUM_ARGS as u32, span);
+
+                            for _ in 0..lowered_arguments.capacity() {
+                                let lowered = self.expression_for_abstract(args.next()?, ctx)?;
+                                ctx.grow_types(lowered)?;
+                                lowered_arguments.push(lowered);
+                            }
 
                             args.finish()?;
+                            let mut lowered_arguments = lowered_arguments.into_inner().unwrap();
+
+                            let fun_overloads = crate::proc::select();
+
+                            #[derive(Debug, Clone, Copy)]
+                            struct Select;
+
+                            impl TryToWgsl for Select {
+                                fn try_to_wgsl(self) -> Option<&'static str> {
+                                    Some("select")
+                                }
+
+                                const DESCRIPTION: &'static str = "`select` built-in";
+                            }
+                            let rule = self.resolve_overloads(
+                                span,
+                                Select,
+                                fun_overloads,
+                                &lowered_arguments,
+                                ctx,
+                            )?;
+
+                            self.apply_automatic_conversions_for_call(
+                                &rule,
+                                &mut lowered_arguments,
+                                ctx,
+                            )?;
+
+                            // If this function returns a predeclared type, register it
+                            // in `Module::special_types`. The typifier will expect to
+                            // be able to find it there.
+                            if let crate::proc::Conclusion::Predeclared(predeclared) =
+                                rule.conclusion
+                            {
+                                ctx.module.generate_predeclared_type(predeclared);
+                            }
+
+                            let [reject, accept, condition] = lowered_arguments;
 
                             crate::Expression::Select {
                                 reject,

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/mathfunction.rs

@@ -5,7 +5,7 @@ use crate::proc::overloads::list::List;
 use crate::proc::overloads::regular::regular;
 use crate::proc::overloads::utils::{
     concrete_int_scalars, float_scalars, float_scalars_unimplemented_abstract, list, pairs, rule,
-    scalar_or_vecn, triples, vector_sizes,
+    scalar_or_vecn, scalars, triples, vector_sizes,
 };
 use crate::proc::overloads::OverloadSet;
 
@@ -187,6 +187,23 @@ fn transpose() -> List {
     )
 }
 
+pub fn select() -> impl OverloadSet {
+    let bool_arg = |input| match input {
+        ir::TypeInner::Scalar(_) => ir::TypeInner::Scalar(ir::Scalar::BOOL),
+        ir::TypeInner::Vector { size, scalar: _ } => ir::TypeInner::Vector {
+            size,
+            scalar: ir::Scalar::BOOL,
+        },
+        _ => unreachable!(),
+    };
+    list(scalars().flat_map(|scalar| {
+        scalar_or_vecn(scalar).map(|input| {
+            let bool_arg = bool_arg(input.clone());
+            rule([input.clone(), input.clone(), bool_arg], input)
+        })
+    }))
+}
+
 fn extract_bits() -> List {
     list(concrete_int_scalars().flat_map(|scalar| {
         scalar_or_vecn(scalar).map(|input| {

naga/src/proc/overloads/mod.rs

@@ -235,3 +235,5 @@ pub trait OverloadSet: Clone {
     /// Return an object that can be formatted with [`core::fmt::Debug`].
     fn for_debug(&self, types: &crate::UniqueArena<ir::Type>) -> impl fmt::Debug;
 }
+
+pub use mathfunction::select;

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

@@ -129,6 +129,9 @@ pub enum ExpressionError {
     WrongArgumentCount(crate::MathFunction),
     #[error("Argument [{1}] to {0:?} as expression {2:?} has an invalid type.")]
     InvalidArgumentType(crate::MathFunction, u32, Handle<crate::Expression>),
+    // TODO: dedupe with above
+    #[error("Argument [{0}] to `select` as expression {1:?} has an invalid type.")]
+    InvalidArgumentTypeSelect(u32, Handle<crate::Expression>),
     #[error(
         "workgroupUniformLoad result type can't be {0:?}. It can only be a constructible type."
     )]
@@ -926,47 +929,32 @@ 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,
-                        }
+                // TODO: dedupe with math functions
+
+                let mut overloads = crate::proc::select();
+                log::debug!(
+                    "initial overloads for `select`: {:#?}",
+                    overloads.for_debug(&module.types)
+                );
+
+                for (i, (expr, ty)) in [reject, accept, condition]
+                    .iter()
+                    .copied()
+                    .map(|arg| (arg, &resolver[arg]))
+                    .enumerate()
+                {
+                    overloads = overloads.arg(i, ty, &module.types);
+                    log::debug!(
+                        "overloads after arg {i}: {:#?}",
+                        overloads.for_debug(&module.types)
+                    );
+
+                    if overloads.is_empty() {
+                        log::debug!("all overloads eliminated");
+                        return Err(ExpressionError::InvalidArgumentTypeSelect(i as u32, expr));
                     }
-                    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(),
-                    });
                 }
+
                 ShaderStages::all()
             }
             E::Derivative { expr, .. } => {