From 5e418eeb5efac5b0714b7d6d872c33046d17a789 Mon Sep 17 00:00:00 2001
From: Shahar Papini <43779613+spapinistarkware@users.noreply.github.com>
Date: Mon, 25 Mar 2024 15:05:35 +0200
Subject: [PATCH] Prepare commitmentschem per size (#537)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

<!-- Reviewable:start -->
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/starkware-libs/stwo/537)
<!-- Reviewable:end -->
---
 src/core/fields/secure_column.rs    |  5 +++++
 src/core/fri.rs                     |  9 +++++++++
 src/core/poly/circle/mod.rs         |  2 +-
 src/core/poly/circle/secure_poly.rs | 23 +++++++++++++++++++++--
 4 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/src/core/fields/secure_column.rs b/src/core/fields/secure_column.rs
index 99b6ffdb9..831891b53 100644
--- a/src/core/fields/secure_column.rs
+++ b/src/core/fields/secure_column.rs
@@ -23,6 +23,11 @@ impl SecureColumn<CPUBackend> {
             .map(|c| &mut c[index])
             .assign(value.to_m31_array());
     }
+
+    // TODO(spapini): Remove when we no longer use CircleEvaluation<SecureField>.
+    pub fn to_cpu(&self) -> Vec<SecureField> {
+        (0..self.len()).map(|i| self.at(i)).collect()
+    }
 }
 impl<B: Backend> SecureColumn<B> {
     pub fn zeros(len: usize) -> Self {
diff --git a/src/core/fri.rs b/src/core/fri.rs
index 15422f5a4..52210749a 100644
--- a/src/core/fri.rs
+++ b/src/core/fri.rs
@@ -881,6 +881,15 @@ impl<F: ExtensionOf<BaseField>> SparseCircleEvaluation<F> {
     }
 }
 
+impl<'a, F: ExtensionOf<BaseField>> IntoIterator for &'a mut SparseCircleEvaluation<F> {
+    type Item = &'a mut CircleEvaluation<CPUBackend, F, BitReversedOrder>;
+    type IntoIter = std::slice::IterMut<'a, CircleEvaluation<CPUBackend, F, BitReversedOrder>>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.subcircle_evals.iter_mut()
+    }
+}
+
 /// Holds a small foldable subset of univariate SecureField polynomial evaluations.
 /// Evaluation is held at the CPU backend.
 #[derive(Debug, Clone)]
diff --git a/src/core/poly/circle/mod.rs b/src/core/poly/circle/mod.rs
index 903a1dcc5..76973d3d1 100644
--- a/src/core/poly/circle/mod.rs
+++ b/src/core/poly/circle/mod.rs
@@ -10,7 +10,7 @@ pub use domain::{CircleDomain, MAX_CIRCLE_DOMAIN_LOG_SIZE};
 pub use evaluation::{CircleEvaluation, CosetSubEvaluation};
 pub use ops::PolyOps;
 pub use poly::CirclePoly;
-pub use secure_poly::SecureCirclePoly;
+pub use secure_poly::{SecureCirclePoly, SecureEvaluation};
 
 #[cfg(test)]
 mod tests {
diff --git a/src/core/poly/circle/secure_poly.rs b/src/core/poly/circle/secure_poly.rs
index c2783afb2..a6d0060e8 100644
--- a/src/core/poly/circle/secure_poly.rs
+++ b/src/core/poly/circle/secure_poly.rs
@@ -1,9 +1,12 @@
 use std::ops::Deref;
 
-use crate::core::backend::cpu::CPUCirclePoly;
+use super::CircleDomain;
+use crate::core::backend::cpu::{CPUCircleEvaluation, CPUCirclePoly};
+use crate::core::backend::{Backend, CPUBackend};
 use crate::core::circle::CirclePoint;
 use crate::core::fields::qm31::SecureField;
-use crate::core::fields::secure_column::SECURE_EXTENSION_DEGREE;
+use crate::core::fields::secure_column::{SecureColumn, SECURE_EXTENSION_DEGREE};
+use crate::core::poly::BitReversedOrder;
 
 pub struct SecureCirclePoly(pub [CPUCirclePoly; SECURE_EXTENSION_DEGREE]);
 
@@ -24,6 +27,10 @@ impl SecureCirclePoly {
         ]
     }
 
+    pub fn log_size(&self) -> u32 {
+        self[0].log_size()
+    }
+
     /// Evaluates the polynomial at a point, given evaluations of its composing base field
     /// polynomials at that point.
     pub fn eval_from_partial_evals(evals: [SecureField; SECURE_EXTENSION_DEGREE]) -> SecureField {
@@ -42,3 +49,15 @@ impl Deref for SecureCirclePoly {
         &self.0
     }
 }
+
+pub struct SecureEvaluation<B: Backend> {
+    pub domain: CircleDomain,
+    pub values: SecureColumn<B>,
+}
+
+impl SecureEvaluation<CPUBackend> {
+    // TODO(spapini): Remove when we no longer use CircleEvaluation<SecureField>.
+    pub fn to_cpu(self) -> CPUCircleEvaluation<SecureField, BitReversedOrder> {
+        CPUCircleEvaluation::new(self.domain, self.values.to_cpu())
+    }
+}