use crate::{
convert::{ToBitsGadget, ToBytesGadget, ToConstraintFieldGadget},
fields::{fp::FpVar, FieldOpsBounds, FieldVar},
prelude::*,
Vec,
};
use ark_ff::{
fields::{CubicExtField, Field},
CubicExtConfig, Zero,
};
use ark_relations::r1cs::{ConstraintSystemRef, Namespace, SynthesisError};
use core::{borrow::Borrow, marker::PhantomData};
use educe::Educe;
#[derive(Educe)]
#[educe(Debug, Clone)]
#[must_use]
pub struct CubicExtVar<BF: FieldVar<P::BaseField, P::BasePrimeField>, P: CubicExtVarConfig<BF>>
where
for<'a> &'a BF: FieldOpsBounds<'a, P::BaseField, BF>,
{
pub c0: BF,
pub c1: BF,
pub c2: BF,
#[educe(Debug(ignore))]
_params: PhantomData<P>,
}
pub trait CubicExtVarConfig<BF: FieldVar<Self::BaseField, Self::BasePrimeField>>:
CubicExtConfig
where
for<'a> &'a BF: FieldOpsBounds<'a, Self::BaseField, BF>,
{
fn mul_base_field_vars_by_frob_coeff(c1: &mut BF, c2: &mut BF, power: usize);
}
impl<BF: FieldVar<P::BaseField, P::BasePrimeField>, P: CubicExtVarConfig<BF>> CubicExtVar<BF, P>
where
for<'a> &'a BF: FieldOpsBounds<'a, P::BaseField, BF>,
{
#[inline]
pub fn new(c0: BF, c1: BF, c2: BF) -> Self {
let _params = PhantomData;
Self {
c0,
c1,
c2,
_params,
}
}
#[inline]
pub fn mul_base_field_by_nonresidue(fe: &BF) -> Result<BF, SynthesisError> {
Ok(fe * P::NONRESIDUE)
}
#[inline]
pub fn mul_by_base_field_constant(&self, fe: P::BaseField) -> Self {
let c0 = &self.c0 * fe;
let c1 = &self.c1 * fe;
let c2 = &self.c2 * fe;
Self::new(c0, c1, c2)
}
#[inline]
pub fn mul_assign_by_base_field_constant(&mut self, fe: P::BaseField) {
*self = (&*self).mul_by_base_field_constant(fe);
}
}
impl<BF, P> R1CSVar<P::BasePrimeField> for CubicExtVar<BF, P>
where
BF: FieldVar<P::BaseField, P::BasePrimeField>,
for<'a> &'a BF: FieldOpsBounds<'a, P::BaseField, BF>,
P: CubicExtVarConfig<BF>,
{
type Value = CubicExtField<P>;
fn cs(&self) -> ConstraintSystemRef<P::BasePrimeField> {
[&self.c0, &self.c1, &self.c2].cs()
}
#[inline]
fn value(&self) -> Result<Self::Value, SynthesisError> {
match (self.c0.value(), self.c1.value(), self.c2.value()) {
(Ok(c0), Ok(c1), Ok(c2)) => Ok(CubicExtField::new(c0, c1, c2)),
(..) => Err(SynthesisError::AssignmentMissing),
}
}
}
impl<BF, P> From<Boolean<P::BasePrimeField>> for CubicExtVar<BF, P>
where
BF: FieldVar<P::BaseField, P::BasePrimeField>,
for<'a> &'a BF: FieldOpsBounds<'a, P::BaseField, BF>,
P: CubicExtVarConfig<BF>,
{
fn from(other: Boolean<P::BasePrimeField>) -> Self {
let c0 = BF::from(other);
let c1 = BF::zero();
let c2 = BF::zero();
Self::new(c0, c1, c2)
}
}
impl<'a, BF, P> FieldOpsBounds<'a, CubicExtField<P>, CubicExtVar<BF, P>> for CubicExtVar<BF, P>
where
BF: FieldVar<P::BaseField, P::BasePrimeField>,
for<'b> &'b BF: FieldOpsBounds<'b, P::BaseField, BF>,
P: CubicExtVarConfig<BF>,
{
}
impl<'a, BF, P> FieldOpsBounds<'a, CubicExtField<P>, CubicExtVar<BF, P>> for &'a CubicExtVar<BF, P>
where
BF: FieldVar<P::BaseField, P::BasePrimeField>,
for<'b> &'b BF: FieldOpsBounds<'b, P::BaseField, BF>,
P: CubicExtVarConfig<BF>,
{
}
impl<BF, P> FieldVar<CubicExtField<P>, P::BasePrimeField> for CubicExtVar<BF, P>
where
BF: FieldVar<P::BaseField, P::BasePrimeField>,
for<'a> &'a BF: FieldOpsBounds<'a, P::BaseField, BF>,
P: CubicExtVarConfig<BF>,
{
fn constant(other: CubicExtField<P>) -> Self {
let c0 = BF::constant(other.c0);
let c1 = BF::constant(other.c1);
let c2 = BF::constant(other.c2);
Self::new(c0, c1, c2)
}
fn zero() -> Self {
let c0 = BF::zero();
let c1 = BF::zero();
let c2 = BF::zero();
Self::new(c0, c1, c2)
}
fn one() -> Self {
let c0 = BF::one();
let c1 = BF::zero();
let c2 = BF::zero();
Self::new(c0, c1, c2)
}
#[inline]
#[tracing::instrument(target = "r1cs")]
fn double(&self) -> Result<Self, SynthesisError> {
let c0 = self.c0.double()?;
let c1 = self.c1.double()?;
let c2 = self.c2.double()?;
Ok(Self::new(c0, c1, c2))
}
#[inline]
#[tracing::instrument(target = "r1cs")]
fn negate(&self) -> Result<Self, SynthesisError> {
let mut result = self.clone();
result.c0.negate_in_place()?;
result.c1.negate_in_place()?;
result.c2.negate_in_place()?;
Ok(result)
}
#[inline]
#[tracing::instrument(target = "r1cs")]
fn square(&self) -> Result<Self, SynthesisError> {
let a = self.c0.clone();
let b = self.c1.clone();
let c = self.c2.clone();
let s0 = a.square()?;
let ab = &a * &b;
let s1 = ab.double()?;
let s2 = (&a - &b + &c).square()?;
let s3 = (&b * &c).double()?;
let s4 = c.square()?;
let c0 = Self::mul_base_field_by_nonresidue(&s3)? + &s0;
let c1 = Self::mul_base_field_by_nonresidue(&s4)? + &s1;
let c2 = s1 + &s2 + &s3 - &s0 - &s4;
Ok(Self::new(c0, c1, c2))
}
#[tracing::instrument(target = "r1cs")]
fn mul_equals(&self, other: &Self, result: &Self) -> Result<(), SynthesisError> {
let v0 = &self.c0 * &other.c0;
let v1 = &self.c1 * &other.c1;
let v2 = &self.c2 * &other.c2;
let nr_a1_plus_a2 = (&self.c1 + &self.c2) * P::NONRESIDUE;
let b1_plus_b2 = &other.c1 + &other.c2;
let nr_v1 = &v1 * P::NONRESIDUE;
let nr_v2 = &v2 * P::NONRESIDUE;
let to_check = &result.c0 - &v0 + &nr_v1 + &nr_v2;
nr_a1_plus_a2.mul_equals(&b1_plus_b2, &to_check)?;
let a0_plus_a1 = &self.c0 + &self.c1;
let b0_plus_b1 = &other.c0 + &other.c1;
let to_check = &result.c1 - &nr_v2 + &v0 + &v1;
a0_plus_a1.mul_equals(&b0_plus_b1, &to_check)?;
let a0_plus_a2 = &self.c0 + &self.c2;
let b0_plus_b2 = &other.c0 + &other.c2;
let to_check = &result.c2 + &v0 - &v1 + &v2;
a0_plus_a2.mul_equals(&b0_plus_b2, &to_check)?;
Ok(())
}
#[tracing::instrument(target = "r1cs")]
fn frobenius_map(&self, power: usize) -> Result<Self, SynthesisError> {
let mut result = self.clone();
result.c0.frobenius_map_in_place(power)?;
result.c1.frobenius_map_in_place(power)?;
result.c2.frobenius_map_in_place(power)?;
P::mul_base_field_vars_by_frob_coeff(&mut result.c1, &mut result.c2, power);
Ok(result)
}
#[tracing::instrument(target = "r1cs")]
fn inverse(&self) -> Result<Self, SynthesisError> {
let mode = if self.is_constant() {
AllocationMode::Constant
} else {
AllocationMode::Witness
};
let inverse = Self::new_variable(
self.cs(),
|| {
self.value()
.map(|f| f.inverse().unwrap_or_else(CubicExtField::zero))
},
mode,
)?;
self.mul_equals(&inverse, &Self::one())?;
Ok(inverse)
}
}
impl_bounded_ops!(
CubicExtVar<BF, P>,
CubicExtField<P>,
Add,
add,
AddAssign,
add_assign,
|this: &'a CubicExtVar<BF, P>, other: &'a CubicExtVar<BF, P>| {
let c0 = &this.c0 + &other.c0;
let c1 = &this.c1 + &other.c1;
let c2 = &this.c2 + &other.c2;
CubicExtVar::new(c0, c1, c2)
},
|this: &'a CubicExtVar<BF, P>, other: CubicExtField<P>| {
this + CubicExtVar::constant(other)
},
(BF: FieldVar<P::BaseField, P::BasePrimeField>, P: CubicExtVarConfig<BF>),
for<'b> &'b BF: FieldOpsBounds<'b, P::BaseField, BF>,
);
impl_bounded_ops!(
CubicExtVar<BF, P>,
CubicExtField<P>,
Sub,
sub,
SubAssign,
sub_assign,
|this: &'a CubicExtVar<BF, P>, other: &'a CubicExtVar<BF, P>| {
let c0 = &this.c0 - &other.c0;
let c1 = &this.c1 - &other.c1;
let c2 = &this.c2 - &other.c2;
CubicExtVar::new(c0, c1, c2)
},
|this: &'a CubicExtVar<BF, P>, other: CubicExtField<P>| {
this - CubicExtVar::constant(other)
},
(BF: FieldVar<P::BaseField, P::BasePrimeField>, P: CubicExtVarConfig<BF>),
for<'b> &'b BF: FieldOpsBounds<'b, P::BaseField, BF>,
);
impl_bounded_ops!(
CubicExtVar<BF, P>,
CubicExtField<P>,
Mul,
mul,
MulAssign,
mul_assign,
|this: &'a CubicExtVar<BF, P>, other: &'a CubicExtVar<BF, P>| {
let v0 = &this.c0 * &other.c0;
let v1 = &this.c1 * &other.c1;
let v2 = &this.c2 * &other.c2;
let c0 =
(((&this.c1 + &this.c2) * (&other.c1 + &other.c2) - &v1 - &v2) * P::NONRESIDUE) + &v0 ;
let c1 =
(&this.c0 + &this.c1) * (&other.c0 + &other.c1) - &v0 - &v1 + (&v2 * P::NONRESIDUE);
let c2 =
(&this.c0 + &this.c2) * (&other.c0 + &other.c2) - &v0 + &v1 - &v2;
CubicExtVar::new(c0, c1, c2)
},
|this: &'a CubicExtVar<BF, P>, other: CubicExtField<P>| {
this * CubicExtVar::constant(other)
},
(BF: FieldVar<P::BaseField, P::BasePrimeField>, P: CubicExtVarConfig<BF>),
for<'b> &'b BF: FieldOpsBounds<'b, P::BaseField, BF>,
);
impl<BF, P> EqGadget<P::BasePrimeField> for CubicExtVar<BF, P>
where
BF: FieldVar<P::BaseField, P::BasePrimeField>,
for<'a> &'a BF: FieldOpsBounds<'a, P::BaseField, BF>,
P: CubicExtVarConfig<BF>,
{
#[tracing::instrument(target = "r1cs")]
fn is_eq(&self, other: &Self) -> Result<Boolean<P::BasePrimeField>, SynthesisError> {
let b0 = self.c0.is_eq(&other.c0)?;
let b1 = self.c1.is_eq(&other.c1)?;
let b2 = self.c2.is_eq(&other.c2)?;
Ok(b0 & b1 & b2)
}
#[inline]
#[tracing::instrument(target = "r1cs")]
fn conditional_enforce_equal(
&self,
other: &Self,
condition: &Boolean<P::BasePrimeField>,
) -> Result<(), SynthesisError> {
self.c0.conditional_enforce_equal(&other.c0, condition)?;
self.c1.conditional_enforce_equal(&other.c1, condition)?;
self.c2.conditional_enforce_equal(&other.c2, condition)?;
Ok(())
}
#[inline]
#[tracing::instrument(target = "r1cs")]
fn conditional_enforce_not_equal(
&self,
other: &Self,
condition: &Boolean<P::BasePrimeField>,
) -> Result<(), SynthesisError> {
let is_equal = self.is_eq(other)?;
(is_equal & condition).enforce_equal(&Boolean::FALSE)
}
}
impl<BF, P> ToBitsGadget<P::BasePrimeField> for CubicExtVar<BF, P>
where
BF: FieldVar<P::BaseField, P::BasePrimeField>,
for<'a> &'a BF: FieldOpsBounds<'a, P::BaseField, BF>,
P: CubicExtVarConfig<BF>,
{
#[tracing::instrument(target = "r1cs")]
fn to_bits_le(&self) -> Result<Vec<Boolean<P::BasePrimeField>>, SynthesisError> {
let mut c0 = self.c0.to_bits_le()?;
let mut c1 = self.c1.to_bits_le()?;
let mut c2 = self.c2.to_bits_le()?;
c0.append(&mut c1);
c0.append(&mut c2);
Ok(c0)
}
#[tracing::instrument(target = "r1cs")]
fn to_non_unique_bits_le(&self) -> Result<Vec<Boolean<P::BasePrimeField>>, SynthesisError> {
let mut c0 = self.c0.to_non_unique_bits_le()?;
let mut c1 = self.c1.to_non_unique_bits_le()?;
let mut c2 = self.c2.to_non_unique_bits_le()?;
c0.append(&mut c1);
c0.append(&mut c2);
Ok(c0)
}
}
impl<BF, P> ToBytesGadget<P::BasePrimeField> for CubicExtVar<BF, P>
where
BF: FieldVar<P::BaseField, P::BasePrimeField>,
for<'a> &'a BF: FieldOpsBounds<'a, P::BaseField, BF>,
P: CubicExtVarConfig<BF>,
{
#[tracing::instrument(target = "r1cs")]
fn to_bytes_le(&self) -> Result<Vec<UInt8<P::BasePrimeField>>, SynthesisError> {
let mut c0 = self.c0.to_bytes_le()?;
let mut c1 = self.c1.to_bytes_le()?;
let mut c2 = self.c2.to_bytes_le()?;
c0.append(&mut c1);
c0.append(&mut c2);
Ok(c0)
}
#[tracing::instrument(target = "r1cs")]
fn to_non_unique_bytes_le(&self) -> Result<Vec<UInt8<P::BasePrimeField>>, SynthesisError> {
let mut c0 = self.c0.to_non_unique_bytes_le()?;
let mut c1 = self.c1.to_non_unique_bytes_le()?;
let mut c2 = self.c2.to_non_unique_bytes_le()?;
c0.append(&mut c1);
c0.append(&mut c2);
Ok(c0)
}
}
impl<BF, P> ToConstraintFieldGadget<P::BasePrimeField> for CubicExtVar<BF, P>
where
BF: FieldVar<P::BaseField, P::BasePrimeField>,
for<'a> &'a BF: FieldOpsBounds<'a, P::BaseField, BF>,
P: CubicExtVarConfig<BF>,
BF: ToConstraintFieldGadget<P::BasePrimeField>,
{
#[tracing::instrument(target = "r1cs")]
fn to_constraint_field(&self) -> Result<Vec<FpVar<P::BasePrimeField>>, SynthesisError> {
let mut res = Vec::new();
res.extend_from_slice(&self.c0.to_constraint_field()?);
res.extend_from_slice(&self.c1.to_constraint_field()?);
res.extend_from_slice(&self.c2.to_constraint_field()?);
Ok(res)
}
}
impl<BF, P> CondSelectGadget<P::BasePrimeField> for CubicExtVar<BF, P>
where
BF: FieldVar<P::BaseField, P::BasePrimeField>,
for<'a> &'a BF: FieldOpsBounds<'a, P::BaseField, BF>,
P: CubicExtVarConfig<BF>,
{
#[inline]
#[tracing::instrument(target = "r1cs")]
fn conditionally_select(
cond: &Boolean<P::BasePrimeField>,
true_value: &Self,
false_value: &Self,
) -> Result<Self, SynthesisError> {
let c0 = BF::conditionally_select(cond, &true_value.c0, &false_value.c0)?;
let c1 = BF::conditionally_select(cond, &true_value.c1, &false_value.c1)?;
let c2 = BF::conditionally_select(cond, &true_value.c2, &false_value.c2)?;
Ok(Self::new(c0, c1, c2))
}
}
impl<BF, P> TwoBitLookupGadget<P::BasePrimeField> for CubicExtVar<BF, P>
where
BF: FieldVar<P::BaseField, P::BasePrimeField>
+ TwoBitLookupGadget<P::BasePrimeField, TableConstant = P::BaseField>,
for<'a> &'a BF: FieldOpsBounds<'a, P::BaseField, BF>,
P: CubicExtVarConfig<BF>,
{
type TableConstant = CubicExtField<P>;
#[tracing::instrument(target = "r1cs")]
fn two_bit_lookup(
b: &[Boolean<P::BasePrimeField>],
c: &[Self::TableConstant],
) -> Result<Self, SynthesisError> {
let c0s = c.iter().map(|f| f.c0).collect::<Vec<_>>();
let c1s = c.iter().map(|f| f.c1).collect::<Vec<_>>();
let c2s = c.iter().map(|f| f.c2).collect::<Vec<_>>();
let c0 = BF::two_bit_lookup(b, &c0s)?;
let c1 = BF::two_bit_lookup(b, &c1s)?;
let c2 = BF::two_bit_lookup(b, &c2s)?;
Ok(Self::new(c0, c1, c2))
}
}
impl<BF, P> ThreeBitCondNegLookupGadget<P::BasePrimeField> for CubicExtVar<BF, P>
where
BF: FieldVar<P::BaseField, P::BasePrimeField>
+ ThreeBitCondNegLookupGadget<P::BasePrimeField, TableConstant = P::BaseField>,
for<'a> &'a BF: FieldOpsBounds<'a, P::BaseField, BF>,
P: CubicExtVarConfig<BF>,
{
type TableConstant = CubicExtField<P>;
#[tracing::instrument(target = "r1cs")]
fn three_bit_cond_neg_lookup(
b: &[Boolean<P::BasePrimeField>],
b0b1: &Boolean<P::BasePrimeField>,
c: &[Self::TableConstant],
) -> Result<Self, SynthesisError> {
let c0s = c.iter().map(|f| f.c0).collect::<Vec<_>>();
let c1s = c.iter().map(|f| f.c1).collect::<Vec<_>>();
let c2s = c.iter().map(|f| f.c2).collect::<Vec<_>>();
let c0 = BF::three_bit_cond_neg_lookup(b, b0b1, &c0s)?;
let c1 = BF::three_bit_cond_neg_lookup(b, b0b1, &c1s)?;
let c2 = BF::three_bit_cond_neg_lookup(b, b0b1, &c2s)?;
Ok(Self::new(c0, c1, c2))
}
}
impl<BF, P> AllocVar<CubicExtField<P>, P::BasePrimeField> for CubicExtVar<BF, P>
where
BF: FieldVar<P::BaseField, P::BasePrimeField>,
for<'a> &'a BF: FieldOpsBounds<'a, P::BaseField, BF>,
P: CubicExtVarConfig<BF>,
{
fn new_variable<T: Borrow<CubicExtField<P>>>(
cs: impl Into<Namespace<P::BasePrimeField>>,
f: impl FnOnce() -> Result<T, SynthesisError>,
mode: AllocationMode,
) -> Result<Self, SynthesisError> {
let ns = cs.into();
let cs = ns.cs();
use SynthesisError::*;
let (c0, c1, c2) = match f() {
Ok(fe) => (Ok(fe.borrow().c0), Ok(fe.borrow().c1), Ok(fe.borrow().c2)),
Err(_) => (
Err(AssignmentMissing),
Err(AssignmentMissing),
Err(AssignmentMissing),
),
};
let c0 = BF::new_variable(ark_relations::ns!(cs, "c0"), || c0, mode)?;
let c1 = BF::new_variable(ark_relations::ns!(cs, "c1"), || c1, mode)?;
let c2 = BF::new_variable(ark_relations::ns!(cs, "c2"), || c2, mode)?;
Ok(Self::new(c0, c1, c2))
}
}