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