From 0bf151247489a2b8ba0c6cbfa6bcd9607157519f Mon Sep 17 00:00:00 2001
From: Tamir Hemo <tamirhmo@gmail.com>
Date: Tue, 6 Feb 2024 15:45:10 -0800
Subject: [PATCH] refactor: move `assert_zero_ext` to `ExtensionBuilder` (#232)

refactor ext logic to ExtensionBuilder trait and make PermutationBuilder inherit from it
---
 air/src/air.rs | 96 +++++++++++++++++++++++++-------------------------
 1 file changed, 48 insertions(+), 48 deletions(-)

diff --git a/air/src/air.rs b/air/src/air.rs
index d5b075ed2..8e29b7d9b 100644
--- a/air/src/air.rs
+++ b/air/src/air.rs
@@ -106,51 +106,40 @@ pub trait AirBuilder: Sized {
         let x = x.into();
         self.assert_zero(x.clone() * (x - Self::Expr::one()));
     }
-
-    fn assert_zero_ext<ExprExt, I>(&mut self, x: I)
-    where
-        ExprExt: AbstractExtensionField<Self::Expr>,
-        I: Into<ExprExt>,
-    {
-        for xb in x.into().as_base_slice().iter().cloned() {
-            self.assert_zero(xb);
-        }
-    }
-
-    fn assert_eq_ext<ExprExt, I1, I2>(&mut self, x: I1, y: I2)
-    where
-        ExprExt: AbstractExtensionField<Self::Expr>,
-        I1: Into<ExprExt>,
-        I2: Into<ExprExt>,
-    {
-        self.assert_zero_ext::<ExprExt, ExprExt>(x.into() - y.into());
-    }
-
-    fn assert_one_ext<ExprExt, I>(&mut self, x: I)
-    where
-        ExprExt: AbstractExtensionField<Self::Expr>,
-        I: Into<ExprExt>,
-    {
-        let xe: ExprExt = x.into();
-        let parts = xe.as_base_slice();
-        self.assert_one(parts[0].clone());
-        for part in &parts[1..] {
-            self.assert_zero(part.clone());
-        }
-    }
 }
 
 pub trait PairBuilder: AirBuilder {
     fn preprocessed(&self) -> Self::M;
 }
 
-pub trait PermutationAirBuilder: AirBuilder {
+pub trait ExtensionBuilder: AirBuilder {
     type EF: ExtensionField<Self::F>;
 
     type ExprEF: AbstractExtensionField<Self::Expr, F = Self::EF>;
 
     type VarEF: Into<Self::ExprEF> + Copy;
 
+    fn assert_zero_ext<I>(&mut self, x: I)
+    where
+        I: Into<Self::ExprEF>;
+
+    fn assert_eq_ext<I1, I2>(&mut self, x: I1, y: I2)
+    where
+        I1: Into<Self::ExprEF>,
+        I2: Into<Self::ExprEF>,
+    {
+        self.assert_zero_ext(x.into() - y.into());
+    }
+
+    fn assert_one_ext<I>(&mut self, x: I)
+    where
+        I: Into<Self::ExprEF>,
+    {
+        self.assert_eq_ext(x, Self::ExprEF::one())
+    }
+}
+
+pub trait PermutationAirBuilder: ExtensionBuilder {
     type MP: MatrixRowSlices<Self::VarEF>;
 
     fn permutation(&self) -> Self::MP;
@@ -165,21 +154,6 @@ pub struct FilteredAirBuilder<'a, AB: AirBuilder> {
     condition: AB::Expr,
 }
 
-impl<'a, AB: PermutationAirBuilder> PermutationAirBuilder for FilteredAirBuilder<'a, AB> {
-    type EF = AB::EF;
-    type VarEF = AB::VarEF;
-    type ExprEF = AB::ExprEF;
-    type MP = AB::MP;
-
-    fn permutation(&self) -> Self::MP {
-        self.inner.permutation()
-    }
-
-    fn permutation_randomness(&self) -> &[Self::EF] {
-        self.inner.permutation_randomness()
-    }
-}
-
 impl<'a, AB: AirBuilder> AirBuilder for FilteredAirBuilder<'a, AB> {
     type F = AB::F;
     type Expr = AB::Expr;
@@ -207,6 +181,32 @@ impl<'a, AB: AirBuilder> AirBuilder for FilteredAirBuilder<'a, AB> {
     }
 }
 
+impl<'a, AB: ExtensionBuilder> ExtensionBuilder for FilteredAirBuilder<'a, AB> {
+    type EF = AB::EF;
+    type VarEF = AB::VarEF;
+    type ExprEF = AB::ExprEF;
+
+    fn assert_zero_ext<I>(&mut self, x: I)
+    where
+        I: Into<Self::ExprEF>,
+    {
+        self.inner
+            .assert_zero_ext(x.into() * self.condition.clone());
+    }
+}
+
+impl<'a, AB: PermutationAirBuilder> PermutationAirBuilder for FilteredAirBuilder<'a, AB> {
+    type MP = AB::MP;
+
+    fn permutation(&self) -> Self::MP {
+        self.inner.permutation()
+    }
+
+    fn permutation_randomness(&self) -> &[Self::EF] {
+        self.inner.permutation_randomness()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use p3_matrix::MatrixRowSlices;