1  // Copyright (c) 2019-2025 Alpha-Delta Network Inc.
  2  // This file is part of the alphavm library.
  3  
  4  // Licensed under the Apache License, Version 2.0 (the "License");
  5  // you may not use this file except in compliance with the License.
  6  // You may obtain a copy of the License at:
  7  
  8  // http://www.apache.org/licenses/LICENSE-2.0
  9  
 10  // Unless required by applicable law or agreed to in writing, software
 11  // distributed under the License is distributed on an "AS IS" BASIS,
 12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 13  // See the License for the specific language governing permissions and
 14  // limitations under the License.
 15  
 16  use crate::{Field, LegendreSymbol, One, PrimeField, SquareRootField, Zero};
 17  use alphavm_utilities::{
 18      FromBytes,
 19      ToBits,
 20      ToBytes,
 21      rand::Uniform,
 22      serialize::{SerializationError, *},
 23  };
 24  
 25  use rand::{
 26      Rng,
 27      distributions::{Distribution, Standard},
 28  };
 29  use serde::{Deserialize, Serialize};
 30  use std::{
 31      cmp::{Ord, Ordering, PartialOrd},
 32      fmt::Debug,
 33      hash::Hash,
 34      io::{Read, Result as IoResult, Write},
 35      ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign},
 36  };
 37  
 38  pub trait Fp2Parameters:
 39      'static + Copy + Clone + Default + Debug + PartialEq + Eq + Hash + Serialize + for<'a> Deserialize<'a> + Send + Sync
 40  {
 41      type Fp: PrimeField;
 42  
 43      /// Coefficients for the Frobenius automorphism.
 44      const FROBENIUS_COEFF_FP2_C1: [Self::Fp; 2];
 45  
 46      const NONRESIDUE: Self::Fp;
 47  
 48      const QUADRATIC_NONRESIDUE: (Self::Fp, Self::Fp);
 49  
 50      #[inline(always)]
 51      fn mul_fp_by_nonresidue(fe: &Self::Fp) -> Self::Fp {
 52          Self::NONRESIDUE * fe
 53      }
 54  }
 55  
 56  #[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
 57  pub struct Fp2<P: Fp2Parameters> {
 58      pub c0: P::Fp,
 59      pub c1: P::Fp,
 60  }
 61  
 62  impl<P: Fp2Parameters> Fp2<P> {
 63      /// Initializes a `Fp2` element from two `Fp` elements.
 64      pub const fn new(c0: P::Fp, c1: P::Fp) -> Self {
 65          Fp2 { c0, c1 }
 66      }
 67  }
 68  
 69  impl<P: Fp2Parameters> Fp2<P> {
 70      /// Norm of Fp2 over Fp: Norm(a) = a.x^2 - beta * a.y^2
 71      pub fn norm(&self) -> P::Fp {
 72          let t0 = self.c0.square();
 73          let mut t1 = self.c1.square();
 74          t1 = -P::mul_fp_by_nonresidue(&t1);
 75          t1.add_assign(t0);
 76          t1
 77      }
 78  
 79      pub fn mul_by_fp(&mut self, element: &P::Fp) {
 80          self.c0.mul_assign(element);
 81          self.c1.mul_assign(element);
 82      }
 83  }
 84  
 85  impl<P: Fp2Parameters> Zero for Fp2<P> {
 86      fn zero() -> Self {
 87          Fp2::new(P::Fp::zero(), P::Fp::zero())
 88      }
 89  
 90      fn is_zero(&self) -> bool {
 91          self.c0.is_zero() && self.c1.is_zero()
 92      }
 93  }
 94  impl<P: Fp2Parameters> One for Fp2<P> {
 95      fn one() -> Self {
 96          Fp2::new(P::Fp::one(), P::Fp::zero())
 97      }
 98  
 99      fn is_one(&self) -> bool {
100          self.c0.is_one() && self.c1.is_zero()
101      }
102  }
103  
104  impl<P: Fp2Parameters> Field for Fp2<P> {
105      type BasePrimeField = P::Fp;
106  
107      fn from_base_prime_field(other: Self::BasePrimeField) -> Self {
108          Self::new(other, P::Fp::zero())
109      }
110  
111      #[inline]
112      fn characteristic<'a>() -> &'a [u64] {
113          P::Fp::characteristic()
114      }
115  
116      fn double(&self) -> Self {
117          let mut result = *self;
118          result.double_in_place();
119          result
120      }
121  
122      fn double_in_place(&mut self) {
123          self.c0.double_in_place();
124          self.c1.double_in_place();
125      }
126  
127      fn square(&self) -> Self {
128          let mut result = *self;
129          result.square_in_place();
130          result
131      }
132  
133      #[inline]
134      fn from_random_bytes_with_flags<F: Flags>(bytes: &[u8]) -> Option<(Self, F)> {
135          let split_at = bytes.len() / 2;
136          if let Some(c0) = P::Fp::from_random_bytes(&bytes[..split_at]) {
137              if let Some((c1, flags)) = P::Fp::from_random_bytes_with_flags::<F>(&bytes[split_at..]) {
138                  return Some((Fp2::new(c0, c1), flags));
139              }
140          }
141          None
142      }
143  
144      #[inline]
145      fn from_random_bytes(bytes: &[u8]) -> Option<Self> {
146          Self::from_random_bytes_with_flags::<EmptyFlags>(bytes).map(|f| f.0)
147      }
148  
149      fn square_in_place(&mut self) -> &mut Self {
150          // v0 = c0 - c1
151          let mut v0 = self.c0 - self.c1;
152          // v3 = c0 - beta * c1
153          let v3 = self.c0 - P::mul_fp_by_nonresidue(&self.c1);
154          // v2 = c0 * c1
155          let v2 = self.c0 * self.c1;
156  
157          // v0 = (v0 * v3) + v2
158          v0 *= &v3;
159          v0 += &v2;
160  
161          self.c1 = v2.double();
162          self.c0 = v0 + P::mul_fp_by_nonresidue(&v2);
163  
164          self
165      }
166  
167      fn inverse(&self) -> Option<Self> {
168          if self.is_zero() {
169              None
170          } else {
171              // Guide to Pairing-based Cryptography, Algorithm 5.19.
172              // v0 = c0.square()
173              let mut v0 = self.c0.square();
174              // v1 = c1.square()
175              let v1 = self.c1.square();
176              // v0 = v0 - beta * v1
177              v0 -= &P::mul_fp_by_nonresidue(&v1);
178              v0.inverse().map(|v1| {
179                  let c0 = self.c0 * v1;
180                  let c1 = -(self.c1 * v1);
181                  Self::new(c0, c1)
182              })
183          }
184      }
185  
186      fn inverse_in_place(&mut self) -> Option<&mut Self> {
187          if let Some(inverse) = self.inverse() {
188              *self = inverse;
189              Some(self)
190          } else {
191              None
192          }
193      }
194  
195      fn frobenius_map(&mut self, power: usize) {
196          self.c1.mul_assign(&P::FROBENIUS_COEFF_FP2_C1[power % 2]);
197      }
198  }
199  
200  impl<P: Fp2Parameters> SquareRootField for Fp2<P>
201  where
202      P::Fp: SquareRootField,
203  {
204      fn legendre(&self) -> LegendreSymbol {
205          self.norm().legendre()
206      }
207  
208      fn sqrt(&self) -> Option<Self> {
209          use crate::LegendreSymbol::*;
210          if self.c1.is_zero() {
211              return self.c0.sqrt().map(|c0| Self::new(c0, P::Fp::zero()));
212          }
213          match self.legendre() {
214              // Square root based on the complex method. See
215              // https://eprint.iacr.org/2012/685.pdf (page 15, algorithm 8)
216              Zero => Some(*self),
217              QuadraticNonResidue => None,
218              QuadraticResidue => {
219                  let two_inv = P::Fp::half();
220                  let alpha = self.norm().sqrt().expect("We are in the QR case, the norm should have a square root");
221                  let mut delta = (alpha + self.c0) * two_inv;
222                  if delta.legendre().is_qnr() {
223                      delta -= &alpha;
224                  }
225                  let c0 = delta.sqrt().expect("Delta must have a square root");
226                  let c0_inv = c0.inverse().expect("c0 must have an inverse");
227                  Some(Self::new(c0, self.c1 * two_inv * c0_inv))
228              }
229          }
230      }
231  
232      fn sqrt_in_place(&mut self) -> Option<&mut Self> {
233          (*self).sqrt().map(|sqrt| {
234              *self = sqrt;
235              self
236          })
237      }
238  }
239  
240  /// `Fp2` elements are ordered lexicographically.
241  impl<P: Fp2Parameters> Ord for Fp2<P> {
242      #[inline(always)]
243      fn cmp(&self, other: &Self) -> Ordering {
244          match self.c1.cmp(&other.c1) {
245              Ordering::Greater => Ordering::Greater,
246              Ordering::Less => Ordering::Less,
247              Ordering::Equal => self.c0.cmp(&other.c0),
248          }
249      }
250  }
251  
252  impl<P: Fp2Parameters> PartialOrd for Fp2<P> {
253      #[inline(always)]
254      fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
255          Some(self.cmp(other))
256      }
257  }
258  
259  impl<P: Fp2Parameters> From<u128> for Fp2<P> {
260      fn from(other: u128) -> Self {
261          Self::new(other.into(), P::Fp::zero())
262      }
263  }
264  
265  impl<P: Fp2Parameters> From<u64> for Fp2<P> {
266      fn from(other: u64) -> Self {
267          Self::new(other.into(), P::Fp::zero())
268      }
269  }
270  
271  impl<P: Fp2Parameters> From<u32> for Fp2<P> {
272      fn from(other: u32) -> Self {
273          Self::new(other.into(), P::Fp::zero())
274      }
275  }
276  
277  impl<P: Fp2Parameters> From<u16> for Fp2<P> {
278      fn from(other: u16) -> Self {
279          Self::new(other.into(), P::Fp::zero())
280      }
281  }
282  
283  impl<P: Fp2Parameters> From<u8> for Fp2<P> {
284      fn from(other: u8) -> Self {
285          Self::new(other.into(), P::Fp::zero())
286      }
287  }
288  
289  impl<P: Fp2Parameters> ToBits for Fp2<P> {
290      fn write_bits_le(&self, vec: &mut Vec<bool>) {
291          self.c0.write_bits_le(vec);
292          self.c1.write_bits_le(vec);
293      }
294  
295      fn write_bits_be(&self, vec: &mut Vec<bool>) {
296          self.c0.write_bits_be(vec);
297          self.c1.write_bits_be(vec);
298      }
299  }
300  
301  impl<P: Fp2Parameters> ToBytes for Fp2<P> {
302      #[inline]
303      fn write_le<W: Write>(&self, mut writer: W) -> IoResult<()> {
304          self.c0.write_le(&mut writer)?;
305          self.c1.write_le(writer)
306      }
307  }
308  
309  impl<P: Fp2Parameters> FromBytes for Fp2<P> {
310      #[inline]
311      fn read_le<R: Read>(mut reader: R) -> IoResult<Self> {
312          let c0 = P::Fp::read_le(&mut reader)?;
313          let c1 = P::Fp::read_le(reader)?;
314          Ok(Fp2::new(c0, c1))
315      }
316  }
317  
318  impl<P: Fp2Parameters> Neg for Fp2<P> {
319      type Output = Self;
320  
321      #[inline]
322      fn neg(self) -> Self {
323          let mut res = self;
324          res.c0 = res.c0.neg();
325          res.c1 = res.c1.neg();
326          res
327      }
328  }
329  
330  impl<P: Fp2Parameters> Distribution<Fp2<P>> for Standard {
331      #[inline]
332      fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Fp2<P> {
333          Fp2::new(Uniform::rand(rng), Uniform::rand(rng))
334      }
335  }
336  
337  impl_add_sub_from_field_ref!(Fp2, Fp2Parameters);
338  impl_mul_div_from_field_ref!(Fp2, Fp2Parameters);
339  
340  impl<P: Fp2Parameters> Add<&'_ Fp2<P>> for Fp2<P> {
341      type Output = Self;
342  
343      #[inline]
344      fn add(self, other: &Self) -> Self {
345          let mut result = self;
346          result.add_assign(other);
347          result
348      }
349  }
350  
351  impl<P: Fp2Parameters> Sub<&'_ Fp2<P>> for Fp2<P> {
352      type Output = Self;
353  
354      #[inline]
355      fn sub(self, other: &Self) -> Self {
356          let mut result = self;
357          result.sub_assign(&other);
358          result
359      }
360  }
361  
362  impl<P: Fp2Parameters> Mul<&'_ Fp2<P>> for Fp2<P> {
363      type Output = Self;
364  
365      #[inline]
366      fn mul(self, other: &Self) -> Self {
367          let mut result = self;
368          result.mul_assign(&other);
369          result
370      }
371  }
372  
373  impl<P: Fp2Parameters> Div<&'_ Fp2<P>> for Fp2<P> {
374      type Output = Self;
375  
376      #[inline]
377      fn div(self, other: &Self) -> Self {
378          let mut result = self;
379          result.mul_assign(&other.inverse().unwrap());
380          result
381      }
382  }
383  
384  impl<P: Fp2Parameters> AddAssign<&'_ Self> for Fp2<P> {
385      #[inline]
386      fn add_assign(&mut self, other: &Self) {
387          self.c0.add_assign(other.c0);
388          self.c1.add_assign(other.c1);
389      }
390  }
391  
392  impl<P: Fp2Parameters> SubAssign<&'_ Self> for Fp2<P> {
393      #[inline]
394      fn sub_assign(&mut self, other: &Self) {
395          self.c0.sub_assign(&other.c0);
396          self.c1.sub_assign(&other.c1);
397      }
398  }
399  
400  impl<P: Fp2Parameters> MulAssign<&'_ Self> for Fp2<P> {
401      #[inline]
402      #[allow(clippy::suspicious_op_assign_impl)]
403      fn mul_assign(&mut self, other: &Self) {
404          *self = Self::new(
405              P::Fp::sum_of_products([self.c0, P::mul_fp_by_nonresidue(&self.c1)].iter(), [other.c0, other.c1].iter()),
406              P::Fp::sum_of_products([self.c0, self.c1].iter(), [other.c1, other.c0].iter()),
407          )
408      }
409  }
410  
411  impl<P: Fp2Parameters> DivAssign<&'_ Self> for Fp2<P> {
412      #[inline]
413      fn div_assign(&mut self, other: &Self) {
414          self.mul_assign(&other.inverse().unwrap());
415      }
416  }
417  
418  impl<P: Fp2Parameters> std::fmt::Display for Fp2<P> {
419      fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
420          write!(f, "Fp2({} + {} * u)", self.c0, self.c1)
421      }
422  }
423  
424  impl<P: Fp2Parameters> CanonicalSerializeWithFlags for Fp2<P> {
425      #[inline]
426      fn serialize_with_flags<W: Write, F: Flags>(&self, mut writer: W, flags: F) -> Result<(), SerializationError> {
427          CanonicalSerialize::serialize_uncompressed(&self.c0, &mut writer)?;
428          self.c1.serialize_with_flags(&mut writer, flags)?;
429          Ok(())
430      }
431  
432      fn serialized_size_with_flags<F: Flags>(&self) -> usize {
433          self.c0.uncompressed_size() + self.c1.serialized_size_with_flags::<F>()
434      }
435  }
436  
437  impl<P: Fp2Parameters> CanonicalSerialize for Fp2<P> {
438      #[inline]
439      fn serialize_with_mode<W: Write>(&self, writer: W, _compress: Compress) -> Result<(), SerializationError> {
440          self.serialize_with_flags(writer, EmptyFlags)
441      }
442  
443      #[inline]
444      fn serialized_size(&self, compress: Compress) -> usize {
445          self.c0.serialized_size(compress) + self.c1.serialized_size(compress)
446      }
447  }
448  
449  impl<P: Fp2Parameters> CanonicalDeserializeWithFlags for Fp2<P> {
450      #[inline]
451      fn deserialize_with_flags<R: Read, F: Flags>(mut reader: R) -> Result<(Self, F), SerializationError> {
452          let c0: P::Fp = CanonicalDeserialize::deserialize_uncompressed(&mut reader)?;
453          let (c1, flags): (P::Fp, _) = CanonicalDeserializeWithFlags::deserialize_with_flags(&mut reader)?;
454          Ok((Fp2::new(c0, c1), flags))
455      }
456  }
457  impl<P: Fp2Parameters> Valid for Fp2<P> {
458      fn check(&self) -> Result<(), alphavm_utilities::SerializationError> {
459          Ok(())
460      }
461  
462      fn batch_check<'a>(
463          _batch: impl Iterator<Item = &'a Self> + Send,
464      ) -> Result<(), alphavm_utilities::SerializationError>
465      where
466          Self: 'a,
467      {
468          Ok(())
469      }
470  }
471  
472  impl<P: Fp2Parameters> CanonicalDeserialize for Fp2<P> {
473      #[inline]
474      fn deserialize_with_mode<R: Read>(
475          mut reader: R,
476          compress: Compress,
477          validate: Validate,
478      ) -> Result<Self, SerializationError> {
479          let c0 = P::Fp::deserialize_with_mode(&mut reader, compress, validate)?;
480          let c1 = P::Fp::deserialize_with_mode(&mut reader, compress, validate)?;
481          Ok(Fp2::new(c0, c1))
482      }
483  }