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