1use ark_ff::PrimeField;
  2use ark_r1cs_std::{
  3    alloc::{AllocVar, AllocationMode},
  4    eq::EqGadget,
  5    fields::{fp::FpVar, FieldVar},
  6    R1CSVar,
  7};
  8use ark_relations::r1cs::{Namespace, SynthesisError};
  9use core::borrow::Borrow;
 10
 11use crate::utils::vec::SparseMatrix;
 12
 13/// `EquivalenceGadget` enforces that two in-circuit variables are equivalent,
 14/// where the equivalence relation is parameterized by `M`:
 15/// - For `FpVar`, it is simply a equality relation, and `M` is unused.
 16/// - For `NonNativeUintVar`, we consider equivalence as a congruence relation,
 17///   in terms of modular arithmetic, so `M` specifies the modulus.
 18pub trait EquivalenceGadget<M> {
 19    fn enforce_equivalent(&self, other: &Self) -> Result<(), SynthesisError>;
 20}
 21impl<M, F: PrimeField> EquivalenceGadget<M> for FpVar<F> {
 22    fn enforce_equivalent(&self, other: &Self) -> Result<(), SynthesisError> {
 23        self.enforce_equal(other)
 24    }
 25}
 26impl<M, T: EquivalenceGadget<M>> EquivalenceGadget<M> for [T] {
 27    fn enforce_equivalent(&self, other: &Self) -> Result<(), SynthesisError> {
 28        self.iter()
 29            .zip(other)
 30            .try_for_each(|(a, b)| a.enforce_equivalent(b))
 31    }
 32}
 33
 34pub trait MatrixGadget<FV> {
 35    fn mul_vector(&self, v: &[FV]) -> Result<Vec<FV>, SynthesisError>;
 36}
 37
 38pub trait VectorGadget<FV> {
 39    fn add(&self, other: &Self) -> Result<Vec<FV>, SynthesisError>;
 40
 41    fn mul_scalar(&self, other: &FV) -> Result<Vec<FV>, SynthesisError>;
 42
 43    fn hadamard(&self, other: &Self) -> Result<Vec<FV>, SynthesisError>;
 44}
 45
 46impl<F: PrimeField> VectorGadget<FpVar<F>> for [FpVar<F>] {
 47    fn add(&self, other: &Self) -> Result<Vec<FpVar<F>>, SynthesisError> {
 48        if self.len() != other.len() {
 49            return Err(SynthesisError::Unsatisfiable);
 50        }
 51        Ok(self.iter().zip(other.iter()).map(|(a, b)| a + b).collect())
 52    }
 53
 54    fn mul_scalar(&self, c: &FpVar<F>) -> Result<Vec<FpVar<F>>, SynthesisError> {
 55        Ok(self.iter().map(|a| a * c).collect())
 56    }
 57
 58    fn hadamard(&self, other: &Self) -> Result<Vec<FpVar<F>>, SynthesisError> {
 59        if self.len() != other.len() {
 60            return Err(SynthesisError::Unsatisfiable);
 61        }
 62        Ok(self.iter().zip(other.iter()).map(|(a, b)| a * b).collect())
 63    }
 64}
 65
 66#[derive(Debug, Clone)]
 67pub struct SparseMatrixVar<FV> {
 68    pub n_rows: usize,
 69    pub n_cols: usize,
 70    // same format as the native SparseMatrix (which follows ark_relations::r1cs::Matrix format
 71    pub coeffs: Vec<Vec<(FV, usize)>>,
 72}
 73
 74impl<F, CF, FV> AllocVar<SparseMatrix<F>, CF> for SparseMatrixVar<FV>
 75where
 76    F: PrimeField,
 77    CF: PrimeField,
 78    FV: AllocVar<F, CF>,
 79{
 80    fn new_variable<T: Borrow<SparseMatrix<F>>>(
 81        cs: impl Into<Namespace<CF>>,
 82        f: impl FnOnce() -> Result<T, SynthesisError>,
 83        mode: AllocationMode,
 84    ) -> Result<Self, SynthesisError> {
 85        f().and_then(|val| {
 86            let cs = cs.into();
 87
 88            let mut coeffs: Vec<Vec<(FV, usize)>> = Vec::new();
 89            for row in val.borrow().coeffs.iter() {
 90                let mut rowVar: Vec<(FV, usize)> = Vec::new();
 91                for &(value, col_i) in row.iter() {
 92                    let coeffVar = FV::new_variable(cs.clone(), || Ok(value), mode)?;
 93                    rowVar.push((coeffVar, col_i));
 94                }
 95                coeffs.push(rowVar);
 96            }
 97
 98            Ok(Self {
 99                n_rows: val.borrow().n_rows,
100                n_cols: val.borrow().n_cols,
101                coeffs,
102            })
103        })
104    }
105}
106
107impl<F: PrimeField> MatrixGadget<FpVar<F>> for SparseMatrixVar<FpVar<F>> {
108    fn mul_vector(&self, v: &[FpVar<F>]) -> Result<Vec<FpVar<F>>, SynthesisError> {
109        Ok(self
110            .coeffs
111            .iter()
112            .map(|row| {
113                let products = row
114                    .iter()
115                    .map(|(value, col_i)| value * &v[*col_i])
116                    .collect::<Vec<_>>();
117                if products.is_constant() {
118                    FpVar::constant(products.value().unwrap_or_default().into_iter().sum())
119                } else {
120                    products.iter().sum()
121                }
122            })
123            .collect())
124    }
125}
126
127/// Interprets the given vector v as the evaluations of a dense multilinear extension of n_vars,
128/// and evaluates it at the given point. This method mimics the behavior of
129/// `utils/mle.rs#dense_vec_to_dense_mle` + `DenseMultilinearExtension::evaluate` but in R1CS
130/// constraints, since dense multilinear extensions are not supported in ark_r1cs_std.
131pub fn eval_mle<F: PrimeField>(
132    // n_vars indicates the number of variables in the MLE
133    n_vars: usize,
134    // v is the vector of the evaluations of the dense multilinear extension (MLE)
135    v: Vec<FpVar<F>>,
136    // point is the point at which we want to evaluate the MLE
137    point: Vec<FpVar<F>>,
138) -> FpVar<F> {
139    // pad to 2^n_vars
140    let mut poly: Vec<FpVar<F>> = [
141        v.to_owned(),
142        std::iter::repeat(FpVar::zero())
143            .take((1 << n_vars) - v.len())
144            .collect(),
145    ]
146    .concat();
147
148    for i in 1..n_vars + 1 {
149        let r = point[i - 1].clone();
150        for b in 0..(1 << (n_vars - 1)) {
151            let left = poly[b << 1].clone();
152            let right = poly[(b << 1) + 1].clone();
153            poly[b] = left.clone() + r.clone() * (right - left);
154        }
155    }
156    poly[0].clone()
157}