From 91bbc9fa5f870bf5ee15ce54ca61b434e00255df Mon Sep 17 00:00:00 2001
From: Kai Schmidt <kaikaliischmidt@gmail.com>
Date: Wed, 12 Jun 2024 08:10:35 -0700
Subject: [PATCH] properly emit a warning for some cases of repeat in an array

---
 src/check.rs       | 12 ++++++------
 src/compile/mod.rs | 35 +++++++++++++++++++++++++++++++++--
 tests/error.ua     |  2 ++
 todo.md            |  1 -
 4 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/src/check.rs b/src/check.rs
index 268922cb4..835bbf199 100644
--- a/src/check.rs
+++ b/src/check.rs
@@ -91,7 +91,7 @@ pub enum SigCheckErrorKind {
     Incorrect,
     Ambiguous,
     LoopOverreach,
-    LoopExcess { sig: Signature, inf: bool },
+    LoopVariable { sig: Signature, inf: bool },
 }
 
 impl SigCheckError {
@@ -107,9 +107,9 @@ impl SigCheckError {
             ..self
         }
     }
-    pub fn loop_excess(self, sig: Signature, inf: bool) -> Self {
+    pub fn loop_variable(self, sig: Signature, inf: bool) -> Self {
         Self {
-            kind: SigCheckErrorKind::LoopExcess { sig, inf },
+            kind: SigCheckErrorKind::LoopVariable { sig, inf },
             ..self
         }
     }
@@ -405,7 +405,7 @@ impl<'a> VirtualEnv<'a> {
                                     return Err(SigCheckError::from(format!(
                                         "repeat with infinity and a function with signature {sig}"
                                     ))
-                                    .loop_excess(sig, true));
+                                    .loop_variable(sig, true));
                                 }
                                 _ => self.handle_sig(sig)?,
                             }
@@ -427,7 +427,7 @@ impl<'a> VirtualEnv<'a> {
                                 return Err(SigCheckError::from(format!(
                                     "repeat with no number and a function with signature {sig}"
                                 ))
-                                .loop_excess(sig, false));
+                                .loop_variable(sig, false));
                             }
                             Ordering::Less => self.handle_sig(sig)?,
                         }
@@ -450,7 +450,7 @@ impl<'a> VirtualEnv<'a> {
                         return Err(SigCheckError::from(format!(
                             "do with a function with signature {comp_sig}"
                         ))
-                        .loop_excess(comp_sig, false));
+                        .loop_variable(comp_sig, false));
                     }
                     self.handle_args_outputs(
                         comp_sig.args,
diff --git a/src/compile/mod.rs b/src/compile/mod.rs
index d88206162..d861e9da8 100644
--- a/src/compile/mod.rs
+++ b/src/compile/mod.rs
@@ -1310,12 +1310,43 @@ code:
         }
         Ok(())
     }
+    /// Emit a warning if a loop inside an array could
+    /// potentially pull in a variable number of values
     fn validate_array_loop_sig(&mut self, instrs: &[Instr], span: &CodeSpan) -> Option<Signature> {
         let inner_sig = instrs_signature(instrs);
         if self.current_bindings.is_empty() && self.scope.kind != ScopeKind::Temp {
             return inner_sig.ok();
         }
         let Err(e) = &inner_sig else {
+            // Case where repeat's function has a balanced signature
+            // This is fine in other contexts, so an error is not returned
+            // from the signature check, but it is not okay in an array.
+            if let Some(i) = instrs
+                .iter()
+                .position(|instr| matches!(instr, Instr::Prim(Primitive::Repeat, _)))
+                .filter(|&i| i > 0)
+            {
+                if let Instr::PushFunc(f) = &instrs[i - 1] {
+                    let body_sig = f.signature();
+                    let before_sig = instrs_signature(&instrs[..i - 1]).ok()?;
+                    let after_sig = instrs_signature(&instrs[i + 1..]).ok()?;
+                    if body_sig.args == body_sig.outputs
+                        && before_sig.args < body_sig.args
+                        && after_sig.args.saturating_sub(body_sig.outputs) < body_sig.args
+                    {
+                        let replacement: String =
+                            repeat('⊙').take(body_sig.args).chain(['∘']).collect();
+                        let message = format!(
+                            "This array contains a loop with an equal number \
+                            of arguments and outputs. This may result in a \
+                            variable number of values being pulled into the \
+                            array. To fix this, insert `{replacement}` on the \
+                            right side of the array.",
+                        );
+                        self.emit_diagnostic(message, DiagnosticKind::Warning, span.clone());
+                    }
+                }
+            }
             return inner_sig.ok();
         };
         let before_sig = (0..instrs.len())
@@ -1326,8 +1357,8 @@ code:
             .find_map(|i| instrs_signature(&instrs[i..]).ok())
             .unwrap();
         match e.kind {
-            SigCheckErrorKind::LoopExcess { sig: body_sig, inf } => {
-                let positive_body_sig = body_sig.outputs > body_sig.args;
+            SigCheckErrorKind::LoopVariable { sig: body_sig, inf } => {
+                let positive_body_sig = body_sig.outputs >= body_sig.args;
                 let balanced_after_args = body_sig.args.saturating_sub(before_sig.outputs);
                 if body_sig.args > 0
                     && (positive_body_sig && after_sig.args != balanced_after_args && !inf)
diff --git a/tests/error.ua b/tests/error.ua
index 570a3849d..8a135a38f 100644
--- a/tests/error.ua
+++ b/tests/error.ua
@@ -41,6 +41,8 @@ F ← ⌊×10[◌◌⍥gen]⧻:0
 
 F ← [+⍥..]
 
+F ← {⍥∘}
+
 ⊕⊢ [0 2 2] [1 2 3]
 
 + ¤[1 2 3] [10_20_30_40 50_60_70_80]
diff --git a/todo.md b/todo.md
index ed8900b1f..63cf8b6da 100644
--- a/todo.md
+++ b/todo.md
@@ -1,7 +1,6 @@
 # Uiua Todo
 
 - 0.12
-  - Fix `{⍥∘}` not giving a warning
   - by inverses
   - `setinv` on and by inverses
   - Document on and by inverses