make excess fold behavior more inline with repeat and do

Author: kaikalii

src/algorithm/reduce.rs

@@ -6,7 +6,7 @@ use std::{convert::identity, iter::repeat};
 use ecow::{eco_vec, EcoVec};
 
 use crate::{
-    algorithm::{get_ops, loops::flip, multi_output, pervade::*},
+    algorithm::{get_ops, loops::flip, pervade::*},
     check::{nodes_clean_sig, nodes_sig},
     cowslice::cowslice,
     Array, ArrayValue, Complex, ImplPrimitive, Node, Ops, Primitive, Shape, SigNode, Uiua,
@@ -1049,7 +1049,7 @@ pub fn fold(ops: Ops, env: &mut Uiua) -> UiuaResult {
     crate::profile_function!();
     let [f] = get_ops(ops, env)?;
     let sig = f.sig;
-    let (iterable_count, acc_count, collect_count) = if sig.args() > sig.outputs() {
+    let (iterable_count, acc_count, excess_count) = if sig.args() > sig.outputs() {
         (sig.args() - sig.outputs(), sig.outputs(), 0)
     } else {
         let iter = sig.args().min(1);
@@ -1095,7 +1095,7 @@ pub fn fold(ops: Ops, env: &mut Uiua) -> UiuaResult {
     if row_count == 0 && arrays.iter().all(Result::is_err) {
         row_count = 1;
     }
-    let mut collect = multi_output(collect_count, Vec::with_capacity(row_count));
+    let mut excess_rows = vec![Vec::new(); excess_count];
     for _ in 0..row_count {
         for array in arrays.iter_mut().rev() {
             env.push(match array {
@@ -1104,17 +1104,23 @@ pub fn fold(ops: Ops, env: &mut Uiua) -> UiuaResult {
             });
         }
         env.exec(f.clone())?;
-        for collected in &mut collect {
-            collected.push(env.remove_nth_back(acc_count)?);
+        if excess_count > 0 {
+            for (i, row) in env
+                .remove_n(excess_count, acc_count + excess_count)?
+                .enumerate()
+            {
+                excess_rows[i].push(row);
+            }
         }
     }
-    let accs = env.pop_n(acc_count)?;
-    for collected in collect.into_iter().rev() {
-        let val = Value::from_row_values(collected, env)?;
-        env.push(val);
+    // Remove preserved/excess values
+    if excess_count > 0 {
+        _ = env.remove_n(acc_count, acc_count)?;
     }
-    for acc in accs {
-        env.push(acc);
+    // Collect excess values
+    for rows in excess_rows.into_iter().rev() {
+        let new_val = Value::from_row_values(rows, env)?;
+        env.push(new_val);
     }
     Ok(())
 }

src/compile/modifier.rs

@@ -1221,10 +1221,10 @@ impl Compiler {
             }
             Fold => {
                 let (sn, _) = self.monadic_modifier_op(modified)?;
-                if sn.sig.args() <= sn.sig.outputs() {
+                if sn.sig.args() < sn.sig.outputs() {
                     self.experimental_error(&modified.modifier.span, || {
                         format!(
-                            "{} with arguments ≤ outputs is experimental. To use it, \
+                            "{} with arguments < outputs is experimental. To use it, \
                             add `# Experimental!` to the top of the file.",
                             prim.format()
                         )

tests/loops.ua

@@ -169,6 +169,8 @@ A ← ↯2_3_4⇡24
 ⍤⤙≍ [2222 0] [∧(⊃(+/+)⋅⋅∘) ♭₂ [1_1]_[10_10]_[100_100]_[1000_1000] 0 0]
 ⍤⤙≍ 60 ∧(+×) 10 1_2_3 0
 ⍤⤙≍ [] ∧(+) [1] []
+⍤⤙≍ [0 1 3 6] ∧(⊸+) [1 2 3 4] 0
+⍤⤙≍ [1 3 6 10] ∧(.+) [1 2 3 4] 0
 
 # Each pervasive
 ⍤⤙≍ ¯5 ∵¯ 5