From 8eae72b6f8911d772b06aff5f528f0b899fb2a24 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Matthias=20G=C3=B6rgens?= <matthias.goergens@gmail.com>
Date: Mon, 6 Jan 2025 15:39:17 +0800
Subject: [PATCH] Suggestions for 'Frontend and backend for building circuits'
 (#801)

Suggestions for https://github.com/scroll-tech/ceno/pull/799

Feel free to pick and choose from the suggestions. I talk about most of
them on your PR.

---------

Co-authored-by: dreamATD <tianyi.liu.08@gmail.com>
---
 gkr_iop/examples/multi_layer_logup.rs | 76 ++++++++++-----------------
 gkr_iop/src/lib.rs                    |  2 +-
 2 files changed, 29 insertions(+), 49 deletions(-)

diff --git a/gkr_iop/examples/multi_layer_logup.rs b/gkr_iop/examples/multi_layer_logup.rs
index 0988f210e..056bcbbe0 100644
--- a/gkr_iop/examples/multi_layer_logup.rs
+++ b/gkr_iop/examples/multi_layer_logup.rs
@@ -1,4 +1,4 @@
-use std::{marker::PhantomData, sync::Arc};
+use std::{marker::PhantomData, mem, sync::Arc};
 
 use ff_ext::ExtensionField;
 use gkr_iop::{
@@ -149,20 +149,10 @@ where
 {
     type Trace = TowerChipTrace;
 
-    fn phase1_witness(&self, phase1: &Self::Trace) -> Vec<Vec<E::BaseField>> {
+    fn phase1_witness(&self, phase1: Self::Trace) -> Vec<Vec<E::BaseField>> {
         let mut res = vec![vec![]; 2];
-        res[self.committed_table] = phase1
-            .table
-            .iter()
-            .cloned()
-            .map(E::BaseField::from)
-            .collect();
-        res[self.committed_count] = phase1
-            .count
-            .iter()
-            .cloned()
-            .map(E::BaseField::from)
-            .collect();
+        res[self.committed_table] = phase1.table.into_iter().map(E::BaseField::from).collect();
+        res[self.committed_count] = phase1.count.into_iter().map(E::BaseField::from).collect();
         res
     }
 
@@ -182,39 +172,29 @@ where
 
         // Compute den_0, den_1, num_0, num_1 for each layer.
         let updated_table = table.iter().map(|x| beta + x).collect_vec();
-        let mut last_den = vec![];
-        let mut last_num = vec![];
-
-        (0..self.params.height).for_each(|i| {
-            if i == 0 {
-                let (num_0, num_1): (Vec<E::BaseField>, Vec<E::BaseField>) =
-                    count.chunks(2).map(|chunk| (chunk[0], chunk[1])).unzip();
-                let (den_0, den_1): (Vec<E>, Vec<E>) = updated_table
-                    .chunks(2)
-                    .map(|chunk| (chunk[0], chunk[1]))
-                    .unzip();
-                (last_den, last_num) = izip!(&den_0, &den_1, &num_0, &num_1)
-                    .map(|(den_0, den_1, num_0, num_1)| {
-                        (*den_0 * den_1, *den_0 * num_1 + *den_1 * num_0)
-                    })
-                    .unzip();
-
-                layer_wits.push(LayerWitness::new(vec![num_0, num_1], vec![den_0, den_1]));
-            } else {
-                let (den_0, den_1): (Vec<E>, Vec<E>) =
-                    last_den.chunks(2).map(|chunk| (chunk[0], chunk[1])).unzip();
-                let (num_0, num_1): (Vec<E>, Vec<E>) =
-                    last_num.chunks(2).map(|chunk| (chunk[0], chunk[1])).unzip();
-
-                (last_den, last_num) = izip!(&den_0, &den_1, &num_0, &num_1)
-                    .map(|(den_0, den_1, num_0, num_1)| {
-                        (*den_0 * den_1, *den_0 * num_1 + *den_1 * num_0)
-                    })
-                    .unzip();
-
-                layer_wits.push(LayerWitness::new(vec![], vec![den_0, den_1, num_0, num_1]));
-            }
-        });
+
+        let (num_0, num_1): (Vec<E::BaseField>, Vec<E::BaseField>) = count.iter().tuples().unzip();
+        let (den_0, den_1): (Vec<E>, Vec<E>) = updated_table.into_iter().tuples().unzip();
+        let (mut last_den, mut last_num): (Vec<_>, Vec<_>) = izip!(&den_0, &den_1, &num_0, &num_1)
+            .map(|(&den_0, &den_1, &num_0, &num_1)| (den_0 * den_1, den_0 * num_1 + den_1 * num_0))
+            .unzip();
+
+        layer_wits.push(LayerWitness::new(vec![num_0, num_1], vec![den_0, den_1]));
+
+        layer_wits.extend((1..self.params.height).map(|_i| {
+            let (den_0, den_1): (Vec<E>, Vec<E>) =
+                mem::take(&mut last_den).into_iter().tuples().unzip();
+            let (num_0, num_1): (Vec<E>, Vec<E>) =
+                mem::take(&mut last_num).into_iter().tuples().unzip();
+
+            (last_den, last_num) = izip!(&den_0, &den_1, &num_0, &num_1)
+                .map(|(&den_0, &den_1, &num_0, &num_1)| {
+                    (den_0 * den_1, den_0 * num_1 + den_1 * num_0)
+                })
+                .unzip();
+
+            LayerWitness::new(vec![], vec![den_0, den_1, num_0, num_1])
+        }));
         layer_wits.reverse();
 
         GKRCircuitWitness { layers: layer_wits }
@@ -234,7 +214,7 @@ fn main() {
         let count = (0..1 << log_size)
             .map(|_| OsRng.gen_range(0..1 << log_size as u64))
             .collect_vec();
-        let phase1_witness = layout.phase1_witness(&TowerChipTrace { table, count });
+        let phase1_witness = layout.phase1_witness(TowerChipTrace { table, count });
 
         let mut prover_transcript = BasicTranscript::<E>::new(b"protocol");
 
diff --git a/gkr_iop/src/lib.rs b/gkr_iop/src/lib.rs
index e1c39ea4a..c5dcf203d 100644
--- a/gkr_iop/src/lib.rs
+++ b/gkr_iop/src/lib.rs
@@ -44,7 +44,7 @@ where
     type Trace;
 
     /// The vectors to be committed in the phase1.
-    fn phase1_witness(&self, phase1: &Self::Trace) -> Vec<Vec<E::BaseField>>;
+    fn phase1_witness(&self, phase1: Self::Trace) -> Vec<Vec<E::BaseField>>;
 
     /// GKR witness.
     fn gkr_witness(&self, phase1: &[Vec<E::BaseField>], challenges: &[E]) -> GKRCircuitWitness<E>