1use ark_ec::{short_weierstrass::SWFlags, AffineRepr, CurveGroup};
  2use ark_ff::{Field, PrimeField};
  3use ark_r1cs_std::{
  4    alloc::{AllocVar, AllocationMode},
  5    eq::EqGadget,
  6    fields::fp::FpVar,
  7    prelude::Boolean,
  8    R1CSVar, ToConstraintFieldGadget,
  9};
 10use ark_relations::r1cs::{ConstraintSystemRef, Namespace, SynthesisError};
 11use ark_serialize::{CanonicalSerialize, CanonicalSerializeWithFlags};
 12use ark_std::Zero;
 13use core::borrow::Borrow;
 14
 15use crate::{
 16    folding::traits::Inputize,
 17    transcript::{AbsorbNonNative, AbsorbNonNativeGadget},
 18};
 19
 20use super::uint::{nonnative_field_to_field_elements, NonNativeUintVar};
 21
 22/// NonNativeAffineVar represents an elliptic curve point in Affine representation in the non-native
 23/// field, over the constraint field. It is not intended to perform operations, but just to contain
 24/// the affine coordinates in order to perform hash operations of the point.
 25#[derive(Debug, Clone)]
 26pub struct NonNativeAffineVar<C: CurveGroup> {
 27    pub x: NonNativeUintVar<C::ScalarField>,
 28    pub y: NonNativeUintVar<C::ScalarField>,
 29}
 30
 31impl<C> AllocVar<C, C::ScalarField> for NonNativeAffineVar<C>
 32where
 33    C: CurveGroup,
 34{
 35    fn new_variable<T: Borrow<C>>(
 36        cs: impl Into<Namespace<C::ScalarField>>,
 37        f: impl FnOnce() -> Result<T, SynthesisError>,
 38        mode: AllocationMode,
 39    ) -> Result<Self, SynthesisError> {
 40        f().and_then(|val| {
 41            let cs = cs.into();
 42
 43            let affine = val.borrow().into_affine();
 44            let zero = (&C::BaseField::zero(), &C::BaseField::zero());
 45            let (x, y) = affine.xy().unwrap_or(zero);
 46
 47            let x = NonNativeUintVar::new_variable(cs.clone(), || Ok(*x), mode)?;
 48            let y = NonNativeUintVar::new_variable(cs.clone(), || Ok(*y), mode)?;
 49
 50            Ok(Self { x, y })
 51        })
 52    }
 53}
 54
 55impl<C: CurveGroup> R1CSVar<C::ScalarField> for NonNativeAffineVar<C> {
 56    type Value = C;
 57
 58    fn cs(&self) -> ConstraintSystemRef<C::ScalarField> {
 59        self.x.cs().or(self.y.cs())
 60    }
 61
 62    fn value(&self) -> Result<Self::Value, SynthesisError> {
 63        debug_assert_eq!(C::BaseField::extension_degree(), 1);
 64
 65        let x = <C::BaseField as Field>::BasePrimeField::from_le_bytes_mod_order(
 66            &self.x.value()?.to_bytes_le(),
 67        );
 68        let y = <C::BaseField as Field>::BasePrimeField::from_le_bytes_mod_order(
 69            &self.y.value()?.to_bytes_le(),
 70        );
 71        // Below is a workaround to convert the `x` and `y` coordinates to a
 72        // point. This is because the `CurveGroup` trait does not provide a
 73        // method to construct a point from `BaseField` elements.
 74        let mut bytes = vec![];
 75        // `unwrap` below is safe because serialization of a `PrimeField` value
 76        // only fails if the serialization flag has more than 8 bits, but here
 77        // we call `serialize_uncompressed` which uses an empty flag.
 78        x.serialize_uncompressed(&mut bytes).unwrap();
 79        // `unwrap` below is also safe, because the bit size of `SWFlags` is 2.
 80        y.serialize_with_flags(
 81            &mut bytes,
 82            if x.is_zero() && y.is_zero() {
 83                SWFlags::PointAtInfinity
 84            } else if y <= -y {
 85                SWFlags::YIsPositive
 86            } else {
 87                SWFlags::YIsNegative
 88            },
 89        )
 90        .unwrap();
 91        // `unwrap` below is safe because `bytes` is constructed from the `x`
 92        // and `y` coordinates of a valid point, and these coordinates are
 93        // serialized in the same way as the `CurveGroup` implementation.
 94        Ok(C::deserialize_uncompressed_unchecked(&bytes[..]).unwrap())
 95    }
 96}
 97
 98impl<C: CurveGroup> ToConstraintFieldGadget<C::ScalarField> for NonNativeAffineVar<C> {
 99    // Used for converting `NonNativeAffineVar` to a vector of `FpVar` with minimum length in
100    // the circuit.
101    fn to_constraint_field(&self) -> Result<Vec<FpVar<C::ScalarField>>, SynthesisError> {
102        let x = self.x.to_constraint_field()?;
103        let y = self.y.to_constraint_field()?;
104        Ok([x, y].concat())
105    }
106}
107
108impl<C: CurveGroup> EqGadget<C::ScalarField> for NonNativeAffineVar<C> {
109    fn is_eq(&self, other: &Self) -> Result<Boolean<C::ScalarField>, SynthesisError> {
110        let mut result = Boolean::TRUE;
111        if self.x.0.len() != other.x.0.len() {
112            return Err(SynthesisError::Unsatisfiable);
113        }
114        if self.y.0.len() != other.y.0.len() {
115            return Err(SynthesisError::Unsatisfiable);
116        }
117        for (l, r) in self
118            .x
119            .0
120            .iter()
121            .chain(&self.y.0)
122            .zip(other.x.0.iter().chain(&other.y.0))
123        {
124            if l.ub != r.ub {
125                return Err(SynthesisError::Unsatisfiable);
126            }
127            result = result.and(&l.v.is_eq(&r.v)?)?;
128        }
129        Ok(result)
130    }
131
132    fn enforce_equal(&self, other: &Self) -> Result<(), SynthesisError> {
133        if self.x.0.len() != other.x.0.len() {
134            return Err(SynthesisError::Unsatisfiable);
135        }
136        if self.y.0.len() != other.y.0.len() {
137            return Err(SynthesisError::Unsatisfiable);
138        }
139        for (l, r) in self
140            .x
141            .0
142            .iter()
143            .chain(&self.y.0)
144            .zip(other.x.0.iter().chain(&other.y.0))
145        {
146            if l.ub != r.ub {
147                return Err(SynthesisError::Unsatisfiable);
148            }
149            l.v.enforce_equal(&r.v)?;
150        }
151        Ok(())
152    }
153}
154
155/// The out-circuit counterpart of `NonNativeAffineVar::to_constraint_field`
156#[allow(clippy::type_complexity)]
157pub(crate) fn nonnative_affine_to_field_elements<C: CurveGroup>(
158    p: C,
159) -> (Vec<C::ScalarField>, Vec<C::ScalarField>) {
160    let affine = p.into_affine();
161    let zero = (&C::BaseField::zero(), &C::BaseField::zero());
162    let (x, y) = affine.xy().unwrap_or(zero);
163
164    let x = nonnative_field_to_field_elements(x);
165    let y = nonnative_field_to_field_elements(y);
166    (x, y)
167}
168
169impl<C: CurveGroup> Inputize<C::ScalarField, NonNativeAffineVar<C>> for C {
170    fn inputize(&self) -> Vec<C::ScalarField> {
171        let affine = self.into_affine();
172        let zero = (&C::BaseField::zero(), &C::BaseField::zero());
173        let (x, y) = affine.xy().unwrap_or(zero);
174
175        let x = x.inputize();
176        let y = y.inputize();
177        [x, y].concat()
178    }
179}
180
181impl<C: CurveGroup> NonNativeAffineVar<C> {
182    pub fn zero() -> Self {
183        // `unwrap` below is safe because we are allocating a constant value,
184        // which is guaranteed to succeed.
185        Self::new_constant(ConstraintSystemRef::None, C::zero()).unwrap()
186    }
187}
188
189impl<C: CurveGroup> AbsorbNonNative<C::ScalarField> for C {
190    fn to_native_sponge_field_elements(&self, dest: &mut Vec<C::ScalarField>) {
191        let (x, y) = nonnative_affine_to_field_elements(*self);
192        dest.extend(x);
193        dest.extend(y);
194    }
195}
196
197impl<C: CurveGroup> AbsorbNonNativeGadget<C::ScalarField> for NonNativeAffineVar<C> {
198    fn to_native_sponge_field_elements(
199        &self,
200    ) -> Result<Vec<FpVar<C::ScalarField>>, SynthesisError> {
201        self.to_constraint_field()
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208    use ark_pallas::{Fr, Projective};
209    use ark_relations::r1cs::ConstraintSystem;
210    use ark_std::UniformRand;
211
212    #[test]
213    fn test_alloc_zero() {
214        let cs = ConstraintSystem::<Fr>::new_ref();
215
216        // dealing with the 'zero' point should not panic when doing the unwrap
217        let p = Projective::zero();
218        assert!(NonNativeAffineVar::<Projective>::new_witness(cs.clone(), || Ok(p)).is_ok());
219    }
220
221    #[test]
222    fn test_improved_to_constraint_field() {
223        let cs = ConstraintSystem::<Fr>::new_ref();
224
225        // check that point_to_nonnative_limbs returns the expected values
226        let mut rng = ark_std::test_rng();
227        let p = Projective::rand(&mut rng);
228        let pVar = NonNativeAffineVar::<Projective>::new_witness(cs.clone(), || Ok(p)).unwrap();
229        let (x, y) = nonnative_affine_to_field_elements(p);
230        assert_eq!(
231            pVar.to_constraint_field().unwrap().value().unwrap(),
232            [x, y].concat()
233        );
234    }
235
236    #[test]
237    fn test_inputize() {
238        let cs = ConstraintSystem::<Fr>::new_ref();
239
240        // check that point_to_nonnative_limbs returns the expected values
241        let mut rng = ark_std::test_rng();
242        let p = Projective::rand(&mut rng);
243        let pVar = NonNativeAffineVar::<Projective>::new_witness(cs.clone(), || Ok(p)).unwrap();
244        let xy = p.inputize();
245
246        assert_eq!(
247            [pVar.x.0.value().unwrap(), pVar.y.0.value().unwrap()].concat(),
248            xy
249        );
250    }
251}