From 6806fe4cdfef58e5c6a85c64240b2d22132443d5 Mon Sep 17 00:00:00 2001
From: Nikita Masych <nikitamasich152@gmail.com>
Date: Wed, 4 Sep 2024 14:36:54 +0300
Subject: [PATCH] feat(boojum): added WitnessVarLengthEncodable implementation
 to tower of extensions

---
 .../implementations/implementation_u16.rs     | 14 ++++++++++-
 .../src/gadgets/tower_extension/fq12.rs       | 22 ++++++++++++++++-
 .../boojum/src/gadgets/tower_extension/fq2.rs | 21 +++++++++++++++-
 .../boojum/src/gadgets/tower_extension/fq6.rs | 24 ++++++++++++++++++-
 4 files changed, 77 insertions(+), 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
index 5425e28..a56fa86 100644
--- 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
index bc9f68a..998b4f8 100644
--- 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
index 93b0e18..24b5955 100644
--- 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
index e68ca98..5409247 100644
--- 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,