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, Fp2, Fp2Parameters, One, 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::Ordering,
 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 Fp6Parameters: 'static + Copy + Clone + Default + Debug + PartialEq + Eq + Hash + Send + Sync {
 39      type Fp2Params: Fp2Parameters;
 40  
 41      /// Coefficients for the Frobenius automorphism.
 42      const FROBENIUS_COEFF_FP6_C1: [Fp2<Self::Fp2Params>; 6];
 43      const FROBENIUS_COEFF_FP6_C2: [Fp2<Self::Fp2Params>; 6];
 44  
 45      const NONRESIDUE: Fp2<Self::Fp2Params>;
 46  
 47      #[inline(always)]
 48      fn mul_fp2_by_nonresidue(fe: &Fp2<Self::Fp2Params>) -> Fp2<Self::Fp2Params> {
 49          Self::NONRESIDUE * fe
 50      }
 51  }
 52  
 53  /// An element of Fp6, represented by c0 + c1 * v + c2 * v^(2).
 54  #[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
 55  pub struct Fp6<P: Fp6Parameters> {
 56      pub c0: Fp2<P::Fp2Params>,
 57      pub c1: Fp2<P::Fp2Params>,
 58      pub c2: Fp2<P::Fp2Params>,
 59  }
 60  
 61  impl<P: Fp6Parameters> Fp6<P> {
 62      /// Initializes an element of `Fp6` from 3 `Fp2` elements.
 63      pub const fn new(c0: Fp2<P::Fp2Params>, c1: Fp2<P::Fp2Params>, c2: Fp2<P::Fp2Params>) -> Self {
 64          Self { c0, c1, c2 }
 65      }
 66  
 67      pub fn mul_by_fp(&mut self, element: &<P::Fp2Params as Fp2Parameters>::Fp) {
 68          self.c0.mul_by_fp(element);
 69          self.c1.mul_by_fp(element);
 70          self.c2.mul_by_fp(element);
 71      }
 72  
 73      pub fn mul_by_fp2(&mut self, element: &Fp2<P::Fp2Params>) {
 74          self.c0.mul_assign(element);
 75          self.c1.mul_assign(element);
 76          self.c2.mul_assign(element);
 77      }
 78  
 79      pub fn mul_by_1(&mut self, c1: &Fp2<P::Fp2Params>) {
 80          let mut b_b = self.c1;
 81          b_b.mul_assign(c1);
 82  
 83          let mut t1 = *c1;
 84          {
 85              let mut tmp = self.c1;
 86              tmp.add_assign(self.c2);
 87  
 88              t1.mul_assign(&tmp);
 89              t1.sub_assign(&b_b);
 90              t1 = P::mul_fp2_by_nonresidue(&t1);
 91          }
 92  
 93          let mut t2 = *c1;
 94          {
 95              let mut tmp = self.c0;
 96              tmp.add_assign(self.c1);
 97  
 98              t2.mul_assign(&tmp);
 99              t2.sub_assign(&b_b);
100          }
101  
102          self.c0 = t1;
103          self.c1 = t2;
104          self.c2 = b_b;
105      }
106  
107      pub fn mul_by_01(&mut self, c0: &Fp2<P::Fp2Params>, c1: &Fp2<P::Fp2Params>) {
108          let mut a_a = self.c0;
109          let mut b_b = self.c1;
110          a_a.mul_assign(c0);
111          b_b.mul_assign(c1);
112  
113          let mut t1 = *c1;
114          {
115              let mut tmp = self.c1;
116              tmp.add_assign(self.c2);
117  
118              t1.mul_assign(&tmp);
119              t1.sub_assign(&b_b);
120              t1 = P::mul_fp2_by_nonresidue(&t1);
121              t1.add_assign(a_a);
122          }
123  
124          let mut t3 = *c0;
125          {
126              let mut tmp = self.c0;
127              tmp.add_assign(self.c2);
128  
129              t3.mul_assign(&tmp);
130              t3.sub_assign(&a_a);
131              t3.add_assign(b_b);
132          }
133  
134          let mut t2 = *c0;
135          t2.add_assign(c1);
136          {
137              let mut tmp = self.c0;
138              tmp.add_assign(self.c1);
139  
140              t2.mul_assign(&tmp);
141              t2.sub_assign(&a_a);
142              t2.sub_assign(&b_b);
143          }
144  
145          self.c0 = t1;
146          self.c1 = t2;
147          self.c2 = t3;
148      }
149  }
150  
151  impl<P: Fp6Parameters> Zero for Fp6<P> {
152      fn zero() -> Self {
153          Self::new(Fp2::zero(), Fp2::zero(), Fp2::zero())
154      }
155  
156      fn is_zero(&self) -> bool {
157          self.c0.is_zero() && self.c1.is_zero() && self.c2.is_zero()
158      }
159  }
160  
161  impl<P: Fp6Parameters> One for Fp6<P> {
162      fn one() -> Self {
163          Self::new(Fp2::one(), Fp2::zero(), Fp2::zero())
164      }
165  
166      fn is_one(&self) -> bool {
167          self.c0.is_one() && self.c1.is_zero() && self.c2.is_zero()
168      }
169  }
170  
171  impl<P: Fp6Parameters> Field for Fp6<P> {
172      type BasePrimeField = <Fp2<P::Fp2Params> as Field>::BasePrimeField;
173  
174      fn from_base_prime_field(other: Self::BasePrimeField) -> Self {
175          Self::new(Fp2::from_base_prime_field(other), Fp2::zero(), Fp2::zero())
176      }
177  
178      #[inline]
179      fn characteristic<'a>() -> &'a [u64] {
180          Fp2::<P::Fp2Params>::characteristic()
181      }
182  
183      fn double(&self) -> Self {
184          let mut result = *self;
185          result.double_in_place();
186          result
187      }
188  
189      fn double_in_place(&mut self) {
190          self.c0.double_in_place();
191          self.c1.double_in_place();
192          self.c2.double_in_place();
193      }
194  
195      #[inline]
196      fn from_random_bytes_with_flags<F: Flags>(bytes: &[u8]) -> Option<(Self, F)> {
197          let split_at = bytes.len() / 3;
198          if let Some(c0) = Fp2::<P::Fp2Params>::from_random_bytes(&bytes[..split_at]) {
199              if let Some(c1) = Fp2::<P::Fp2Params>::from_random_bytes(&bytes[split_at..2 * split_at]) {
200                  if let Some((c2, flags)) =
201                      Fp2::<P::Fp2Params>::from_random_bytes_with_flags::<F>(&bytes[2 * split_at..])
202                  {
203                      return Some((Fp6::new(c0, c1, c2), flags));
204                  }
205              }
206          }
207          None
208      }
209  
210      #[inline]
211      fn from_random_bytes(bytes: &[u8]) -> Option<Self> {
212          Self::from_random_bytes_with_flags::<EmptyFlags>(bytes).map(|f| f.0)
213      }
214  
215      fn square(&self) -> Self {
216          let mut result = *self;
217          result.square_in_place();
218          result
219      }
220  
221      fn square_in_place(&mut self) -> &mut Self {
222          let s0 = self.c0.square();
223          let s1 = (self.c0 * self.c1).double();
224          let s2 = (self.c0 - self.c1 + self.c2).square();
225          let s3 = (self.c1 * self.c2).double();
226          let s4 = self.c2.square();
227  
228          self.c0 = s0 + P::mul_fp2_by_nonresidue(&s3);
229          self.c1 = s1 + P::mul_fp2_by_nonresidue(&s4);
230          self.c2 = s1 + s2 + s3 - s0 - s4;
231  
232          self
233      }
234  
235      fn inverse(&self) -> Option<Self> {
236          if self.is_zero() {
237              None
238          } else {
239              let mut c0 = self.c2;
240              c0 = P::mul_fp2_by_nonresidue(&c0);
241              c0.mul_assign(&self.c1);
242              c0 = c0.neg();
243              {
244                  let mut c0s = self.c0;
245                  c0s.square_in_place();
246                  c0.add_assign(c0s);
247              }
248              let mut c1 = self.c2;
249              c1.square_in_place();
250              c1 = P::mul_fp2_by_nonresidue(&c1);
251              {
252                  let mut c01 = self.c0;
253                  c01.mul_assign(&self.c1);
254                  c1.sub_assign(&c01);
255              }
256              let mut c2 = self.c1;
257              c2.square_in_place();
258              {
259                  let mut c02 = self.c0;
260                  c02.mul_assign(&self.c2);
261                  c2.sub_assign(&c02);
262              }
263  
264              let mut tmp1 = self.c2;
265              tmp1.mul_assign(&c1);
266              let mut tmp2 = self.c1;
267              tmp2.mul_assign(&c2);
268              tmp1.add_assign(tmp2);
269              tmp1 = P::mul_fp2_by_nonresidue(&tmp1);
270              tmp2 = self.c0;
271              tmp2.mul_assign(&c0);
272              tmp1.add_assign(tmp2);
273  
274              tmp1.inverse().map(|t| Self::new(t * c0, t * c1, t * c2))
275          }
276      }
277  
278      fn inverse_in_place(&mut self) -> Option<&mut Self> {
279          if let Some(inverse) = self.inverse() {
280              *self = inverse;
281              Some(self)
282          } else {
283              None
284          }
285      }
286  
287      fn frobenius_map(&mut self, power: usize) {
288          self.c0.frobenius_map(power);
289          self.c1.frobenius_map(power);
290          self.c2.frobenius_map(power);
291  
292          self.c1.mul_assign(&P::FROBENIUS_COEFF_FP6_C1[power % 6]);
293          self.c2.mul_assign(&P::FROBENIUS_COEFF_FP6_C2[power % 6]);
294      }
295  }
296  
297  impl<P: Fp6Parameters> std::fmt::Display for Fp6<P> {
298      fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
299          write!(f, "Fq6_3over2({} + {} * v, {} * v^2)", self.c0, self.c1, self.c2)
300      }
301  }
302  
303  impl<P: Fp6Parameters> Distribution<Fp6<P>> for Standard {
304      #[inline]
305      fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Fp6<P> {
306          Fp6::new(Uniform::rand(rng), Uniform::rand(rng), Uniform::rand(rng))
307      }
308  }
309  
310  impl<P: Fp6Parameters> Neg for Fp6<P> {
311      type Output = Self;
312  
313      #[inline]
314      fn neg(self) -> Self {
315          let mut copy = Self::zero();
316          copy.c0 = self.c0.neg();
317          copy.c1 = self.c1.neg();
318          copy.c2 = self.c2.neg();
319          copy
320      }
321  }
322  
323  impl_add_sub_from_field_ref!(Fp6, Fp6Parameters);
324  impl_mul_div_from_field_ref!(Fp6, Fp6Parameters);
325  
326  impl<P: Fp6Parameters> Add<&'_ Self> for Fp6<P> {
327      type Output = Self;
328  
329      #[inline]
330      fn add(self, other: &Self) -> Self {
331          let mut result = self;
332          result.add_assign(other);
333          result
334      }
335  }
336  
337  impl<P: Fp6Parameters> Sub<&'_ Self> for Fp6<P> {
338      type Output = Self;
339  
340      #[inline]
341      fn sub(self, other: &Self) -> Self {
342          let mut result = self;
343          result.sub_assign(&other);
344          result
345      }
346  }
347  
348  impl<P: Fp6Parameters> Mul<&'_ Self> for Fp6<P> {
349      type Output = Self;
350  
351      #[inline]
352      fn mul(self, other: &Self) -> Self {
353          let mut result = self;
354          result.mul_assign(&other);
355          result
356      }
357  }
358  
359  impl<P: Fp6Parameters> Div<&'_ Self> for Fp6<P> {
360      type Output = Self;
361  
362      #[inline]
363      fn div(self, other: &Self) -> Self {
364          let mut result = self;
365          result.mul_assign(&other.inverse().unwrap());
366          result
367      }
368  }
369  
370  impl<P: Fp6Parameters> AddAssign<&'_ Self> for Fp6<P> {
371      #[inline]
372      fn add_assign(&mut self, other: &Self) {
373          self.c0 += &other.c0;
374          self.c1 += &other.c1;
375          self.c2 += &other.c2;
376      }
377  }
378  
379  impl<P: Fp6Parameters> SubAssign<&'_ Self> for Fp6<P> {
380      #[inline]
381      fn sub_assign(&mut self, other: &Self) {
382          self.c0 -= &other.c0;
383          self.c1 -= &other.c1;
384          self.c2 -= &other.c2;
385      }
386  }
387  
388  impl<P: Fp6Parameters> MulAssign<&'_ Self> for Fp6<P> {
389      #[inline]
390      fn mul_assign(&mut self, other: &Self) {
391          let v0 = self.c0 * other.c0;
392          let v1 = self.c1 * other.c1;
393          let v2 = self.c2 * other.c2;
394  
395          let c0 = P::mul_fp2_by_nonresidue(&((self.c1 + self.c2) * (other.c1 + other.c2) - v1 - v2)) + v0;
396          let c1 = (self.c0 + self.c1) * (other.c0 + other.c1) - v0 - v1 + P::mul_fp2_by_nonresidue(&v2);
397          let c2 = (self.c0 + self.c2) * (other.c0 + other.c2) - v0 - v2 + v1;
398  
399          self.c0 = c0;
400          self.c1 = c1;
401          self.c2 = c2;
402      }
403  }
404  
405  impl<P: Fp6Parameters> DivAssign<&'_ Self> for Fp6<P> {
406      #[inline]
407      fn div_assign(&mut self, other: &Self) {
408          self.mul_assign(&other.inverse().unwrap());
409      }
410  }
411  
412  impl<P: Fp6Parameters> From<&'_ [bool]> for Fp6<P> {
413      fn from(_bits: &[bool]) -> Self {
414          unimplemented!()
415      }
416  }
417  
418  /// `Fp3` elements are ordered lexicographically.
419  impl<P: Fp6Parameters> Ord for Fp6<P> {
420      #[inline(always)]
421      fn cmp(&self, other: &Self) -> Ordering {
422          let c2_cmp = self.c2.cmp(&other.c2);
423          let c1_cmp = self.c1.cmp(&other.c1);
424          let c0_cmp = self.c0.cmp(&other.c0);
425          match c2_cmp {
426              Ordering::Equal => match c1_cmp {
427                  Ordering::Equal => c0_cmp,
428                  _ => c1_cmp,
429              },
430              _ => c2_cmp,
431          }
432      }
433  }
434  
435  impl<P: Fp6Parameters> PartialOrd for Fp6<P> {
436      #[inline(always)]
437      fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
438          Some(self.cmp(other))
439      }
440  }
441  
442  impl<P: Fp6Parameters> From<u128> for Fp6<P> {
443      fn from(other: u128) -> Self {
444          Self::new(other.into(), Fp2::zero(), Fp2::zero())
445      }
446  }
447  
448  impl<P: Fp6Parameters> From<u64> for Fp6<P> {
449      fn from(other: u64) -> Self {
450          Self::new(other.into(), Fp2::zero(), Fp2::zero())
451      }
452  }
453  
454  impl<P: Fp6Parameters> From<u32> for Fp6<P> {
455      fn from(other: u32) -> Self {
456          Self::new(other.into(), Fp2::zero(), Fp2::zero())
457      }
458  }
459  
460  impl<P: Fp6Parameters> From<u16> for Fp6<P> {
461      fn from(other: u16) -> Self {
462          Self::new(other.into(), Fp2::zero(), Fp2::zero())
463      }
464  }
465  
466  impl<P: Fp6Parameters> From<u8> for Fp6<P> {
467      fn from(other: u8) -> Self {
468          Self::new(other.into(), Fp2::zero(), Fp2::zero())
469      }
470  }
471  
472  impl<P: Fp6Parameters> ToBits for Fp6<P> {
473      fn write_bits_le(&self, vec: &mut Vec<bool>) {
474          self.c0.write_bits_le(vec);
475          self.c1.write_bits_le(vec);
476          self.c2.write_bits_le(vec);
477      }
478  
479      fn write_bits_be(&self, vec: &mut Vec<bool>) {
480          self.c0.write_bits_be(vec);
481          self.c1.write_bits_be(vec);
482          self.c2.write_bits_be(vec);
483      }
484  }
485  
486  impl<P: Fp6Parameters> ToBytes for Fp6<P> {
487      #[inline]
488      fn write_le<W: Write>(&self, mut writer: W) -> IoResult<()> {
489          self.c0.write_le(&mut writer)?;
490          self.c1.write_le(&mut writer)?;
491          self.c2.write_le(&mut writer)
492      }
493  }
494  
495  impl<P: Fp6Parameters> FromBytes for Fp6<P> {
496      #[inline]
497      fn read_le<R: Read>(mut reader: R) -> IoResult<Self> {
498          let c0 = Fp2::read_le(&mut reader)?;
499          let c1 = Fp2::read_le(&mut reader)?;
500          let c2 = Fp2::read_le(&mut reader)?;
501          Ok(Fp6::new(c0, c1, c2))
502      }
503  }
504  
505  impl<P: Fp6Parameters> CanonicalSerializeWithFlags for Fp6<P> {
506      #[inline]
507      fn serialize_with_flags<W: Write, F: Flags>(&self, mut writer: W, flags: F) -> Result<(), SerializationError> {
508          CanonicalSerialize::serialize_uncompressed(&self.c0, &mut writer)?;
509          CanonicalSerialize::serialize_uncompressed(&self.c1, &mut writer)?;
510          self.c2.serialize_with_flags(writer, flags)?;
511          Ok(())
512      }
513  
514      fn serialized_size_with_flags<F: Flags>(&self) -> usize {
515          self.c0.uncompressed_size() + self.c1.uncompressed_size() + self.c2.serialized_size_with_flags::<F>()
516      }
517  }
518  
519  impl<P: Fp6Parameters> CanonicalSerialize for Fp6<P> {
520      #[inline]
521      fn serialize_with_mode<W: Write>(&self, writer: W, _compress: Compress) -> Result<(), SerializationError> {
522          self.serialize_with_flags(writer, EmptyFlags)
523      }
524  
525      #[inline]
526      fn serialized_size(&self, compress: Compress) -> usize {
527          self.c0.serialized_size(compress) + self.c1.serialized_size(compress) + self.c2.serialized_size(compress)
528      }
529  }
530  
531  impl<P: Fp6Parameters> Valid for Fp6<P> {
532      fn check(&self) -> Result<(), alphavm_utilities::SerializationError> {
533          Ok(())
534      }
535  
536      fn batch_check<'a>(_batch: impl Iterator<Item = &'a Self>) -> Result<(), alphavm_utilities::SerializationError>
537      where
538          Self: 'a,
539      {
540          Ok(())
541      }
542  }
543  
544  impl<P: Fp6Parameters> CanonicalDeserializeWithFlags for Fp6<P> {
545      #[inline]
546      fn deserialize_with_flags<R: Read, F: Flags>(mut reader: R) -> Result<(Self, F), SerializationError> {
547          let c0 = CanonicalDeserialize::deserialize_uncompressed(&mut reader)?;
548          let c1 = CanonicalDeserialize::deserialize_uncompressed(&mut reader)?;
549          let (c2, flags): (_, _) = CanonicalDeserializeWithFlags::deserialize_with_flags(&mut reader)?;
550          Ok((Fp6::new(c0, c1, c2), flags))
551      }
552  }
553  
554  impl<P: Fp6Parameters> CanonicalDeserialize for Fp6<P> {
555      #[inline]
556      fn deserialize_with_mode<R: Read>(
557          mut reader: R,
558          compress: Compress,
559          validate: Validate,
560      ) -> Result<Self, SerializationError> {
561          let c0 = CanonicalDeserialize::deserialize_with_mode(&mut reader, compress, validate)?;
562          let c1 = CanonicalDeserialize::deserialize_with_mode(&mut reader, compress, validate)?;
563          let c2 = CanonicalDeserialize::deserialize_with_mode(&mut reader, compress, validate)?;
564          Ok(Fp6::new(c0, c1, c2))
565      }
566  }