1use ark_ff::PrimeField;
  2use ark_r1cs_std::{
  3    alloc::{AllocVar, AllocationMode},
  4    fields::FieldVar,
  5};
  6use ark_relations::r1cs::{Namespace, SynthesisError};
  7use core::{borrow::Borrow, marker::PhantomData};
  8
  9use crate::utils::vec::SparseMatrix;
 10
 11pub fn mat_vec_mul_sparse<F: PrimeField, CF: PrimeField, FV: FieldVar<F, CF>>(
 12    m: SparseMatrixVar<F, CF, FV>,
 13    v: Vec<FV>,
 14) -> Vec<FV> {
 15    let mut res = vec![FV::zero(); m.n_rows];
 16    for (row_i, row) in m.coeffs.iter().enumerate() {
 17        for (value, col_i) in row.iter() {
 18            res[row_i] += value.clone().mul(&v[*col_i].clone());
 19        }
 20    }
 21    res
 22}
 23pub fn vec_add<F: PrimeField, CF: PrimeField, FV: FieldVar<F, CF>>(
 24    a: &Vec<FV>,
 25    b: &Vec<FV>,
 26) -> Result<Vec<FV>, SynthesisError> {
 27    if a.len() != b.len() {
 28        return Err(SynthesisError::Unsatisfiable);
 29    }
 30    let mut r: Vec<FV> = vec![FV::zero(); a.len()];
 31    for i in 0..a.len() {
 32        r[i] = a[i].clone() + b[i].clone();
 33    }
 34    Ok(r)
 35}
 36pub fn vec_scalar_mul<F: PrimeField, CF: PrimeField, FV: FieldVar<F, CF>>(
 37    vec: &Vec<FV>,
 38    c: &FV,
 39) -> Vec<FV> {
 40    let mut result = vec![FV::zero(); vec.len()];
 41    for (i, a) in vec.iter().enumerate() {
 42        result[i] = a.clone() * c;
 43    }
 44    result
 45}
 46pub fn hadamard<F: PrimeField, CF: PrimeField, FV: FieldVar<F, CF>>(
 47    a: &Vec<FV>,
 48    b: &Vec<FV>,
 49) -> Result<Vec<FV>, SynthesisError> {
 50    if a.len() != b.len() {
 51        return Err(SynthesisError::Unsatisfiable);
 52    }
 53    let mut r: Vec<FV> = vec![FV::zero(); a.len()];
 54    for i in 0..a.len() {
 55        r[i] = a[i].clone() * b[i].clone();
 56    }
 57    Ok(r)
 58}
 59
 60#[derive(Debug, Clone)]
 61pub struct SparseMatrixVar<F: PrimeField, CF: PrimeField, FV: FieldVar<F, CF>> {
 62    _f: PhantomData<F>,
 63    _cf: PhantomData<CF>,
 64    _fv: PhantomData<FV>,
 65    pub n_rows: usize,
 66    pub n_cols: usize,
 67    // same format as the native SparseMatrix (which follows ark_relations::r1cs::Matrix format
 68    pub coeffs: Vec<Vec<(FV, usize)>>,
 69}
 70
 71impl<F, CF, FV> AllocVar<SparseMatrix<F>, CF> for SparseMatrixVar<F, CF, FV>
 72where
 73    F: PrimeField,
 74    CF: PrimeField,
 75    FV: FieldVar<F, CF>,
 76{
 77    fn new_variable<T: Borrow<SparseMatrix<F>>>(
 78        cs: impl Into<Namespace<CF>>,
 79        f: impl FnOnce() -> Result<T, SynthesisError>,
 80        mode: AllocationMode,
 81    ) -> Result<Self, SynthesisError> {
 82        f().and_then(|val| {
 83            let cs = cs.into();
 84
 85            let mut coeffs: Vec<Vec<(FV, usize)>> = Vec::new();
 86            for row in val.borrow().coeffs.iter() {
 87                let mut rowVar: Vec<(FV, usize)> = Vec::new();
 88                for &(value, col_i) in row.iter() {
 89                    let coeffVar = FV::new_variable(cs.clone(), || Ok(value), mode)?;
 90                    rowVar.push((coeffVar, col_i));
 91                }
 92                coeffs.push(rowVar);
 93            }
 94
 95            Ok(Self {
 96                _f: PhantomData,
 97                _cf: PhantomData,
 98                _fv: PhantomData,
 99                n_rows: val.borrow().n_rows,
100                n_cols: val.borrow().n_cols,
101                coeffs,
102            })
103        })
104    }
105}