feat(boojum): added WitnessVarLengthEncodable implementation to tower…

… of extensions
matter-labs · Sep 5, 2024 · 6806fe4 · 6806fe4
1 parent ef7db18
commit 6806fe4
Show file tree

Hide file tree

Showing 4 changed files with 77 additions and 4 deletions.
diff --git a/crates/boojum/src/gadgets/non_native_field/implementations/implementation_u16.rs b/crates/boojum/src/gadgets/non_native_field/implementations/implementation_u16.rs
@@ -1088,6 +1088,18 @@ impl<F: SmallField, T: pairing::ff::PrimeField, const N: usize> CircuitVarLength
     }
 }
 
+impl<F: SmallField, T: pairing::ff::PrimeField, const N: usize> WitnessVarLengthEncodable<F>
+    for NonNativeFieldOverU16<F, T, N>
+{
+    fn witness_encoding_length(_witness: &Self::Witness) -> usize {
+        N
+    }
+
+    fn encode_witness_to_buffer(_witness: &Self::Witness, _dst: &mut Vec<F>) {
+        unimplemented!("need to cast_into_source");
+    }
+}
+
 // We need this to ensure no conflicting implementations without negative impls
 
 #[derive(Derivative, Serialize, PartialEq)]
@@ -1199,7 +1211,7 @@ impl<F: SmallField, T: pairing::ff::PrimeField, const N: usize> WitnessCastable<
 }
 
 use crate::gadgets::traits::castable::Convertor;
-use crate::gadgets::traits::encodable::CircuitVarLengthEncodable;
+use crate::gadgets::traits::encodable::{CircuitVarLengthEncodable, WitnessVarLengthEncodable};
 
 impl<F: SmallField, T: pairing::ff::PrimeField, const N: usize> CSWitnessable<F, N>
     for NonNativeFieldOverU16<F, T, N>

diff --git a/crates/boojum/src/gadgets/tower_extension/fq12.rs b/crates/boojum/src/gadgets/tower_extension/fq12.rs
@@ -12,7 +12,7 @@ use super::{
 };
 
 use crate::gadgets::traits::allocatable::CSPlaceholder;
-use crate::gadgets::traits::encodable::CircuitVarLengthEncodable;
+use crate::gadgets::traits::encodable::{CircuitVarLengthEncodable, WitnessVarLengthEncodable};
 use crate::{
     cs::traits::cs::ConstraintSystem,
     field::SmallField,
@@ -507,6 +507,26 @@ where
     }
 }
 
+impl<F, T, NN, P> WitnessVarLengthEncodable<F> for Fq12<F, T, NN, P>
+where
+    F: SmallField,
+    T: PrimeField,
+    NN: NonNativeField<F, T> + WitnessVarLengthEncodable<F>,
+    P: Extension12Params<T>,
+{
+    fn witness_encoding_length(witness: &Self::Witness) -> usize {
+        let (c0, c1) = witness;
+        Fq6::<F, T, NN, P::Ex6>::witness_encoding_length(c0)
+            + Fq6::<F, T, NN, P::Ex6>::witness_encoding_length(c1)
+    }
+
+    fn encode_witness_to_buffer(witness: &Self::Witness, dst: &mut Vec<F>) {
+        let (c0, c1) = witness;
+        Fq6::<F, T, NN, P::Ex6>::encode_witness_to_buffer(c0, dst);
+        Fq6::<F, T, NN, P::Ex6>::encode_witness_to_buffer(c1, dst);
+    }
+}
+
 impl<F, T, NN, P> NonNativeField<F, T> for Fq12<F, T, NN, P>
 where
     F: SmallField,

diff --git a/crates/boojum/src/gadgets/tower_extension/fq2.rs b/crates/boojum/src/gadgets/tower_extension/fq2.rs
@@ -9,7 +9,7 @@ use super::params::{bn256::BN256Extension2Params, Extension2Params};
 
 use crate::cs::Variable;
 use crate::gadgets::traits::allocatable::CSPlaceholder;
-use crate::gadgets::traits::encodable::CircuitVarLengthEncodable;
+use crate::gadgets::traits::encodable::{CircuitVarLengthEncodable, WitnessVarLengthEncodable};
 use crate::{
     cs::traits::cs::ConstraintSystem,
     field::SmallField,
@@ -403,6 +403,25 @@ where
     }
 }
 
+impl<F, T, NN, P> WitnessVarLengthEncodable<F> for Fq2<F, T, NN, P>
+where
+    F: SmallField,
+    T: PrimeField,
+    NN: NonNativeField<F, T> + WitnessVarLengthEncodable<F>,
+    P: Extension2Params<T>,
+{
+    fn witness_encoding_length(witness: &Self::Witness) -> usize {
+        let (c0, c1) = witness;
+        NN::witness_encoding_length(c0) + NN::witness_encoding_length(c1)
+    }
+
+    fn encode_witness_to_buffer(witness: &Self::Witness, dst: &mut Vec<F>) {
+        let (c0, c1) = witness;
+        NN::encode_witness_to_buffer(c0, dst);
+        NN::encode_witness_to_buffer(c1, dst);
+    }
+}
+
 impl<F, T, NN, P> NonNativeField<F, T> for Fq2<F, T, NN, P>
 where
     F: SmallField,

diff --git a/crates/boojum/src/gadgets/tower_extension/fq6.rs b/crates/boojum/src/gadgets/tower_extension/fq6.rs
@@ -12,7 +12,7 @@ use super::{
 
 use crate::cs::Variable;
 use crate::gadgets::traits::allocatable::CSPlaceholder;
-use crate::gadgets::traits::encodable::CircuitVarLengthEncodable;
+use crate::gadgets::traits::encodable::{CircuitVarLengthEncodable, WitnessVarLengthEncodable};
 use crate::{
     cs::traits::cs::ConstraintSystem,
     field::SmallField,
@@ -594,6 +594,28 @@ where
     }
 }
 
+impl<F, T, NN, P> WitnessVarLengthEncodable<F> for Fq6<F, T, NN, P>
+where
+    F: SmallField,
+    T: PrimeField,
+    NN: NonNativeField<F, T> + WitnessVarLengthEncodable<F>,
+    P: Extension6Params<T>,
+{
+    fn witness_encoding_length(witness: &Self::Witness) -> usize {
+        let (c0, c1, c2) = witness;
+        Fq2::<F, T, NN, P::Ex2>::witness_encoding_length(c0)
+            + Fq2::<F, T, NN, P::Ex2>::witness_encoding_length(c1)
+            + Fq2::<F, T, NN, P::Ex2>::witness_encoding_length(c2)
+    }
+
+    fn encode_witness_to_buffer(witness: &Self::Witness, dst: &mut Vec<F>) {
+        let (c0, c1, c2) = witness;
+        Fq2::<F, T, NN, P::Ex2>::encode_witness_to_buffer(c0, dst);
+        Fq2::<F, T, NN, P::Ex2>::encode_witness_to_buffer(c1, dst);
+        Fq2::<F, T, NN, P::Ex2>::encode_witness_to_buffer(c2, dst);
+    }
+}
+
 impl<F, T, NN, P> NonNativeField<F, T> for Fq6<F, T, NN, P>
 where
     F: SmallField,