Merge branch 'kevjue/field_extension' into kevjue/recursive_verifier_…

…take_two

core/src/air/extension.rs

@@ -0,0 +1,109 @@
+use std::{
+    mem::size_of,
+    ops::{Add, Mul},
+};
+
+use core::borrow::{Borrow, BorrowMut};
+use itertools::Itertools;
+use p3_field::{
+    extension::{BinomialExtensionField, BinomiallyExtendable},
+    field_to_array, AbstractExtensionField, AbstractField, Field,
+};
+use sp1_derive::AlignedBorrow;
+
+use super::SP1AirBuilder;
+
+pub const DEGREE: usize = 4;
+
+#[derive(AlignedBorrow, Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
+#[repr(C)]
+pub struct Extension<T>(pub [T; DEGREE]); // Degree 4 is hard coded for now.  TODO:  Change to a const generic
+
+impl<E: AbstractField> Extension<E> {
+    // Returns the one element of the extension field
+    pub fn one<AB: SP1AirBuilder<Expr = E>>() -> Extension<AB::Expr>
+    where
+        AB::Expr: AbstractField,
+    {
+        let one = AB::Expr::one();
+        Extension(field_to_array(one))
+    }
+
+    // Converts a field element to extension element
+    pub fn from<AB: SP1AirBuilder<Expr = E>>(x: E) -> Extension<AB::Expr> {
+        Extension(field_to_array(x))
+    }
+
+    // Negates an extension field Element
+    pub fn neg<AB: SP1AirBuilder<Expr = E>>(self) -> Extension<AB::Expr> {
+        Extension(self.0.map(|x| AB::Expr::zero() - x))
+    }
+
+    // Adds an extension field element
+    pub fn add<AB: SP1AirBuilder<Expr = E>>(self, rhs: &Self) -> Extension<AB::Expr>
+    where
+        E: Add<E, Output = AB::Expr>,
+    {
+        let mut elements = Vec::new();
+
+        for (e1, e2) in self.0.into_iter().zip_eq(rhs.0.clone().into_iter()) {
+            elements.push(e1 + e2);
+        }
+
+        Extension(elements.try_into().unwrap())
+    }
+
+    // Subtracts an extension field element
+    pub fn sub<AB: SP1AirBuilder<Expr = E>>(self, rhs: &Self) -> Extension<AB::Expr>
+    where
+        E: Add<E, Output = AB::Expr>,
+    {
+        let mut elements = Vec::new();
+
+        for (e1, e2) in self.0.into_iter().zip_eq(rhs.0.clone().into_iter()) {
+            elements.push(e1 - e2);
+        }
+
+        Extension(elements.try_into().unwrap())
+    }
+
+    // Multiplies an extension field element
+    pub fn mul<AB: SP1AirBuilder<Expr = E>>(self, rhs: &Self) -> Extension<AB::Expr>
+    where
+        E: Mul<E, Output = AB::Expr>,
+    {
+        let mut elements = Vec::new();
+
+        for (e1, e2) in self.0.into_iter().zip_eq(rhs.0.clone().into_iter()) {
+            elements.push(e1 * e2);
+        }
+
+        Extension(elements.try_into().unwrap())
+    }
+
+    pub fn as_base_slice(&self) -> &[E] {
+        &self.0
+    }
+}
+
+impl<V> Extension<V> {
+    // Converts a field element with var base elements to one with expr base elements.
+    pub fn from_var<AB: SP1AirBuilder<Var = V>>(self) -> Extension<AB::Expr>
+    where
+        V: Into<AB::Expr>,
+    {
+        Extension(self.0.map(|x| x.into()))
+    }
+}
+
+impl<F> From<BinomialExtensionField<F, 4>> for Extension<F>
+where
+    F: Field,
+    F::F: BinomiallyExtendable<4>,
+{
+    fn from(value: BinomialExtensionField<F, 4>) -> Self {
+        let base_slice = value.as_base_slice();
+
+        Self(base_slice.try_into().unwrap())
+    }
+}

core/src/air/mod.rs

@@ -1,11 +1,13 @@
 mod builder;
+mod extension;
 mod interaction;
 mod machine;
 mod polynomial;
 mod sub_builder;
 mod word;
 
 pub use builder::*;
+pub use extension::*;
 pub use interaction::*;
 pub use machine::*;
 pub use polynomial::*;

core/src/operations/div_extension.rs

@@ -0,0 +1,55 @@
+//! An operation to performce div on the inputs.
+//!
+use core::borrow::Borrow;
+use core::borrow::BorrowMut;
+use p3_field::extension::BinomialExtensionField;
+use p3_field::extension::BinomiallyExtendable;
+use sp1_derive::AlignedBorrow;
+use std::mem::size_of;
+
+use crate::air::Extension;
+use crate::air::SP1AirBuilder;
+use crate::air::DEGREE;
+
+use super::IsEqualExtOperation;
+
+/// A set of columns needed to compute whether the given word is 0.
+#[derive(AlignedBorrow, Default, Debug, Clone, Copy)]
+#[repr(C)]
+pub struct DivExtOperation<T> {
+    pub is_equal: IsEqualExtOperation<T>,
+
+    /// Result is the quotient
+    pub result: Extension<T>,
+}
+
+impl<F: BinomiallyExtendable<DEGREE>> DivExtOperation<F> {
+    pub fn populate(
+        &mut self,
+        a: BinomialExtensionField<F, DEGREE>,
+        b: BinomialExtensionField<F, DEGREE>,
+    ) -> BinomialExtensionField<F, DEGREE> {
+        let result = a / b;
+        self.result = result.into();
+
+        let product = b * result;
+        self.is_equal.populate(a, product);
+
+        result
+    }
+
+    pub fn eval<AB: SP1AirBuilder>(
+        builder: &mut AB,
+        a: Extension<AB::Expr>,
+        b: Extension<AB::Expr>,
+        cols: DivExtOperation<AB::Var>,
+        is_real: AB::Expr,
+    ) where
+        AB::F: BinomiallyExtendable<DEGREE>,
+    {
+        builder.assert_bool(is_real.clone());
+
+        let product = b.mul::<AB>(&cols.result.from_var::<AB>());
+        IsEqualExtOperation::<AB::F>::eval(builder, a, product, cols.is_equal, is_real.clone());
+    }
+}

core/src/operations/is_equal_extension.rs

@@ -0,0 +1,56 @@
+use core::borrow::Borrow;
+use core::borrow::BorrowMut;
+use p3_field::extension::BinomialExtensionField;
+use p3_field::extension::BinomiallyExtendable;
+use sp1_derive::AlignedBorrow;
+use std::mem::size_of;
+
+use crate::air::Extension;
+use crate::air::SP1AirBuilder;
+use crate::air::DEGREE;
+
+use super::IsZeroExtOperation;
+
+/// A set of columns needed to compute the equality of two field extension elements.
+#[derive(AlignedBorrow, Default, Debug, Clone, Copy)]
+#[repr(C)]
+pub struct IsEqualExtOperation<T> {
+    /// An operation to check whether the differences in field extension elements is zero.
+    pub is_diff_zero: IsZeroExtOperation<T>,
+}
+
+impl<F: BinomiallyExtendable<DEGREE>> IsEqualExtOperation<F> {
+    pub fn populate(
+        &mut self,
+        a: BinomialExtensionField<F, DEGREE>,
+        b: BinomialExtensionField<F, DEGREE>,
+    ) -> u32 {
+        let diff = a - b;
+        self.is_diff_zero.populate(diff);
+        (a == b) as u32
+    }
+
+    pub fn eval<AB: SP1AirBuilder>(
+        builder: &mut AB,
+        a: Extension<AB::Expr>,
+        b: Extension<AB::Expr>,
+        cols: IsEqualExtOperation<AB::Var>,
+        is_real: AB::Expr,
+    ) where
+        AB::F: BinomiallyExtendable<DEGREE>,
+    {
+        builder.assert_bool(is_real.clone());
+
+        // Calculate differences.
+        let diff = a.sub::<AB>(&b);
+
+        // Check if the difference is 0.
+        IsZeroExtOperation::<AB::F>::eval(builder, diff, cols.is_diff_zero, is_real.clone());
+
+        // Degree 3 constraint to avoid "OodEvaluationMismatch".
+        builder.assert_zero(
+            is_real.clone() * is_real.clone() * is_real.clone()
+                - is_real.clone() * is_real.clone() * is_real.clone(),
+        );
+    }
+}

core/src/operations/is_zero_extension.rs

@@ -0,0 +1,93 @@
+//! An operation to check if the input word is 0.
+//!
+//! This is bijective (i.e., returns 1 if and only if the input is 0). It is also worth noting that
+//! this operation doesn't do a range check.
+use core::borrow::Borrow;
+use core::borrow::BorrowMut;
+use p3_air::AirBuilder;
+use p3_field::extension::BinomialExtensionField;
+use p3_field::extension::BinomiallyExtendable;
+use p3_field::AbstractExtensionField;
+use sp1_derive::AlignedBorrow;
+use std::mem::size_of;
+
+use crate::air::Extension;
+use crate::air::SP1AirBuilder;
+use crate::air::DEGREE;
+
+use super::IsZeroOperation;
+
+/// A set of columns needed to compute whether the given field ext element is 0.
+#[derive(AlignedBorrow, Default, Debug, Clone, Copy)]
+#[repr(C)]
+pub struct IsZeroExtOperation<T> {
+    /// `IsZeroOperation` to check if each base field element in the input field ext element is zero.
+    pub is_zero_base_element: [IsZeroOperation<T>; DEGREE],
+
+    /// A boolean flag indicating whether the first and second base field elements are 0.
+    /// This equals `is_zero_byte[0] * is_zero_byte[1]`.
+    pub is_lower_half_zero: T,
+
+    /// A boolean flag indicating whether the third and fourth base field elements are 0.
+    pub is_upper_half_zero: T,
+
+    /// A boolean flag indicating whether the field ext element is zero. This equals `is_zero_byte[0] * ... *
+    /// is_zero_byte[DEGREE - 1]`.
+    pub result: T,
+}
+
+impl<F: BinomiallyExtendable<DEGREE>> IsZeroExtOperation<F> {
+    pub fn populate(&mut self, a: BinomialExtensionField<F, DEGREE>) -> u32 {
+        let mut is_zero = true;
+        let base_slice = a.as_base_slice();
+        for i in 0..DEGREE {
+            is_zero &= self.is_zero_base_element[i].populate_from_field_element(base_slice[i]) == 1;
+        }
+        self.is_lower_half_zero =
+            self.is_zero_base_element[0].result * self.is_zero_base_element[1].result;
+        self.is_upper_half_zero =
+            self.is_zero_base_element[2].result * self.is_zero_base_element[3].result;
+        self.result = F::from_bool(is_zero);
+        is_zero as u32
+    }
+
+    pub fn eval<AB: SP1AirBuilder>(
+        builder: &mut AB,
+        a: Extension<AB::Expr>,
+        cols: IsZeroExtOperation<AB::Var>,
+        is_real: AB::Expr,
+    ) {
+        let base_slice = a.as_base_slice();
+
+        // Calculate whether each byte is 0.
+        for i in 0..DEGREE {
+            IsZeroOperation::<AB::F>::eval(
+                builder,
+                base_slice[i].clone(),
+                cols.is_zero_base_element[i],
+                is_real.clone(),
+            );
+        }
+
+        // From here, we only assert when is_real is true.
+        builder.assert_bool(is_real.clone());
+        let mut builder_is_real = builder.when(is_real.clone());
+
+        // Calculate is_upper_half_zero and is_lower_half_zero and finally the result.
+        builder_is_real.assert_bool(cols.is_lower_half_zero);
+        builder_is_real.assert_bool(cols.is_upper_half_zero);
+        builder_is_real.assert_bool(cols.result);
+        builder_is_real.assert_eq(
+            cols.is_lower_half_zero,
+            cols.is_zero_base_element[0].result * cols.is_zero_base_element[1].result,
+        );
+        builder_is_real.assert_eq(
+            cols.is_upper_half_zero,
+            cols.is_zero_base_element[2].result * cols.is_zero_base_element[3].result,
+        );
+        builder_is_real.assert_eq(
+            cols.result,
+            cols.is_lower_half_zero * cols.is_upper_half_zero,
+        );
+    }
+}

core/src/operations/mod.rs

@@ -8,11 +8,14 @@ mod add;
 mod add4;
 mod add5;
 mod and;
+mod div_extension;
 pub mod field;
 mod fixed_rotate_right;
 mod fixed_shift_right;
+mod is_equal_extension;
 mod is_equal_word;
 mod is_zero;
+mod is_zero_extension;
 mod is_zero_word;
 mod not;
 mod or;
@@ -24,8 +27,10 @@ pub use add5::*;
 pub use and::*;
 pub use fixed_rotate_right::*;
 pub use fixed_shift_right::*;
+pub use is_equal_extension::*;
 pub use is_equal_word::*;
 pub use is_zero::*;
+pub use is_zero_extension::*;
 pub use is_zero_word::*;
 pub use not::*;
 pub use or::*;

core/src/syscall/precompiles/sha256/compress/air.rs

@@ -26,7 +26,7 @@ where
         let local: &ShaCompressCols<AB::Var> = main.row_slice(0).borrow();
         let next: &ShaCompressCols<AB::Var> = main.row_slice(1).borrow();
 
-        self.contrain_control_flow_flags(builder, local, next);
+        self.constrain_control_flow_flags(builder, local, next);
 
         self.constrain_memory(builder, local);
 
@@ -37,7 +37,7 @@ where
 }
 
 impl ShaCompressChip {
-    fn contrain_control_flow_flags<AB: SP1AirBuilder>(
+    fn constrain_control_flow_flags<AB: SP1AirBuilder>(
         &self,
         builder: &mut AB,
         local: &ShaCompressCols<AB::Var>,