Add multilinear extension (#565)

Cargo.toml

@@ -8,6 +8,7 @@ edition = "2021"
 [dependencies]
 blake2 = "0.10.6"
 blake3 = "1.5.0"
+derivative = "2.2.0"
 hex = "0.4.3"
 itertools = "0.12.0"
 num-traits = "0.2.17"

src/core/backend/cpu/lookups/mle.rs

@@ -0,0 +1,38 @@
+use std::iter::zip;
+
+use crate::core::backend::CPUBackend;
+use crate::core::fields::m31::BaseField;
+use crate::core::fields::qm31::SecureField;
+use crate::core::lookups::mle::{Mle, MleOps};
+
+impl MleOps<BaseField> for CPUBackend {
+    fn fix_first(mle: Mle<Self, BaseField>, assignment: SecureField) -> Mle<Self, SecureField> {
+        let midpoint = mle.len() / 2;
+        let (lhs_evals, rhs_evals) = mle.split_at(midpoint);
+
+        let res = zip(lhs_evals, rhs_evals)
+            // Equivalent to `eq(0, assignment) * lhs_eval + eq(1, assignment) * rhs_eval`.
+            .map(|(&lhs_eval, &rhs_eval)| assignment * (rhs_eval - lhs_eval) + lhs_eval)
+            .collect();
+
+        Mle::new(res)
+    }
+}
+
+impl MleOps<SecureField> for CPUBackend {
+    fn fix_first(mle: Mle<Self, SecureField>, assignment: SecureField) -> Mle<Self, SecureField> {
+        let midpoint = mle.len() / 2;
+        let mut evals = mle.into_evals();
+
+        for i in 0..midpoint {
+            let lhs_eval = evals[i];
+            let rhs_eval = evals[i + midpoint];
+            // Equivalent to `eq(0, assignment) * lhs_eval + eq(1, assignment) * rhs_eval`.
+            evals[i] = lhs_eval + assignment * (rhs_eval - lhs_eval);
+        }
+
+        evals.truncate(midpoint);
+
+        Mle::new(evals)
+    }
+}

src/core/backend/cpu/lookups/mod.rs

@@ -0,0 +1 @@
+mod mle;

src/core/backend/cpu/mod.rs

@@ -2,12 +2,14 @@ mod accumulation;
 mod blake2s;
 mod circle;
 mod fri;
+mod lookups;
 pub mod quotients;
 
 use std::fmt::Debug;
 
 use super::{Backend, Column, ColumnOps, FieldOps};
 use crate::core::fields::Field;
+use crate::core::lookups::mle::Mle;
 use crate::core::poly::circle::{CircleEvaluation, CirclePoly};
 use crate::core::utils::bit_reverse;
 
@@ -49,6 +51,7 @@ impl<T: Debug + Clone + Default> Column<T> for Vec<T> {
 
 pub type CPUCirclePoly = CirclePoly<CPUBackend>;
 pub type CPUCircleEvaluation<F, EvalOrder> = CircleEvaluation<CPUBackend, F, EvalOrder>;
+pub type CPUMle<F> = Mle<CPUBackend, F>;
 
 #[cfg(test)]
 mod tests {

src/core/lookups/mle.rs

@@ -0,0 +1,60 @@
+use std::ops::Deref;
+
+use derivative::Derivative;
+
+use crate::core::backend::{Col, Column, ColumnOps};
+use crate::core::fields::qm31::SecureField;
+use crate::core::fields::Field;
+
+/// TODO
+pub trait MleOps<F: Field>: ColumnOps<F> + Sized {
+    /// Returns a transformed [`Mle`] where the first variable is fixed to `assignment`.
+    fn fix_first(mle: Mle<Self, F>, assignment: SecureField) -> Mle<Self, SecureField>
+    where
+        Self: MleOps<SecureField>;
+}
+
+/// Multilinear Extension stored as evaluations of a multilinear polynomial over the boolean
+/// hypercube in bit-reversed order.
+#[derive(Derivative)]
+#[derivative(Debug(bound = ""), Clone(bound = ""))]
+pub struct Mle<B: ColumnOps<F>, F: Field> {
+    evals: Col<B, F>,
+}
+
+impl<B: MleOps<F>, F: Field> Mle<B, F> {
+    /// Creates a [`Mle`] from evaluations of a multilinear polynomial on the boolean hypercube.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the number of evaluations is not a power of two.
+    pub fn new(evals: Col<B, F>) -> Self {
+        assert!(evals.len().is_power_of_two());
+        Self { evals }
+    }
+
+    pub fn into_evals(self) -> Col<B, F> {
+        self.evals
+    }
+
+    /// Returns a transformed polynomial where the first variable is fixed to `assignment`.
+    pub fn fix_first(self, assignment: SecureField) -> Mle<B, SecureField>
+    where
+        B: MleOps<SecureField>,
+    {
+        B::fix_first(self, assignment)
+    }
+
+    /// Returns the number of variables in the polynomial.
+    pub fn n_variables(&self) -> usize {
+        self.evals.len().ilog2() as usize
+    }
+}
+
+impl<B: ColumnOps<F>, F: Field> Deref for Mle<B, F> {
+    type Target = Col<B, F>;
+
+    fn deref(&self) -> &Col<B, F> {
+        &self.evals
+    }
+}

src/core/lookups/mod.rs

@@ -0,0 +1 @@
+pub mod mle;

src/core/mod.rs

@@ -9,6 +9,7 @@ pub mod constraints;
 pub mod fft;
 pub mod fields;
 pub mod fri;
+pub mod lookups;
 pub mod poly;
 pub mod proof_of_work;
 pub mod prover;