1use super::{
  2    params::{get_params, OptimizationType},
  3    reduce::{bigint_to_basefield, limbs_to_bigint, Reducer},
  4    AllocatedEmulatedFpVar,
  5};
  6use crate::{fields::fp::FpVar, prelude::*};
  7use ark_ff::PrimeField;
  8use ark_relations::{
  9    ns,
 10    r1cs::{ConstraintSystemRef, OptimizationGoal, Result as R1CSResult},
 11};
 12use ark_std::{marker::PhantomData, vec::Vec};
 13use num_bigint::BigUint;
 14
 15/// The allocated form of `MulResultVar` (introduced below)
 16#[derive(Debug)]
 17#[must_use]
 18pub struct AllocatedMulResultVar<TargetF: PrimeField, BaseF: PrimeField> {
 19    /// Constraint system reference
 20    pub cs: ConstraintSystemRef<BaseF>,
 21    /// Limbs of the intermediate representations
 22    pub limbs: Vec<FpVar<BaseF>>,
 23    /// The cumulative num of additions
 24    pub prod_of_num_of_additions: BaseF,
 25    #[doc(hidden)]
 26    pub target_phantom: PhantomData<TargetF>,
 27}
 28
 29impl<TargetF: PrimeField, BaseF: PrimeField> From<&AllocatedEmulatedFpVar<TargetF, BaseF>>
 30    for AllocatedMulResultVar<TargetF, BaseF>
 31{
 32    fn from(src: &AllocatedEmulatedFpVar<TargetF, BaseF>) -> Self {
 33        let params = get_params(
 34            TargetF::MODULUS_BIT_SIZE as usize,
 35            BaseF::MODULUS_BIT_SIZE as usize,
 36            src.get_optimization_type(),
 37        );
 38
 39        let mut limbs = src.limbs.clone();
 40        limbs.reverse();
 41        limbs.resize(2 * params.num_limbs - 1, FpVar::<BaseF>::zero());
 42        limbs.reverse();
 43
 44        let prod_of_num_of_additions = src.num_of_additions_over_normal_form + &BaseF::one();
 45
 46        Self {
 47            cs: src.cs(),
 48            limbs,
 49            prod_of_num_of_additions,
 50            target_phantom: PhantomData,
 51        }
 52    }
 53}
 54
 55impl<TargetF: PrimeField, BaseF: PrimeField> AllocatedMulResultVar<TargetF, BaseF> {
 56    /// Get the CS
 57    pub fn cs(&self) -> ConstraintSystemRef<BaseF> {
 58        self.cs.clone()
 59    }
 60
 61    /// Get the value of the multiplication result
 62    pub fn value(&self) -> R1CSResult<TargetF> {
 63        let params = get_params(
 64            TargetF::MODULUS_BIT_SIZE as usize,
 65            BaseF::MODULUS_BIT_SIZE as usize,
 66            self.get_optimization_type(),
 67        );
 68
 69        let p_representations =
 70            AllocatedEmulatedFpVar::<TargetF, BaseF>::get_limbs_representations_from_big_integer(
 71                &<TargetF as PrimeField>::MODULUS,
 72                self.get_optimization_type(),
 73            )?;
 74        let p_bigint = limbs_to_bigint(params.bits_per_limb, &p_representations);
 75
 76        let mut limbs_values = Vec::<BaseF>::new();
 77        for limb in self.limbs.iter() {
 78            limbs_values.push(limb.value().unwrap_or_default());
 79        }
 80        let value_bigint = limbs_to_bigint(params.bits_per_limb, &limbs_values);
 81
 82        let res = bigint_to_basefield::<TargetF>(&(value_bigint % p_bigint));
 83        Ok(res)
 84    }
 85
 86    /// Constraints for reducing the result of a multiplication mod p, to get an
 87    /// original representation.
 88    pub fn reduce(&self) -> R1CSResult<AllocatedEmulatedFpVar<TargetF, BaseF>> {
 89        let params = get_params(
 90            TargetF::MODULUS_BIT_SIZE as usize,
 91            BaseF::MODULUS_BIT_SIZE as usize,
 92            self.get_optimization_type(),
 93        );
 94
 95        // Step 1: get p
 96        let p_representations =
 97            AllocatedEmulatedFpVar::<TargetF, BaseF>::get_limbs_representations_from_big_integer(
 98                &<TargetF as PrimeField>::MODULUS,
 99                self.get_optimization_type(),
100            )?;
101        let p_bigint = limbs_to_bigint(params.bits_per_limb, &p_representations);
102
103        let mut p_gadget_limbs = Vec::new();
104        for limb in p_representations.iter() {
105            p_gadget_limbs.push(FpVar::<BaseF>::new_constant(self.cs(), limb)?);
106        }
107        let p_gadget = AllocatedEmulatedFpVar::<TargetF, BaseF> {
108            cs: self.cs(),
109            limbs: p_gadget_limbs,
110            num_of_additions_over_normal_form: BaseF::one(),
111            is_in_the_normal_form: false,
112            target_phantom: PhantomData,
113        };
114
115        // Step 2: compute surfeit
116        let surfeit = overhead!(self.prod_of_num_of_additions + BaseF::one()) + 1 + 1;
117
118        // Step 3: allocate k
119        let k_bits = {
120            let mut res = Vec::new();
121
122            let mut limbs_values = Vec::<BaseF>::new();
123            for limb in self.limbs.iter() {
124                limbs_values.push(limb.value().unwrap_or_default());
125            }
126
127            let value_bigint = limbs_to_bigint(params.bits_per_limb, &limbs_values);
128            let mut k_cur = value_bigint / p_bigint;
129
130            let total_len = TargetF::MODULUS_BIT_SIZE as usize + surfeit;
131
132            for _ in 0..total_len {
133                res.push(Boolean::<BaseF>::new_witness(self.cs(), || {
134                    Ok(&k_cur % 2u64 == BigUint::from(1u64))
135                })?);
136                k_cur /= 2u64;
137            }
138            res
139        };
140
141        let k_limbs = {
142            let zero = FpVar::Constant(BaseF::zero());
143            let mut limbs = Vec::new();
144
145            let mut k_bits_cur = k_bits.clone();
146
147            for i in 0..params.num_limbs {
148                let this_limb_size = if i != params.num_limbs - 1 {
149                    params.bits_per_limb
150                } else {
151                    k_bits.len() - (params.num_limbs - 1) * params.bits_per_limb
152                };
153
154                let this_limb_bits = k_bits_cur[0..this_limb_size].to_vec();
155                k_bits_cur = k_bits_cur[this_limb_size..].to_vec();
156
157                let mut limb = zero.clone();
158                let mut cur = BaseF::one();
159
160                for bit in this_limb_bits.iter() {
161                    limb += &(FpVar::<BaseF>::from(bit.clone()) * cur);
162                    cur.double_in_place();
163                }
164                limbs.push(limb);
165            }
166
167            limbs.reverse();
168            limbs
169        };
170
171        let k_gadget = AllocatedEmulatedFpVar::<TargetF, BaseF> {
172            cs: self.cs(),
173            limbs: k_limbs,
174            num_of_additions_over_normal_form: self.prod_of_num_of_additions,
175            is_in_the_normal_form: false,
176            target_phantom: PhantomData,
177        };
178
179        let cs = self.cs();
180
181        let r_gadget = AllocatedEmulatedFpVar::<TargetF, BaseF>::new_witness(ns!(cs, "r"), || {
182            Ok(self.value()?)
183        })?;
184
185        let params = get_params(
186            TargetF::MODULUS_BIT_SIZE as usize,
187            BaseF::MODULUS_BIT_SIZE as usize,
188            self.get_optimization_type(),
189        );
190
191        // Step 1: reduce `self` and `other` if neceessary
192        let mut prod_limbs = Vec::new();
193        let zero = FpVar::<BaseF>::zero();
194
195        for _ in 0..2 * params.num_limbs - 1 {
196            prod_limbs.push(zero.clone());
197        }
198
199        for i in 0..params.num_limbs {
200            for j in 0..params.num_limbs {
201                prod_limbs[i + j] = &prod_limbs[i + j] + (&p_gadget.limbs[i] * &k_gadget.limbs[j]);
202            }
203        }
204
205        let mut kp_plus_r_gadget = Self {
206            cs,
207            limbs: prod_limbs,
208            prod_of_num_of_additions: (p_gadget.num_of_additions_over_normal_form + BaseF::one())
209                * (k_gadget.num_of_additions_over_normal_form + BaseF::one()),
210            target_phantom: PhantomData,
211        };
212
213        let kp_plus_r_limbs_len = kp_plus_r_gadget.limbs.len();
214        for (i, limb) in r_gadget.limbs.iter().rev().enumerate() {
215            kp_plus_r_gadget.limbs[kp_plus_r_limbs_len - 1 - i] += limb;
216        }
217
218        Reducer::<TargetF, BaseF>::group_and_check_equality(
219            surfeit,
220            2 * params.bits_per_limb,
221            params.bits_per_limb,
222            &self.limbs,
223            &kp_plus_r_gadget.limbs,
224        )?;
225
226        Ok(r_gadget)
227    }
228
229    /// Add unreduced elements.
230    #[tracing::instrument(target = "r1cs")]
231    pub fn add(&self, other: &Self) -> R1CSResult<Self> {
232        assert_eq!(self.get_optimization_type(), other.get_optimization_type());
233
234        let mut new_limbs = Vec::new();
235
236        for (l1, l2) in self.limbs.iter().zip(other.limbs.iter()) {
237            let new_limb = l1 + l2;
238            new_limbs.push(new_limb);
239        }
240
241        Ok(Self {
242            cs: self.cs(),
243            limbs: new_limbs,
244            prod_of_num_of_additions: self.prod_of_num_of_additions
245                + other.prod_of_num_of_additions,
246            target_phantom: PhantomData,
247        })
248    }
249
250    /// Add native constant elem
251    #[tracing::instrument(target = "r1cs")]
252    pub fn add_constant(&self, other: &TargetF) -> R1CSResult<Self> {
253        let mut other_limbs = AllocatedEmulatedFpVar::<TargetF, BaseF>::get_limbs_representations(
254            other,
255            self.get_optimization_type(),
256        )?;
257        other_limbs.reverse();
258
259        let mut new_limbs = Vec::new();
260
261        for (i, limb) in self.limbs.iter().rev().enumerate() {
262            if i < other_limbs.len() {
263                new_limbs.push(limb + other_limbs[i]);
264            } else {
265                new_limbs.push((*limb).clone());
266            }
267        }
268
269        new_limbs.reverse();
270
271        Ok(Self {
272            cs: self.cs(),
273            limbs: new_limbs,
274            prod_of_num_of_additions: self.prod_of_num_of_additions + BaseF::one(),
275            target_phantom: PhantomData,
276        })
277    }
278
279    pub(crate) fn get_optimization_type(&self) -> OptimizationType {
280        match self.cs().optimization_goal() {
281            OptimizationGoal::None => OptimizationType::Constraints,
282            OptimizationGoal::Constraints => OptimizationType::Constraints,
283            OptimizationGoal::Weight => OptimizationType::Weight,
284        }
285    }
286}