From 99cca29841e3ae3efa61f7fadef6fa7f6eb4466d Mon Sep 17 00:00:00 2001
From: MercysJest <wdnguyen@stanford.edu>
Date: Mon, 8 Jan 2024 11:03:41 -0800
Subject: [PATCH] Update Absorb impl to remove requirement that SWAffine and
 TEAffine's basefields are prime fields.

---
 src/sponge/absorb.rs | 36 +++++++++++++++++++++++++++---------
 1 file changed, 27 insertions(+), 9 deletions(-)

diff --git a/src/sponge/absorb.rs b/src/sponge/absorb.rs
index 5588f979..5a8e49df 100644
--- a/src/sponge/absorb.rs
+++ b/src/sponge/absorb.rs
@@ -5,7 +5,7 @@ use ark_ec::{
     twisted_edwards::TECurveConfig as TEModelParameters,
 };
 use ark_ff::models::{Fp, FpConfig};
-use ark_ff::{BigInteger, PrimeField, ToConstraintField};
+use ark_ff::{BigInteger, Field, PrimeField, ToConstraintField};
 use ark_serialize::CanonicalSerialize;
 use ark_std::vec::Vec;
 
@@ -227,26 +227,44 @@ impl Absorb for isize {
     }
 }
 
-impl<CF: PrimeField, P: TEModelParameters<BaseField = CF>> Absorb for TEAffine<P> {
+impl<P: TEModelParameters> Absorb for TEAffine<P>
+where
+    P::BaseField: ToConstraintField<<P::BaseField as Field>::BasePrimeField>,
+{
     fn to_sponge_bytes(&self, dest: &mut Vec<u8>) {
-        dest.append(&mut self.x.into_bigint().to_bytes_le());
-        dest.append(&mut self.y.into_bigint().to_bytes_le());
+        self.x
+            .to_field_elements()
+            .unwrap()
+            .into_iter()
+            .chain(self.y.to_field_elements().unwrap())
+            .for_each(|elem| {
+                dest.append(&mut elem.into_bigint().to_bytes_le());
+            });
     }
 
     fn to_sponge_field_elements<F: PrimeField>(&self, dest: &mut Vec<F>) {
-        field_cast::<P::BaseField, _>(&self.to_field_elements().unwrap(), dest).unwrap();
+        field_cast(&self.to_field_elements().unwrap(), dest).unwrap();
     }
 }
 
-impl<CF: PrimeField, P: SWModelParameters<BaseField = CF>> Absorb for SWAffine<P> {
+impl<P: SWModelParameters> Absorb for SWAffine<P>
+where
+    P::BaseField: ToConstraintField<<P::BaseField as Field>::BasePrimeField>,
+{
     fn to_sponge_bytes(&self, dest: &mut Vec<u8>) {
-        dest.append(&mut self.x.into_bigint().to_bytes_le());
-        dest.append(&mut self.y.into_bigint().to_bytes_le());
+        self.x
+            .to_field_elements()
+            .unwrap()
+            .into_iter()
+            .chain(self.y.to_field_elements().unwrap())
+            .for_each(|elem| {
+                dest.append(&mut elem.into_bigint().to_bytes_le());
+            });
         dest.push(self.infinity.into());
     }
 
     fn to_sponge_field_elements<F: PrimeField>(&self, dest: &mut Vec<F>) {
-        field_cast::<P::BaseField, _>(&self.to_field_elements().unwrap(), dest).unwrap();
+        field_cast(&self.to_field_elements().unwrap(), dest).unwrap();
     }
 }