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};
@@ -2481,13 +2482,46 @@ 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 = proc::select::overloads();
+
+                            let rule = self.resolve_overloads(
+                                span,
+                                proc::select::WgslSymbol,
+                                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 proc::Conclusion::Predeclared(predeclared) = rule.conclusion {
+                                ctx.module.generate_predeclared_type(predeclared);
+                            }
+
+                            let [reject, accept, condition] = lowered_arguments;
 
                             ir::Expression::Select {
                                 reject,
@@ -2988,6 +3022,70 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
         })
     }
 
+    fn resolve_arr_things<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.resolve_const_things(
+            span,
+            fun,
+            |args, ctx| {
+                let args = args.try_into().unwrap();
+                resolve_overloads(args, ctx)
+            },
+            ast_arguments.into(),
+            ctx,
+        )
+        .map(|arr| arr.into_inner().unwrap())
+    }
+
+    fn resolve_const_things<const NUM_ARGS: usize, F, O, R>(
+        &mut self,
+        span: Span,
+        fun: F,
+        resolve_overloads: R,
+        ast_arguments: ArrayVec<Handle<ast::Expression<'source>>, { NUM_ARGS }>,
+        ctx: &mut ExpressionContext<'source, '_, '_>,
+    ) -> Result<'source, ArrayVec<Handle<ir::Expression>, { NUM_ARGS }>>
+    where
+        F: TryToWgsl + core::fmt::Debug + Copy,
+        O: proc::OverloadSet,
+        R: FnOnce(&[Handle<ir::Expression>], &mut ExpressionContext<'source, '_, '_>) -> O,
+    {
+        let mut lowered_arguments = ArrayVec::<_, { NUM_ARGS }>::new();
+
+        for &arg in ast_arguments.iter() {
+            let lowered = self.expression_for_abstract(arg, ctx)?;
+            ctx.grow_types(lowered)?;
+            lowered_arguments.push(lowered);
+        }
+
+        let fun_overloads = resolve_overloads(&lowered_arguments, ctx);
+        let rule = self.resolve_overloads(span, fun, 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 proc::Conclusion::Predeclared(predeclared) = rule.conclusion {
+            ctx.module.generate_predeclared_type(predeclared);
+        }
+
+        Ok(lowered_arguments)
+    }
+
     /// Choose the right overload for a function call.
     ///
     /// Return a [`Rule`] representing the most preferred overload in
@@ -3826,6 +3924,9 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
     }
 }
 
+// TODO: Surely this already exists! Find it.
+const MAX_BUILTIN_ARGS: usize = 4;
+
 impl ir::AtomicFunction {
     pub fn map(word: &str) -> Option<Self> {
         Some(match word {

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,37 @@
+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(Debug, 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";
+}

naga/src/valid/expression.rs

@@ -1,7 +1,8 @@
-use alloc::{format, string::String};
+use alloc::{borrow::ToOwned, format, string::String};
 
 use super::{compose::validate_compose, FunctionInfo, ModuleInfo, ShaderStages, TypeFlags};
 use crate::arena::UniqueArena;
+use crate::common::wgsl::TryToWgsl as _;
 use crate::{
     arena::Handle,
     proc,
@@ -929,47 +930,36 @@ 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 = proc::select::overloads();
+                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::InvalidArgumentType(
+                            proc::select::WgslSymbol::DESCRIPTION.to_owned(),
+                            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, .. } => {