refactor(naga): make math_function_helper a more generic function_helper

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};
@@ -2465,7 +2466,21 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
 
                     ir::Expression::Derivative { axis, ctrl, expr }
                 } else if let Some(fun) = conv::map_standard_fun(function.name) {
-                    self.math_function_helper(span, fun, arguments, ctx)?
+                    let lowered_arguments = self.function_helper(
+                        span,
+                        fun,
+                        |_args, _ctx| fun.overloads(),
+                        ArrayVec::<_, 4>::from_iter(arguments.iter().copied()),
+                        ctx,
+                    )?;
+
+                    ir::Expression::Math {
+                        fun,
+                        arg: lowered_arguments[0],
+                        arg1: lowered_arguments.get(1).cloned(),
+                        arg2: lowered_arguments.get(2).cloned(),
+                        arg3: lowered_arguments.get(3).cloned(),
+                    }
                 } else if let Some(fun) = Texture::map(function.name) {
                     self.texture_sample_helper(fun, arguments, span, ctx)?
                 } else if let Some((op, cop)) = conv::map_subgroup_operation(function.name) {
@@ -2952,31 +2967,33 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
         }
     }
 
-    /// Generate a Naga IR [`Math`] expression.
-    ///
-    /// Generate Naga IR for a call to the [`MathFunction`] `fun`, whose
+    /// Generate a Naga IR expression for a call to `fun`, whose
     /// unlowered arguments are `ast_arguments`.
     ///
     /// The `span` argument should give the span of the function name in the
     /// call expression.
-    ///
-    /// [`Math`]: ir::Expression::Math
-    /// [`MathFunction`]: ir::MathFunction
-    fn math_function_helper(
+    fn function_helper<const NUM_ARGS: usize, F, O, R>(
         &mut self,
         span: Span,
-        fun: ir::MathFunction,
-        ast_arguments: &[Handle<ast::Expression<'source>>],
+        fun: F,
+        resolve_overloads: R,
+        ast_arguments: ArrayVec<Handle<ast::Expression<'source>>, { NUM_ARGS }>,
         ctx: &mut ExpressionContext<'source, '_, '_>,
-    ) -> Result<'source, ir::Expression> {
-        let mut lowered_arguments = Vec::with_capacity(ast_arguments.len());
-        for &arg in ast_arguments {
+    ) -> 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 = fun.overloads();
+        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)?;
 
@@ -2987,13 +3004,7 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
             ctx.module.generate_predeclared_type(predeclared);
         }
 
-        Ok(ir::Expression::Math {
-            fun,
-            arg: lowered_arguments[0],
-            arg1: lowered_arguments.get(1).cloned(),
-            arg2: lowered_arguments.get(2).cloned(),
-            arg3: lowered_arguments.get(3).cloned(),
-        })
+        Ok(lowered_arguments)
     }
 
     /// Choose the right overload for a function call.