core: add np_argmax/np_argmin functions

nac3core/src/toplevel/builtins.rs

@@ -510,7 +510,10 @@ impl<'a> BuiltinBuilder<'a> {
 
             PrimDef::FunMin | PrimDef::FunMax => self.build_min_max_function(prim),
 
-            PrimDef::FunNpMin | PrimDef::FunNpMax => self.build_np_min_max_function(prim),
+            PrimDef::FunNpArgmin 
+            | PrimDef::FunNpArgmax
+            | PrimDef::FunNpMin 
+            | PrimDef::FunNpMax => self.build_np_max_min_function(prim),
 
             PrimDef::FunNpMinimum | PrimDef::FunNpMaximum => {
                 self.build_np_minimum_maximum_function(prim)
@@ -1555,39 +1558,42 @@ impl<'a> BuiltinBuilder<'a> {
         }
     }
 
-    /// Build the functions `np_min()` and `np_max()`.
+    /// Build the functions `np_max()`, `np_min()`, `np_argmax()` and `np_argmin()`
-    fn build_np_min_max_function(&mut self, prim: PrimDef) -> TopLevelDef {
+    /// Calls `call_numpy_max_min` with the function name
-        debug_assert_prim_is_allowed(prim, &[PrimDef::FunNpMin, PrimDef::FunNpMax]);
+    fn build_np_max_min_function(&mut self, prim: PrimDef) -> TopLevelDef {
+        debug_assert_prim_is_allowed(prim, &[PrimDef::FunNpArgmin, PrimDef::FunNpArgmax, PrimDef::FunNpMin, PrimDef::FunNpMax]);
 
-        let ret_ty = self.unifier.get_fresh_var(Some("R".into()), None);
+        let (var_map, ret_ty) = match prim {
-        let var_map = self
+            PrimDef::FunNpArgmax | PrimDef::FunNpArgmin => {
-            .num_or_ndarray_var_map
+                (self.num_or_ndarray_var_map.clone(), self.primitives.int64)
-            .clone()
+            },
-            .into_iter()
+            PrimDef::FunNpMax | PrimDef::FunNpMin => {
-            .chain(once((ret_ty.id, ret_ty.ty)))
+                let ret_ty = self.unifier.get_fresh_var(Some("R".into()), None);
-            .collect::<IndexMap<_, _>>();
+                let var_map = self
+                    .num_or_ndarray_var_map
+                    .clone()
+                    .into_iter()
+                    .chain(once((ret_ty.id, ret_ty.ty)))
+                    .collect::<IndexMap<_, _>>();
+                (var_map, ret_ty.ty)
+            },
+            _ => unreachable!()
+        };
 
         create_fn_by_codegen(
             self.unifier,
             &var_map,
             prim.name(),
-            ret_ty.ty,
+            ret_ty,
-            &[(self.float_or_ndarray_ty.ty, "a")],
+            &[(self.num_or_ndarray_ty.ty, "a")],
             Box::new(move |ctx, _, fun, args, generator| {
                 let a_ty = fun.0.args[0].ty;
                 let a = args[0].1.clone().to_basic_value_enum(ctx, generator, a_ty)?;
-
+                
-                let func = match prim {
+                Ok(Some(builtin_fns::call_numpy_max_min(generator, ctx, (a_ty, a), &prim.name())?))
-                    PrimDef::FunNpMin => builtin_fns::call_numpy_min,
-                    PrimDef::FunNpMax => builtin_fns::call_numpy_max,
-                    _ => unreachable!(),
-                };
-
-                Ok(Some(func(generator, ctx, (a_ty, a))?))
             }),
         )
     }
-
     /// Build the functions `np_minimum()` and `np_maximum()`.
     fn build_np_minimum_maximum_function(&mut self, prim: PrimDef) -> TopLevelDef {
         debug_assert_prim_is_allowed(prim, &[PrimDef::FunNpMinimum, PrimDef::FunNpMaximum]);

nac3core/src/toplevel/helper.rs

@@ -62,9 +62,11 @@ pub enum PrimDef {
     FunMin,
     FunNpMin,
     FunNpMinimum,
+    FunNpArgmin,
     FunMax,
     FunNpMax,
     FunNpMaximum,
+    FunNpArgmax,
     FunAbs,
     FunNpIsNan,
     FunNpIsInf,
@@ -216,9 +218,11 @@ impl PrimDef {
             PrimDef::FunMin => fun("min", None),
             PrimDef::FunNpMin => fun("np_min", None),
             PrimDef::FunNpMinimum => fun("np_minimum", None),
+            PrimDef::FunNpArgmin => fun("np_argmin", None),
             PrimDef::FunMax => fun("max", None),
             PrimDef::FunNpMax => fun("np_max", None),
             PrimDef::FunNpMaximum => fun("np_maximum", None),
+            PrimDef::FunNpArgmax => fun("np_argmax", None),
             PrimDef::FunAbs => fun("abs", None),
             PrimDef::FunNpIsNan => fun("np_isnan", None),
             PrimDef::FunNpIsInf => fun("np_isinf", None),