1use ark_ec::{AffineRepr, CurveGroup};
  2use ark_ff::{BigInteger, PrimeField};
  3use ark_r1cs_std::{
  4    alloc::{AllocVar, AllocationMode},
  5    boolean::Boolean,
  6    fields::{
  7        fp::FpVar,
  8        nonnative::{
  9            params::{get_params, OptimizationType},
 10            AllocatedNonNativeFieldVar, NonNativeFieldVar,
 11        },
 12        FieldVar,
 13    },
 14    ToBitsGadget, ToConstraintFieldGadget,
 15};
 16use ark_relations::r1cs::{ConstraintSystemRef, Namespace, OptimizationGoal, SynthesisError};
 17use ark_std::Zero;
 18use core::borrow::Borrow;
 19use std::marker::PhantomData;
 20
 21/// Compose a vector boolean into a `NonNativeFieldVar`
 22pub fn nonnative_field_var_from_le_bits<TargetField: PrimeField, BaseField: PrimeField>(
 23    cs: ConstraintSystemRef<BaseField>,
 24    bits: &[Boolean<BaseField>],
 25) -> Result<NonNativeFieldVar<TargetField, BaseField>, SynthesisError> {
 26    let params = get_params(
 27        TargetField::MODULUS_BIT_SIZE as usize,
 28        BaseField::MODULUS_BIT_SIZE as usize,
 29        match cs.optimization_goal() {
 30            OptimizationGoal::None => OptimizationType::Constraints,
 31            OptimizationGoal::Constraints => OptimizationType::Constraints,
 32            OptimizationGoal::Weight => OptimizationType::Weight,
 33        },
 34    );
 35
 36    // push the lower limbs first
 37    let mut limbs = bits
 38        .chunks(params.bits_per_limb)
 39        .map(Boolean::le_bits_to_fp_var)
 40        .collect::<Result<Vec<_>, _>>()?;
 41    limbs.resize(params.num_limbs, FpVar::zero());
 42    limbs.reverse();
 43
 44    Ok(AllocatedNonNativeFieldVar {
 45        cs,
 46        limbs,
 47        num_of_additions_over_normal_form: BaseField::one(),
 48        is_in_the_normal_form: false,
 49        target_phantom: PhantomData,
 50    }
 51    .into())
 52}
 53
 54/// A more efficient version of `NonNativeFieldVar::to_constraint_field`
 55pub fn nonnative_field_var_to_constraint_field<TargetField: PrimeField, BaseField: PrimeField>(
 56    f: &NonNativeFieldVar<TargetField, BaseField>,
 57) -> Result<Vec<FpVar<BaseField>>, SynthesisError> {
 58    let bits = f.to_bits_le()?;
 59
 60    let bits_per_limb = BaseField::MODULUS_BIT_SIZE as usize - 1;
 61    let num_limbs = (TargetField::MODULUS_BIT_SIZE as usize).div_ceil(bits_per_limb);
 62
 63    let mut limbs = bits
 64        .chunks(bits_per_limb)
 65        .map(|chunk| {
 66            let mut limb = FpVar::<BaseField>::zero();
 67            let mut w = BaseField::one();
 68            for b in chunk.iter() {
 69                limb += FpVar::from(b.clone()) * w;
 70                w.double_in_place();
 71            }
 72            limb
 73        })
 74        .collect::<Vec<FpVar<BaseField>>>();
 75    limbs.resize(num_limbs, FpVar::zero());
 76    limbs.reverse();
 77
 78    Ok(limbs)
 79}
 80
 81/// The out-circuit counterpart of `nonnative_field_var_to_constraint_field`
 82pub fn nonnative_field_to_field_elements<TargetField: PrimeField, BaseField: PrimeField>(
 83    f: &TargetField,
 84) -> Vec<BaseField> {
 85    let bits = f.into_bigint().to_bits_le();
 86
 87    let bits_per_limb = BaseField::MODULUS_BIT_SIZE as usize - 1;
 88    let num_limbs = (TargetField::MODULUS_BIT_SIZE as usize).div_ceil(bits_per_limb);
 89
 90    let mut limbs = bits
 91        .chunks(bits_per_limb)
 92        .map(|chunk| {
 93            let mut limb = BaseField::zero();
 94            let mut w = BaseField::one();
 95            for &b in chunk.iter() {
 96                limb += BaseField::from(b) * w;
 97                w.double_in_place();
 98            }
 99            limb
100        })
101        .collect::<Vec<BaseField>>();
102    limbs.resize(num_limbs, BaseField::zero());
103    limbs.reverse();
104
105    limbs
106}
107
108/// NonNativeAffineVar represents an elliptic curve point in Affine representation in the non-native
109/// field, over the constraint field. It is not intended to perform operations, but just to contain
110/// the affine coordinates in order to perform hash operations of the point.
111#[derive(Debug, Clone)]
112pub struct NonNativeAffineVar<C: CurveGroup>
113where
114    <C as ark_ec::CurveGroup>::BaseField: ark_ff::PrimeField,
115{
116    pub x: NonNativeFieldVar<C::BaseField, C::ScalarField>,
117    pub y: NonNativeFieldVar<C::BaseField, C::ScalarField>,
118}
119
120impl<C> AllocVar<C, C::ScalarField> for NonNativeAffineVar<C>
121where
122    C: CurveGroup,
123    <C as ark_ec::CurveGroup>::BaseField: ark_ff::PrimeField,
124{
125    fn new_variable<T: Borrow<C>>(
126        cs: impl Into<Namespace<C::ScalarField>>,
127        f: impl FnOnce() -> Result<T, SynthesisError>,
128        mode: AllocationMode,
129    ) -> Result<Self, SynthesisError> {
130        f().and_then(|val| {
131            let cs = cs.into();
132
133            let affine = val.borrow().into_affine();
134            let zero_point = (&C::BaseField::zero(), &C::BaseField::zero());
135            let xy = affine.xy().unwrap_or(zero_point);
136
137            let x = NonNativeFieldVar::<C::BaseField, C::ScalarField>::new_variable(
138                cs.clone(),
139                || Ok(xy.0),
140                mode,
141            )?;
142            let y = NonNativeFieldVar::<C::BaseField, C::ScalarField>::new_variable(
143                cs.clone(),
144                || Ok(xy.1),
145                mode,
146            )?;
147
148            Ok(Self { x, y })
149        })
150    }
151}
152
153impl<C: CurveGroup> ToConstraintFieldGadget<C::ScalarField> for NonNativeAffineVar<C>
154where
155    <C as ark_ec::CurveGroup>::BaseField: ark_ff::PrimeField,
156{
157    // A more efficient version of `point_to_nonnative_limbs_custom_opt`.
158    // Used for converting `NonNativeAffineVar` to a vector of `FpVar` with minimum length in
159    // the circuit.
160    fn to_constraint_field(&self) -> Result<Vec<FpVar<C::ScalarField>>, SynthesisError> {
161        let x = nonnative_field_var_to_constraint_field(&self.x)?;
162        let y = nonnative_field_var_to_constraint_field(&self.y)?;
163        Ok([x, y].concat())
164    }
165}
166
167/// The out-circuit counterpart of `NonNativeAffineVar::to_constraint_field`
168#[allow(clippy::type_complexity)]
169pub fn nonnative_affine_to_field_elements<C: CurveGroup>(
170    p: C,
171) -> Result<(Vec<C::ScalarField>, Vec<C::ScalarField>), SynthesisError>
172where
173    <C as ark_ec::CurveGroup>::BaseField: ark_ff::PrimeField,
174{
175    let affine = p.into_affine();
176    if affine.is_zero() {
177        let x = nonnative_field_to_field_elements(&C::BaseField::zero());
178        let y = nonnative_field_to_field_elements(&C::BaseField::zero());
179        return Ok((x, y));
180    }
181
182    let (x, y) = affine.xy().unwrap();
183    let x = nonnative_field_to_field_elements(x);
184    let y = nonnative_field_to_field_elements(y);
185    Ok((x, y))
186}
187
188impl<C: CurveGroup> NonNativeAffineVar<C>
189where
190    <C as ark_ec::CurveGroup>::BaseField: ark_ff::PrimeField,
191{
192    // A wrapper of `point_to_nonnative_limbs_custom_opt` with constraints-focused optimization
193    // type (which is the default optimization type for arkworks' Groth16).
194    // Used for extracting a list of field elements of type `C::ScalarField` from the public input
195    // `p`, in exactly the same way as how `NonNativeAffineVar` is represented as limbs of type
196    // `FpVar` in-circuit.
197    #[allow(clippy::type_complexity)]
198    pub fn inputize(p: C) -> Result<(Vec<C::ScalarField>, Vec<C::ScalarField>), SynthesisError> {
199        point_to_nonnative_limbs_custom_opt(p, OptimizationType::Constraints)
200    }
201}
202
203// Used to compute (outside the circuit) the limbs representation of a point.
204// For `OptimizationType::Constraints`, the result matches the one used in-circuit.
205// For `OptimizationType::Weight`, the result vector is more dense and is suitable for hashing.
206// It is possible to further optimize the length of the result vector (see
207// `nonnative_affine_to_field_elements`)
208#[allow(clippy::type_complexity)]
209fn point_to_nonnative_limbs_custom_opt<C: CurveGroup>(
210    p: C,
211    optimization_type: OptimizationType,
212) -> Result<(Vec<C::ScalarField>, Vec<C::ScalarField>), SynthesisError>
213where
214    <C as ark_ec::CurveGroup>::BaseField: ark_ff::PrimeField,
215{
216    let affine = p.into_affine();
217    if affine.is_zero() {
218        let x =
219            AllocatedNonNativeFieldVar::<C::BaseField, C::ScalarField>::get_limbs_representations(
220                &C::BaseField::zero(),
221                optimization_type,
222            )?;
223        let y =
224            AllocatedNonNativeFieldVar::<C::BaseField, C::ScalarField>::get_limbs_representations(
225                &C::BaseField::zero(),
226                optimization_type,
227            )?;
228        return Ok((x, y));
229    }
230
231    let (x, y) = affine.xy().unwrap();
232    let x = AllocatedNonNativeFieldVar::<C::BaseField, C::ScalarField>::get_limbs_representations(
233        x,
234        optimization_type,
235    )?;
236    let y = AllocatedNonNativeFieldVar::<C::BaseField, C::ScalarField>::get_limbs_representations(
237        y,
238        optimization_type,
239    )?;
240    Ok((x, y))
241}
242
243#[cfg(test)]
244mod tests {
245    use super::*;
246    use ark_pallas::{Fr, Projective};
247    use ark_r1cs_std::R1CSVar;
248    use ark_relations::r1cs::ConstraintSystem;
249    use ark_std::UniformRand;
250
251    #[test]
252    fn test_alloc_zero() {
253        let cs = ConstraintSystem::<Fr>::new_ref();
254
255        // dealing with the 'zero' point should not panic when doing the unwrap
256        let p = Projective::zero();
257        assert!(NonNativeAffineVar::<Projective>::new_witness(cs.clone(), || Ok(p)).is_ok());
258    }
259
260    #[test]
261    fn test_arkworks_to_constraint_field() {
262        let cs = ConstraintSystem::<Fr>::new_ref();
263
264        // check that point_to_nonnative_limbs returns the expected values
265        let mut rng = ark_std::test_rng();
266        let p = Projective::rand(&mut rng);
267        let pVar = NonNativeAffineVar::<Projective>::new_witness(cs.clone(), || Ok(p)).unwrap();
268        let (x, y) = point_to_nonnative_limbs_custom_opt(p, OptimizationType::Weight).unwrap();
269        assert_eq!(pVar.x.to_constraint_field().unwrap().value().unwrap(), x);
270        assert_eq!(pVar.y.to_constraint_field().unwrap().value().unwrap(), y);
271    }
272
273    #[test]
274    fn test_improved_to_constraint_field() {
275        let cs = ConstraintSystem::<Fr>::new_ref();
276
277        // check that point_to_nonnative_limbs returns the expected values
278        let mut rng = ark_std::test_rng();
279        let p = Projective::rand(&mut rng);
280        let pVar = NonNativeAffineVar::<Projective>::new_witness(cs.clone(), || Ok(p)).unwrap();
281        let (x, y) = nonnative_affine_to_field_elements(p).unwrap();
282        assert_eq!(
283            pVar.to_constraint_field().unwrap().value().unwrap(),
284            [x, y].concat()
285        );
286    }
287
288    #[test]
289    fn test_inputize() {
290        let cs = ConstraintSystem::<Fr>::new_ref();
291
292        // check that point_to_nonnative_limbs returns the expected values
293        let mut rng = ark_std::test_rng();
294        let p = Projective::rand(&mut rng);
295        let pVar = NonNativeAffineVar::<Projective>::new_witness(cs.clone(), || Ok(p)).unwrap();
296        let (x, y) = NonNativeAffineVar::inputize(p).unwrap();
297
298        match (pVar.x, pVar.y) {
299            (NonNativeFieldVar::Var(p_x), NonNativeFieldVar::Var(p_y)) => {
300                assert_eq!(p_x.limbs.value().unwrap(), x);
301                assert_eq!(p_y.limbs.value().unwrap(), y);
302            }
303            _ => unreachable!(),
304        }
305    }
306}