From 39e058228bd99aca5f3f4c230c14da57cbb48548 Mon Sep 17 00:00:00 2001
From: Gabriel Barreto <gabriel.aquino.barreto@gmail.com>
Date: Sat, 27 Jul 2024 08:19:40 -0300
Subject: [PATCH] REPL more info on iterations

---
 src/lair/execute.rs   | 20 ++++++++++++++++++++
 src/lair/func_chip.rs |  2 +-
 src/lurk/cli/repl.rs  | 14 +++++++++-----
 3 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/src/lair/execute.rs b/src/lair/execute.rs
index 63f530b6..5cf64b3a 100644
--- a/src/lair/execute.rs
+++ b/src/lair/execute.rs
@@ -370,6 +370,26 @@ impl<F: PrimeField32> QueryRecord<F> {
     pub fn expect_public_values(&self) -> &[F] {
         self.public_values.as_ref().expect("Public values not set")
     }
+
+    pub fn stats<C1: Chipset<F>, C2: Chipset<F>>(
+        &self,
+        toplevel: &Toplevel<F, C1, C2>,
+    ) -> (usize, usize) {
+        let mut steps = 0;
+        let mut width = 0;
+        for (i, queries) in self.func_queries.iter().enumerate() {
+            steps += queries.len();
+            let func = toplevel.func_by_index(i);
+            let func_width = func.compute_layout_sizes(toplevel).total();
+            width += queries.len() * func_width;
+        }
+        for (i, queries) in self.mem_queries.iter().enumerate() {
+            steps += queries.len();
+            let mem_width = 4 + mem_index_to_len(i);
+            width += queries.len() * mem_width;
+        }
+        (steps, width)
+    }
 }
 
 impl<F: PrimeField32, C1: Chipset<F>, C2: Chipset<F>> Toplevel<F, C1, C2> {
diff --git a/src/lair/func_chip.rs b/src/lair/func_chip.rs
index 9b000215..350f359c 100644
--- a/src/lair/func_chip.rs
+++ b/src/lair/func_chip.rs
@@ -20,7 +20,7 @@ pub type LayoutSizes = ColumnLayout<usize, usize>;
 
 impl LayoutSizes {
     #[inline]
-    fn total(&self) -> usize {
+    pub fn total(&self) -> usize {
         self.nonce + self.input + self.aux + self.sel + self.output
     }
 }
diff --git a/src/lurk/cli/repl.rs b/src/lurk/cli/repl.rs
index 6768cfb8..9dd301f1 100644
--- a/src/lurk/cli/repl.rs
+++ b/src/lurk/cli/repl.rs
@@ -206,11 +206,15 @@ impl<C1: Chipset<BabyBear>, C2: Chipset<BabyBear>> Repl<BabyBear, C1, C2> {
     }
 }
 
-fn pretty_iterations_display(iterations: usize) -> String {
+fn pretty_iterations_display(iterations: usize, width: usize) -> String {
     if iterations != 1 {
-        format!("{iterations} iterations")
+        let average_width = width as f64 / iterations as f64;
+        format!(
+            "{iterations} iterations, {width} total width, average step size {:.2}",
+            average_width,
+        )
     } else {
-        "1 iteration".into()
+        format!("1 iteration, width {width}")
     }
 }
 
@@ -492,10 +496,10 @@ impl<F: PrimeField32, C1: Chipset<F>, C2: Chipset<F>> Repl<F, C1, C2> {
         let env = env.unwrap_or(self.env);
         let result = self.reduce_with_env(expr, &env)?;
         self.memoize_dag(result.tag, &result.digest);
-        let iterations = self.queries.func_queries[self.func_indices.eval].len();
+        let (iterations, width) = self.queries.stats(&self.toplevel);
         println!(
             "[{}] => {}",
-            pretty_iterations_display(iterations),
+            pretty_iterations_display(iterations, width),
             self.fmt(&result)
         );
         Ok(result)