From 64a3154751be6cc51e595babaacf29b5d4c9d562 Mon Sep 17 00:00:00 2001
From: Kai Schmidt <kaikaliischmidt@gmail.com>
Date: Sun, 16 Jun 2024 12:36:22 -0700
Subject: [PATCH] make repeat use pervasion rules

---
 src/algorithm/loops.rs | 112 +++++++++++++++++++++++++++--------------
 tests/loops.ua         |   3 ++
 todo.md                |   1 -
 3 files changed, 78 insertions(+), 38 deletions(-)

diff --git a/src/algorithm/loops.rs b/src/algorithm/loops.rs
index 854b1fda7..c98ae1c21 100644
--- a/src/algorithm/loops.rs
+++ b/src/algorithm/loops.rs
@@ -10,6 +10,7 @@ use std::{
 use ecow::{eco_vec, EcoVec};
 
 use crate::{
+    algorithm::{fixed_rows, FixedRowsData},
     array::{Array, ArrayValue},
     cowslice::CowSlice,
     value::Value,
@@ -26,58 +27,95 @@ pub fn repeat(env: &mut Uiua) -> UiuaResult {
     crate::profile_function!();
     let f = env.pop_function()?;
     let n = env.pop("repetition count")?;
-    let n = match n {
-        Value::Num(n) => n,
-        Value::Byte(n) => n.convert(),
-        val => {
-            return Err(env.error(format!(
-                "Repetitions must be a scalar or list of \
-                natural numbers or infinity, \
-                but it is {}",
-                val.type_name_plural()
-            )))
-        }
-    };
+    fn rep_count(value: Value, env: &Uiua) -> UiuaResult<Array<f64>> {
+        Ok(match value {
+            Value::Num(n) => n,
+            Value::Byte(n) => n.convert(),
+            val => {
+                return Err(env.error(format!(
+                    "Repetitions must be a scalar or list of \
+                    natural numbers or infinity, \
+                    but it is {}",
+                    val.type_name_plural()
+                )))
+            }
+        })
+    }
     if n.rank() == 0 {
         // Scalar repeat
+        let n = rep_count(n, env)?;
         repeat_impl(f, n.data[0], env)
     } else {
-        // Non-scalar repeat
+        // Array
         let sig = f.signature();
         // Collect arguments
-        let mut args_rows: Vec<_> = Vec::with_capacity(sig.args);
+        let mut args = Vec::with_capacity(sig.args + 1);
+        let new_shape = n.shape().clone();
+        args.push(n);
         for i in 0..sig.args {
             let arg = env.pop(i + 1)?;
-            if !arg.shape().starts_with(n.shape()) {
-                return Err(env.error(format!(
-                    "The repetitions' shape {} is not compatible \
-                    with the argument {}'s shape {}",
-                    n.shape(),
-                    i + 1,
-                    arg.shape(),
-                )));
-            }
-            let row_shape = Shape::from(&arg.shape()[n.rank()..]);
-            args_rows.push(arg.into_row_shaped_slices(row_shape));
-        }
-        args_rows.reverse();
-        let mut outputs = multi_output(sig.outputs, Vec::with_capacity(n.data.len()));
-        // Repeate with each repetition count
-        for n in n.data {
-            for arg in args_rows.iter_mut().rev() {
-                env.push(arg.next().unwrap());
+            args.push(arg);
+        }
+        args[1..].reverse();
+        let FixedRowsData {
+            mut rows,
+            row_count,
+            ..
+        } = fixed_rows(Primitive::Repeat.format(), sig.outputs, args, env)?;
+
+        // Switch with each selector element
+        let mut outputs = multi_output(sig.outputs, Vec::new());
+        let mut rows_to_sel = Vec::with_capacity(sig.args);
+        for _ in 0..row_count {
+            let n = rep_count(
+                match &mut rows[0] {
+                    Ok(n) => n.next().unwrap(),
+                    Err(n) => n.clone(),
+                },
+                env,
+            )?;
+            // println!("ns: {} {:?}", n.shape, n.data);
+            rows_to_sel.clear();
+            for row in rows[1..].iter_mut() {
+                let row = match row {
+                    Ok(row) => row.next().unwrap(),
+                    Err(row) => row.clone(),
+                };
+                // println!("row: {:?}", row);
+                if n.rank() > row.rank() {
+                    rows_to_sel.push(Err(row));
+                } else {
+                    let row_shape = row.shape()[n.rank()..].into();
+                    rows_to_sel.push(Ok(row.into_row_shaped_slices(row_shape)));
+                }
             }
-            repeat_impl(f.clone(), n, env)?;
-            for i in 0..sig.outputs {
-                outputs[i].push(env.pop("repeat's output")?);
+            for sel_row_slice in n.row_slices() {
+                for &elem in sel_row_slice {
+                    // println!("  elem: {}", elem);
+                    for row in &mut rows_to_sel {
+                        let row = match row {
+                            Ok(row) => row.next().unwrap(),
+                            Err(row) => row.clone(),
+                        };
+                        // println!("  row: {:?}", row);
+                        env.push(row);
+                    }
+                    repeat_impl(f.clone(), elem, env)?;
+                    for i in 0..sig.outputs {
+                        let res = env.pop("repeat output")?;
+                        // println!("    res: {:?}", res);
+                        outputs[i].push(res);
+                    }
+                }
             }
         }
-        // Collect outputs
+        // Collect output
         for output in outputs.into_iter().rev() {
             let mut new_value = Value::from_row_values(output, env)?;
-            let mut new_shape = n.shape.clone();
+            let mut new_shape = new_shape.clone();
             new_shape.extend_from_slice(&new_value.shape()[1..]);
             *new_value.shape_mut() = new_shape;
+            new_value.validate_shape();
             env.push(new_value);
         }
         Ok(())
diff --git a/tests/loops.ua b/tests/loops.ua
index 4c78b91d8..f6e4597ef 100644
--- a/tests/loops.ua
+++ b/tests/loops.ua
@@ -224,6 +224,9 @@ A ← ↯2_3_4⇡24
 ⍤⟜≍: [10_20 40_80] ⍥(×2) [1_2 3_4] [5_5 5_5]
 ⍤⟜≍: [10_10 20_20] ⍥(×2) [1 2] [5_5 5_5]
 ⍤⟜≍: [] ⊙◌[⍥(×2.)] 0 1
+⍤⟜≍: [+1.+1. 6_7_8] ⍥(+1) [+1.+1. 1_2_3] 5
+⍤⟜≍: [+1.+1. 2_4_6] ⍥(+1) [+1.+1. 1_2_3] ¤[1 2 3]
+⍤⟜≍: [..4_5_6] ⍥(+1) [+1.+1. 1_2_3] [1 2 3]
 
 # Do
 ⍤⟜≍: 1024 ⍢(×2)(<1000) 1
diff --git a/todo.md b/todo.md
index 54237db66..d030283bb 100644
--- a/todo.md
+++ b/todo.md
@@ -1,7 +1,6 @@
 # Uiua Todo
 
 - 0.12
-  - Pervasive `repeat`
   - `# Track caller!`
   - Stack-source locality tutorial
   - Axis swizzles?