check more loop modifier signatures properly

uiua-lang · Oct 13, 2023 · 35a6b2a · 35a6b2a
1 parent d2b58e0
commit 35a6b2a
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 86 deletions.
diff --git a/src/algorithm/loops.rs b/src/algorithm/loops.rs
@@ -274,10 +274,9 @@ where
 
 fn generic_scan(f: Value, xs: Value, env: &mut Uiua) -> UiuaResult {
     let sig = f.signature();
-    if sig.outputs != 1 {
+    if sig != (2, 1) {
         return Err(env.error(format!(
-            "Scan's function must return 1 value, but {} returns {}",
-            f, sig.outputs
+            "Scan's function's signature must be |2.1, but it is {sig}"
         )));
     }
     if xs.row_count() == 0 {
@@ -834,10 +833,9 @@ fn fast_table_join_or_couple<T: ArrayValue>(a: Array<T>, b: Array<T>) -> Array<T
 
 fn generic_table(f: Value, xs: Value, ys: Value, env: &mut Uiua) -> UiuaResult {
     let sig = f.signature();
-    if sig.outputs != 1 {
+    if sig != (2, 1) {
         return Err(env.error(format!(
-            "Table's function must return 1 value, but {} returns {}",
-            f, sig.outputs
+            "Table's function's signature must be |2.1, but it is {sig}"
         )));
     }
     let mut new_shape = Shape::from(xs.shape());
@@ -868,10 +866,9 @@ pub fn cross(env: &mut Uiua) -> UiuaResult {
     let xs = env.pop(ArrayArg(1))?;
     let ys = env.pop(ArrayArg(2))?;
     let sig = f.signature();
-    if sig.outputs != 1 {
+    if sig != (2, 1) {
         return Err(env.error(format!(
-            "Cross's function must return 1 value, but {} returns {}",
-            f, sig.outputs
+            "Cross's function's signature must be |2.1, but it is {sig}"
         )));
     }
     let mut new_shape = tiny_vec![xs.row_count(), ys.row_count()];

diff --git a/src/check.rs b/src/check.rs
@@ -146,8 +146,24 @@ impl<'a> VirtualEnv<'a> {
                     };
                     self.handle_args_outputs(1, outputs)?;
                 }
-                Each | Rows => self.handle_variadic_mod(prim)?,
-                Table | Cross => self.handle_mod(prim, Some(2), Some(1), 2, None)?,
+                Each | Rows => {
+                    let sig = self.pop()?.expect_function(|| prim)?;
+                    if sig.outputs != 1 {
+                        return Err(format!(
+                            "{prim}'s function must have 1 output, but its signature is {sig}"
+                        ));
+                    }
+                    self.handle_sig(sig)?
+                }
+                Table | Cross => {
+                    let sig = self.pop()?.expect_function(|| prim)?;
+                    if sig != (2, 1) {
+                        return Err(format!(
+                            "{prim}'s function's signature must be |2.1, but it is {sig}"
+                        ));
+                    }
+                    self.handle_args_outputs(2, 1)?;
+                }
                 Distribute => {
                     let sig = self.pop()?.expect_function(|| prim)?;
                     self.handle_sig(sig)?
@@ -167,7 +183,10 @@ impl<'a> VirtualEnv<'a> {
                     };
                     self.handle_args_outputs(args, outputs)?;
                 }
-                Spawn => self.handle_mod(prim, None, None, 1, Some(1))?,
+                Spawn => {
+                    let sig = self.pop()?.expect_function(|| prim)?;
+                    self.handle_args_outputs(sig.args, 1)?;
+                }
                 Repeat => {
                     let f = self.pop()?;
                     let n = self.pop()?;
@@ -182,27 +201,22 @@ impl<'a> VirtualEnv<'a> {
                         if n.fract() == 0.0 && n >= 0.0 {
                             let n = n as usize;
                             if n > 0 {
-                                if let BasicValue::Func(f) = f {
-                                    let sig = f.signature();
-                                    let (args, outputs) = match sig.args.cmp(&sig.outputs) {
-                                        Ordering::Equal => (sig.args, sig.outputs),
-                                        Ordering::Less => {
-                                            (sig.args, n * (sig.outputs - sig.args) + sig.args)
-                                        }
-                                        Ordering::Greater => (
-                                            (n - 1) * (sig.args - sig.outputs) + sig.args,
-                                            sig.outputs,
-                                        ),
-                                    };
-                                    for _ in 0..args {
-                                        self.pop()?;
+                                let sig = f.signature();
+                                let (args, outputs) = match sig.args.cmp(&sig.outputs) {
+                                    Ordering::Equal => (sig.args, sig.outputs),
+                                    Ordering::Less => {
+                                        (sig.args, n * (sig.outputs - sig.args) + sig.args)
                                     }
-                                    self.set_min_height();
-                                    for _ in 0..outputs {
-                                        self.stack.push(BasicValue::Other);
+                                    Ordering::Greater => {
+                                        ((n - 1) * (sig.args - sig.outputs) + sig.args, sig.outputs)
                                     }
-                                } else {
-                                    self.handle_mod(prim, Some(0), Some(1), n, None)?
+                                };
+                                for _ in 0..args {
+                                    self.pop()?;
+                                }
+                                self.set_min_height();
+                                for _ in 0..outputs {
+                                    self.stack.push(BasicValue::Other);
                                 }
                             }
                         } else {
@@ -514,61 +528,6 @@ impl<'a> VirtualEnv<'a> {
         }
         Ok(())
     }
-    fn handle_mod(
-        &mut self,
-        prim: &Primitive,
-        f_args: Option<usize>,
-        f_outputs: Option<usize>,
-        m_args: usize,
-        m_outputs: Option<usize>,
-    ) -> Result<(), String> {
-        if let BasicValue::Func(f) = self.pop()? {
-            let sig = f.signature();
-            if let Some(f_args) = f_args {
-                if sig.args != f_args {
-                    return Err(format!(
-                        "{prim}'s function {f:?} has {} args, expected {}",
-                        sig.args, f_args
-                    ));
-                }
-            }
-            if let Some(f_outputs) = f_outputs {
-                if sig.outputs != f_outputs {
-                    return Err(format!(
-                        "{prim}'s function {f:?} has {} outputs, expected {}",
-                        sig.outputs, f_outputs
-                    ));
-                }
-            }
-            for _ in 0..m_args {
-                self.pop()?;
-            }
-            self.set_min_height();
-            let outputs = m_outputs.unwrap_or(f_outputs.unwrap_or(sig.outputs));
-            for _ in 0..outputs {
-                self.stack.push(BasicValue::Other);
-            }
-            Ok(())
-        } else {
-            Err(format!("{prim} without function"))
-        }
-    }
-    fn handle_variadic_mod(&mut self, prim: &Primitive) -> Result<(), String> {
-        if let BasicValue::Func(f) = self.pop()? {
-            let sig = f.signature();
-            if sig.outputs != 1 {
-                return Err(format!("{prim}'s function {f:?} did not return 1 value",));
-            }
-            for _ in 0..sig.args {
-                self.pop()?;
-            }
-            self.set_min_height();
-            self.stack.push(BasicValue::Other);
-            Ok(())
-        } else {
-            Err(format!("{prim} without function"))
-        }
-    }
 }
 
 fn instrs_contain_break_at_non_end(instrs: &[Instr]) -> bool {