From 07f9121858132edd27ddb793738d59135427da3e Mon Sep 17 00:00:00 2001
From: Yuncong Zhang <yczhangsjtu@gmail.com>
Date: Wed, 6 Nov 2024 08:39:03 +0800
Subject: [PATCH 1/3] Add and reimplement some utility functions.

---
 mpcs/src/util.rs | 55 +++++++++++++++++++++++++++---------------------
 1 file changed, 31 insertions(+), 24 deletions(-)

diff --git a/mpcs/src/util.rs b/mpcs/src/util.rs
index 80ecf2535..c70b269a3 100644
--- a/mpcs/src/util.rs
+++ b/mpcs/src/util.rs
@@ -6,7 +6,7 @@ pub mod plonky2_util;
 use ff::{Field, PrimeField};
 use ff_ext::ExtensionField;
 use goldilocks::SmallField;
-use itertools::{Itertools, izip};
+use itertools::{Either, Itertools, izip};
 use multilinear_extensions::mle::{DenseMultilinearExtension, FieldType};
 use serde::{Deserialize, Serialize, de::DeserializeOwned};
 pub mod merkle_tree;
@@ -148,38 +148,45 @@ pub fn field_type_index_set_ext<E: ExtensionField>(
     }
 }
 
-pub struct FieldTypeIterExt<'a, E: ExtensionField> {
-    inner: &'a FieldType<E>,
-    index: usize,
+pub fn poly_iter_ext<E: ExtensionField>(
+    poly: &DenseMultilinearExtension<E>,
+) -> impl Iterator<Item = E> + '_ {
+    field_type_iter_ext(&poly.evaluations)
 }
 
-impl<'a, E: ExtensionField> Iterator for FieldTypeIterExt<'a, E> {
-    type Item = E;
+pub fn field_type_iter_ext<E: ExtensionField>(
+    evaluations: &FieldType<E>,
+) -> impl Iterator<Item = E> + '_ {
+    match evaluations {
+        FieldType::Ext(coeffs) => Either::Left(coeffs.iter().copied()),
+        FieldType::Base(coeffs) => Either::Right(coeffs.iter().map(|x| (*x).into())),
+        _ => unreachable!(),
+    }
+}
 
-    fn next(&mut self) -> Option<Self::Item> {
-        if self.index >= self.inner.len() {
-            None
-        } else {
-            let res = field_type_index_ext(self.inner, self.index);
-            self.index += 1;
-            Some(res)
-        }
+pub fn field_type_to_ext_vec<E: ExtensionField>(evaluations: &FieldType<E>) -> Vec<E> {
+    match evaluations {
+        FieldType::Ext(coeffs) => coeffs.to_vec(),
+        FieldType::Base(coeffs) => coeffs.iter().map(|x| (*x).into()).collect(),
+        _ => unreachable!(),
     }
 }
 
-pub fn poly_iter_ext<E: ExtensionField>(
-    poly: &DenseMultilinearExtension<E>,
-) -> FieldTypeIterExt<E> {
-    FieldTypeIterExt {
-        inner: &poly.evaluations,
-        index: 0,
+pub fn field_type_as_ext<E: ExtensionField>(values: &FieldType<E>) -> &Vec<E> {
+    match values {
+        FieldType::Ext(coeffs) => coeffs,
+        FieldType::Base(_) => panic!("Expected base field"),
+        _ => unreachable!(),
     }
 }
 
-pub fn field_type_iter_ext<E: ExtensionField>(evaluations: &FieldType<E>) -> FieldTypeIterExt<E> {
-    FieldTypeIterExt {
-        inner: evaluations,
-        index: 0,
+pub fn field_type_iter_base<E: ExtensionField>(
+    values: &FieldType<E>,
+) -> impl Iterator<Item = &E::BaseField> + '_ {
+    match values {
+        FieldType::Ext(coeffs) => Either::Left(coeffs.iter().flat_map(|x| x.as_bases())),
+        FieldType::Base(coeffs) => Either::Right(coeffs.iter()),
+        _ => unreachable!(),
     }
 }
 

From a394572260986f05b24a61eb41fa67c0790432f4 Mon Sep 17 00:00:00 2001
From: Yuncong Zhang <yczhangsjtu@gmail.com>
Date: Wed, 6 Nov 2024 08:43:23 +0800
Subject: [PATCH 2/3] Fix typo

---
 mpcs/src/util.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mpcs/src/util.rs b/mpcs/src/util.rs
index c70b269a3..1c2b2ea53 100644
--- a/mpcs/src/util.rs
+++ b/mpcs/src/util.rs
@@ -175,7 +175,7 @@ pub fn field_type_to_ext_vec<E: ExtensionField>(evaluations: &FieldType<E>) -> V
 pub fn field_type_as_ext<E: ExtensionField>(values: &FieldType<E>) -> &Vec<E> {
     match values {
         FieldType::Ext(coeffs) => coeffs,
-        FieldType::Base(_) => panic!("Expected base field"),
+        FieldType::Base(_) => panic!("Expected ext field"),
         _ => unreachable!(),
     }
 }

From c5364cd06a5b454901b2e4a829954af9dd09e6a6 Mon Sep 17 00:00:00 2001
From: Yuncong Zhang <yczhangsjtu@gmail.com>
Date: Wed, 6 Nov 2024 17:48:20 +0800
Subject: [PATCH 3/3] Update according to suggestion.

---
 mpcs/src/util.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mpcs/src/util.rs b/mpcs/src/util.rs
index 1c2b2ea53..7688b53ec 100644
--- a/mpcs/src/util.rs
+++ b/mpcs/src/util.rs
@@ -167,7 +167,7 @@ pub fn field_type_iter_ext<E: ExtensionField>(
 pub fn field_type_to_ext_vec<E: ExtensionField>(evaluations: &FieldType<E>) -> Vec<E> {
     match evaluations {
         FieldType::Ext(coeffs) => coeffs.to_vec(),
-        FieldType::Base(coeffs) => coeffs.iter().map(|x| (*x).into()).collect(),
+        FieldType::Base(coeffs) => coeffs.iter().map(|&x| x.into()).collect(),
         _ => unreachable!(),
     }
 }